def realspace_benchmark_new_algorithm_comparison(model, ea=None, num_inplane_angles=360, rad=0.6, save_animation=False, animation_name=None): N = model.shape[0] ea, euler_angles = gen_EAs_randomly(ea, num_inplane_angles) sl = projector.project(model, ea, rad=rad, truncate=True) rot_slice = projector.project(M, euler_angles, rad=rad, truncate=True) full_slice = projector.trunc_to_full(sl, N, rad) full_rot_slices = projector.trunc_to_full(rot_slice, N, rad) proj = density.fspace_to_real(full_slice) full_rot_projs = np.zeros_like(full_rot_slices, dtype=density.real_t) for i, fs in enumerate(full_rot_slices): full_rot_projs[i] = density.fspace_to_real(fs) trunc_proj = projector.full_to_trunc(proj, rad) trunc_rot_projs = projector.full_to_trunc(full_rot_projs, rad) corr_trunc_proj = correlation.calc_angular_correlation(trunc_proj, N, rad) corr_trunc_rot_projs = correlation.calc_angular_correlation(trunc_rot_projs, N, rad) diff = calc_difference(trunc_proj, trunc_rot_projs, only_real=True) corr_diff = calc_difference(corr_trunc_proj, corr_trunc_rot_projs, only_real=True) vis_real_space_comparison( full_rot_projs, projector.trunc_to_full(corr_trunc_rot_projs, N, rad), diff, corr_diff, original_img=proj, original_corr_img=projector.trunc_to_full(corr_trunc_proj, N, rad), save_animation=save_animation, animation_name=animation_name)
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 apply(self,M,normalize=True,kernel='sinc',kernelsize=3,rad=2): if self.symclass == '': return M if n.iscomplexobj(M): symRs = n.array(self.get_rotations(exclude_dsym=True), dtype=n.float32) N = M.shape[0] if rad*self.get_order()*N < 500: # VERY heuristic symM = sincint.symmetrize_fspace_volume(M,rad,kernel,kernelsize,symRs=symRs) if self.symclass == 'd': symR = n.array([ [ [ -1.0, 0, 0 ], [ 0, 1.0, 0 ], [ 0, 0, -1.0 ] ] ], dtype=n.float32) symM = sincint.symmetrize_fspace_volume(symM,rad,kernel,kernelsize,symRs=symR) else: return density.real_to_fspace(self.apply(density.fspace_to_real(M),normalize=normalize)) else: symRs = n.array(self.get_rotations(), dtype=n.float32) symM = sincint.symmetrize_volume(n.require(M,dtype=n.float32),symRs) # symRs = n.array(self.get_rotations(exclude_dsym=True), dtype=n.float32) # # symM = sincint.symmetrize_volume_z(M,symRs) # # if self.symclass == 'd': # symR = n.array([ [ [ -1.0, 0, 0 ], # [ 0, -1.0, 0 ], # [ 0, 0, 1.0 ] ] ], dtype=n.float32) # symM = n.swapaxes(sincint.symmetrize_volume_z(n.swapaxes(symM,1,2),symR),1,2) if normalize: symM /= self.get_order() return symM
def eval(self, M=None, compute_gradient=True, fM=None, **kwargs): tic_start = time.time() if self.kernel.slice_premult is not None: pfM = density.real_to_fspace(self.kernel.slice_premult * M) else: pfM = density.real_to_fspace(M) pmtime = time.time() - tic_start ret = self.kernel.eval(fM=pfM, M=None, compute_gradient=compute_gradient) if not self.minibatch['test_batch'] and not kwargs.get( 'intermediate', False): tic_record = time.time() curr_var = ret[-1]['sigma2_est'] assert n.all(n.isfinite(curr_var)) if self.error_history.N_sum != self.cryodata.N_batches: self.error_history.setup(curr_var, self.cryodata.N_batches, allow_decay=False) self.error_history.set_value(self.minibatch['id'], curr_var) curr_corr = ret[-1]['correlation'] assert n.all(n.isfinite(curr_corr)) if self.correlation_history.N_sum != self.cryodata.N_batches: self.correlation_history.setup(curr_corr, self.cryodata.N_batches, allow_decay=False) self.correlation_history.set_value(self.minibatch['id'], curr_corr) curr_power = ret[-1]['power'] assert n.all(n.isfinite(curr_power)) if self.power_history.N_sum != self.cryodata.N_batches: self.power_history.setup(curr_power, self.cryodata.N_batches, allow_decay=False) self.power_history.set_value(self.minibatch['id'], curr_power) curr_mask = self.kernel.truncmask if self.mask_history.N_sum != self.cryodata.N_batches: self.mask_history.setup(n.require(curr_mask, dtype=n.float32), self.cryodata.N_batches, allow_decay=False) self.mask_history.set_value(self.minibatch['id'], curr_mask) ret[-1]['like_timing']['record'] = time.time() - tic_record if compute_gradient and self.kernel.slice_premult is not None: tic_record = time.time() ret = (ret[0], self.kernel.slice_premult * density.fspace_to_real(ret[1]), ret[2]) ret[-1]['like_timing']['premult'] = pmtime + time.time( ) - tic_record ret[-1]['like_timing']['total'] = time.time() - tic_start return ret
def convert_parameter(self,x,comp_real=False,comp_fspace=False): is_x0 = x is self.x0 if is_x0: M, fM = self.M0, self.fM0 else: M, fM = param2density(x, self.xtype, self.M0.shape, \ precond=self.precond) if comp_real and M is None: M = density.fspace_to_real(fM) if comp_fspace and fM is None: fM = density.real_to_fspace(M) return M, fM
def eval(self, M=None, compute_gradient=True, fM=None, **kwargs): tic_start = time.time() if self.kernel.slice_premult is not None: pfM = density.real_to_fspace(self.kernel.slice_premult * M) else: pfM = density.real_to_fspace(M) pmtime = time.time() - tic_start ret = self.kernel.eval(fM=pfM, M=None, compute_gradient=compute_gradient) if not self.minibatch['test_batch'] and not kwargs.get('intermediate', False): tic_record = time.time() curr_var = ret[-1]['sigma2_est'] assert n.all(n.isfinite(curr_var)) if self.error_history.N_sum != self.cryodata.N_batches: self.error_history.setup(curr_var, self.cryodata.N_batches, allow_decay=False) self.error_history.set_value(self.minibatch['id'], curr_var) curr_corr = ret[-1]['correlation'] assert n.all(n.isfinite(curr_corr)) if self.correlation_history.N_sum != self.cryodata.N_batches: self.correlation_history.setup(curr_corr, self.cryodata.N_batches, allow_decay=False) self.correlation_history.set_value(self.minibatch['id'], curr_corr) curr_power = ret[-1]['power'] assert n.all(n.isfinite(curr_power)) if self.power_history.N_sum != self.cryodata.N_batches: self.power_history.setup(curr_power, self.cryodata.N_batches, allow_decay=False) self.power_history.set_value(self.minibatch['id'], curr_power) curr_mask = self.kernel.truncmask if self.mask_history.N_sum != self.cryodata.N_batches: self.mask_history.setup(n.require(curr_mask, dtype=n.float32), self.cryodata.N_batches, allow_decay=False) self.mask_history.set_value(self.minibatch['id'], curr_mask) ret[-1]['like_timing']['record'] = time.time() - tic_record if compute_gradient and self.kernel.slice_premult is not None: tic_record = time.time() ret = (ret[0], self.kernel.slice_premult * density.fspace_to_real(ret[1]), ret[2]) ret[-1]['like_timing']['premult'] = pmtime + time.time() - tic_record ret[-1]['like_timing']['total'] = time.time() - tic_start return ret
def apply(self, M, normalize=True, kernel='sinc', kernelsize=3, rad=2): if self.symclass == '': return M if np.iscomplexobj(M): symRs = np.array(self.get_rotations(exclude_dsym=True), dtype=np.float32) N = M.shape[0] if rad * self.get_order() * N < 500: # VERY heuristic symM = sincint.symmetrize_fspace_volume(M, rad, kernel, kernelsize, symRs=symRs) if self.symclass == 'd': symR = np.array( [[[-1.0, 0, 0], [0, 1.0, 0], [0, 0, -1.0]]], dtype=np.float32) symM = sincint.symmetrize_fspace_volume(symM, rad, kernel, kernelsize, symRs=symR) else: return density.real_to_fspace( self.apply(density.fspace_to_real(M), normalize=normalize)) else: symRs = np.array(self.get_rotations(), dtype=np.float32) symM = sincint.symmetrize_volume(np.require(M, dtype=np.float32), symRs) # symRs = n.array(self.get_rotations(exclude_dsym=True), dtype=n.float32) # symM = sincint.symmetrize_volume_z(M,symRs) # if self.symclass == 'd': # symR = n.array([ [ [ -1.0, 0, 0 ], # [ 0, -1.0, 0 ], # [ 0, 0, 1.0 ] ] ], dtype=n.float32) # symM = n.swapaxes(sincint.symmetrize_volume_z(n.swapaxes(symM,1,2),symR),1,2) if normalize: symM /= self.get_order() return symM
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 learn_params(self, params, cparams, M=None, fM=None): anyfspace = any([obj.fspace for obj in self.objs]) anyrspace = any([not obj.fspace for obj in self.objs]) N = None if fM is None and anyfspace: assert M is not None, 'M or fM must be set!' N = M.shape[0] fM = density.real_to_fspace(M) elif fM is not None: N = fM.shape[0] if M is None and anyrspace: assert fM is not None, 'M or fM must be set!' N = fM.shape[0] M = density.fspace_to_real(fM) elif M is not None: assert N is None or N == M.shape[0] N = M.shape[0] assert N is not None for obj in self.objs: obj.learn_params(params, cparams, M=M, fM=fM)
def learn_params(self, params, cparams, M=None, fM=None): anyfspace = any([obj.fspace for obj in self.objs]) anyrspace = any([not obj.fspace for obj in self.objs]) N = None if fM is None and anyfspace: assert M is not None, 'M or fM must be set!' N = M.shape[0] fM = density.real_to_fspace(M) elif fM is not None: N = fM.shape[0] if M is None and anyrspace: assert fM is not None, 'M or fM must be set!' N = fM.shape[0] M = density.fspace_to_real(fM) elif M is not None: assert N is None or N == M.shape[0] N = M.shape[0] assert N is not None for obj in self.objs: obj.learn_params(params,cparams,M=M,fM=fM)
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
def eval(self, M=None, fM=None, compute_gradient=True, all_grads=False, **kwargs): anyfspace = any([obj.fspace for obj in self.objs]) anyrspace = any([not obj.fspace for obj in self.objs]) N = None if fM is None and anyfspace: assert M is not None, 'M or fM must be set!' N = M.shape[0] fM = density.real_to_fspace(M) elif fM is not None: N = fM.shape[0] if M is None and anyrspace: assert fM is not None, 'M or fM must be set!' N = fM.shape[0] M = density.fspace_to_real(fM) elif M is not None: assert N is None or N == M.shape[0] N = M.shape[0] assert N is not None logP = 0 logPs = [] if compute_gradient: if all_grads: dlogP = density.zeros_like( fM) if self.fspace else density.zeros_like(M) dlogPs = [] else: if (not self.fspace) or anyrspace: dlogPdM = density.zeros_like(M) if self.fspace or anyfspace: dlogPdfM = density.zeros_like(fM) outputs = {} for w, obj in zip(self.ws, self.objs): if compute_gradient: clogP, cdlogP, coutputs = obj.eval( M=M, fM=fM, compute_gradient=compute_gradient, **kwargs) if w is not None and w != 1: clogP *= w cdlogP *= w if all_grads: if obj.fspace == self.fspace: dlogPs.append(cdlogP) elif self.fspace: dlogPs.append(density.real_to_fspace(cdlogP)) else: dlogPs.append(density.fspace_to_real(cdlogP)) dlogP += dlogPs[-1] else: if obj.fspace: dlogPdfM += cdlogP else: dlogPdM += cdlogP else: clogP, coutputs = obj.eval(M=M, fM=fM, compute_gradient=compute_gradient, **kwargs) if w is not None and w != 1: clogP *= w logP += clogP logPs.append(clogP) outputs.update(coutputs) if compute_gradient and not all_grads: if self.fspace: dlogP = dlogPdfM if anyrspace: dlogP += density.real_to_fspace(dlogPdM) else: dlogP = dlogPdM if anyfspace: dlogP += density.fspace_to_real(dlogPdfM) outputs['all_logPs'] = logPs if compute_gradient and all_grads: outputs['all_dlogPs'] = dlogPs if compute_gradient: return logP, dlogP, outputs else: return logP, outputs
def calc_angular_correlation( trunc_slices, N, rad, beamstop_rad=None, pixel_size=1.0, interpolation='nearest', sort_theta=True, clip=True, outside=False, ): """compute angular correlation for input array outside: True or False (default: False) calculate angular correlation in radius or outside of radius sort_theta: True or False (default: True) sort theta when slicing the same rho in trunc array """ # 1. get a input (single: N_T or multi: N_R x N_T) with normal sequence. # 2. sort truncation array by rho value of polar coordinates # 3. apply angular correlation function to sorted slice for both real part and imaginary part # 4. deal with outlier beyond 3 sigma (no enough points to do sampling via fft) # (oversampling is unavailable, hence dropout points beyond 3 sigma) # 5. return angluar correlation slice with normal sequence. # 1. iscomplex = np.iscomplexobj(trunc_slices) if outside: trunc_xy = gencoords_outside(N, 2, rad) else: if beamstop_rad is None: trunc_xy = geometry.gencoords(N, 2, rad) else: trunc_xy = geometry.gencoords_centermask(N, 2, rad, beamstop_rad) if trunc_slices.ndim < 2: assert trunc_xy.shape[0] == trunc_slices.shape[ 0], "wrong length of trunc slice or wrong radius" else: assert trunc_xy.shape[0] == trunc_slices.shape[ 1], "wrong length of trunc slice or wrong radius" # 2. pol_trunc_xy = cart2pol(trunc_xy) if sort_theta: # lexsort; first, sort rho; second, sort theta sorted_idx = np.lexsort((pol_trunc_xy[:, 1], pol_trunc_xy[:, 0])) else: sorted_idx = np.argsort(pol_trunc_xy[:, 0]) axis = trunc_slices.ndim - 1 sorted_rho = np.take(pol_trunc_xy[:, 0], sorted_idx) sorted_slice = np.take(trunc_slices, sorted_idx, axis=axis) # 3. if 'none' in interpolation: pass elif 'nearest' in interpolation: sorted_rho = np.round(sorted_rho) elif 'linear' in interpolation: raise NotImplementedError() else: raise ValueError('unsupported method for interpolation') # sorted_rho_freqs = sorted_rho / (N * pixel_size) resolution = 1.0 / (N * pixel_size) _, unique_idx, unique_counts = np.unique(sorted_rho, return_index=True, return_counts=True) indices = [slice(None)] * trunc_slices.ndim angular_correlation = np.zeros_like(trunc_slices, dtype=trunc_slices.dtype) for i, count in enumerate(unique_counts): indices[axis] = slice(unique_idx[i], unique_idx[i] + count) # minimum points to do fft (2 or 4 times than Nyquist frequency) # minimum_sample_points = (4 / count) / resolution minimum_sample_points = 2000 same_rho = np.copy(sorted_slice[indices]) if count < minimum_sample_points: for shift in range(count): curr_delta_phi = same_rho * np.roll(same_rho, shift, axis=axis) indices[axis] = unique_idx[i] + shift angular_correlation[indices] = np.mean(curr_delta_phi, axis=axis) else: # use view (slicing) or copy (fancy indexing, np.take(), np.put())? fpcimg_real = density.real_to_fspace( same_rho.real, axes=(axis, )) # polar image in fourier sapce angular_correlation[indices].real = density.fspace_to_real( fpcimg_real * fpcimg_real.conjugate(), axes=(axis, )).real if iscomplex: # FIXME: stupid way. optimize this fpcimg_fourier = density.real_to_fspace( same_rho.imag, axes=(axis, )) # polar image in fourier sapce angular_correlation[indices].imag = density.fspace_to_real( fpcimg_fourier * fpcimg_fourier.conjugate(), axes=(axis, )).real # check inf and nan if np.any(np.isinf(angular_correlation)): warnings.warn( "Some values in angular correlation occur inf. These values have been set to zeros." ) angular_correlation.real[np.isinf(angular_correlation.real)] = 0 if iscomplex: angular_correlation.imag[np.isinf(angular_correlation.imag)] = 0 if np.any(np.isnan(angular_correlation)): warnings.warn( "Some values in angular correlation occur inf. These values have been set to zeros." ) angular_correlation.real[np.isnan(angular_correlation.real)] = 0 if iscomplex: angular_correlation.imag[np.isnan(angular_correlation.imag)] = 0 # 4. if clip: factor = 3.0 for i, count in enumerate(unique_counts): # minimum_sample_points = (4 / count) / resolution indices[axis] = slice(unique_idx[i], unique_idx[i] + count) mean = np.tile(angular_correlation[indices].mean(axis), (count, 1)).T std = np.tile(angular_correlation[indices].std(axis), (count, 1)).T # print(mean) # print(std) vmin = mean.mean(axis) - factor * std.mean(axis) vmax = mean.mean(axis) + factor * std.mean(axis) # if np.all(std < 1e-16): # # why ??? # warnings.warn("Standard deviation all equal to zero") # vmin = mean.mean(axis) - factor * std.mean(axis) # vmax = mean.mean(axis) + factor * std.mean(axis) # else: # # Normalize to N(0, 1) # angular_correlation[indices] = (angular_correlation[indices] - mean) / std # vmin = -factor # vmax = +factor angular_correlation[indices] = np.clip( angular_correlation[indices].T, vmin, vmax).T # set outlier to nearby boundary # 5. corr_trunc_slices = np.take(angular_correlation, sorted_idx.argsort(), axis=axis) return corr_trunc_slices
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
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] trunc_slice = slices_sampled[0] premult_trunc_slice = premult_slices_sampled[0] premult_trunc_slice_shift = S * premult_slices_sampled[0] fourier_slice = TtoF.dot(trunc_slice).reshape((N, N)) real_proj = density.fspace_to_real(fourier_slice) premult_fourier_slice = TtoF.dot(premult_trunc_slice).reshape((N, N)) premult_fourier_slice_shift = TtoF.dot(premult_trunc_slice_shift).reshape( (N, N)) premult_real_proj = density.fspace_to_real(premult_fourier_slice) premult_real_proj_shift = density.fspace_to_real(premult_fourier_slice_shift) fig, axes = plt.subplots(2, 2) ax = axes.flatten() ax[0].imshow(real_proj) ax[1].imshow(np.log(np.abs(fourier_slice))) ax[2].imshow(premult_real_proj) ax[3].imshow(np.log(np.abs(premult_fourier_slice))) fig, axes = plt.subplots(2, 2) ax = axes.flatten()
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] trunc_slice = slices_sampled[0] premult_trunc_slice = premult_slices_sampled[0] premult_trunc_slice_shift = S * premult_slices_sampled[0] fourier_slice = TtoF.dot(trunc_slice).reshape((N, N)) real_proj = density.fspace_to_real(fourier_slice) premult_fourier_slice = TtoF.dot(premult_trunc_slice).reshape((N, N)) premult_fourier_slice_shift = TtoF.dot(premult_trunc_slice_shift).reshape((N, N)) premult_real_proj = density.fspace_to_real(premult_fourier_slice) premult_real_proj_shift = density.fspace_to_real(premult_fourier_slice_shift) fig, axes = plt.subplots(2, 2) ax = axes.flatten() ax[0].imshow(real_proj) ax[1].imshow(np.log(np.abs(fourier_slice))) ax[2].imshow(premult_real_proj) ax[3].imshow(np.log(np.abs(premult_fourier_slice))) fig, axes = plt.subplots(2, 2) ax = axes.flatten() ax[0].imshow(premult_real_proj)
def eval(self, M=None, fM=None, compute_gradient=True, all_grads=False,**kwargs): anyfspace = any([obj.fspace for obj in self.objs]) anyrspace = any([not obj.fspace for obj in self.objs]) N = None if fM is None and anyfspace: assert M is not None, 'M or fM must be set!' N = M.shape[0] fM = density.real_to_fspace(M) elif fM is not None: N = fM.shape[0] if M is None and anyrspace: assert fM is not None, 'M or fM must be set!' N = fM.shape[0] M = density.fspace_to_real(fM) elif M is not None: assert N is None or N == M.shape[0] N = M.shape[0] assert N is not None logP = 0 logPs = [] if compute_gradient: if all_grads: dlogP = density.zeros_like(fM) if self.fspace else density.zeros_like(M) dlogPs = [] else: if (not self.fspace) or anyrspace: dlogPdM = density.zeros_like(M) if self.fspace or anyfspace: dlogPdfM = density.zeros_like(fM) outputs = {} for w,obj in zip(self.ws,self.objs): if compute_gradient: clogP, cdlogP, coutputs = obj.eval(M = M, fM = fM, compute_gradient = compute_gradient, **kwargs) if w is not None and w != 1: clogP *= w cdlogP *= w if all_grads: if obj.fspace == self.fspace: dlogPs.append(cdlogP) elif self.fspace: dlogPs.append(density.real_to_fspace(cdlogP)) else: dlogPs.append(density.fspace_to_real(cdlogP)) dlogP += dlogPs[-1] else: if obj.fspace: dlogPdfM += cdlogP else: dlogPdM += cdlogP else: clogP, coutputs = obj.eval(M = M, fM = fM, compute_gradient = compute_gradient, **kwargs) if w is not None and w != 1: clogP *= w logP += clogP logPs.append(clogP) outputs.update(coutputs) if compute_gradient and not all_grads: if self.fspace: dlogP = dlogPdfM if anyrspace: dlogP += density.real_to_fspace(dlogPdM) else: dlogP = dlogPdM if anyfspace: dlogP += density.fspace_to_real(dlogPdfM) outputs['all_logPs'] = logPs if compute_gradient and all_grads: outputs['all_dlogPs'] = dlogPs if compute_gradient: return logP, dlogP, outputs else: return logP, outputs