def window(v, func='hanning', params=None): """ applies a windowing function to the 3D volume v (inplace, as reference) """ N = v.shape[0] D = v.ndim if any([d != N for d in list(v.shape)]) or D != 3: raise Exception("Error: Volume is not Cube.") def apply_seperable_window(v, w): v *= np.reshape(w, (-1, 1, 1)) v *= np.reshape(w, (1, -1, 1)) v *= np.reshape(w, (1, 1, -1)) if func == "hanning": w = np.hanning(N) apply_seperable_window(v, w) elif func == 'hamming': w = np.hamming(N) apply_seperable_window(v, w) elif func == 'gaussian': raise Exception('Unimplimented') elif func == 'circle': c = gencoords(N, 3) if params == None: r = N / 2 - 1 else: r = params[0] * (N / 2 * 1) v *= (np.sum(c**2, 1) < (r ** 2)).reshape((N, N, N)) elif func == 'box': v[:, 0, 0] = 0.0 v[0, :, 0] = 0.0 v[0, 0, :] = 0.0 else: raise Exception("Error: Window Type Not Supported")
def correlation_noise(num_images=1000, N=128, rad=0.6, stack_noise=False): noise_stack = np.require(np.random.randn(num_images, N, N), dtype=density.real_t) real_image = noise_stack[np.random.randint(num_images)] corr_real_image = correlation.calc_full_ac(real_image, rad=rad) fourier_noise = density.real_to_fspace(real_image) fourier_corr_image = correlation.calc_full_ac(fourier_noise, rad=rad) _, _, mask = geometry.gencoords(N, 2, rad, True) plot_noise_histogram(corr_real_image, fourier_corr_image, mask, mask) if stack_noise: center = int(N/2) x_shift, y_shift = np.random.randint(-center, center, size=2) fourier_noise_stack = density.real_to_fspace(noise_stack, axes=(1, 2)) corr_noise_stack = np.zeros_like(noise_stack, dtype=density.real_t) fourier_corr_noise_stack = np.zeros_like(fourier_noise_stack, dtype=density.complex_t) for i in range(num_images): corr_noise_stack[i] = correlation.calc_full_ac(noise_stack[i], rad=rad) fourier_corr_noise_stack[i] = correlation.calc_full_ac(fourier_noise_stack[i], rad=rad) noise_zoom = noise_stack[:, x_shift, y_shift] fourier_noise_zoom = fourier_noise_stack[:, x_shift, y_shift] plot_stack_noise(noise_zoom, fourier_noise_zoom)
def generate_phantom_density(N, window, sigma, num_blobs, seed=None): if seed is not None: np.random.seed(seed) M = np.zeros((N, N, N), dtype=np.float32) coords = gencoords(N, 3).reshape((N**3, 3)) inside_window = np.sum(coords**2, axis=1).reshape((N, N, N)) < window**2 curr_c = np.array([0.0, 0.0, 0.0]) curr_n = 0 while curr_n < num_blobs: csigma = sigma * np.exp(0.25 * np.random.randn()) radM = np.sum((coords - curr_c.reshape((1, 3))) ** 2, axis=1).reshape((N, N, N)) inside = np.logical_and(radM < (3 * csigma)**2, inside_window) # M[inside] = 1 M[inside] += np.exp(-0.5 * (radM[inside] / csigma**2)) curr_n += 1 curr_dir = np.random.randn(3) curr_dir /= np.sum(curr_dir**2) curr_c += 2.0 * csigma * curr_dir curr_w = np.sqrt(np.sum(curr_c**2)) while curr_w > window: curr_n_dir = curr_c / curr_w curr_r_dir = (2 * np.dot(curr_dir, curr_n_dir)) * \ curr_n_dir - curr_dir curr_c = curr_n_dir + (curr_w - window) * curr_r_dir curr_w = np.sqrt(np.sum(curr_c**2)) return M
def generate_phantom_density(N, window, sigma, num_blobs, seed=None): if seed is not None: np.random.seed(seed) M = np.zeros((N, N, N), dtype=np.float32) coords = gencoords(N, 3).reshape((N**3, 3)) inside_window = np.sum(coords**2, axis=1).reshape((N, N, N)) < window**2 curr_c = np.array([0.0, 0.0, 0.0]) curr_n = 0 while curr_n < num_blobs: csigma = sigma * np.exp(0.25 * np.random.randn()) radM = np.sum((coords - curr_c.reshape((1, 3)))**2, axis=1).reshape( (N, N, N)) inside = np.logical_and(radM < (3 * csigma)**2, inside_window) # M[inside] = 1 M[inside] += np.exp(-0.5 * (radM[inside] / csigma**2)) curr_n += 1 curr_dir = np.random.randn(3) curr_dir /= np.sum(curr_dir**2) curr_c += 2.0 * csigma * curr_dir curr_w = np.sqrt(np.sum(curr_c**2)) while curr_w > window: curr_n_dir = curr_c / curr_w curr_r_dir = (2 * np.dot(curr_dir, curr_n_dir)) * \ curr_n_dir - curr_dir curr_c = curr_n_dir + (curr_w - window) * curr_r_dir curr_w = np.sqrt(np.sum(curr_c**2)) return M
def window(v, func='hanning', params=None): """ applies a windowing function to the 3D volume v (inplace, as reference) """ N = v.shape[0] D = v.ndim if any([d != N for d in list(v.shape)]) or D != 3: raise Exception("Error: Volume is not Cube.") def apply_seperable_window(v, w): v *= np.reshape(w, (-1, 1, 1)) v *= np.reshape(w, (1, -1, 1)) v *= np.reshape(w, (1, 1, -1)) if func == "hanning": w = np.hanning(N) apply_seperable_window(v, w) elif func == 'hamming': w = np.hamming(N) apply_seperable_window(v, w) elif func == 'gaussian': raise Exception('Unimplimented') elif func == 'circle': c = gencoords(N, 3) if params == None: r = N / 2 - 1 else: r = params[0] * (N / 2 * 1) v *= (np.sum(c**2, 1) < (r**2)).reshape((N, N, N)) elif func == 'box': v[:, 0, 0] = 0.0 v[0, :, 0] = 0.0 v[0, 0, :] = 0.0 else: raise Exception("Error: Window Type Not Supported")
def get_envelope_map(self,sigma2,rho,env_lb=None,env_ub=None,minFreq=None,bfactor=None,rotavg=True): N = self.cryodata.N N_D = float(self.cryodata.N_D_Train) num_batches = float(self.cryodata.num_batches) psize = self.params['pixel_size'] beamstop_freq = self.params.get('beamstop_freq', None) mean_corr = self.correlation_history.get_mean().reshape((N,N)) mean_power = self.power_history.get_mean().reshape((N,N)) mean_mask = self.mask_history.get_mean().reshape((N,N)) mask_w = self.mask_history.get_wsum() * (N_D / num_batches) if rotavg: mean_corr = cryoem.rotational_average(mean_corr,normalize=True,doexpand=True) mean_power = cryoem.rotational_average(mean_power,normalize=True,doexpand=True) mean_mask = cryoem.rotational_average(mean_mask,normalize=False,doexpand=True) if isinstance(sigma2,np.ndarray): sigma2 = sigma2.reshape((N,N)) if bfactor is not None: coords = gencoords(N,2).reshape((N**2,2)) freqs = np.sqrt(np.sum(coords**2,axis=1))/(psize*N) prior_envelope = ctf.envelope_function(freqs,bfactor).reshape((N,N)) else: prior_envelope = 1.0 obsw = (mask_w * mean_mask / sigma2) exp_env = (mean_corr * obsw + prior_envelope*rho) / (mean_power * obsw + rho) if minFreq is not None: # Only consider envelope parameters for frequencies above a threshold minRad = minFreq*2.0*psize if beamstop_freq is None: _, _, minRadMask = gencoords(N, 2, minRad, True) else: maskRad = beamstop_freq * 2.0 * psize _, _, minRadMask = gencoords_centermask(N, 2, minRad, maskRad, True) exp_env[minRadMask.reshape((N,N))] = 1.0 if env_lb is not None or env_ub is not None: np.clip(exp_env,env_lb,env_ub,out=exp_env) return exp_env
def window_images(self, rad=0.99): N = self.get_num_pixels() coords = geometry.gencoords(N, 2).reshape((N**2, 2)) Cs = np.sum(coords**2, axis=1).reshape((N, N) ) > (rad * N / 2.0 - 1.5)**2 for img in self: img[Cs] = 0
def compute_shift_phases(pts, N, rad): xy = geometry.gencoords(N, 2, rad) N_T = xy.shape[0] N_S = pts.shape[0] shifts = np.empty((N_S, N_T), dtype=np.complex64) for (i, (sx, sy)) in enumerate(pts): shifts[i] = np.exp(2.0j * np.pi / N * (xy[:, 0] * sx + xy[:, 1] * sy)) return shifts
def shift_vis(model): N = model.shape[0] rad = 0.8 kernel = 'lanczos' kernsize = 4 xy, trunc_xy, truncmask = geometry.gencoords(N, 2, rad, True) N_T = trunc_xy.shape[0] premult = cryoops.compute_premultiplier(N, kernel=kernel, kernsize=kernsize) TtoF = sincint.gentrunctofull(N=N, rad=rad) fM = density.real_to_fspace(model) prefM = density.real_to_fspace( premult.reshape((1, 1, -1)) * premult.reshape( (1, -1, 1)) * premult.reshape((-1, 1, 1)) * model) pt = np.random.randn(3) pt /= np.linalg.norm(pt) psi = 2 * np.pi * np.random.rand() ea = geometry.genEA(pt)[0] ea[2] = psi print('project model for Euler angel: ({:.2f}, {:.2f}, {:.2f}) degree'. format(*np.rad2deg(ea))) rot_matrix = geometry.rotmat3D_EA(*ea)[:, 0:2] slop = cryoops.compute_projection_matrix([rot_matrix], N, kernel, kernsize, rad, 'rots') # trunc_slice = slop.dot(prefM.reshape((-1,))) trunc_slice = cryoem.getslices(prefM, slop) fourier_slice = TtoF.dot(trunc_slice).reshape(N, N) real_proj = density.fspace_to_real(fourier_slice) fig, axes = plt.subplots(4, 4, figsize=(12.8, 8)) im_real = axes[0, 0].imshow(real_proj) im_fourier = axes[1, 0].imshow(np.log(np.abs(fourier_slice))) for i, ax in enumerate(axes[:, 1:].T): shift = np.random.randn(2) * (N / 4.0) S = cryoops.compute_shift_phases(shift.reshape(1, 2), N, rad)[0] shift_trunc_slice = S * trunc_slice shift_fourier_slice = TtoF.dot(shift_trunc_slice).reshape(N, N) shift_real_proj = density.fspace_to_real(shift_fourier_slice) ax[0].imshow(shift_real_proj) ax[1].imshow(np.log(np.abs(shift_fourier_slice))) ax[2].imshow(np.log(shift_fourier_slice.real)) ax[3].imshow(np.log(shift_fourier_slice.imag)) fig.tight_layout() plt.show()
def get_envelope_map(self,sigma2,rho,env_lb=None,env_ub=None,minFreq=None,bfactor=None,rotavg=True): N = self.cryodata.N N_D = float(self.cryodata.N_D_Train) num_batches = float(self.cryodata.num_batches) psize = self.params['pixel_size'] mean_corr = self.correlation_history.get_mean().reshape((N,N)) mean_power = self.power_history.get_mean().reshape((N,N)) mean_mask = self.mask_history.get_mean().reshape((N,N)) mask_w = self.mask_history.get_wsum() * (N_D / num_batches) if rotavg: mean_corr = cryoem.rotational_average(mean_corr,normalize=True,doexpand=True) mean_power = cryoem.rotational_average(mean_power,normalize=True,doexpand=True) mean_mask = cryoem.rotational_average(mean_mask,normalize=False,doexpand=True) if isinstance(sigma2,np.ndarray): sigma2 = sigma2.reshape((N,N)) if bfactor is not None: coords = gencoords(N,2).reshape((N**2,2)) freqs = np.sqrt(np.sum(coords**2,axis=1))/(psize*N) prior_envelope = ctf.envelope_function(freqs,bfactor).reshape((N,N)) else: prior_envelope = 1.0 obsw = (mask_w * mean_mask / sigma2) exp_env = (mean_corr * obsw + prior_envelope*rho) / (mean_power * obsw + rho) if minFreq is not None: # Only consider envelope parameters for frequencies above a threshold minRad = minFreq*2.0*psize _, _, minRadMask = gencoords(N, 2, minRad, True) exp_env[minRadMask.reshape((N,N))] = 1.0 if env_lb is not None or env_ub is not None: np.clip(exp_env,env_lb,env_ub,out=exp_env) return exp_env
def rotational_expand(vals, N, D, interp_order=1): interp_coords = np.sqrt(np.sum(gencoords(N, D).reshape( (N**D, D))**2, axis=1)).reshape((1,) + D * (N,)) if np.iscomplexobj(vals): rotexp = 1.0j * spinterp.map_coordinates(vals.imag, interp_coords, order=interp_order, mode='nearest') rotexp += spinterp.map_coordinates(vals.real, interp_coords, order=interp_order, mode='nearest') else: rotexp = spinterp.map_coordinates(vals, interp_coords, order=interp_order, mode='nearest') return rotexp
def envelope_vis(N=128, rad=1): trunc_xy = geometry.gencoords(N, 2, rad) psize = 2.8 bfactor = 500 freqs = np.sqrt(np.sum(trunc_xy**2, axis=1)) / (psize * N) envelope = ctf.envelope_function(freqs, bfactor) fig, ax = plt.subplots() ax.plot(freqs, envelope, '.', label='bfactor: %d' % bfactor) ax.legend(frameon=False) ax.set_title('envelope') plt.show()
def compute_density_moments(M, mu=None): N = M.shape[0] absM = (M**2).reshape((N**3, 1)) absM /= np.sum(absM) coords = gencoords(N, 3).reshape((N**3, 3)) if mu == None: wcoords = coords.reshape((N**3, 3)) * absM mu = np.sum(wcoords, axis=0).reshape((1, 3)) wccoords = np.sqrt(absM / N**3) * (coords - mu) covar = np.dot(wccoords.T, wccoords) return mu, covar
def rotational_average(M, maxRadius=None, doexpand=False, normalize=True, return_cnt=False): N = M.shape[0] D = len(M.shape) assert D >= 2, 'Cannot rotationally average a 1D array' pts = gencoords(N, D).reshape((N**D, D)) r = np.sqrt(np.sum(pts**2, axis=1)).reshape(M.shape) ir = np.require(np.floor(r), dtype='uint32') f = r - ir if maxRadius is None: maxRadius = np.ceil(np.sqrt(D) * N / D) if maxRadius < np.max(ir) + 2: valid_ir = ir + 1 < maxRadius ir = ir[valid_ir] f = f[valid_ir] M = M[valid_ir] if np.iscomplexobj(M): raps = 1.0j * np.bincount(ir, weights=(1 - f) * M.imag, minlength=maxRadius) + \ np.bincount(ir + 1, weights=f * M.imag, minlength=maxRadius) raps += np.bincount(ir, weights=(1 - f) * M.real, minlength=maxRadius) + \ np.bincount(ir + 1, weights=f * M.real, minlength=maxRadius) else: raps = np.bincount(ir, weights=(1 - f) * M, minlength=maxRadius) + \ np.bincount(ir + 1, weights=f * M, minlength=maxRadius) raps = raps[0:maxRadius] if normalize or return_cnt: cnt = np.bincount(ir, weights=(1 - f), minlength=maxRadius) + \ np.bincount(ir + 1, weights=f, minlength=maxRadius) cnt = cnt[0:maxRadius] if normalize: raps[cnt <= 0] = 0 raps[cnt > 0] /= cnt[cnt > 0] if doexpand: raps = rotational_expand(raps, N, D) if return_cnt: return raps, cnt else: return raps
def no_correlation(num_images=1000, N=128, rad=0.8): center = int(N/2) x_shift, y_shift = np.random.randint(-center, center, size=2) noise_stack = np.require(np.random.randn(num_images, N, N), dtype=density.real_t) real_image = noise_stack[np.random.randint(num_images)] fourier_noise_stack = density.real_to_fspace(noise_stack, axes=(1, 2)) fourier_noise = density.real_to_fspace(real_image) noise_zoom = noise_stack[:, x_shift, y_shift] fourier_noise_zoom = fourier_noise_stack[:, x_shift, y_shift] _, _, mask = geometry.gencoords(N, 2, rad, True) plot_noise_histogram(real_image, fourier_noise, rmask=mask, fmask=mask) plot_stack_noise(noise_zoom, fourier_noise_zoom)
def float_images(self, rad=0.99): N = self.get_num_pixels() coords = geometry.gencoords(N, 2).reshape((N**2, 2)) Cs = np.sum(coords**2, axis=1).reshape((N, N)) \ > (rad * N / 2.0 - 1.5)**2 vals = [] for img in self: corner_pixels = img[Cs] float_val = np.mean(corner_pixels) img -= float_val vals.append(float_val) return vals
def rotate_density(M, R, t=None, upsamp=1.0): assert len(M.shape) == 3 N = M.shape[0] Nup = int(np.round(N * upsamp)) # print "Upsampling by", upsamp, "to", Nup, "^3" coords = gencoords(Nup, 3).reshape((Nup**3, 3)) / float(upsamp) if t is None: interp_coords = np.transpose(np.dot(coords, R.T)).reshape( (3, Nup, Nup, Nup)) + N / 2 else: interp_coords = np.transpose( np.dot(coords, R.T) + t).reshape((3, Nup, Nup, Nup)) + N / 2 out = spinterp.map_coordinates(M, interp_coords, order=1) return out
def view_rad_range(N=128): fig, axes = plt.subplots(4, 5, figsize=(12.8, 8)) # , sharex=True, sharey=True, squeeze=False) rad_list = np.arange(0.1, 1.1, step=0.2) for i, rad in enumerate(rad_list): TtoF = sincint.gentrunctofull(N, rad) xy, trunc_xy, truncmask = geometry.gencoords(N, 2, rad, True) N_T = trunc_xy.shape[0] trunc = np.arange(0, N_T) image = TtoF.dot(trunc).reshape(N, N) axes[0, i].imshow(image, origin='lower') axes[0, i].set_title('radius: {:.2f}'.format(rad)) # outside of radius xy_outside = xy[~truncmask] image_outside_rad = np.zeros((N, N)) image_outside_rad[~truncmask.reshape(N, N)] = np.arange(xy_outside.shape[0]) axes[1, i].imshow(image_outside_rad, origin='lower') # sort trunc_xy coordinates pol_trunc_xy = correlation.cart2pol(trunc_xy) sorted_idx = np.lexsort((pol_trunc_xy[:, 1], pol_trunc_xy[:, 0])) # lexsort; first, sort rho; second, sort theta pol_trunc = trunc[sorted_idx.argsort()] pol_image = TtoF.dot(pol_trunc).reshape(N, N) axes[2, i].imshow(pol_image, origin='lower') # pol coordinate in outside part pol_xy_outside = correlation.cart2pol(xy_outside) outside_sorted_idx = np.lexsort((pol_xy_outside[:, 1], pol_xy_outside[:, 0])) pol_image_outside = np.zeros((N, N)) pol_image_outside[~truncmask.reshape(N, N)] = np.arange(pol_xy_outside.shape[0])[outside_sorted_idx.argsort()] axes[3, i].imshow(pol_image_outside, origin='lower') for i, ax in enumerate(axes.flat): m, n = np.unravel_index(i, (4, 5)) if m != 3: ax.set_xticks([]) else: ax.set_xticks([0, int(N/4), int(N/2), int(N*3/4), int(N-1)]) if n != 0: ax.set_yticks([]) else: ax.set_yticks([0, int(N/4), int(N/2), int(N*3/4), int(N-1)]) # fig.savefig('cart_coordinate_view_to_polar_coordinate_view.png', dpi=300) plt.show()
def rotate_density(M, R, t=None, upsamp=1.0): assert len(M.shape) == 3 N = M.shape[0] Nup = int(np.round(N * upsamp)) # print "Upsampling by", upsamp, "to", Nup, "^3" coords = gencoords(Nup, 3).reshape((Nup**3, 3)) / float(upsamp) if t is None: interp_coords = np.transpose(np.dot(coords, R.T)).reshape( (3, Nup, Nup, Nup)) + int(N / 2) else: interp_coords = np.transpose(np.dot(coords, R.T) + t).reshape( (3, Nup, Nup, Nup)) + int(N / 2) out = spinterp.map_coordinates(M, interp_coords, order=1) return out
def dataset_loading_test(params, visualize=False): imgpath = params['inpath'] psize = params['resolution'] imgstk = MRCImageStack(imgpath, psize) # if params.get('float_images', True): # imgstk.float_images() ctfpath = params['ctfpath'] mscope_params = params['microscope_params'] ctfstk = CTFStack(ctfpath, mscope_params) cryodata = CryoDataset(imgstk, ctfstk) cryodata.compute_noise_statistics() # if params.get('window_images',True): # imgstk.window_images() cryodata.divide_dataset(params['minisize'], params['test_imgs'], params['partition'], params['num_partitions'], params['random_seed']) # cryodata.set_datasign(params.get('datasign', 'auto')) # if params.get('normalize_data',True): # cryodata.normalize_dataset() # voxel_size = cryodata.pixel_size N = cryodata.imgstack.get_num_pixels() fspace_stack = FourierStack(cryodata.imgstack, caching = True, zeropad=1) premult = cryoops.compute_premultiplier(N + 2 * int(1 * (N/2)), 'lanczos', 8) premult = premult.reshape((-1,1)) * premult.reshape((1,-1)) fspace_stack.set_transform(premult, 1) if visualize: rad = 0.99 coords = geometry.gencoords(N, 2).reshape((N**2, 2)) Cs = np.sum(coords**2, axis=1).reshape((N, N)) > (rad * N / 2.0 - 1.5)**2 idx = np.random.randint(cryodata.imgstack.num_images) normalized = cryodata.imgstack.get_image(1) f_normalized = fspace_stack.get_image(1) plot_noise_histogram(normalized, f_normalized, rmask=~Cs, fmask=None, plot_unmask=False) plt.show() return cryodata, fspace_stack
def compute_symmetry_errs(rs,symop,gfV,gfVnopre,radwn,stream=None): N = gfV[0].shape[0] offRs = geometry.expmaps(rs) Rsyms = symop.get_rotations(include_identity=False) pts = geometry.gencoords(N, 3, radwn).reshape((-1,3)) N_R = offRs.shape[0] N_S = 1 N_T = pts.shape[0] N_T_aligned = int(n.ceil(float(N_T)/(2*cudaworker.cukrns.blocksize)))*(2*cudaworker.cukrns.blocksize) print(' ') print('NR, NS, NT ', N_R, N_S, N_T, N_T_aligned, (N_R*N_T_aligned)/1.e9) with cudaworker.GPUContext(): gts = gpuarray.zeros((1,3),n.float32) gret = gpuarray.zeros((N_R,N_S),n.float32) gphis = gpuarray.empty((N_R,N_S),n.float32) gP = (gpuarray.empty((N_R,N_T_aligned),n.float32), gpuarray.empty((N_R,N_T_aligned),n.float32)) gQ = (gpuarray.empty((1,N_T_aligned),n.float32), gpuarray.empty((1,N_T_aligned),n.float32)) gpts = gpuarray.to_gpu(pts.reshape((-1,3)).astype(n.float32)) gI = gpuarray.to_gpu(n.identity(3,dtype=n.float32).reshape((1,9))) cudaworker.cukrns.slice('linear',2,gQ,gfVnopre,gpts, gI,None,1,stream=stream) for Rsym in Rsyms: cRsym = Rsym.reshape((3,3)) Rs = n.array([cR.reshape((3,3)).dot(cRsym.dot(cR.reshape((3,3)).T)) for cR in offRs],dtype=n.float32) gRs = gpuarray.to_gpu(Rs.reshape((-1,9)).astype(n.float32)) cudaworker.cukrns.slice('linear',2,gP,gfV,gpts,gRs,None,1,stream=stream) cudaworker.cukrns.shifted_squared_error(N_R, N_S, N_T, gphis.reshape((N_R*N_S)), gP, gQ, gts, gpts, 2.0*n.pi/N, None, broadcast_Q=True, stream=stream) if stream is not None: stream.synchronize() gret += gphis ret = gret.get() return ret.reshape((N_R))
def demo(N=128, rad=0.5): TtoF = sincint.gentrunctofull(N=N, rad=rad) xy, trunc_xy, truncmask = geometry.gencoords(N, 2, rad, True) print('shape of TtoF:', TtoF.shape) print('slice shape:', trunc_xy.shape[0]) trunc_slice = np.arange(trunc_xy.shape[0]) sliced_image = TtoF.dot(trunc_slice).reshape(N, N) trunc_xy_idx = np.int_(trunc_xy + int(N/2)) # Compare speed for getting slices in this way new_trunc_slice = sliced_image[trunc_xy_idx[:, 0], trunc_xy_idx[:, 1]] print('error:', sum(trunc_slice - new_trunc_slice)) pol_trunc_xy = correlation.cart2pol(trunc_xy) # inside of rad # sort trunc_xy coordinates sorted_idx = np.lexsort((pol_trunc_xy[:, 1], pol_trunc_xy[:, 0])) # lexsort; first, sort rho; second, sort theta sorted_pol_trunc_xy = pol_trunc_xy[sorted_idx] # reconstuct sorted coordinates into original state reco_pol_trunc_xy = sorted_pol_trunc_xy[sorted_idx.argsort()] print('error for reconstructed coordinates:', sum(correlation.pol2cart(reco_pol_trunc_xy) - trunc_xy)) reco_trunc_slice = trunc_slice[sorted_idx.argsort()] bingo_sliced_image = TtoF.dot(reco_trunc_slice).reshape(N, N) # outside of rad xy_outside = xy[~truncmask] sliced_image_outside_rad = np.zeros((N, N)) sliced_image_outside_rad[~truncmask.reshape(N, N)] = np.arange(xy_outside.shape[0]) pol_xy_outside = correlation.cart2pol(xy_outside) outside_sorted_idx = np.lexsort((pol_xy_outside[:, 1], pol_xy_outside[:, 0])) # lexsort; first, sort rho; second, sort theta sorted_pol_xy_outside = pol_xy_outside[outside_sorted_idx] reco_pol_xy_outside = np.arange(xy_outside.shape[0])[outside_sorted_idx.argsort()] bingo_sliced_image_outside_rad = np.zeros((N, N)) bingo_sliced_image_outside_rad[~truncmask.reshape(N, N)] = reco_pol_xy_outside fig, axes = plt.subplots(2, 2) ax = axes.flatten() ax[0].imshow(sliced_image) ax[1].imshow(bingo_sliced_image) ax[2].imshow(sliced_image_outside_rad) ax[3].imshow(bingo_sliced_image_outside_rad) plt.show()
def compare_interpolation(N=128, rad=1): _, trunc_xy, _ = geometry.gencoords(N, 2, rad, True) pol_trunc_xy = correlation.cart2pol(trunc_xy) sorted_idx = np.lexsort((pol_trunc_xy[:, 1], pol_trunc_xy[:, 0])) # lexsort; first, sort rho; second, sort theta sorted_pol_trunc_xy = pol_trunc_xy[sorted_idx] interpolation = ['none', 'nearest', 'nearest_decimal_1', 'nearest_half'] fig, ax = plt.subplots(nrows=len(interpolation), sharex=True) # fig, ax = plt.subplots() def round_to(n, precision): # correction = 0.5 if n >= 0 else -0.5 correction = np.ones_like(n) * 0.5 correction[n < 0] = -0.5 return np.int_(n / precision + correction) * precision def round_half(n): return round_to(n, 0.5) def get_ip_func(ip_method): if 'none' == ip_method.lower(): return lambda x: x elif 'nearest' == ip_method.lower(): return np.round elif 'nearest_decimal_1' == ip_method.lower(): return lambda x: np.round(x, 1) elif 'nearest_half' == ip_method.lower(): return round_half else: raise ValueError('please input correct interpolation method.') for i, ip in enumerate(interpolation): ip_func = get_ip_func(ip) ip_pol_xy = ip_func(sorted_pol_trunc_xy[:, 0]) unique_value, unique_index, unique_inverse, unique_counts = np.unique(ip_pol_xy, return_index=True, return_inverse=True, return_counts=True) ax[i].plot(unique_value, unique_counts, label='interpolation: {}'.format(ip)) ax[i].legend(frameon=False) ax[i].set_ylabel('counts') ax[-1].set_xlabel('radius') plt.show()
def compute_fsc(VF1, VF2, maxrad, width=1.0, thresholds=[0.143, 0.5]): assert VF1.shape == VF2.shape N = VF1.shape[0] r = np.sqrt(np.sum(gencoords(N, 3).reshape((N, N, N, 3))**2, axis=3)) prev_rad = -np.inf fsc = [] rads = [] resInd = len(thresholds) * [None] for i, rad in enumerate(np.arange(1.5, maxrad * N / 2.0, width)): cxyz = np.logical_and(r >= prev_rad, r < rad) cF1 = VF1[cxyz] cF2 = VF2[cxyz] if len(cF1) == 0: break cCorr = np.vdot(cF1, cF2) / np.sqrt( np.vdot(cF1, cF1) * np.vdot(cF2, cF2)) for j, thr in enumerate(thresholds): if cCorr < thr and resInd[j] is None: resInd[j] = i fsc.append(cCorr.real) rads.append(rad / (N / 2.0)) prev_rad = rad fsc = np.array(fsc) rads = np.array(rads) resolutions = [] for rI, thr in zip(resInd, thresholds): if rI is None: resolutions.append(rads[-1]) elif rI == 0: resolutions.append(np.inf) else: x = (thr - fsc[rI]) / (fsc[rI - 1] - fsc[rI]) resolutions.append(x * rads[rI - 1] + (1 - x) * rads[rI]) return rads, fsc, thresholds, resolutions
def estimate_noise_variance(self, esttype='robust', zerosub=False, rad=1.0): N = self.get_num_pixels() Cs = np.sum(geometry.gencoords(N, 2).reshape((N**2, 2))**2, axis=1).reshape((N, N)) > (rad * N / 2.0 - 1.5)**2 vals = [] for img in self: cvals = img[Cs] vals.append(cvals) if esttype == 'robust': if zerosub: var = ( 1.4826 * np.median(np.abs(np.asarray(vals) - np.median(vals))))**2 else: var = (1.4826 * np.median(np.abs(vals)))**2 elif esttype == 'mle': var = np.mean(np.asarray(vals)**2, dtype=np.float64) if zerosub: var -= np.mean(vals, dtype=np.float64)**2 return var
def compute_fsc(VF1, VF2, maxrad, width=1.0, thresholds=[0.143, 0.5]): assert VF1.shape == VF2.shape N = VF1.shape[0] r = np.sqrt(np.sum(gencoords(N, 3).reshape((N, N, N, 3))**2, axis=3)) prev_rad = -np.inf fsc = [] rads = [] resInd = len(thresholds) * [None] for i, rad in enumerate(np.arange(1.5, maxrad * N / 2.0, width)): cxyz = np.logical_and(r >= prev_rad, r < rad) cF1 = VF1[cxyz] cF2 = VF2[cxyz] if len(cF1) == 0: break cCorr = np.vdot(cF1, cF2) / np.sqrt(np.vdot(cF1, cF1) * np.vdot(cF2, cF2)) for j, thr in enumerate(thresholds): if cCorr < thr and resInd[j] is None: resInd[j] = i fsc.append(cCorr.real) rads.append(rad / (N / 2.0)) prev_rad = rad fsc = np.array(fsc) rads = np.array(rads) resolutions = [] for rI, thr in zip(resInd, thresholds): if rI is None: resolutions.append(rads[-1]) elif rI == 0: resolutions.append(np.inf) else: x = (thr - fsc[rI]) / (fsc[rI - 1] - fsc[rI]) resolutions.append(x * rads[rI - 1] + (1 - x) * rads[rI]) return rads, fsc, thresholds, resolutions
def merge_slices(slices, Rs, N, rad, beamstop_rad=None, res=None): center = int(N / 2) if beamstop_rad is None: coords = gencoords(N, 2, rad) else: coords = gencoords_centermask(N, 2, rad, beamstop_rad) assert slices.shape[1] == coords.shape[0] if res is None: res = np.zeros((N, ) * 3, dtype=np.float32) else: assert res.shape == (N, ) * 3 assert res.dtype == slices.dtype res[:] = 0.0 model_weight = np.zeros((N, ) * 3) for i, R in enumerate(Rs): curr_slices = slices[i, :] for j, xy in enumerate(coords): voxel_intensity = curr_slices[j] rot_coord = R.dot(xy.T).reshape(1, -1)[0] + center rot_coord = np.int_(np.round(rot_coord)) in_x = rot_coord[0] >= 0 and rot_coord[0] < N in_y = rot_coord[1] >= 0 and rot_coord[1] < N in_z = rot_coord[2] >= 0 and rot_coord[2] < N if in_x and in_y and in_z: index_coord = tuple(rot_coord) model_voxel_intensity = res[index_coord] model_weight[index_coord] += 1 voxel_weight = model_weight[index_coord] delta_intensity = voxel_intensity - model_voxel_intensity model_voxel_intensity += delta_intensity / voxel_weight res[index_coord] = model_voxel_intensity return res
def getslices_interp(V, Rs, rad, beamstop_rad=None, res=None): ndim = V.ndim assert ndim > 1 num_slices = len(Rs) # if ndim == 2: # assert Rs.shape[1] == 2 # elif ndim == 3: # assert Rs.shape[1] == 3 # Rs.shape[2] == 2 N = V.shape[0] center = int(N / 2) if beamstop_rad is None: coords = gencoords(N, 2, rad) else: coords = gencoords_centermask(N, 2, rad, beamstop_rad) N_T = coords.shape[0] grid = (np.arange(N), ) * ndim slicing_func = RegularGridInterpolator(grid, V, bounds_error=False, fill_value=0.0) if res is None: res = np.zeros((num_slices, N_T), dtype=V.dtype) else: assert res.shape[0] == Rs.shape[0] assert res.dtype == V.dtype res[:] = 0 for i, R in enumerate(Rs): rotated_coords = R.dot(coords.T).T + center res[i] = slicing_func(rotated_coords) # res[i] = interpn(grid, V, rotated_coords) # res[i] = spinterp.map_coordinates(V, rotated_coords.T) return res
def __init__(self, model, dataset_params, ctf_params, interp_params={'kern': 'lanczos', 'kernsize': 4.0, 'zeropad': 0, 'dopremult': True}, load_cache=True): self.dataset_params = dataset_params if model is not None: # assert False assert isinstance(model, np.ndarray), "Unexpected data type for input model" self.num_pixels = model.shape[0] N = self.num_pixels self.num_images = dataset_params['num_images'] assert self.num_images > 1, "it's better to make num_images larger than 1." self.pixel_size = float(dataset_params['pixel_size']) euler_angles = dataset_params['euler_angles'] self.is_sym = get_symmetryop(dataset_params.get('symmetry', None)) if euler_angles is None and self.is_sym is None: pt = np.random.randn(self.num_images, 3) pt /= np.linalg.norm(pt, axis=1, keepdims=True) euler_angles = geometry.genEA(pt) euler_angles[:, 2] = 2 * np.pi * np.random.rand(self.num_images) elif euler_angles is None and self.is_sym is not None: euler_angles = np.zeros((self.num_images, 3)) for i, ea in enumerate(euler_angles): while True: pt = np.random.randn(3) pt /= np.linalg.norm(pt) if self.is_sym.in_asymunit(pt.reshape(-1, 3)): break ea[0:2] = geometry.genEA(pt)[0][0:2] ea[2] = 2 * np.pi * np.random.rand() self.euler_angles = euler_angles.reshape((-1, 3)) if ctf_params is not None: self.use_ctf = True ctf_map = ctf.compute_full_ctf(None, N, ctf_params['psize'], ctf_params['akv'], ctf_params['cs'], ctf_params['wgh'], ctf_params['df1'], ctf_params['df2'], ctf_params['angast'], ctf_params['dscale'], ctf_params.get('bfactor', 500)) self.ctf_params = copy(ctf_params) if 'bfactor' in self.ctf_params.keys(): self.ctf_params.pop('bfactor') else: self.use_ctf = False ctf_map = np.ones((N**2,), dtype=density.real_t) kernel = 'lanczos' ksize = 6 rad = 0.95 # premult = cryoops.compute_premultiplier(N, kernel, ksize) TtoF = sincint.gentrunctofull(N=N, rad=rad) base_coords = geometry.gencoords(N, 2, rad) # premulter = premult.reshape((1, 1, -1)) \ # * premult.reshape((1, -1, 1)) \ # * premult.reshape((-1, 1, 1)) # fM = density.real_to_fspace(premulter * model) fM = model # if load_cache: # try: print("Generating Dataset ... :") tic = time.time() imgdata = np.empty((self.num_images, N, N), dtype=density.real_t) for i, ea in zip(range(self.num_images), self.euler_angles): R = geometry.rotmat3D_EA(*ea)[:, 0:2] slop = cryoops.compute_projection_matrix( [R], N, kernel, ksize, rad, 'rots') # D = slop.dot(fM.reshape((-1,))) rotated_coords = R.dot(base_coords.T).T + int(N/2) D = interpn((np.arange(N),) * 3, fM, rotated_coords) np.maximum(D, 0.0, out=D) intensity = ctf_map.reshape((N, N)) * TtoF.dot(D).reshape((N, N)) np.maximum(1e-8, intensity, out=intensity) intensity = np.float_( np.random.poisson(intensity) ) imgdata[i] = np.require(intensity, dtype=density.real_t) self.imgdata = imgdata print(" cost {} seconds.".format(time.time()-tic)) self.set_transform(interp_params) # self.prep_processing() else: euler_angles = [] with open(self.dataset_params['gtpath']) as par: par.readline() # 'C PHI THETA PSI SHX SHY FILM DF1 DF2 ANGAST' while True: try: line = par.readline().split() euler_angles.append([float(line[1]), float(line[2]), float(line[3])]) except Exception: break self.euler_angles = np.deg2rad(np.asarray(euler_angles)) num_images = self.dataset_params.get('num_images', 200) imgdata = mrc.readMRCimgs(self.dataset_params['inpath'], 0, num_images) self.imgdata = np.transpose(imgdata, axes=(2, 0, 1)) self.num_images = self.imgdata.shape[0] self.num_pixels = self.imgdata.shape[1] N = self.num_pixels self.pixel_size = self.dataset_params['resolution'] self.is_sym = self.dataset_params.get('symmetry', None) self.use_ctf = False ctf_map = np.ones((N**2,), dtype=density.real_t) self.set_transform(interp_params)
def compute_full_ctf(rots, N, psize, akv, csf, wgh, dfmid1, dfmid2, angastf, dscale, bfactor): freqs = geometry.gencoords(N, 2) / (N * psize) return compute_ctf(freqs, rots, akv, csf, wgh, dfmid1, dfmid2, angastf, dscale, bfactor)
def genphantomdata(N_D, phantompath): mscope_params = { 'akv': 200, 'wgh': 0.07, 'cs': 2.0, 'psize': 3.0, 'bfactor': 500.0 } M = mrc.readMRC(phantompath) N = M.shape[0] rad = 0.95 M_totalmass = 1000000 # M_totalmass = 1500000 kernel = 'lanczos' ksize = 6 tic = time.time() N_D = int(N_D) N = int(N) rad = float(rad) psize = mscope_params['psize'] bfactor = mscope_params['bfactor'] M_totalmass = float(M_totalmass) ctfparfile = 'particle/examplectfs.par' srcctf_stack = CTFStack(ctfparfile, mscope_params) genctf_stack = GeneratedCTFStack( mscope_params, parfields=['PHI', 'THETA', 'PSI', 'SHX', 'SHY']) TtoF = sincint.gentrunctofull(N=N, rad=rad) Cmap = np.sort( np.random.random_integers(0, srcctf_stack.get_num_ctfs() - 1, N_D)) cryoem.window(M, 'circle') M[M < 0] = 0 if M_totalmass is not None: M *= M_totalmass / M.sum() # oversampling oversampling_factor = 3 psize = psize * oversampling_factor V = density.real_to_fspace_with_oversampling(M, oversampling_factor) fM = V.real**2 + V.imag**2 # mrc.writeMRC('particle/EMD6044_fM_totalmass_{}_oversampling_{}.mrc'.format(str(int(M_totalmass)).zfill(5), oversampling_factor), fM, psz=psize) print("Generating data...") sys.stdout.flush() imgdata = np.empty((N_D, N, N), dtype=density.real_t) pardata = {'R': []} prevctfI = None coords = geometry.gencoords(N, 2, rad) slicing_func = RegularGridInterpolator((np.arange(N), ) * 3, fM, bounds_error=False, fill_value=0.0) for i, srcctfI in enumerate(Cmap): ellapse_time = time.time() - tic remain_time = float(N_D - i) * ellapse_time / max(i, 1) print("\r%.2f Percent.. (Elapsed: %s, Remaining: %s)" % (i / float(N_D) * 100.0, format_timedelta(ellapse_time), format_timedelta(remain_time)), end='') sys.stdout.flush() # Get the CTF for this image cCTF = srcctf_stack.get_ctf(srcctfI) if prevctfI != srcctfI: genctfI = genctf_stack.add_ctf(cCTF) C = cCTF.dense_ctf(N, psize, bfactor).reshape((N, N)) prevctfI = srcctfI # Randomly generate the viewing direction/shift pt = np.random.randn(3) pt /= np.linalg.norm(pt) psi = 2 * np.pi * np.random.rand() EA = geometry.genEA(pt)[0] EA[2] = psi # Rotate coordinates and get slice image by interpolation R = geometry.rotmat3D_EA(*EA)[:, 0:2] rotated_coords = R.dot(coords.T).T + int(N / 2) slice_data = slicing_func(rotated_coords) intensity = TtoF.dot(slice_data) np.maximum(intensity, 0.0, out=intensity) # Add poisson noise img = np.float_(np.random.poisson(intensity.reshape(N, N))) np.maximum(1e-8, img, out=img) imgdata[i] = np.require(img, dtype=density.real_t) genctf_stack.add_img(genctfI, PHI=EA[0] * 180.0 / np.pi, THETA=EA[1] * 180.0 / np.pi, PSI=EA[2] * 180.0 / np.pi, SHX=0.0, SHY=0.0) pardata['R'].append(R) print("\n\rDone in ", time.time() - tic, " seconds.") return imgdata, genctf_stack, pardata, mscope_params
import sincint M = mrc.readMRC('./particle/EMD-6044.mrc') # M = mrc.readMRC('./particle/1AON.mrc') # M = M / np.sum(M) M = M[:124, :124, :124] mrc.writeMRC('./particle/EMD-6044-cropped.mrc', M, psz=3.0) N = M.shape[0] print(M.shape) rad = 1 kernel = 'lanczos' ksize = 4 xy, trunc_xy, truncmask = geometry.gencoords(N, 2, rad, True) # premult = cryoops.compute_premultiplier(N, kernel='lanczos', kernsize=6) premult = cryoops.compute_premultiplier(N, kernel, ksize) TtoF = sincint.gentrunctofull(N=N, rad=rad) fM = density.real_to_fspace(M) prefM = density.real_to_fspace(premult.reshape( (1, 1, -1)) * premult.reshape((1, -1, 1)) * premult.reshape((-1, 1, 1)) * M) EAs_grid = healpix.gen_EAs_grid(nside=2, psi_step=360) Rs = [geometry.rotmat3D_EA(*EA)[:, 0:2] for EA in EAs_grid] slice_ops = cryoops.compute_projection_matrix(Rs, N, kern='lanczos', kernsize=ksize, rad=rad, projdirtype='rots') slices_sampled = cryoem.getslices(fM, slice_ops).reshape((EAs_grid.shape[0], trunc_xy.shape[0])) premult_slices_sampled = cryoem.getslices(prefM, slice_ops).reshape((EAs_grid.shape[0], trunc_xy.shape[0]))
cs = 2.0 df1, df2, angast = 44722, 49349, 45.0 * (np.pi / 180.0) dscale = 1.0 v1 = compute_ctf(fcoords, rots, akv, cs, wgh, df1, df2, angast, dscale).reshape((-1,)) v2 = compute_ctf(rotfcoords, None, akv, cs, wgh, df1, df2, angast, dscale).reshape((-1,)) # This being small confirms that using the rots parameter is equivalent to # rotating the coordinates print(np.abs(v1 - v2).max()) N = 512 psz = 5.6 rad = 0.25 fcoords = geometry.gencoords(N, 2, rad) / (N * psz) ctf1_rot = compute_full_ctf( rots, N, psz, akv, cs, wgh, df1, df2, angast, dscale, None) ctf2_full = compute_full_ctf( None, N, psz, akv, cs, wgh, df1, df2, angast, dscale, None) P_rot = coops.compute_inplanerot_matrix(rots, N, 'lanczos', 10, rad) ctf2_rot = P_rot.dot(ctf2_full).reshape((-1,)) P_null = coops.compute_inplanerot_matrix(np.array([0]), N, 'linear', 2, rad) ctf1_rot = P_null.dot(ctf1_rot).reshape((-1,)) roterr = ctf1_rot - ctf2_rot relerr = np.abs(roterr) / np.maximum(np.abs(ctf1_rot), np.abs(ctf2_rot)) # This being small confirms that compute_inplane_rotmatrix and rots use
def set_data(self,cparams,minibatch): self.params = cparams self.minibatch = minibatch factoredRI = cparams.get('likelihood_factored_slicing',True) max_freq = cparams['max_frequency'] psize = cparams['pixel_size'] rad_cutoff = cparams.get('rad_cutoff', 1.0) rad = min(rad_cutoff,max_freq*2.0*psize) self.xy, self.trunc_xy, self.truncmask = gencoords(self.N, 2, rad, True) self.trunc_freq = np.require(self.trunc_xy / (self.N*psize), dtype=np.float32) self.N_T = self.trunc_xy.shape[0] interp_change = self.rad != rad or self.factoredRI != factoredRI if interp_change: print("Iteration {0}: freq = {3}, rad = {1}, N_T = {2}".format(cparams['iteration'], rad, self.N_T, max_freq)) self.rad = rad self.factoredRI = factoredRI # Setup the quadrature schemes if not factoredRI: self.set_proj_quad(rad) else: self.set_slice_quad(rad) self.set_inplane_quad(rad) # Check shift quadrature self.set_shift_quad(rad) # Setup inlier model self.inlier_sigma2 = cparams['sigma']**2 base_sigma2 = self.cryodata.noise_var if isinstance(self.inlier_sigma2,np.ndarray): self.inlier_sigma2 = self.inlier_sigma2.reshape(self.truncmask.shape) self.inlier_sigma2_trunc = self.inlier_sigma2[self.truncmask != 0] self.inlier_const = (self.N_T/2.0)*np.log(2.0*np.pi) + 0.5*np.sum(np.log(self.inlier_sigma2_trunc)) else: self.inlier_sigma2_trunc = self.inlier_sigma2 self.inlier_const = (self.N_T/2.0)*np.log(2.0*np.pi*self.inlier_sigma2) # Compute the likelihood for the image content outside of rad _,_,fspace_truncmask = gencoords(self.fspace_stack.get_num_pixels(), 2, rad*self.fspace_stack.get_num_pixels()/self.N, True) self.imgpower = np.empty((self.minibatch['N_M'],),dtype=density.real_t) self.imgpower_trunc = np.empty((self.minibatch['N_M'],),dtype=density.real_t) for idx,Idx in enumerate(self.minibatch['img_idxs']): Img = self.fspace_stack.get_image(Idx) self.imgpower[idx] = np.sum(Img.real**2) + np.sum(Img.imag**2) Img_trunc = Img[fspace_truncmask.reshape(Img.shape) == 0] self.imgpower_trunc[idx] = np.sum(Img_trunc.real**2) + np.sum(Img_trunc.imag**2) like_trunc = 0.5*self.imgpower_trunc/base_sigma2 self.inlier_like_trunc = like_trunc self.inlier_const += ((self.N**2 - self.N_T)/2.0)*np.log(2.0*np.pi*base_sigma2) # Setup the envelope function envelope = self.params.get('exp_envelope',None) if envelope is not None: envelope = envelope.reshape((-1,)) envelope = envelope[self.truncmask != 0] envelope = np.require(envelope,dtype=np.float32) else: bfactor = self.params.get('learn_like_envelope_bfactor',500.0) if bfactor is not None: freqs = np.sqrt(np.sum(self.trunc_xy**2,axis=1))/(psize*self.N) envelope = ctf.envelope_function(freqs,bfactor) self.envelope = envelope
import sincint M = mrc.readMRC('./particle/EMD-6044.mrc') # M = mrc.readMRC('./particle/1AON.mrc') # M = M / np.sum(M) M = M[:124, :124, :124] mrc.writeMRC('./particle/EMD-6044-cropped.mrc', M, psz=3.0) N = M.shape[0] print(M.shape) rad = 1 kernel = 'lanczos' ksize = 4 xy, trunc_xy, truncmask = geometry.gencoords(N, 2, rad, True) # premult = cryoops.compute_premultiplier(N, kernel='lanczos', kernsize=6) premult = cryoops.compute_premultiplier(N, kernel, ksize) TtoF = sincint.gentrunctofull(N=N, rad=rad) fM = density.real_to_fspace(M) prefM = density.real_to_fspace( premult.reshape((1, 1, -1)) * premult.reshape( (1, -1, 1)) * premult.reshape((-1, 1, 1)) * M) EAs_grid = healpix.gen_EAs_grid(nside=2, psi_step=360) Rs = [geometry.rotmat3D_EA(*EA)[:, 0:2] for EA in EAs_grid] slice_ops = cryoops.compute_projection_matrix(Rs, N, kern='lanczos', kernsize=ksize,
def ctf_vis(): import cryoops as coops fcoords = np.random.randn(10, 2) rots = np.array([np.pi / 3.0]) R = np.array([[np.cos(rots), -np.sin(rots)], [np.sin(rots), np.cos(rots)]]).reshape((2, 2)) rotfcoords = np.dot(fcoords, R.T) akv = 200 wgh = 0.07 cs = 2.0 df1, df2, angast = 44722, 49349, 45.0 * (np.pi / 180.0) dscale = 1.0 v1 = ctf.compute_ctf(fcoords, rots, akv, cs, wgh, df1, df2, angast, dscale).reshape((-1,)) v2 = ctf.compute_ctf(rotfcoords, None, akv, cs, wgh, df1, df2, angast, dscale).reshape((-1,)) # This being small confirms that using the rots parameter is equivalent to # rotating the coordinates print(np.abs(v1 - v2).max()) # N = 512 N = 128 psz = 5.6 rad = 0.25 fcoords = geometry.gencoords(N, 2, rad) / (N * psz) ctf1_rot = ctf.compute_full_ctf( rots, N, psz, akv, cs, wgh, df1, df2, angast, dscale, None) ctf2_full = ctf.compute_full_ctf( None, N, psz, akv, cs, wgh, df1, df2, angast, dscale, None) print(ctf1_rot.shape) fig, ax = plt.subplots(1, 2) ax[0].imshow(ctf1_rot.reshape(N, N)) ax[1].imshow(ctf2_full.reshape(N, N)) ax[0].set_title('inplane-rotation: 0') ax[1].set_title('inplane-rotation: 60 degree') fig, ax = plt.subplots() x = np.arange(int(-N/2.0), int(N/2.0)) ax.plot(x, ctf1_rot.reshape(N, N)[int(N/2), :], label='0 degree') ax.plot(x, ctf2_full.reshape(N, N)[int(N/2), :], label='60 degree') ax.set_xlabel('pixel') ax.set_ylabel('CTF value') ax.legend() plt.show() P_rot = coops.compute_inplanerot_matrix(rots, N, 'lanczos', 10, rad) ctf2_rot = P_rot.dot(ctf2_full).reshape((-1,)) P_null = coops.compute_inplanerot_matrix(np.array([0]), N, 'linear', 2, rad) ctf1_rot = P_null.dot(ctf1_rot).reshape((-1,)) print(fcoords.shape) print(P_rot.shape) print(ctf2_rot.shape) roterr = ctf1_rot - ctf2_rot relerr = np.abs(roterr) / np.maximum(np.abs(ctf1_rot), np.abs(ctf2_rot)) # This being small confirms that compute_inplane_rotmatrix and rots use # the same rotation convention print(relerr.max(), relerr.mean())