def fourier_ring_correlation(obj,
ref,
step_size=1,
save_path='frc',
save_mask=False):
if not os.path.exists(save_path):
os.makedirs(save_path)
radius_max = int(min(obj.shape) / 2)
f_obj = np_fftshift(fft2(obj))
f_ref = np_fftshift(fft2(ref))
f_prod = f_obj * np.conjugate(f_ref)
f_obj_2 = np.real(f_obj * np.conjugate(f_obj))
f_ref_2 = np.real(f_ref * np.conjugate(f_ref))
radius_ls = np.arange(1, radius_max, step_size)
fsc_ls = []
np.save(os.path.join(save_path, 'radii.npy'), radius_ls)
for rad in radius_ls:
print(rad)
if os.path.exists(
os.path.join(save_path,
'mask_rad_{:04d}.tiff'.format(int(rad)))):
mask = dxchange.read_tiff(
os.path.join(save_path,
'mask_rad_{:04d}.tiff'.format(int(rad))))
else:
mask = generate_ring(obj.shape, rad, anti_aliasing=2)
if save_mask:
dxchange.write_tiff(
mask,
os.path.join(save_path,
'mask_rad_{:04d}.tiff'.format(int(rad))),
dtype='float32',
overwrite=True)
fsc = abs(np.sum(f_prod * mask))
fsc /= np.sqrt(np.sum(f_obj_2 * mask) * np.sum(f_ref_2 * mask))
fsc_ls.append(fsc)
np.save(os.path.join(save_path, 'fsc.npy'), fsc_ls)
matplotlib.rcParams['pdf.fonttype'] = 'truetype'
fontProperties = {
'family': 'serif',
'serif': ['Times New Roman'],
'weight': 'normal',
'size': 12
}
plt.rc('font', **fontProperties)
plt.plot(radius_ls.astype(float) / radius_ls[-1], fsc_ls)
plt.xlabel('Spatial frequency (1 / Nyquist)')
plt.ylabel('FRC')
plt.savefig(os.path.join(save_path, 'frc.pdf'), format='pdf')
def multislice_propagate_batch_numpy(grid_delta_batch, grid_beta_batch, probe_real, probe_imag, energy_ev, psize_cm, free_prop_cm=None, obj_batch_shape=None):
minibatch_size = obj_batch_shape[0]
grid_shape = obj_batch_shape[1:]
voxel_nm = np.array([psize_cm] * 3) * 1.e7
wavefront = np.zeros([minibatch_size, obj_batch_shape[1], obj_batch_shape[2]], dtype='complex64')
wavefront += (probe_real + 1j * probe_imag)
lmbda_nm = 1240. / energy_ev
mean_voxel_nm = np.prod(voxel_nm) ** (1. / 3)
size_nm = np.array(grid_shape) * voxel_nm
n_slice = obj_batch_shape[-1]
delta_nm = voxel_nm[-1]
# h = get_kernel_ir(delta_nm, lmbda_nm, voxel_nm, grid_shape)
h = get_kernel(delta_nm, lmbda_nm, voxel_nm, grid_shape)
k = 2. * PI * delta_nm / lmbda_nm
probe_array = []
for i in range(n_slice):
delta_slice = grid_delta_batch[:, :, :, i]
beta_slice = grid_beta_batch[:, :, :, i]
c = np.exp(1j * k * delta_slice) * np.exp(-k * beta_slice)
wavefront = wavefront * c
if i < n_slice - 1:
wavefront = ifft2(np_ifftshift(np_fftshift(fft2(wavefront), axes=[1, 2]) * h, axes=[1, 2]))
probe_array.append(wavefront)
if free_prop_cm is not None:
#1dxchange.write_tiff(abs(wavefront), '2d_1024/monitor_output/wv', dtype='float32', overwrite=True)
if free_prop_cm == 'inf':
wavefront = np_fftshift(fft2(wavefront), axes=[1, 2])
else:
dist_nm = free_prop_cm * 1e7
l = np.prod(size_nm)**(1. / 3)
crit_samp = lmbda_nm * dist_nm / l
algorithm = 'TF' if mean_voxel_nm > crit_samp else 'IR'
# print(algorithm)
algorithm = 'TF'
if algorithm == 'TF':
h = get_kernel(dist_nm, lmbda_nm, voxel_nm, grid_shape)
wavefront = ifft2(np_ifftshift(np_fftshift(fft2(wavefront), axes=[1, 2]) * h, axes=[1, 2]))
else:
h = get_kernel_ir(dist_nm, lmbda_nm, voxel_nm, grid_shape)
wavefront = ifft2(np_ifftshift(np_fftshift(fft2(wavefront), axes=[1, 2]) * h, axes=[1, 2]))
# dxchange.write_tiff(abs(wavefront), '2d_512/monitor_output/wv', dtype='float32', overwrite=True)
# dxchange.write_tiff(np.angle(h), '2d_512/monitor_output/h', dtype='float32', overwrite=True)
return wavefront, np.array(probe_array)
def get_kernel_ir(dist_nm, lmbda_nm, voxel_nm, grid_shape):
"""
Get Fresnel propagation kernel for IR algorithm.
Parameters:
-----------
simulator : :class:`acquisition.Simulator`
The Simulator object.
dist : float
Propagation distance in cm.
"""
size_nm = np.array(voxel_nm) * np.array(grid_shape)
k = 2 * PI / lmbda_nm
ymin, xmin = np.array(size_nm)[:2] / -2.
dy, dx = voxel_nm[0:2]
x = np.arange(xmin, xmin + size_nm[1], dx)
y = np.arange(ymin, ymin + size_nm[0], dy)
x, y = np.meshgrid(x, y)
try:
h = np.exp(1j * k * dist_nm) / (1j * lmbda_nm * dist_nm) * np.exp(
1j * k / (2 * dist_nm) * (x**2 + y**2))
H = np_fftshift(fft2(h)) * voxel_nm[0] * voxel_nm[1]
dxchange.write_tiff(x,
'2d_512/monitor_output/x',
dtype='float32',
overwrite=True)
except:
h = tf.exp(1j * k * dist_nm) / (1j * lmbda_nm * dist_nm) * tf.exp(
1j * k / (2 * dist_nm) * (x**2 + y**2))
# h = tf.convert_to_tensor(h, dtype='complex64')
H = fftshift(tf.fft2d(h)) * voxel_nm[0] * voxel_nm[1]
return H
def fresnel_propagate_numpy(wavefront, energy_ev, psize_cm, dist_cm):
lmbda_nm = 1240. / energy_ev
lmbda_cm = 0.000124 / energy_ev
psize_nm = psize_cm * 1e7
dist_nm = dist_cm * 1e7
if dist_cm == 'inf':
wavefront = np_fftshift(fft2(wavefront))
else:
n = np.mean(wavefront.shape)
z_crit_cm = (psize_cm * n)**2 / (lmbda_cm * n)
algorithm = 'TF' if dist_cm < z_crit_cm else 'IR'
if algorithm == 'TF':
h = get_kernel(dist_nm, lmbda_nm, [psize_nm, psize_nm],
wavefront.shape)
wavefront = ifft2(np_ifftshift(np_fftshift(fft2(wavefront)) * h))
else:
h = get_kernel_ir(dist_nm, lmbda_nm, [psize_nm, psize_nm],
wavefront.shape)
wavefront = np_ifftshift(ifft2(np_fftshift(fft2(wavefront)) * h))
return wavefront
def fourier_ring_correlation(obj, ref, step_size=1, save_path=None, save_mask=True, save_fname='fsc', threshold_curve=False):
if not os.path.exists(save_path):
os.makedirs(save_path)
if obj.ndim == 2:
fft_func = fft2
gen_mask = generate_ring
gen_kwargs = {'anti_aliasing: 2'}
elif obj.ndim == 3:
fft_func = fftn
gen_mask = generate_shell
gen_kwargs = {}
radius_max = min(obj.shape) // 2
f_obj = np_fftshift(fft_func(obj))
f_ref = np_fftshift(fft_func(ref))
f_prod = f_obj * np.conjugate(f_ref)
f_obj_2 = np.real(f_obj * np.conjugate(f_obj))
f_ref_2 = np.real(f_ref * np.conjugate(f_ref))
radius_ls = np.arange(1, radius_max, step_size)
fsc_ls = []
if save_path is not None:
np.save(os.path.join(save_path, 'radii.npy'), radius_ls)
for rad in radius_ls:
if os.path.exists(os.path.join(save_path, 'mask_rad_{:04d}.tiff'.format(int(rad)))):
mask = dxchange.read_tiff(os.path.join(save_path, 'mask_rad_{:04d}.tiff'.format(int(rad))))
else:
mask = gen_mask(obj.shape, rad, **gen_kwargs)
if save_mask:
dxchange.write_tiff(mask, os.path.join(save_path, 'mask_rad_{:04d}.tiff'.format(int(rad))),
dtype='float32', overwrite=True)
fsc = abs(np.sum(f_prod * mask))
fsc /= np.sqrt(np.sum(f_obj_2 * mask) * np.sum(f_ref_2 * mask))
fsc_ls.append(fsc)
if save_path is not None:
np.save(os.path.join(save_path, '{}.npy'.format(save_fname)), fsc_ls)
return np.array(fsc_ls)
def multidistance_ctf(prj_ls,
dist_cm_ls,
psize_cm,
energy_kev,
kappa=50,
sigma_cut=0.01,
alpha_1=5e-4,
alpha_2=1e-16):
prj_ls = np.array(prj_ls)
dist_cm_ls = np.array(dist_cm_ls)
dist_nm_ls = dist_cm_ls * 1.e7
lmbda_nm = 1.24 / energy_kev
psize_nm = psize_cm * 1.e7
prj_shape = prj_ls.shape[1:]
u_max = 1. / (2. * psize_nm)
v_max = 1. / (2. * psize_nm)
u, v = gen_mesh([v_max, u_max], prj_shape)
xi_mesh = PI * lmbda_nm * (u**2 + v**2)
xi_ls = np.zeros([len(dist_cm_ls), *prj_shape])
for i in range(len(dist_cm_ls)):
xi_ls[i] = xi_mesh * dist_nm_ls[i]
abs_nu = np.sqrt(u**2 + v**2)
nu_cut = 0.6 * u_max
f = 0.5 * (1 - erf((abs_nu - nu_cut) / sigma_cut))
alpha = alpha_1 * f + alpha_2 * (1 - f)
# plt.imshow(abs(np.log(np_fftshift(fft2(prj_ls[0] - 1, axes=(-2, -1)), axes=(-2, -1)))))
# plt.imshow(alpha)
# plt.show()
# alpha = 0
phase = np.sum(
np_fftshift(fft2(prj_ls - 1, axes=(-2, -1)), axes=(-2, -1)) *
(np.sin(xi_ls) + 1. / kappa * np.cos(xi_ls)),
axis=0)
phase /= (
np.sum(2 * (np.sin(xi_ls) + 1. / kappa * np.cos(xi_ls))**2, axis=0) +
alpha)
phase = ifft2(np_ifftshift(phase, axes=(-2, -1)), axes=(-2, -1))
return np.abs(phase)
def free_propagate_spherical_numpy(wavefront,
dist_cm,
r_cm,
wavelen_nm,
probe_size,
theta_max=PI / 18,
phi_max=PI / 18):
dist_nm = dist_cm * 1.e7
r_nm = r_cm * 1.e7
k_theta = PI / theta_max * (np.arange(probe_size[0]) -
float(probe_size[0] - 1) / 2)
k_phi = PI / phi_max * (np.arange(probe_size[1]) -
float(probe_size[1] - 1) / 2)
k_phi, k_theta = np.meshgrid(k_phi, k_theta)
k = 2 * PI / wavelen_nm
wavefront = np_fftshift(fft2(wavefront), axes=[-1, -2])
wavefront *= np.exp(-1j / (2 * k) * (k_theta**2 + k_phi**2) *
(1. / (r_nm + dist_nm) - 1. / r_nm))
wavefront = ifft2(np_ifftshift(wavefront, axes=[-1, -2]))
return wavefront
def get_kernel_ir(dist_nm, lmbda_nm, voxel_nm, grid_shape):
"""
Get Fresnel propagation kernel for IR algorithm.
Parameters:
-----------
simulator : :class:`acquisition.Simulator`
The Simulator object.
dist : float
Propagation distance in cm.
"""
size_nm = np.array(voxel_nm) * np.array(grid_shape)
k = 2 * PI / lmbda_nm
ymin, xmin = np.array(size_nm)[:2] / -2.
x = np.linspace(xmin, -xmin, grid_shape[1])
y = np.linspace(ymin, -ymin, grid_shape[0])
x, y = np.meshgrid(x, y)
h = np.exp(1j * k * dist_nm) / (1j * lmbda_nm * dist_nm) * np.exp(
1j * k / (2. * dist_nm) * (x**2 + y**2))
H = np_fftshift(fft2(h)) * voxel_nm[0] * voxel_nm[1]
return H
