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