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 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 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
