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 premult_test(model, kernel='lanczos', kernsize=6): if isinstance(model, str): M = mrc.readMRC(model) elif isinstance(model, np.ndarray): M = model shape = np.asarray(M.shape) assert (shape - shape.mean()).sum() == 0 N = M.shape[0] rad = 0.6 premult = cryoops.compute_premultiplier(N, kernel, kernsize) TtoF = sincint.gentrunctofull(N=N, rad=rad) premulter = premult.reshape((1, 1, -1)) \ * premult.reshape((1, -1, 1)) \ * premult.reshape((-1, 1, 1)) fM = density.real_to_fspace(M) prefM = density.real_to_fspace(premulter * M) 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(fM.reshape((-1, ))) premult_trunc_slice = slop.dot(prefM.reshape((-1, ))) proj = density.fspace_to_real(TtoF.dot(trunc_slice).reshape(N, N)) premult_proj = density.fspace_to_real( TtoF.dot(premult_trunc_slice).reshape(N, N)) fig, ax = plt.subplots(1, 3, figsize=(14.4, 4.8)) im_proj = ax[0].imshow(proj, origin='lower') fig.colorbar(im_proj, ax=ax[0]) ax[0].set_title('no premulter') im_pre = ax[1].imshow(premult_proj, origin='lower') fig.colorbar(im_pre, ax=ax[1]) ax[1].set_title('with premulter') im_diff = ax[2].imshow(proj - premult_proj, origin='lower') fig.colorbar(im_diff, ax=ax[2]) ax[2].set_title('difference of two image') fig.tight_layout() plt.show()
def premult_vis(N=128): kernel_function = ['lanczos', 'sinc'] for kernel in kernel_function: fig, ax = plt.subplots() kernel_size = np.arange(3, 11) for ksize in kernel_size: premult = cryoops.compute_premultiplier(N, kernel, ksize) ax.plot(range(premult.shape[0]), premult, label='kernel size: {}'.format(ksize)) ax.set_title('kernel function: {}'.format(kernel)) ax.set_xlabel('N') ax.legend(frameon=False) plt.show()
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 set_transform(self, interp_params, caching=True): self.caching = caching self.transformed = {} self.interp_params = interp_params zeropad = interp_params.get('zeropad', 0) kernel = interp_params.get('kern', 'lanczos') kernsize = interp_params.get('kernsize', 4) self.zeropad = int(zeropad * (self.get_num_pixels() / 2)) Nzp = 2 * self.zeropad + self.num_pixels self.zpimg = np.zeros((Nzp, Nzp), dtype=density.real_t) if interp_params.get('dopremult', True): premult = cryoops.compute_premultiplier(Nzp, kernel, kernsize) reshape = ((-1, 1), (1, -1)) self.premult = np.prod([premult.reshape(rs) for rs in reshape]) else: self.premult = None if self.premult is not None: assert self.premult.shape[0] == Nzp assert self.premult.shape[1] == Nzp
def project(model, euler_angles, rad=0.95, truncate=False): if isinstance(model, str): M = mrc.readMRC(model) elif isinstance(model, np.ndarray): M = model N = M.shape[0] kernel = 'lanczos' ksize = 6 premult = cryoops.compute_premultiplier(N, kernel, ksize) TtoF = sincint.gentrunctofull(N=N, rad=rad) premulter = premult.reshape((1, 1, -1)) \ * premult.reshape((1, -1, 1)) \ * premult.reshape((-1, 1, 1)) # premulter = 1 fM = density.real_to_fspace(premulter * M) euler_angles = euler_angles.reshape((-1, 3)) num_projs = euler_angles.shape[0] if truncate: projs = np.zeros((num_projs, TtoF.shape[1]), dtype=fM.dtype) else: projs = np.zeros((num_projs, N, N), dtype=M.dtype) for i, ea in enumerate(euler_angles): rot_matrix = geometry.rotmat3D_EA(*ea)[:, 0:2] slop = cryoops.compute_projection_matrix([rot_matrix], N, kernel, ksize, rad, 'rots') trunc_slice = slop.dot(fM.reshape((-1,))) if truncate: projs[i, :] = trunc_slice else: projs[i, :, :] = density.fspace_to_real(TtoF.dot(trunc_slice).reshape(N, N)) if num_projs == 1 and not truncate: projs = projs.reshape((N, N)) return projs
def set_inplane_quad(self,rad): # Get (and generate if needed) the quadrature scheme for inplane rotation params = self.params tic = time.time() degree_I = params.get('quad_degree_I','auto') usFactor_I = params.get('quad_undersample_I',params.get('quad_undersample',1.0)) kern_I = params.get('interp_kernel_I',params.get('interp_kernel',None)) kernsize_I = params.get('interp_kernel_size_I',params.get('interp_kernel_size',None)) zeropad_I = params.get('interp_zeropad_I',params.get('interp_zeropad',0)) dopremult_I = params.get('interp_premult_I',params.get('interp_premult',True)) maxAngle = quadrature.compute_max_angle(self.N,rad,usFactor_I) if degree_I == 'auto': degree_I = np.uint32(np.ceil(2.0 * np.pi / maxAngle)) resolution_I = max(0.5*quadrature.compute_max_angle(self.N,rad), 2.0*np.pi / degree_I) inplane_params = { 'degree':degree_I } interp_params_I = { 'N':self.N, 'kern':kern_I, 'kernsize':kernsize_I, 'rad':rad, 'zeropad':zeropad_I, 'dopremult':dopremult_I } domain_change_I = self.inplane_params != inplane_params interp_change_I = self.inplane_interp != interp_params_I transform_change = self.inplane_interp is None or \ self.inplane_interp['kern'] != interp_params_I['kern'] or \ self.inplane_interp['kernsize'] != interp_params_I['kernsize'] or \ self.inplane_interp['zeropad'] != interp_params_I['zeropad'] if domain_change_I: inplane_quad = {} inplane_quad['resolution'] = resolution_I inplane_quad['thetas'] = np.linspace(0, 2.0*np.pi, degree_I, endpoint=False) inplane_quad['thetas'] += inplane_quad['thetas'][1]/2.0 inplane_quad['W'] = np.require((2.0*np.pi/float(degree_I))*np.ones((degree_I,)),dtype=np.float32) self.quad_domain_I = quadrature.FixedCircleDomain(inplane_quad['thetas'], inplane_quad['resolution']) self.N_I = len(self.quad_domain_I) self.sampler_I.setup(params, self.N_D, self.N_D_Train, self.quad_domain_I) self.inplane_quad = inplane_quad self.inplane_params = inplane_params if domain_change_I or interp_change_I: print(" Inplane Ops: %d, " % self.N_I); sys.stdout.flush() if self.N_I < self.otf_thresh_I: self.using_precomp_inplane = True print("generated in", end=''); sys.stdout.flush() self.inplane_ops = self.quad_domain_I.compute_operator(interp_params_I) print(" {0} secs.".format(time.time() - tic)) else: self.using_precomp_inplane = False print("generating OTF.") self.inplane_ops = None self.inplane_interp = interp_params_I if transform_change: if dopremult_I: premult = cryoops.compute_premultiplier(self.N + 2*int(interp_params_I['zeropad']*(self.N/2)), interp_params_I['kern'],interp_params_I['kernsize']) premult = premult.reshape((-1,1)) * premult.reshape((1,-1)) else: premult = None self.fspace_stack.set_transform(premult,interp_params_I['zeropad'])
def set_slice_quad(self,rad): # Get (and generate if needed) the quadrature scheme for slicing params = self.params tic = time.time() N = self.N degree_R = params.get('quad_degree_R','auto') quad_scheme_R = params.get('quad_type_R','sk97') sym = get_symmetryop(params.get('symmetry',None)) if params.get('perfect_symmetry',True) else None usFactor_R = params.get('quad_undersample_R',params.get('quad_undersample',1.0)) kern_R = params.get('interp_kernel_R',params.get('interp_kernel',None)) kernsize_R = params.get('interp_kernel_size_R',params.get('interp_kernel_size',None)) zeropad_R = params.get('interp_zeropad_R',params.get('interp_zeropad',0)) dopremult_R = params.get('interp_premult_R',params.get('interp_premult',True)) quad_R = quadrature.quad_schemes[('dir',quad_scheme_R)] if degree_R == 'auto': degree_R,resolution_R = quad_R.compute_degree(N,rad,usFactor_R) resolution_R = max(0.5*quadrature.compute_max_angle(self.N,rad), resolution_R) slice_params = { 'quad_type':quad_scheme_R, 'degree':degree_R, 'sym':sym } interp_params_R = { 'N':self.N, 'kern':kern_R, 'kernsize':kernsize_R, 'rad':rad, 'zeropad':zeropad_R, 'dopremult':dopremult_R } domain_change_R = slice_params != self.slice_params interp_change_R = self.slice_interp != interp_params_R transform_change = self.slice_interp is None or \ self.slice_interp['kern'] != interp_params_R['kern'] or \ self.slice_interp['kernsize'] != interp_params_R['kernsize'] or \ self.slice_interp['zeropad'] != interp_params_R['zeropad'] if domain_change_R: slice_quad = {} slice_quad['resolution'] = resolution_R slice_quad['degree'] = degree_R slice_quad['symop'] = sym slice_quad['dir'],slice_quad['W'] = quad_R.get_quad_points(degree_R,slice_quad['symop']) slice_quad['W'] = np.require(slice_quad['W'], dtype=np.float32) self.quad_domain_R = quadrature.FixedSphereDomain(slice_quad['dir'], slice_quad['resolution'],\ sym=sym) self.N_R = len(self.quad_domain_R) self.sampler_R.setup(params, self.N_D, self.N_D_Train, self.quad_domain_R) self.slice_quad = slice_quad self.slice_params = slice_params if domain_change_R or interp_change_R: symorder = 1 if self.slice_quad['symop'] is None else self.slice_quad['symop'].get_order() print(" Slice Ops: %d, " % self.N_R); sys.stdout.flush() if self.N_R*symorder < self.otf_thresh_R: self.using_precomp_slicing = True print("generated in", end=''); sys.stdout.flush() self.slice_ops = self.quad_domain_R.compute_operator(interp_params_R) print(" {0} secs.".format(time.time() - tic)) Gsz = (self.N_R,self.N_T) self.G = np.empty(Gsz, dtype=self.G_datatype) self.slices = np.empty(np.prod(Gsz), dtype=np.complex64) else: self.using_precomp_slicing = False print("generating OTF.") self.slice_ops = None self.G = np.empty((self.N,self.N,self.N),dtype=self.G_datatype) self.slices = None self.slice_interp = interp_params_R if transform_change: if dopremult_R: premult = cryoops.compute_premultiplier(self.N + 2*int(interp_params_R['zeropad']*(self.N/2)), interp_params_R['kern'],interp_params_R['kernsize']) premult = premult.reshape((-1,1,1)) * premult.reshape((1,-1,1)) * premult.reshape((1,1,-1)) else: premult = None self.slice_premult = premult self.slice_zeropad = interp_params_R['zeropad'] assert interp_params_R['zeropad'] == 0,'Zero padding for slicing not yet implemented'
def set_proj_quad(self,rad): # Get (and generate if needed) the quadrature scheme for slicing params = self.params tic = time.time() N = self.N quad_scheme_R = params.get('quad_type_R','sk97') quad_R = quadrature.quad_schemes[('dir',quad_scheme_R)] degree_R = params.get('quad_degree_R','auto') degree_I = params.get('quad_degree_I','auto') usFactor_R = params.get('quad_undersample_R',params.get('quad_undersample',1.0)) usFactor_I = params.get('quad_undersample_I',params.get('quad_undersample',1.0)) kern_R = params.get('interp_kernel_R',params.get('interp_kernel',None)) kernsize_R = params.get('interp_kernel_size_R',params.get('interp_kernel_size',None)) zeropad_R = params.get('interp_zeropad_R',params.get('interp_zeropad',0)) dopremult_R = params.get('interp_premult_R',params.get('interp_premult',True)) sym = get_symmetryop(params.get('symmetry',None)) if params.get('perfect_symmetry',True) else None maxAngle = quadrature.compute_max_angle(self.N,rad,usFactor_I) if degree_I == 'auto': degree_I = np.uint32(np.ceil(2.0 * np.pi / maxAngle)) if degree_R == 'auto': degree_R,resolution_R = quad_R.compute_degree(N,rad,usFactor_R) resolution_R = max(0.5*quadrature.compute_max_angle(self.N,rad), resolution_R) resolution_I = max(0.5*quadrature.compute_max_angle(self.N,rad), 2.0*np.pi / degree_I) slice_params = { 'quad_type':quad_scheme_R, 'degree':degree_R, 'sym':sym } inplane_params = { 'degree':degree_I } proj_params = { 'quad_type_R':quad_scheme_R, 'degree_R':degree_R, 'sym':sym, 'degree_I':degree_I } interp_params_RI = { 'N':self.N, 'kern':kern_R, 'kernsize':kernsize_R, 'rad':rad, 'zeropad':zeropad_R, 'dopremult':dopremult_R } interp_change_RI = self.proj_interp != interp_params_RI transform_change = self.slice_interp is None or \ self.slice_interp['kern'] != interp_params_RI['kern'] or \ self.slice_interp['kernsize'] != interp_params_RI['kernsize'] or \ self.slice_interp['zeropad'] != interp_params_RI['zeropad'] domain_change_R = self.slice_params != slice_params domain_change_I = self.inplane_params != inplane_params domain_change_RI = self.proj_params != proj_params if domain_change_RI: proj_quad = {} proj_quad['resolution'] = max(resolution_R,resolution_I) proj_quad['degree_R'] = degree_R proj_quad['degree_I'] = degree_I proj_quad['symop'] = sym proj_quad['dir'],proj_quad['W_R'] = quad_R.get_quad_points(degree_R,proj_quad['symop']) proj_quad['W_R'] = np.require(proj_quad['W_R'], dtype=np.float32) proj_quad['thetas'] = np.linspace(0, 2.0*np.pi, degree_I, endpoint=False) proj_quad['thetas'] += proj_quad['thetas'][1]/2.0 proj_quad['W_I'] = np.require((2.0*np.pi/float(degree_I))*np.ones((degree_I,)),dtype=np.float32) self.quad_domain_RI = quadrature.FixedSO3Domain( proj_quad['dir'], -proj_quad['thetas'], proj_quad['resolution'], sym=sym) self.N_RI = len(self.quad_domain_RI) self.proj_quad = proj_quad self.proj_params = proj_params if domain_change_R: self.quad_domain_R = quadrature.FixedSphereDomain(proj_quad['dir'], proj_quad['resolution'], sym=sym) self.N_R = len(self.quad_domain_R) self.sampler_R.setup(params, self.N_D, self.N_D_Train, self.quad_domain_R) self.slice_params = slice_params if domain_change_I: self.quad_domain_I = quadrature.FixedCircleDomain(proj_quad['thetas'], proj_quad['resolution']) self.N_I = len(self.quad_domain_I) self.sampler_I.setup(params, self.N_D, self.N_D_Train, self.quad_domain_I) self.inplane_params = inplane_params if domain_change_RI or interp_change_RI: symorder = 1 if self.proj_quad['symop'] is None else self.proj_quad['symop'].get_order() print(" Projection Ops: %d (%d slice, %d inplane), " % (self.N_RI, self.N_R, self.N_I)); sys.stdout.flush() if self.N_RI*symorder < self.otf_thresh_RI: self.using_precomp_slicing = True print("generated in", end=''); sys.stdout.flush() self.slice_ops = self.quad_domain_RI.compute_operator(interp_params_RI) print(" {0} secs.".format(time.time() - tic)) Gsz = (self.N_RI,self.N_T) self.G = np.empty(Gsz, dtype=self.G_datatype) self.slices = np.empty(np.prod(Gsz), dtype=np.complex64) else: self.using_precomp_slicing = False print("generating OTF.") self.slice_ops = None self.G = np.empty((N,N,N),dtype=self.G_datatype) self.slices = None self.using_precomp_inplane = False self.inplane_ops = None self.proj_interp = interp_params_RI if transform_change: if dopremult_R: premult = cryoops.compute_premultiplier(self.N + 2*int(interp_params_RI['zeropad']*(self.N/2)), interp_params_RI['kern'],interp_params_RI['kernsize']) premult = premult.reshape((-1,1,1)) * premult.reshape((1,-1,1)) * premult.reshape((1,1,-1)) else: premult = None self.slice_premult = premult self.slice_zeropad = interp_params_RI['zeropad'] assert interp_params_RI['zeropad'] == 0,'Zero padding for slicing not yet implemented'
def genphantomdata(N_D, phantompath, ctfparfile): mscope_params = { 'akv': 200, 'wgh': 0.07, 'cs': 2.0, 'psize': 2.8, 'bfactor': 500.0 } N = 128 rad = 0.95 shift_sigma = 3.0 sigma_noise = 25.0 M_totalmass = 80000 kernel = 'lanczos' ksize = 6 premult = cryoops.compute_premultiplier(N, kernel, ksize) tic = time.time() N_D = int(N_D) N = int(N) rad = float(rad) psize = mscope_params['psize'] bfactor = mscope_params['bfactor'] shift_sigma = float(shift_sigma) sigma_noise = float(sigma_noise) M_totalmass = float(M_totalmass) 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 = n.sort( n.random.random_integers(0, srcctf_stack.get_num_ctfs() - 1, N_D)) M = mrc.readMRC(phantompath) cryoem.window(M, 'circle') M[M < 0] = 0 if M_totalmass is not None: M *= M_totalmass / M.sum() V = density.real_to_fspace( premult.reshape((1, 1, -1)) * premult.reshape( (1, -1, 1)) * premult.reshape((-1, 1, 1)) * M) print "Generating data..." sys.stdout.flush() imgdata = n.empty((N_D, N, N), dtype=density.real_t) pardata = {'R': [], 't': []} prevctfI = None 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)), 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**2, )) prevctfI = srcctfI # Randomly generate the viewing direction/shift pt = n.random.randn(3) pt /= n.linalg.norm(pt) psi = 2 * n.pi * n.random.rand() EA = geom.genEA(pt)[0] EA[2] = psi shift = n.random.randn(2) * shift_sigma R = geom.rotmat3D_EA(*EA)[:, 0:2] slop = cryoops.compute_projection_matrix([R], N, kernel, ksize, rad, 'rots') S = cryoops.compute_shift_phases(shift.reshape((1, 2)), N, rad)[0] D = slop.dot(V.reshape((-1, ))) D *= S imgdata[i] = density.fspace_to_real((C * TtoF.dot(D)).reshape( (N, N))) + n.require(n.random.randn(N, N) * sigma_noise, dtype=density.real_t) genctf_stack.add_img(genctfI, PHI=EA[0] * 180.0 / n.pi, THETA=EA[1] * 180.0 / n.pi, PSI=EA[2] * 180.0 / n.pi, SHX=shift[0], SHY=shift[1]) pardata['R'].append(R) pardata['t'].append(shift) print "\rDone in ", time.time() - tic, " seconds." return imgdata, genctf_stack, pardata, mscope_params
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')
def sagd_init(data_dir, model_file, use_angular_correlation=False): data_params = { 'dataset_name': "1AON", 'inpath': os.path.join(data_dir, 'imgdata.mrc'), 'ctfpath': os.path.join(data_dir, 'defocus.txt'), 'microscope_params': { 'akv': 200, 'wgh': 0.07, 'cs': 2.0 }, 'resolution': 2.8, 'sigma': 'noise_std', 'sigma_out': 'data_std', 'minisize': 20, 'test_imgs': 20, 'partition': 0, 'num_partitions': 0, 'random_seed': 1, # 'symmetry': 'C7' } # Setup dataset print("Loading dataset %s" % data_dir) cryodata, _ = dataset_loading_test(data_params) # mleDC, _, mleDC_est_std = cryodata.get_dc_estimate() # modelscale = (np.abs(mleDC) + 2*mleDC_est_std)/cryodata.N modelscale = 1.0 if model_file is not None: print("Loading density map %s" % model_file) M = readMRC(model_file) else: print("Generating random initial density map ...") M = cryoem.generate_phantom_density(cryodata.N, 0.95 * cryodata.N / 2.0, \ 5 * cryodata.N / 128.0, 30, seed=0) M *= modelscale / M.sum() slice_interp = { 'kern': 'lanczos', 'kernsize': 4, 'zeropad': 0, 'dopremult': True } # fM = SimpleKernel.get_fft(M, slice_interp) M_totalmass = 5000 M *= M_totalmass / M.sum() N = M.shape[0] kernel = 'lanczos' ksize = 6 premult = cryoops.compute_premultiplier(N, kernel, ksize) V = density.real_to_fspace( premult.reshape((1, 1, -1)) * premult.reshape( (1, -1, 1)) * premult.reshape((-1, 1, 1)) * M) M = V.real**2 + V.imag**2 freqs_3d = geometry.gencoords_base(N, 3) / (N * data_params['resolution']) freq_radius_3d = np.sqrt((freqs_3d**2).sum(axis=1)) mask_3d_outlier = np.require(np.float_(freq_radius_3d > 0.015).reshape( (N, N, N)), dtype=density.real_t) fM = M * mask_3d_outlier cparams = { 'use_angular_correlation': use_angular_correlation, 'likelihood': 'UnknownRSLikelihood()', 'kernel': 'multicpu', 'prior_name': "'Null'", 'sparsity_lambda': 0.9, 'prior': 'NullPrior()', # 'prior_name': "'CAR'", # 'prior': 'CARPrior()', # 'car_type': 'gauss0.5', # 'car_tau': 75.0, 'iteration': 0, 'pixel_size': cryodata.pixel_size, 'max_frequency': 0.02, 'num_batches': cryodata.N_batches, 'interp_kernel_R': 'lanczos', 'interp_kernel_size_R': 4, 'interp_zeropad_R': 0.0, 'interp_premult_R': True, 'interp_kernel_I': 'lanczos', 'interp_kernel_size_I': 8, 'interp_zeropad_I': 0.0, # 1.0, 'interp_premult_I': True, 'sigma': cryodata.noise_var, 'modelscale': modelscale, # 'symmetry': 'C7' } is_params = { # importance sampling # Ignore the first 50 iterations entirely 'is_prior_prob': max(0.05, 2**(-0.005 * max(0, cparams['iteration'] - 50))), 'is_temperature': max(1.0, 2**(750.0 / max(1, cparams['iteration'] - 50))), 'is_ess_scale': 10, 'is_fisher_chirality_flip': cparams['iteration'] < 2500, 'is_on_R': True, 'is_global_prob_R': 0.9, 'is_on_I': True, 'is_global_prob_I': 1e-10, 'is_on_S': True, 'is_global_prob_S': 0.9, 'is_gaussian_sigmascale_S': 0.67, } cparams.update(is_params) return cryodata, (M, fM), cparams
def genphantomdata(N_D, phantompath, ctfparfile): # mscope_params = {'akv': 200, 'wgh': 0.07, # 'cs': 2.0, 'psize': 2.8, 'bfactor': 500.0} 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 shift_sigma = 3.0 sigma_noise = 25.0 M_totalmass = 80000 kernel = 'lanczos' ksize = 6 premult = cryoops.compute_premultiplier(N, kernel, ksize) tic = time.time() N_D = int(N_D) N = int(N) rad = float(rad) psize = mscope_params['psize'] bfactor = mscope_params['bfactor'] shift_sigma = float(shift_sigma) sigma_noise = float(sigma_noise) M_totalmass = float(M_totalmass) 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() V = density.real_to_fspace( premult.reshape((1, 1, -1)) * premult.reshape((1, -1, 1)) * premult.reshape((-1, 1, 1)) * M) print("Generating data...") sys.stdout.flush() imgdata = np.empty((N_D, N, N), dtype=density.real_t) pardata = {'R': [], 't': []} prevctfI = None 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))) 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**2,)) 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 shift = np.random.randn(2) * shift_sigma R = geometry.rotmat3D_EA(*EA)[:, 0:2] slop = cryoops.compute_projection_matrix( [R], N, kernel, ksize, rad, 'rots') S = cryoops.compute_shift_phases(shift.reshape((1, 2)), N, rad)[0] D = slop.dot(V.reshape((-1,))) D *= S imgdata[i] = density.fspace_to_real((C * TtoF.dot(D)).reshape((N, N))) + np.require( np.random.randn(N, N) * sigma_noise, 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=shift[0], SHY=shift[1]) pardata['R'].append(R) pardata['t'].append(shift) print("\rDone in ", time.time() - tic, " seconds.") return imgdata, genctf_stack, pardata, mscope_params
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])) S = cryoops.compute_shift_phases(np.asarray([100, -20]).reshape((1,2)), N, rad)[0]