def translation_align__given_unshifted_fft(v1f, v2f): cor = fftshift(N.real(ifftn(v1f * N.conj(v2f)))) mid_co = IVU.fft_mid_co(cor.shape) loc = N.unravel_index(cor.argmax(), cor.shape) return {'loc': (loc - mid_co), 'cor': cor[loc[0], loc[1], loc[2]]}
def bandpass_denoising(a, name, save_flag=False): b_time = time.time() grid = GV.grid_displacement_to_center(a.shape, GV.fft_mid_co(a.shape)) rad = GV.grid_distance_to_center(grid) rad = np.round(rad).astype(np.int) # create a mask that only center frequencies components will be left curve = np.zeros(rad.shape) # TODO: change the curve value as desired curve[int(rad.shape[0] / 8) * 3:int(rad.shape[0] / 8) * 5, int(rad.shape[1] / 8) * 3:int(rad.shape[1] / 8) * 5, int(rad.shape[2] / 8) * 3:int(rad.shape[2] / 8) * 5] = 1 #perform FFT and filter the data with the mask and then transform the filtered data back vf = ifftn(ifftshift((fftshift(fftn(a)) * curve))) vf = np.real(vf) end_time = time.time() print('Bandpass de-noise takes', end_time - b_time, 's') if save_flag: img = (vf[:, :, int(vf.shape[2] / 2)]).copy() # TODO: Change the image and tomogram saving path img_path = '/Users/apple/Desktop/Lab/Zach_Project/Denoising_Result/Bandpass/' + str( name) + '_BP.png' plt.imsave(img_path, img, cmap='gray') mrc_path = '/Users/apple/Desktop/Lab/Zach_Project/Denoising_Result/Bandpass/' + str( name) + '_BP.mrc' io_file.put_mrc_data(vf, mrc_path) return img
def tilt_mask(size, tilt_ang1, tilt_ang2=None, tilt_axis=1, light_axis=2, sphere_mask_bool=True): assert tilt_axis != light_axis if tilt_ang2 is None: tilt_ang2 = float(N.abs(tilt_ang1)) tilt_ang1 = -tilt_ang2 else: assert tilt_ang1 < 0 assert tilt_ang2 > 0 tilt_ang1 = (tilt_ang1 / 180.0) * N.pi tilt_ang2 = (tilt_ang2 / 180.0) * N.pi g = IVU.grid_displacement_to_center(size=size, mid_co=IVU.fft_mid_co(siz=size)) plane_axis = set([0,1,2]) plane_axis.difference_update([light_axis,tilt_axis]) assert len(plane_axis) == 1 plane_axis = list(plane_axis)[0] x_light = g[light_axis] x_plane = g[plane_axis] m = N.zeros(size, dtype=float) m[ N.logical_and(x_light <= (N.tan(tilt_ang1) * x_plane), x_light >= (N.tan(tilt_ang2) * x_plane))] = 1.0 m[ N.logical_and(x_light >= (N.tan(tilt_ang1) * x_plane), x_light <= (N.tan(tilt_ang2) * x_plane))] = 1.0 if sphere_mask_bool: m *= MU.sphere_mask(m.shape) return m
def fsc(v1, v2, band_width_radius=1.0): """Fourier Shell correlation between two volumes""" siz = v1.shape assert (siz == v2.shape) origin_co = GV.fft_mid_co(siz) x = N.mgrid[0:siz[0], 0:siz[1], 0:siz[2]] x = x.astype(N.float) for dim_i in range(3): x[dim_i] -= origin_co[dim_i] rad = N.sqrt(N.square(x).sum(axis=0)) vol_rad = int(N.floor(N.min(siz) / 2.0) + 1) v1f = NF.fftshift(NF.fftn(v1)) v2f = NF.fftshift(NF.fftn(v2)) fsc_cors = N.zeros(vol_rad) # the interpolation can also be performed using scipy.ndimage.interpolation.map_coordinates() for r in range(vol_rad): ind = (abs(rad - r) <= band_width_radius) c1 = v1f[ind] c2 = v2f[ind] fsc_cor_t = N.sum(c1 * N.conj(c2)) / N.sqrt(N.sum(N.abs(c1) ** 2) * N.sum(N.abs(c2) ** 2)) fsc_cors[r] = N.real(fsc_cor_t) return fsc_cors
def wedge_mask(size, ang1, ang2=None, tilt_axis=1, sphere_mask=True, verbose=False): # should define both tilt axis and electron beam (missing wedge) direction warnings.warn("The definition of wedge mask is still ambiguous") if ang2 is None: ang2 = float(N.abs(ang1)) ang1 = -ang2 else: assert ang1 < 0 assert ang2 > 0 if verbose: print('image.vol.wedge.util.wedge_mask()', 'ang1', ang1, 'ang2', ang2, 'tilt_axis', tilt_axis, 'sphere_mask', sphere_mask) ang1 = (ang1 / 180.0) * N.pi ang2 = (ang2 / 180.0) * N.pi g = AIVU.grid_displacement_to_center(size=size, mid_co=AIVU.fft_mid_co(siz=size)) if tilt_axis == 0: # y-z plane # y axis x0 = g[1] # z axis x1 = g[2] elif tilt_axis == 1: # x-z plane # x axis x0 = g[0] # z axis x1 = g[2] elif tilt_axis == 2: # x-y plane # x axis x0 = g[0] # y axis x1 = g[1] m = N.zeros(size, dtype=float) m[N.logical_and(x0 >= (N.tan(ang2) * x1), x0 >= (N.tan(ang1) * x1))] = 1.0 m[N.logical_and(x0 <= (N.tan(ang1) * x1), x0 <= (N.tan(ang2) * x1))] = 1.0 if sphere_mask: m *= MU.sphere_mask(m.shape) return m
def filter_given_curve(v, curve): grid = GV.grid_displacement_to_center(v.shape, GV.fft_mid_co(v.shape)) rad = GV.grid_distance_to_center(grid) rad = N.round(rad).astype(N.int) b = N.zeros(rad.shape) for (i, a) in enumerate(curve): b[(rad == i)] = a vf = ifftn(ifftshift((fftshift(fftn(v)) * b))) vf = N.real(vf) return vf
def ssnr__get_rad(siz): grid = GV.grid_displacement_to_center(siz, GV.fft_mid_co(siz)) rad = GV.grid_distance_to_center(grid) return rad
def rotate_mask(v, angle=None, rm=None): c1 = IVU.fft_mid_co(v.shape) c2 = N.copy(c1) vr = rotate(v, angle=angle, rm=rm, c1=c1, c2=c2, default_val=float('NaN')) vr[N.logical_not(N.isfinite(vr))] = 0.0 return vr