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