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 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 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 fourierspace_benchmark_comparison(model, ea=None, num_inplane_angles=360, modulus=False, save_animation=False, animation_name=None): N = model.shape[0] ea, euler_angles = gen_EAs_randomly(ea, num_inplane_angles) proj = projector.project(model, ea) one_slice = density.real_to_fspace(proj) rot_projs = projector.project(M, euler_angles) rot_slices = np.zeros_like(rot_projs, dtype=density.complex_t) for i, pj in enumerate(rot_projs): rot_slices[i] = density.real_to_fspace(pj) if modulus: one_slice = np.abs(one_slice) rot_slices = np.abs(rot_slices) if modulus: corr_slice = correlation.get_corr_img(one_slice) corr_rot_slices = correlation.get_corr_imgs(rot_slices) diff = calc_difference(one_slice, rot_slices, only_real=True) corr_diff = calc_difference(corr_slice, corr_rot_slices, only_real=True) vis_real_space_comparison( np.log(rot_slices), np.log(corr_rot_slices), diff, corr_diff, original_img=np.log(one_slice), original_corr_img=np.log(corr_slice), save_animation=save_animation, animation_name=animation_name ) else: corr_slice = np.zeros((int(N/2.0), 360), dtype=density.complex_t) corr_slice.imag = correlation.get_corr_img(one_slice.imag) corr_slice.real = correlation.get_corr_img(one_slice.real) corr_rot_slices = np.zeros((num_inplane_angles, int(N/2.0), 360), dtype=density.complex_t) corr_rot_slices.real = correlation.get_corr_imgs(rot_slices.real) corr_rot_slices.imag = correlation.get_corr_imgs(rot_slices.imag) diff_real, diff_imag = calc_difference(one_slice, rot_slices) corr_diff_real, corr_diff_imag = calc_difference(corr_slice, corr_rot_slices) vis_fourier_space_comparison( rot_slices, corr_rot_slices, diff_real, diff_imag, corr_diff_real, corr_diff_imag, original_img=one_slice, original_corr_img=corr_slice, save_animation=save_animation, animation_name=animation_name )
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 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 get_image(self, idx): if not self.caching: self.transformed = {} if idx not in self.transformed: self.fft_lock.acquire() if self.zeropad: N = self.stack.get_num_pixels() img = self.zpimg img[self.zeropad:( N + self.zeropad), self.zeropad:(N + self.zeropad)] = self.stack.get_image(idx) else: img = self.stack.get_image(idx) if self.premult is not None: img = self.premult * img self.transformed[idx] = density.real_to_fspace(img) self.fft_lock.release() self.fspacesum += self.transformed[idx] self.powersum += self.transformed[idx].real**2 + \ self.transformed[idx].imag**2 self.nsum += 1 return self.transformed[idx]
def get_image(self, idx): if not self.caching: self.transformed = {} if idx not in self.transformed: self.fft_lock.acquire() if self.zeropad: N = self.stack.get_num_pixels() img = self.zpimg img[self.zeropad:(N + self.zeropad), self.zeropad:(N + self.zeropad)] = self.stack.get_image(idx) else: img = self.stack.get_image(idx) if self.premult is not None: img = self.premult * img self.transformed[idx] = density.real_to_fspace(img) self.fft_lock.release() self.fspacesum += self.transformed[idx] self.powersum += self.transformed[idx].real**2 + self.transformed[ idx].imag**2 self.nsum += 1 return self.transformed[idx]
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 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)
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] trunc_slice = slices_sampled[0] premult_trunc_slice = premult_slices_sampled[0]
def dowork(self): """Do one atom of work. I.E. Execute one minibatch""" timing = {} # Time each minibatch tic_mini = time.time() self.iteration += 1 # Fetch the current batches trainbatch = self.cryodata.get_next_minibatch(self.cparams.get('shuffle_minibatches',True)) # Get the current epoch cepoch = self.cryodata.get_epoch(frac=True) epoch = self.cryodata.get_epoch() num_data = self.cryodata.N_D_Train # Evaluate the parameters self.eval_params() timing['setup'] = time.time() - tic_mini # Do hyperparameter learning if self.cparams.get('learn_params',False): tic_learn = time.time() if self.cparams.get('learn_prior_params',True): tic_learn_prior = time.time() self.prior_func.learn_params(self.params, self.cparams, M=self.M, fM=self.fM) timing['learn_prior'] = time.time() - tic_learn_prior if self.cparams.get('learn_likelihood_params',True): tic_learn_like = time.time() self.like_func.learn_params(self.params, self.cparams, M=self.M, fM=self.fM) timing['learn_like'] = time.time() - tic_learn_like if self.cparams.get('learn_prior_params',True) or self.cparams.get('learn_likelihood_params',True): timing['learn_total'] = time.time() - tic_learn # Time each epoch if self.tic_epoch == None: self.ostream("Epoch: %d" % epoch) self.tic_epoch = (tic_mini,epoch) elif self.tic_epoch[1] != epoch: self.ostream("Epoch Total - %.6f seconds " % \ (tic_mini - self.tic_epoch[0])) self.tic_epoch = (tic_mini,epoch) sym = get_symmetryop(self.cparams.get('symmetry',None)) if sym is not None: self.obj.ws[1] = 1.0/sym.get_order() tic_mstats = time.time() self.ostream(self.cparams['name']," Iteration:", self.iteration,\ " Epoch:", epoch, " Host:", socket.gethostname()) # Compute density statistics N = self.cryodata.N M_sum = self.M.sum(dtype=n.float64) M_zeros = (self.M == 0).sum() M_mean = M_sum/N**3 M_max = self.M.max() M_min = self.M.min() # self.ostream(" Density (min/max/avg/sum/zeros): " + # "%.2e / %.2e / %.2e / %.2e / %g " % # (M_min, M_max, M_mean, M_sum, M_zeros)) self.statout.output(total_density=[M_sum], avg_density=[M_mean], nonzero_density=[M_zeros], max_density=[M_max], min_density=[M_min]) timing['density_stats'] = time.time() - tic_mstats # evaluate test batch if requested if self.iteration <= 1 or self.cparams.get('evaluate_test_set',self.iteration%5): tic_test = time.time() testbatch = self.cryodata.get_testbatch() self.obj.set_data(self.cparams,testbatch) testLogP, res_test = self.obj.eval(M=self.M, fM=self.fM, compute_gradient=False) self.outputbatchinfo(testbatch, res_test, testLogP, 'test', 'Test') timing['test_batch'] = time.time() - tic_test else: testLogP, res_test = None, None # setup the wrapper for the objective function tic_objsetup = time.time() self.obj.set_data(self.cparams,trainbatch) self.obj_wrapper.set_objective(self.obj) x0 = self.obj_wrapper.set_density(self.M,self.fM) evalobj = self.obj_wrapper.eval_obj timing['obj_setup'] = time.time() - tic_objsetup # Get step size self.num_data_evals += trainbatch['N_M'] # at least one gradient tic_objstep = time.time() trainLogP, dlogP, v, res_train, extra_num_data = self.step.do_step(x0, self.cparams, self.cryodata, evalobj, batch=trainbatch) # Apply the step x = x0 + v timing['step'] = time.time() - tic_objstep # Convert from parameters to value tic_stepfinalize = time.time() prevM = n.copy(self.M) self.M, self.fM = self.obj_wrapper.convert_parameter(x,comp_real=True) apply_sym = sym is not None and self.cparams.get('perfect_symmetry',True) and self.cparams.get('apply_symmetry',True) if apply_sym: self.M = sym.apply(self.M) # Truncate the density to bounds if they exist if self.cparams['density_lb'] is not None: n.maximum(self.M,self.cparams['density_lb']*self.cparams['modelscale'],out=self.M) if self.cparams['density_ub'] is not None: n.minimum(self.M,self.cparams['density_ub']*self.cparams['modelscale'],out=self.M) # Compute net change self.dM = prevM - self.M # Convert to fourier space (may not be required) if self.fM is None or apply_sym \ or self.cparams['density_lb'] != None \ or self.cparams['density_ub'] != None: self.fM = density.real_to_fspace(self.M) timing['step_finalize'] = time.time() - tic_stepfinalize # Compute step statistics tic_stepstats = time.time() step_size = n.linalg.norm(self.dM) grad_size = n.linalg.norm(dlogP) M_norm = n.linalg.norm(self.M) self.num_data_evals += extra_num_data inc_ratio = step_size / M_norm self.statout.output(step_size=[step_size], inc_ratio=[inc_ratio], grad_size=[grad_size], norm_density=[M_norm]) timing['step_stats'] = time.time() - tic_stepstats # Update import sampling distributions tic_isupdate = time.time() self.sampler_R.perform_update() self.sampler_I.perform_update() self.sampler_S.perform_update() self.diagout.output(global_phi_R=self.sampler_R.get_global_dist()) self.diagout.output(global_phi_I=self.sampler_I.get_global_dist()) self.diagout.output(global_phi_S=self.sampler_S.get_global_dist()) timing['is_update'] = time.time() - tic_isupdate # Output basic diagnostics tic_diagnostics = time.time() self.diagout.output(iteration=self.iteration, epoch=epoch, cepoch=cepoch) if self.logpost_history.N_sum != self.cryodata.N_batches: self.logpost_history.setup(trainLogP,self.cryodata.N_batches) self.logpost_history.set_value(trainbatch['id'],trainLogP) if self.like_history.N_sum != self.cryodata.N_batches: self.like_history.setup(res_train['L'],self.cryodata.N_batches) self.like_history.set_value(trainbatch['id'],res_train['L']) self.outputbatchinfo(trainbatch, res_train, trainLogP, 'train', 'Train') # Dump parameters here to catch the defaults used in evaluation self.diagout.output(params=self.cparams, envelope_mle=self.like_func.get_envelope_mle(), sigma2_mle=self.like_func.get_sigma2_mle(), hostname=socket.gethostname()) self.statout.output(num_data=[num_data], num_data_evals=[self.num_data_evals], iteration=[self.iteration], epoch=[epoch], cepoch=[cepoch], logp=[self.logpost_history.get_mean()], like=[self.like_history.get_mean()], sigma=[self.like_func.get_rmse()], time=[time.time()]) timing['diagnostics'] = time.time() - tic_diagnostics checkpoint_it = self.iteration % self.cparams.get('checkpoint_frequency',50) == 0 save_it = checkpoint_it or self.cparams['save_iteration'] or \ time.time() - self.last_save > self.cparams.get('save_time',n.inf) if save_it: tic_save = time.time() self.last_save = tic_save if self.io_queue.qsize(): print "Warning: IO queue has become backlogged with {0} remaining, waiting for it to clear".format(self.io_queue.qsize()) self.io_queue.join() self.io_queue.put(( 'pkl', self.statout.fname, copy(self.statout.outdict) )) self.io_queue.put(( 'pkl', self.diagout.fname, deepcopy(self.diagout.outdict) )) self.io_queue.put(( 'pkl', self.likeout.fname, deepcopy(self.likeout.outdict) )) self.io_queue.put(( 'mrc', opj(self.outbase,'model.mrc'), \ (n.require(self.M,dtype=density.real_t),self.voxel_size) )) self.io_queue.put(( 'mrc', opj(self.outbase,'dmodel.mrc'), \ (n.require(self.dM,dtype=density.real_t),self.voxel_size) )) if checkpoint_it: self.io_queue.put(( 'cp', self.diagout.fname, self.diagout.fname+'-{0:06}'.format(self.iteration) )) self.io_queue.put(( 'cp', self.likeout.fname, self.likeout.fname+'-{0:06}'.format(self.iteration) )) self.io_queue.put(( 'cp', opj(self.outbase,'model.mrc'), opj(self.outbase,'model-{0:06}.mrc'.format(self.iteration)) )) timing['save'] = time.time() - tic_save time_total = time.time() - tic_mini self.ostream(" Minibatch Total - %.2f seconds Total Runtime - %s" % (time_total, format_timedelta(datetime.now() - self.startdatetime) )) return self.iteration < self.cparams.get('max_iterations',n.inf) and \ cepoch < self.cparams.get('max_epochs',n.inf)
def __init__(self, expbase, cmdparams=None): """cryodata is a CryoData instance. expbase is a path to the base of folder where this experiment's files will be stored. The folder above expbase will also be searched for .params files. These will be loaded first.""" BackgroundWorker.__init__(self) # Create a background thread which handles IO self.io_queue = Queue() self.io_thread = Thread(target=self.ioworker) self.io_thread.daemon = True self.io_thread.start() # General setup ---------------------------------------------------- self.expbase = expbase self.outbase = None # Paramter setup --------------------------------------------------- # search above expbase for params files _,_,filenames = os.walk(opj(expbase,'../')).next() self.paramfiles = [opj(opj(expbase,'../'), fname) \ for fname in filenames if fname.endswith('.params')] # search expbase for params files _,_,filenames = os.walk(opj(expbase)).next() self.paramfiles += [opj(expbase,fname) \ for fname in filenames if fname.endswith('.params')] if 'local.params' in filenames: self.paramfiles += [opj(expbase,'local.params')] # load parameter files self.params = Params(self.paramfiles) self.cparams = None if cmdparams is not None: # Set parameter specified on the command line for k,v in cmdparams.iteritems(): self.params[k] = v # Dataset setup ------------------------------------------------------- self.imgpath = self.params['inpath'] psize = self.params['resolution'] if not isinstance(self.imgpath,list): imgstk = MRCImageStack(self.imgpath,psize) else: imgstk = CombinedImageStack([MRCImageStack(cimgpath,psize) for cimgpath in self.imgpath]) if self.params.get('float_images',True): imgstk.float_images() self.ctfpath = self.params['ctfpath'] mscope_params = self.params['microscope_params'] if not isinstance(self.ctfpath,list): ctfstk = CTFStack(self.ctfpath,mscope_params) else: ctfstk = CombinedCTFStack([CTFStack(cctfpath,mscope_params) for cctfpath in self.ctfpath]) self.cryodata = CryoDataset(imgstk,ctfstk) self.cryodata.compute_noise_statistics() if self.params.get('window_images',True): imgstk.window_images() minibatch_size = self.params['minisize'] testset_size = self.params['test_imgs'] partition = self.params.get('partition',0) num_partitions = self.params.get('num_partitions',1) seed = self.params['random_seed'] if isinstance(partition,str): partition = eval(partition) if isinstance(num_partitions,str): num_partitions = eval(num_partitions) if isinstance(seed,str): seed = eval(seed) self.cryodata.divide_dataset(minibatch_size,testset_size,partition,num_partitions,seed) self.cryodata.set_datasign(self.params.get('datasign','auto')) if self.params.get('normalize_data',True): self.cryodata.normalize_dataset() self.voxel_size = self.cryodata.pixel_size # Iterations setup ------------------------------------------------- self.iteration = 0 self.tic_epoch = None self.num_data_evals = 0 self.eval_params() outdir = self.cparams.get('outdir',None) if outdir is None: if self.cparams.get('num_partitions',1) > 1: outdir = 'partition{0}'.format(self.cparams['partition']) else: outdir = '' self.outbase = opj(self.expbase,outdir) if not os.path.isdir(self.outbase): os.makedirs(self.outbase) # Output setup ----------------------------------------------------- self.ostream = OutputStream(opj(self.outbase,'stdout')) self.ostream(80*"=") self.ostream("Experiment: " + expbase + \ " Kernel: " + self.params['kernel']) self.ostream("Started on " + socket.gethostname() + \ " At: " + time.strftime('%B %d %Y: %I:%M:%S %p')) self.ostream("Git SHA1: " + gitutil.git_get_SHA1()) self.ostream(80*"=") gitutil.git_info_dump(opj(self.outbase, 'gitinfo')) self.startdatetime = datetime.now() # for diagnostics and parameters self.diagout = Output(opj(self.outbase, 'diag'),runningout=False) # for stats (per image etc) self.statout = Output(opj(self.outbase, 'stat'),runningout=True) # for likelihoods of individual images self.likeout = Output(opj(self.outbase, 'like'),runningout=False) self.img_likes = n.empty(self.cryodata.N_D) self.img_likes[:] = n.inf # optimization state vars ------------------------------------------ init_model = self.cparams.get('init_model',None) if init_model is not None: filename = init_model if filename.upper().endswith('.MRC'): M = readMRC(filename) else: with open(filename) as fp: M = cPickle.load(fp) if type(M)==list: M = M[-1]['M'] if M.shape != 3*(self.cryodata.N,): M = cryoem.resize_ndarray(M,3*(self.cryodata.N,),axes=(0,1,2)) else: init_seed = self.cparams.get('init_random_seed',0) + self.cparams.get('partition',0) print "Randomly generating initial density (init_random_seed = {0})...".format(init_seed), ; sys.stdout.flush() tic = time.time() M = cryoem.generate_phantom_density(self.cryodata.N, 0.95*self.cryodata.N/2.0, \ 5*self.cryodata.N/128.0, 30, seed=init_seed) print "done in {0}s".format(time.time() - tic) tic = time.time() print "Windowing and aligning initial density...", ; sys.stdout.flush() # window the initial density wfunc = self.cparams.get('init_window','circle') cryoem.window(M,wfunc) # Center and orient the initial density cryoem.align_density(M) print "done in {0:.2f}s".format(time.time() - tic) # apply the symmetry operator init_sym = get_symmetryop(self.cparams.get('init_symmetry',self.cparams.get('symmetry',None))) if init_sym is not None: tic = time.time() print "Applying symmetry operator...", ; sys.stdout.flush() M = init_sym.apply(M) print "done in {0:.2f}s".format(time.time() - tic) tic = time.time() print "Scaling initial model...", ; sys.stdout.flush() modelscale = self.cparams.get('modelscale','auto') mleDC, _, mleDC_est_std = self.cryodata.get_dc_estimate() if modelscale == 'auto': # Err on the side of a weaker prior by using a larger value for modelscale modelscale = (n.abs(mleDC) + 2*mleDC_est_std)/self.cryodata.N print "estimated modelscale = {0:.3g}...".format(modelscale), ; sys.stdout.flush() self.params['modelscale'] = modelscale self.cparams['modelscale'] = modelscale M *= modelscale/M.sum() print "done in {0:.2f}s".format(time.time() - tic) if mleDC_est_std/n.abs(mleDC) > 0.05: print " WARNING: the DC component estimate has a high relative variance, it may be inaccurate!" if ((modelscale*self.cryodata.N - n.abs(mleDC)) / mleDC_est_std) > 3: print " WARNING: the selected modelscale value is more than 3 std devs different than the estimated one. Be sure this is correct." self.M = n.require(M,dtype=density.real_t) self.fM = density.real_to_fspace(M) self.dM = density.zeros_like(self.M) self.step = eval(self.cparams['optim_algo']) self.step.setup(self.cparams, self.diagout, self.statout, self.ostream) # Objective function setup -------------------------------------------- param_type = self.cparams.get('parameterization','real') cplx_param = param_type in ['complex','complex_coeff','complex_herm_coeff'] self.like_func = eval_objective(self.cparams['likelihood']) self.prior_func = eval_objective(self.cparams['prior']) if self.cparams.get('penalty',None) is not None: self.penalty_func = eval_objective(self.cparams['penalty']) prior_func = SumObjectives(self.prior_func.fspace, \ [self.penalty_func,self.prior_func], None) else: prior_func = self.prior_func self.obj = SumObjectives(cplx_param, [self.like_func,prior_func], [None,None]) self.obj.setup(self.cparams, self.diagout, self.statout, self.ostream) self.obj.set_dataset(self.cryodata) self.obj_wrapper = ObjectiveWrapper(param_type) self.last_save = time.time() self.logpost_history = FiniteRunningSum() self.like_history = FiniteRunningSum() # Importance Samplers ------------------------------------------------- self.is_sym = get_symmetryop(self.cparams.get('is_symmetry',self.cparams.get('symmetry',None))) self.sampler_R = FixedFisherImportanceSampler('_R',self.is_sym) self.sampler_I = FixedFisherImportanceSampler('_I') self.sampler_S = FixedGaussianImportanceSampler('_S') self.like_func.set_samplers(sampler_R=self.sampler_R,sampler_I=self.sampler_I,sampler_S=self.sampler_S)
def updatevis(self, levels=[0.2,0.5,0.8]): if self.M is None or self.diag is None or self.stat is None: return cdiag = self.diag cparams = cdiag['params'] sym = get_symmetryop(cparams.get('symmetry',None)) quad_sym = sym if cparams.get('perfect_symmetry',True) else None resolution = cparams['voxel_size'] name = cparams['name'] maxfreq = cparams['max_frequency'] N = self.M.shape[0] rad_cutoff = cparams.get('rad_cutoff', 1.0) rad = min(rad_cutoff,maxfreq*2.0*resolution) # Show objective function self.show_objective_plot(self.get_figure('stats')) # Show information about noise and error self.show_error_plot(self.get_figure('error')) self.show_noise_plot(self.get_figure('noise')) # Plot the envelope function if we have the info if 'envelope_mle' in cdiag: self.show_envelope_plot(self.get_figure('envelope')) else: self.close_figure('envelope') if sym is None: assert quad_sym is None alignedM,R = c.align_density(self.M) if self.show_grad: aligneddM = c.rotate_density(self.dM,R) else: aligneddM = None else: alignedM, aligneddM = self.M, self.dM R = n.identity(3) self.alignedM,self.aligneddM,self.alignedR = alignedM,aligneddM,R self.fM = density.real_to_fspace(self.M) self.figMslices.set_data(alignedM) glbl_phi_R = n.array([cdiag['global_phi_R']]).ravel() if len(glbl_phi_R) == 1: glbl_phi_R = None glbl_phi_I = cdiag['global_phi_I'] glbl_phi_S = cdiag['global_phi_S'] # Get direction quadrature quad_R = quadrature.quad_schemes[('dir',cparams.get('quad_type_R','sk97'))] quad_degree_R = cparams.get('quad_degree_R','auto') if quad_degree_R == 'auto': usFactor_R = cparams.get('quad_undersample_R', cparams.get('quad_undersample',1.0)) quad_degree_R,_ = quad_R.compute_degree(N,rad,usFactor_R) origlebDirs,_ = quad_R.get_quad_points(quad_degree_R,quad_sym) lebDirs = n.dot(origlebDirs,R) # Get shift quadrature quad_S = quadrature.quad_schemes[('shift',cparams.get('quad_type_S','hermite'))] quad_degree_S = cparams.get('quad_degree_S','auto') if quad_degree_S == 'auto': usFactor_S = cparams.get('quad_undersample_S', cparams.get('quad_undersample',1.0)) quad_degree_S = quad_S.get_degree(N,rad, cparams['quad_shiftsigma']/resolution, cparams['quad_shiftextent']/resolution, usFactor_S) pts_S,_ = quad_S.get_quad_points(quad_degree_S, cparams['quad_shiftsigma']/resolution, cparams['quad_shiftextent']/resolution, cparams.get('quad_shifttrunc','circ')) vmax_R = 5.0/len(glbl_phi_R) vmax_S = 5.0/len(glbl_phi_S) # Density visualization mlab.figure(self.fig1) mlab.clf() self.curr_contours = plot_density(alignedM, self.contours, levels) # dispPhiR = glbl_phi_R # dispDirs = lebDirs # plot_directions(alignedM.shape[0]*dispDirs + alignedM.shape[0]/2.0, # dispPhiR, # 0, vmax_R) mlab.view(focalpoint=[alignedM.shape[0]/2.0,alignedM.shape[0]/2.0,alignedM.shape[0]/2.0],distance=1.5*alignedM.shape[0]) if glbl_phi_R is not None: plt.figure(self.get_figure('global_is_dists').number) plt.clf() plot_importance_dists(name,lebDirs,pts_S*resolution,glbl_phi_R,glbl_phi_I,glbl_phi_S,vmax_R,vmax_S) if self.show_grad: # Statistics of dM self.figdMslices.set_data(aligneddM) plt.figure(self.get_figure('step_stats').number) plt.clf() plt.suptitle(name + ' Step Statistics') plt.subplot(1,2,1) plt.hist(self.dM.reshape((-1,)),bins=0.5*self.dM.shape[0],log=True) plt.title('Voxel Histogram') (fs,raps) = rot_power_spectra(self.dM,resolution=resolution) plt.subplot(1,2,2) plt.plot(fs/(N/2.0)/(2.0*resolution),raps,label='RAPS') plt.plot((rad/(2.0*resolution))*n.ones((2,)), n.array([raps[raps > 0].min(),raps.max()])) plt.yscale('log') plt.title('RAPS Step') if not self.extra_plots: self.close_figure('density_stats') return # Statistics of M self.show_density_plot(self.get_figure('density_stats'))
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 = 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]))
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
zeropad_size = int(zeropad * (N / 2)) zp_N = zeropad_size * 2 + N zp_M_shape = (zp_N,) * 3 ZP_M = np.zeros(zp_M_shape, dtype=density.real_t) zp_M_slicer = (slice( zeropad_size, (N + zeropad_size) ),) * 3 M_totalmass = 5000 M *= M_totalmass / M.sum() ZP_M[zp_M_slicer] = M N = M.shape[0] kernel = 'lanczos' ksize = 6 premult = cryoops.compute_premultiplier(zp_N, kernel, ksize) V = density.real_to_fspace(premult.reshape((1, 1, -1)) * premult.reshape((1, -1, 1)) * premult.reshape((-1, 1, 1)) * ZP_M) # V = density.real_to_fspace(ZP_M) ZP_fM = V.real ** 2 + V.imag ** 2 fM = ZP_fM[zp_M_slicer] # mask_3d_outlier = geometry.gen_dense_beamstop_mask(N, 3, 0.015, psize=2.8) # fM *= mask_3d_outlier # fM = mrc.readMRC('particle/1AON_fM_totalmass_5000.mrc') * mask_3d_outlier imgdata = mrc.readMRCimgs('data/1AON_xfel_5000_totalmass_05000/imgdata.mrc', 420, 1) curr_img = imgdata[:, :, 0] zp_img = np.zeros((zp_N,)*2, dtype=density.real_t) zp_img[zp_M_slicer[0:2]] = curr_img # curr_img = zp_img
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 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
def updatevis(self, levels=[0.2,0.5,0.8]): if self.M is None or self.diag is None or self.stat is None: return cdiag = self.diag cparams = cdiag['params'] sym = get_symmetryop(cparams.get('symmetry',None)) quad_sym = sym if cparams.get('perfect_symmetry',True) else None resolution = cparams['voxel_size'] name = cparams['name'] maxfreq = cparams['max_frequency'] N = self.M.shape[0] rad_cutoff = cparams.get('rad_cutoff', 1.0) rad = min(rad_cutoff,maxfreq*2.0*resolution) # Show objective function self.show_objective_plot(self.get_figure('stats')) # Show information about noise and error self.show_error_plot(self.get_figure('error')) self.show_noise_plot(self.get_figure('noise')) # Plot the envelope function if we have the info if 'envelope_mle' in cdiag: self.show_envelope_plot(self.get_figure('envelope')) else: self.close_figure('envelope') if sym is None: assert quad_sym is None alignedM,R = cryoem.align_density(self.M) if self.show_grad: aligneddM = cryoem.rotate_density(self.dM,R) else: aligneddM = None else: alignedM, aligneddM = self.M, self.dM R = np.identity(3) self.alignedM,self.aligneddM,self.alignedR = alignedM,aligneddM,R self.fM = density.real_to_fspace(self.M) self.figMslices.set_data(alignedM) glbl_phi_R = np.array([cdiag['global_phi_R']]).ravel() if len(glbl_phi_R) == 1: glbl_phi_R = None glbl_phi_I = cdiag['global_phi_I'] if 'global_phi_S' in cdiag: glbl_phi_S = cdiag['global_phi_S'] else: glbl_phi_S = None # Get direction quadrature quad_R = quadrature.quad_schemes[('dir',cparams.get('quad_type_R','sk97'))] quad_degree_R = cparams.get('quad_degree_R','auto') if quad_degree_R == 'auto': usFactor_R = cparams.get('quad_undersample_R', cparams.get('quad_undersample',1.0)) quad_degree_R,_ = quad_R.compute_degree(N,rad,usFactor_R) origlebDirs,_ = quad_R.get_quad_points(quad_degree_R,quad_sym) lebDirs = np.dot(origlebDirs,R) vmax_R = 5.0/len(glbl_phi_R) # Get shift quadrature if 'global_phi_S' in cdiag: quad_S = quadrature.quad_schemes[('shift',cparams.get('quad_type_S','hermite'))] quad_degree_S = cparams.get('quad_degree_S','auto') if quad_degree_S == 'auto': usFactor_S = cparams.get('quad_undersample_S', cparams.get('quad_undersample',1.0)) quad_degree_S = quad_S.get_degree(N,rad, cparams['quad_shiftsigma']/resolution, cparams['quad_shiftextent']/resolution, usFactor_S) pts_S,_ = quad_S.get_quad_points(quad_degree_S, cparams['quad_shiftsigma']/resolution, cparams['quad_shiftextent']/resolution, cparams.get('quad_shifttrunc','circ')) vmax_S = 5.0/len(glbl_phi_S) else: pts_S = np.zeros_like([0]) vmax_S = None # Density visualization mlab.figure(self.fig1) mlab.clf() self.curr_contours = plot_density(alignedM, self.contours, levels) # dispPhiR = glbl_phi_R # dispDirs = lebDirs # plot_directions(alignedM.shape[0]*dispDirs + alignedM.shape[0]/2.0, # dispPhiR, # 0, vmax_R) mlab.view(focalpoint=[alignedM.shape[0]/2.0,alignedM.shape[0]/2.0,alignedM.shape[0]/2.0],distance=1.5*alignedM.shape[0]) if glbl_phi_R is not None: plt.figure(self.get_figure('global_is_dists').number) plt.clf() plot_importance_dists(name,lebDirs,pts_S*resolution,glbl_phi_R,glbl_phi_I,glbl_phi_S,vmax_R,vmax_S) if self.show_grad: # Statistics of dM self.figdMslices.set_data(aligneddM) plt.figure(self.get_figure('step_stats').number) plt.clf() plt.suptitle(name + ' Step Statistics') plt.subplot(1,2,1) plt.hist(self.dM.reshape((-1,)),bins=0.5*self.dM.shape[0],log=True) plt.title('Voxel Histogram') (fs,raps) = rot_power_spectra(self.dM,resolution=resolution) plt.subplot(1,2,2) plt.plot(fs/(N/2.0)/(2.0*resolution),raps,label='RAPS') plt.plot((rad/(2.0*resolution))*np.ones((2,)), np.array([raps[raps > 0].min(),raps.max()])) plt.yscale('log') plt.title('RAPS Step') if not self.extra_plots: self.close_figure('density_stats') return # Statistics of M self.show_density_plot(self.get_figure('density_stats'))
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