def plot_rmsd(self): processed_idxs = self.cached_workspace.keys() quad_domain_R = self.quad_domain_R num_idxs = len(processed_idxs) top1_rmsd = np.zeros(num_idxs) topN_rmsd = np.zeros(num_idxs) cutoff_R = 7 topN_weight = np.exp(-np.arange(cutoff_R) / 2) topN_weight /= topN_weight.sum() for i, idx in enumerate(processed_idxs): ea = self.cryodata.euler_angles[idx][0:2] tiled_ea = np.tile(ea, (cutoff_R, 1)) cphi_R = self.cached_workspace[idx]['cphi_R'] sorted_indices_R = (-cphi_R).argsort() potential_R = quad_domain_R.dirs[sorted_indices_R[0:cutoff_R]] eas_of_dirs = geometry.genEA(potential_R)[:, 0:2] top1_rmsd[i] = np.sqrt(((ea[0:2] - eas_of_dirs[0])**2).mean()) topN_rmsd[i] = (np.sqrt(((tiled_ea - eas_of_dirs) ** 2).mean(axis=1)) * topN_weight).sum() print("Slicing quadrature scheme, resolution {}, num of points {}".format( quad_domain_R.resolution, len(quad_domain_R.dirs))) print("Top 1 RMSD:", top1_rmsd.mean()) print("Top {} RMSD: {}".format(cutoff_R, topN_rmsd.mean()))
def gen_EAs_randomly(ea=None, num_inplane_angles=360): if ea is None: pt = np.random.randn(3) pt /= np.linalg.norm(pt) ea = geometry.genEA(pt)[0] euler_angles = np.vstack([np.repeat(ea[0], num_inplane_angles), np.repeat(ea[1], num_inplane_angles), np.linspace(0, 2*np.pi, num_inplane_angles, endpoint=False)]).T return ea, euler_angles
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 gen_euler_angles(num_EAs): EAs = list() for i in range(num_EAs): # 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 EAs.append(np.rad2deg(EA)) # print(EA) return EAs
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 gen_exp_samples(num_EAs, phantompath, data_dst): EAs = list() for i in range(num_EAs): # 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 EAs.append(np.rad2deg(EA)) # EAs_grid = gen_ref_EAs_grid() # grid_size = EAs_grid.shape[0] # EAs = EAs_grid[np.random.randint(0, grid_size, size=num_EAs)] ang_star = data_dst + '_gen.star' star.easy_writeSTAR_relion(ang_star, EAs=EAs) projs_star, projs_mrcs = relion.project(phantompath, data_dst, ang=ang_star) return projs_star, projs_mrcs
def __init__(self, model, dataset_params, ctf_params, interp_params={'kern': 'lanczos', 'kernsize': 4.0, 'zeropad': 0, 'dopremult': True}, load_cache=True): self.dataset_params = dataset_params if model is not None: # assert False assert isinstance(model, np.ndarray), "Unexpected data type for input model" self.num_pixels = model.shape[0] N = self.num_pixels self.num_images = dataset_params['num_images'] assert self.num_images > 1, "it's better to make num_images larger than 1." self.pixel_size = float(dataset_params['pixel_size']) euler_angles = dataset_params['euler_angles'] self.is_sym = get_symmetryop(dataset_params.get('symmetry', None)) if euler_angles is None and self.is_sym is None: pt = np.random.randn(self.num_images, 3) pt /= np.linalg.norm(pt, axis=1, keepdims=True) euler_angles = geometry.genEA(pt) euler_angles[:, 2] = 2 * np.pi * np.random.rand(self.num_images) elif euler_angles is None and self.is_sym is not None: euler_angles = np.zeros((self.num_images, 3)) for i, ea in enumerate(euler_angles): while True: pt = np.random.randn(3) pt /= np.linalg.norm(pt) if self.is_sym.in_asymunit(pt.reshape(-1, 3)): break ea[0:2] = geometry.genEA(pt)[0][0:2] ea[2] = 2 * np.pi * np.random.rand() self.euler_angles = euler_angles.reshape((-1, 3)) if ctf_params is not None: self.use_ctf = True ctf_map = ctf.compute_full_ctf(None, N, ctf_params['psize'], ctf_params['akv'], ctf_params['cs'], ctf_params['wgh'], ctf_params['df1'], ctf_params['df2'], ctf_params['angast'], ctf_params['dscale'], ctf_params.get('bfactor', 500)) self.ctf_params = copy(ctf_params) if 'bfactor' in self.ctf_params.keys(): self.ctf_params.pop('bfactor') else: self.use_ctf = False ctf_map = np.ones((N**2,), dtype=density.real_t) kernel = 'lanczos' ksize = 6 rad = 0.95 # premult = cryoops.compute_premultiplier(N, kernel, ksize) TtoF = sincint.gentrunctofull(N=N, rad=rad) base_coords = geometry.gencoords(N, 2, rad) # premulter = premult.reshape((1, 1, -1)) \ # * premult.reshape((1, -1, 1)) \ # * premult.reshape((-1, 1, 1)) # fM = density.real_to_fspace(premulter * model) fM = model # if load_cache: # try: print("Generating Dataset ... :") tic = time.time() imgdata = np.empty((self.num_images, N, N), dtype=density.real_t) for i, ea in zip(range(self.num_images), self.euler_angles): R = geometry.rotmat3D_EA(*ea)[:, 0:2] slop = cryoops.compute_projection_matrix( [R], N, kernel, ksize, rad, 'rots') # D = slop.dot(fM.reshape((-1,))) rotated_coords = R.dot(base_coords.T).T + int(N/2) D = interpn((np.arange(N),) * 3, fM, rotated_coords) np.maximum(D, 0.0, out=D) intensity = ctf_map.reshape((N, N)) * TtoF.dot(D).reshape((N, N)) np.maximum(1e-8, intensity, out=intensity) intensity = np.float_( np.random.poisson(intensity) ) imgdata[i] = np.require(intensity, dtype=density.real_t) self.imgdata = imgdata print(" cost {} seconds.".format(time.time()-tic)) self.set_transform(interp_params) # self.prep_processing() else: euler_angles = [] with open(self.dataset_params['gtpath']) as par: par.readline() # 'C PHI THETA PSI SHX SHY FILM DF1 DF2 ANGAST' while True: try: line = par.readline().split() euler_angles.append([float(line[1]), float(line[2]), float(line[3])]) except Exception: break self.euler_angles = np.deg2rad(np.asarray(euler_angles)) num_images = self.dataset_params.get('num_images', 200) imgdata = mrc.readMRCimgs(self.dataset_params['inpath'], 0, num_images) self.imgdata = np.transpose(imgdata, axes=(2, 0, 1)) self.num_images = self.imgdata.shape[0] self.num_pixels = self.imgdata.shape[1] N = self.num_pixels self.pixel_size = self.dataset_params['resolution'] self.is_sym = self.dataset_params.get('symmetry', None) self.use_ctf = False ctf_map = np.ones((N**2,), dtype=density.real_t) self.set_transform(interp_params)
def genphantomdata(N_D, phantompath): mscope_params = { 'akv': 200, 'wgh': 0.07, 'cs': 2.0, 'psize': 3.0, 'bfactor': 500.0 } M = mrc.readMRC(phantompath) N = M.shape[0] rad = 0.95 M_totalmass = 1000000 # M_totalmass = 1500000 kernel = 'lanczos' ksize = 6 tic = time.time() N_D = int(N_D) N = int(N) rad = float(rad) psize = mscope_params['psize'] bfactor = mscope_params['bfactor'] M_totalmass = float(M_totalmass) ctfparfile = 'particle/examplectfs.par' srcctf_stack = CTFStack(ctfparfile, mscope_params) genctf_stack = GeneratedCTFStack( mscope_params, parfields=['PHI', 'THETA', 'PSI', 'SHX', 'SHY']) TtoF = sincint.gentrunctofull(N=N, rad=rad) Cmap = np.sort( np.random.random_integers(0, srcctf_stack.get_num_ctfs() - 1, N_D)) cryoem.window(M, 'circle') M[M < 0] = 0 if M_totalmass is not None: M *= M_totalmass / M.sum() # oversampling oversampling_factor = 3 psize = psize * oversampling_factor V = density.real_to_fspace_with_oversampling(M, oversampling_factor) fM = V.real**2 + V.imag**2 # mrc.writeMRC('particle/EMD6044_fM_totalmass_{}_oversampling_{}.mrc'.format(str(int(M_totalmass)).zfill(5), oversampling_factor), fM, psz=psize) print("Generating data...") sys.stdout.flush() imgdata = np.empty((N_D, N, N), dtype=density.real_t) pardata = {'R': []} prevctfI = None coords = geometry.gencoords(N, 2, rad) slicing_func = RegularGridInterpolator((np.arange(N), ) * 3, fM, bounds_error=False, fill_value=0.0) for i, srcctfI in enumerate(Cmap): ellapse_time = time.time() - tic remain_time = float(N_D - i) * ellapse_time / max(i, 1) print("\r%.2f Percent.. (Elapsed: %s, Remaining: %s)" % (i / float(N_D) * 100.0, format_timedelta(ellapse_time), format_timedelta(remain_time)), end='') sys.stdout.flush() # Get the CTF for this image cCTF = srcctf_stack.get_ctf(srcctfI) if prevctfI != srcctfI: genctfI = genctf_stack.add_ctf(cCTF) C = cCTF.dense_ctf(N, psize, bfactor).reshape((N, N)) prevctfI = srcctfI # Randomly generate the viewing direction/shift pt = np.random.randn(3) pt /= np.linalg.norm(pt) psi = 2 * np.pi * np.random.rand() EA = geometry.genEA(pt)[0] EA[2] = psi # Rotate coordinates and get slice image by interpolation R = geometry.rotmat3D_EA(*EA)[:, 0:2] rotated_coords = R.dot(coords.T).T + int(N / 2) slice_data = slicing_func(rotated_coords) intensity = TtoF.dot(slice_data) np.maximum(intensity, 0.0, out=intensity) # Add poisson noise img = np.float_(np.random.poisson(intensity.reshape(N, N))) np.maximum(1e-8, img, out=img) imgdata[i] = np.require(img, dtype=density.real_t) genctf_stack.add_img(genctfI, PHI=EA[0] * 180.0 / np.pi, THETA=EA[1] * 180.0 / np.pi, PSI=EA[2] * 180.0 / np.pi, SHX=0.0, SHY=0.0) pardata['R'].append(R) print("\n\rDone in ", time.time() - tic, " seconds.") return imgdata, genctf_stack, pardata, mscope_params
def gen_slices(model_files, fspace=False, log_scale=True): for model in model_files: M = mrc.readMRC(model) N = M.shape[0] print('model size: {0}x{0}x{0}'.format(N)) oversampling_factor = 3 zeropad = oversampling_factor - 1 # oversampling factor = zeropad + 1 psize = 3.0 * oversampling_factor beamstop_freq = 0.003 mask = geometry.gen_dense_beamstop_mask(N, 2, beamstop_freq, psize=psize) # mask = None if fspace: fM = M else: M_totalmass = 1000000 M *= M_totalmass / M.sum() V = density.real_to_fspace_with_oversampling(M, oversampling_factor) fM = V.real ** 2 + V.imag ** 2 mask_3D = geometry.gen_dense_beamstop_mask(N, 3, beamstop_freq, psize=psize) fM *= mask_3D mrc.writeMRC('particle/{}_fM_totalmass_{}_oversampling_{}.mrc'.format( os.path.splitext(os.path.basename(model))[0], str(int(M_totalmass)).zfill(5), oversampling_factor ), fM, psz=psize) slicing_func = RegularGridInterpolator([np.arange(N),]*3, fM, bounds_error=False, fill_value=0.0) coords = geometry.gencoords_base(N, 2) fig, axes = plt.subplots(3, 3, figsize=(12.9, 9.6)) for i, ax in enumerate(axes.flat): row, col = np.unravel_index(i, (3, 3)) # 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 R = geometry.rotmat3D_EA(*EA)[:, 0:2] rotated_coords = R.dot(coords.T).T + int(N/2) img = slicing_func(rotated_coords).reshape(N, N) img = np.require(np.random.poisson(img), dtype=np.float32) if log_scale: img = np.log(np.maximum(img, 0)) if mask is not None: img *= mask im = ax.imshow(img, origin='lower') # cmap='Greys' ticks = [0, int(N/4.0), int(N/2.0), int(N*3.0/4.0), int(N-1)] if row == 2: ax.set_xticks(ticks) else: ax.set_xticks([]) if col == 0: ax.set_yticks(ticks) else: ax.set_yticks([]) fig.colorbar(im, ax=ax) # fig.subplots_adjust(right=0.8) # cbar_ax = fig.add_axes([0.85, 0.15, 0.05, 0.7]) # fig.colorbar(im, cax=cbar_ax) fig.suptitle('simulated experimental data of XFEL for {}'.format(model)) # fig.tight_layout() plt.show()
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 calculate_nside_resolution(): NSIDE = [2**i for i in range(11)] print( 'given nside | number of pixels | resolution (pixel size in degree) | Maximum angular distance (degree) | pixel area (in square degrees)' ) for nside in NSIDE: npix = hp.nside2npix(nside) resol = np.rad2deg(hp.nside2resol(nside)) maxrad = np.rad2deg(hp.max_pixrad(nside)) pixarea = hp.nside2pixarea(nside, degrees=True) print( '{0:^11} | {1:^16} | {2:^33.4f} | {3:^33.4f} | {4:^30.6f}'.format( nside, npix, resol, maxrad, pixarea)) if __name__ == '__main__': calculate_nside_resolution() # generate random distribution of Euler angles v = np.random.randn(100, 3) v = v / np.linalg.norm(v, axis=1).repeat(3).reshape(-1, 3) EA = geometry.genEA(v) phi = EA[:, 0] # phi += 2 * np.pi theta = EA[:, 1] # visulization hp.mollview() hp.visufunc.projscatter(theta, phi, 'r.') hp.graticule() plt.show()
# 512 | 3145728 | 0.1145 | 0.1196 | 0.013114 # 1024 | 12582912 | 0.0573 | 0.0598 | 0.003278 def calculate_nside_resolution(): NSIDE = [2**i for i in range(11)] print('given nside | number of pixels | resolution (pixel size in degree) | Maximum angular distance (degree) | pixel area (in square degrees)') for nside in NSIDE: npix = hp.nside2npix(nside) resol = np.rad2deg(hp.nside2resol(nside)) maxrad = np.rad2deg(hp.max_pixrad(nside)) pixarea = hp.nside2pixarea(nside, degrees=True) print('{0:^11} | {1:^16} | {2:^33.4f} | {3:^33.4f} | {4:^30.6f}'.format(nside, npix, resol, maxrad, pixarea)) if __name__ == '__main__': calculate_nside_resolution() # generate random distribution of Euler angles v = np.random.randn(100,3) v = v / np.linalg.norm(v, axis=1).repeat(3).reshape(-1,3) EA = geometry.genEA(v) phi = EA[:, 0] # phi += 2 * np.pi theta = EA[:, 1] # visulization hp.mollview() hp.visufunc.projscatter(theta, phi, 'r.') hp.graticule() plt.show()
def plot_rmsd(cached_cphi, euler_angles, quad_domain_R, ac=0): processed_idxs = cached_cphi.keys() quad_domain_R = quad_domain_R num_idxs = len(processed_idxs) top1_rmsd = np.zeros(num_idxs) topN_rmsd = np.zeros(num_idxs) topN_C_rmsd = np.zeros(num_idxs) cutoff_R = 5 topN_weight = np.exp(-np.arange(cutoff_R) / 2) topN_weight /= topN_weight.sum() for i, idx in enumerate(processed_idxs): ea = euler_angles[idx][0:2] tiled_ea = np.tile(ea, (cutoff_R, 1)) cphi_R = cached_cphi[idx]['cphi_R'] sorted_indices_R = (-cphi_R).argsort() potential_R = quad_domain_R.dirs[sorted_indices_R[0:cutoff_R]] eas_of_dirs = geometry.genEA(potential_R)[:, 0:2] if ac == 0: top1_rmsd[i] = np.sqrt(((ea[0:2] - eas_of_dirs[0])**2).mean()) topN_rmsd[i] = (np.sqrt( ((tiled_ea - eas_of_dirs)**2).mean(axis=1)) * topN_weight).sum() topN_C_rmsd[i] = (np.sqrt( ((tiled_ea - eas_of_dirs)**2).mean(axis=1))).min() else: top1_dir = eas_of_dirs[0] top1_dir_chiral = np.pi + np.asarray( [+1.0, -1.0]) * eas_of_dirs[0] # 手性对称位置的坐标 top1_dir_chiral[0] -= (top1_dir_chiral[0] > 2 * np.pi) * 2 * np.pi # 超过 2pi 的地方减去 2pi first = np.sqrt(((ea[0:2] - top1_dir)**2).mean()) first_chiral = np.sqrt(((ea[0:2] - top1_dir_chiral)**2).mean()) top1_rmsd[i] = np.min([first, first_chiral]) topN_rmsd[i] = (np.sqrt( ((tiled_ea - eas_of_dirs)**2).mean(axis=1)) * topN_weight).sum() eas_of_dirs_chiral = (np.pi + np.tile([+1.0, -1.0], (cutoff_R, 1)) * eas_of_dirs) eas_of_dirs_chiral[:, 0] -= (eas_of_dirs_chiral[:, 0] > 2 * np.pi) * 2 * np.pi topN_C = np.sqrt(((tiled_ea - eas_of_dirs)**2).mean(axis=1)) topN_C_chiral = np.sqrt( ((tiled_ea - eas_of_dirs_chiral)**2).mean(axis=1)) topN_C_rmsd[i] = (np.minimum(topN_C, topN_C_chiral)).min() print("Slicing quadrature scheme, resolution {}, num of points {}".format( quad_domain_R.resolution, len(quad_domain_R.dirs))) print("Top 1 RMSD:", top1_rmsd.mean()) print("Top {} RMSD: {}".format(cutoff_R, topN_rmsd.mean())) print("Top {}:RMSD {}".format(cutoff_R, topN_C_rmsd.mean())) # plt.hist(np.rad2deg(topN_C_rmsd)) # plt.show() return top1_rmsd.mean(), topN_rmsd.mean(), topN_C_rmsd.mean()