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
