def setUp(self): n = 32 L = 8 filters = [ RadialCTFFilter(5, 200, defocus=d, Cs=2.0, alpha=0.1) for d in np.linspace(1.5e4, 2.5e4, 7) ] self.dtype = np.float32 self.noise_var = 0.1848 # Initial noise filter to generate noise images. # Noise variance is set to a value far away that is used to calculate # covariance matrix and CWF coefficients in order to check the function # for rebuilding positive definite covariance matrix. noise_filter = ScalarFilter(dim=2, value=self.noise_var * 0.001) self.src = Simulation(L, n, unique_filters=filters, dtype=self.dtype, noise_filter=noise_filter) self.basis = FFBBasis2D((L, L), dtype=self.dtype) self.coeff = self.basis.evaluate_t(self.src.images(0, self.src.n)) self.ctf_idx = self.src.filter_indices self.ctf_fb = [f.fb_mat(self.basis) for f in self.src.unique_filters] self.cov2d = RotCov2D(self.basis) self.bcov2d = BatchedRotCov2D(self.src, self.basis, batch_size=7)
def testMSE(self): # need larger numbers of images and higher resolution for good MSE dtype = np.float32 img_size = 64 num_imgs = 1024 noise_var = 0.1848 noise_filter = ScalarFilter(dim=2, value=noise_var) filters = [ RadialCTFFilter(5, 200, defocus=d, Cs=2.0, alpha=0.1) for d in np.linspace(1.5e4, 2.5e4, 7) ] # set simulation object sim = Simulation( L=img_size, n=num_imgs, unique_filters=filters, offsets=0.0, amplitudes=1.0, dtype=dtype, noise_filter=noise_filter, ) imgs_clean = sim.projections() # Specify the fast FB basis method for expending the 2D images ffbbasis = FFBBasis2D((img_size, img_size), dtype=dtype) denoiser = DenoiserCov2D(sim, ffbbasis, noise_var) denoised_src = denoiser.denoise(batch_size=64) imgs_denoised = denoised_src.images(0, num_imgs) # Calculate the normalized RMSE of the estimated images. nrmse_ims = (imgs_denoised - imgs_clean).norm() / imgs_clean.norm() self.assertTrue(nrmse_ims < 0.25)
def setUp(self): self.sim = Simulation(n=1024, L=8, filters=[ RadialCTFFilter(defocus=d) for d in np.linspace(1.5e4, 2.5e4, 7) ], seed=0, noise_filter=IdentityFilter(), dtype='single')
def setUp(self): self.dtype = np.float32 L = 8 n = 32 pixel_size = 5.0 * 65 / L voltage = 200 defocus_min = 1.5e4 defocus_max = 2.5e4 defocus_ct = 7 self.noise_var = 1.3957e-4 noise_filter = ScalarFilter(dim=2, value=self.noise_var) unique_filters = [ RadialCTFFilter(pixel_size, voltage, defocus=d, Cs=2.0, alpha=0.1) for d in np.linspace(defocus_min, defocus_max, defocus_ct) ] vols = Volume( np.load(os.path.join(DATA_DIR, "clean70SRibosome_vol.npy")).astype( self.dtype ) ) # RCOPT vols = vols.downsample((L * np.ones(3, dtype=int))) * 1.0e3 # Since FFBBasis2D doesn't yet implement dtype, we'll set this to double to match its built in types. sim = Simulation( n=n, L=L, vols=vols, unique_filters=unique_filters, offsets=0.0, amplitudes=1.0, dtype=self.dtype, noise_filter=noise_filter, ) self.basis = FFBBasis2D((L, L), dtype=self.dtype) self.h_idx = sim.filter_indices self.h_ctf_fb = [filt.fb_mat(self.basis) for filt in unique_filters] self.imgs_clean = sim.projections() self.imgs_ctf_clean = sim.clean_images() self.imgs_ctf_noise = sim.images(start=0, num=n) self.cov2d = RotCov2D(self.basis) self.coeff_clean = self.basis.evaluate_t(self.imgs_clean) self.coeff = self.basis.evaluate_t(self.imgs_ctf_noise)
def setUp(self): L = 32 n = 64 pixel_size = 5 voltage = 200 defocus_min = 1.5e4 defocus_max = 2.5e4 defocus_ct = 7 Cs = 2.0 alpha = 0.1 self.dtype = np.float32 filters = [ RadialCTFFilter(pixel_size, voltage, defocus=d, Cs=Cs, alpha=alpha) for d in np.linspace(defocus_min, defocus_max, defocus_ct) ] vols = Volume( np.load(os.path.join(DATA_DIR, "clean70SRibosome_vol.npy")).astype( self.dtype)) vols = vols.downsample((L * np.ones(3, dtype=int))) sim = Simulation(L=L, n=n, vols=vols, unique_filters=filters, dtype=self.dtype) self.orient_est = CLSyncVoting(sim, L // 2, 36)
def setUpClass(cls): cls.dtype = np.float32 cls.sim = Simulation( n=1024, unique_filters=[ RadialCTFFilter(defocus=d) for d in np.linspace(1.5e4, 2.5e4, 7) ], dtype=cls.dtype, ) basis = FBBasis3D((8, 8, 8), dtype=cls.dtype) cls.noise_variance = 0.0030762743633643615 cls.mean_estimator = MeanEstimator(cls.sim, basis) cls.mean_est = Volume( np.load(os.path.join(DATA_DIR, "mean_8_8_8.npy")).astype(cls.dtype)) # Passing in a mean_kernel argument to the following constructor speeds up some calculations cls.covar_estimator = CovarianceEstimator( cls.sim, basis, mean_kernel=cls.mean_estimator.kernel, preconditioner="none") cls.covar_estimator_with_preconditioner = CovarianceEstimator( cls.sim, basis, mean_kernel=cls.mean_estimator.kernel, preconditioner="circulant", )
def setUp(self): self.sim = Simulation( n=1024, unique_filters=[ RadialCTFFilter(defocus=d) for d in np.linspace(1.5e4, 2.5e4, 7) ], )
def setUp(self): n = 32 L = 8 noise_var = 0.1848 pixel_size = 5 voltage = 200 defocus_min = 1.5e4 defocus_max = 2.5e4 defocus_ct = 7 Cs = 2.0 alpha = 0.1 filters = [ RadialCTFFilter(pixel_size, voltage, defocus=d, Cs=2.0, alpha=0.1) for d in np.linspace(defocus_min, defocus_max, defocus_ct) ] # Since FFBBasis2D doesn't yet implement dtype, we'll set this to double to match its built in types. src = Simulation(L, n, filters=filters, dtype='double') basis = FFBBasis2D((L, L)) unique_filters = list(set(src.filters)) ctf_idx = np.array([unique_filters.index(f) for f in src.filters]) ctf_fb = [f.fb_mat(basis) for f in unique_filters] im = src.images(0, src.n) coeff = basis.evaluate_t(im.data).astype(src.dtype) cov2d = RotCov2D(basis) bcov2d = BatchedRotCov2D(src, basis, batch_size=7) self.src = src self.basis = basis self.ctf_fb = ctf_fb self.ctf_idx = ctf_idx self.cov2d = cov2d self.bcov2d = bcov2d self.coeff = coeff
def setUp(self): sim = Simulation(n=1024, filters=[ RadialCTFFilter(defocus=d) for d in np.linspace(1.5e4, 2.5e4, 7) ]) basis = FBBasis3D((8, 8, 8)) self.estimator = MeanEstimator(sim, basis, preconditioner='none') self.estimator_with_preconditioner = MeanEstimator( sim, basis, preconditioner='circulant')
def setUp(self): self.L = 64 self.n = 128 self.dtype = np.float32 self.noise_filter = FunctionFilter( lambda x, y: np.exp(-(x**2 + y**2) / 2)) self.sim = Simulation( L=self.L, n=self.n, unique_filters=[ RadialCTFFilter(defocus=d) for d in np.linspace(1.5e4, 2.5e4, 7) ], noise_filter=self.noise_filter, dtype=self.dtype, ) self.imgs_org = self.sim.images(start=0, num=self.n)
def setUp(self): self.dtype = np.float32 sim = Simulation( n=1024, unique_filters=[ RadialCTFFilter(defocus=d) for d in np.linspace(1.5e4, 2.5e4, 7) ], dtype=self.dtype, ) basis = FBBasis3D((8, 8, 8), dtype=self.dtype) self.estimator = MeanEstimator(sim, basis, preconditioner="none") self.estimator_with_preconditioner = MeanEstimator( sim, basis, preconditioner="circulant" )
def setUp(self): L = 8 n = 32 C = 1 SNR = 1 pixel_size = 5 voltage = 200 defocus_min = 1.5e4 defocus_max = 2.5e4 defocus_ct = 7 Cs = 2.0 alpha = 0.1 filters = [ RadialCTFFilter(pixel_size, voltage, defocus=d, Cs=2.0, alpha=0.1) for d in np.linspace(defocus_min, defocus_max, defocus_ct) ] # Since FFBBasis2D doesn't yet implement dtype, we'll set this to double to match its built in types. sim = Simulation(n=n, C=C, filters=filters, dtype='double') vols = np.load(os.path.join(DATA_DIR, 'clean70SRibosome_vol.npy')) vols = vols[..., np.newaxis] vols = downsample(vols, (L * np.ones(3, dtype=int))) sim.vols = vols self.basis = FFBBasis2D((L, L)) # use new methods to generate random rotations and clean images sim.rots = qrand_rots(n, seed=0) self.imgs_clean = vol2img(vols[..., 0], sim.rots) self.h_idx = np.array([filters.index(f) for f in sim.filters]) self.filters = filters self.h_ctf_fb = [filt.fb_mat(self.basis) for filt in self.filters] self.imgs_ctf_clean = sim.eval_filters(self.imgs_clean) sim.cache(self.imgs_ctf_clean) power_clean = anorm(self.imgs_ctf_clean)**2 / np.size( self.imgs_ctf_clean) self.noise_var = power_clean / SNR self.imgs_ctf_noise = self.imgs_ctf_clean + np.sqrt( self.noise_var) * randn(L, L, n, seed=0) self.cov2d = RotCov2D(self.basis) self.coeff_clean = self.basis.evaluate_t(self.imgs_clean) self.coeff = self.basis.evaluate_t(self.imgs_ctf_noise)
def testDownsample(self): # generate a 3D map with density decays as Gaussian function g3d = grid_3d(self.L, dtype=self.dtype) coords = np.array( [g3d["x"].flatten(), g3d["y"].flatten(), g3d["z"].flatten()]) sigma = 0.2 vol = np.exp(-0.5 * np.sum(np.abs(coords / sigma)**2, axis=0)).astype( self.dtype) vol = np.reshape(vol, g3d["x"].shape) vols = Volume(vol) # set noise to zero and CFT filters to unity for simulation object noise_var = 0 noise_filter = ScalarFilter(dim=2, value=noise_var) sim = Simulation( L=self.L, n=self.n, vols=vols, offsets=0.0, amplitudes=1.0, unique_filters=[ ScalarFilter(dim=2, value=1) for d in np.linspace(1.5e4, 2.5e4, 7) ], noise_filter=noise_filter, dtype=self.dtype, ) # get images before downsample imgs_org = sim.images(start=0, num=self.n) # get images after downsample max_resolution = 32 sim.downsample(max_resolution) imgs_ds = sim.images(start=0, num=self.n) # Check individual grid points self.assertTrue( np.allclose( imgs_org[:, 32, 32], imgs_ds[:, 16, 16], atol=utest_tolerance(self.dtype), )) # check resolution self.assertTrue(np.allclose(max_resolution, imgs_ds.shape[1])) # check energy conservation after downsample self.assertTrue( np.allclose( anorm(imgs_org.asnumpy(), axes=(1, 2)) / self.L, anorm(imgs_ds.asnumpy(), axes=(1, 2)) / max_resolution, atol=utest_tolerance(self.dtype), ))
def setUpClass(cls): cls.sim = Simulation(n=1024, filters=[ RadialCTFFilter(defocus=d) for d in np.linspace(1.5e4, 2.5e4, 7) ]) basis = FBBasis3D((8, 8, 8)) cls.noise_variance = 0.0030762743633643615 cls.mean_estimator = MeanEstimator(cls.sim, basis) cls.mean_est = np.load(os.path.join(DATA_DIR, 'mean_8_8_8.npy')) # Passing in a mean_kernel argument to the following constructor speeds up some calculations cls.covar_estimator = CovarianceEstimator( cls.sim, basis, mean_kernel=cls.mean_estimator.kernel, preconditioner='none') cls.covar_estimator_with_preconditioner = CovarianceEstimator( cls.sim, basis, mean_kernel=cls.mean_estimator.kernel, preconditioner='circulant')
class SimTestCase(TestCase): def setUp(self): self.sim = Simulation( n=1024, L=8, unique_filters=[ RadialCTFFilter(defocus=d) for d in np.linspace(1.5e4, 2.5e4, 7) ], seed=0, noise_filter=IdentityFilter(), dtype="single", ) def tearDown(self): pass def testGaussianBlob(self): blobs = self.sim.vols.asnumpy() ref = np.load(os.path.join(DATA_DIR, "sim_blobs.npy")) self.assertTrue(np.allclose(blobs, ref)) def testSimulationRots(self): self.assertTrue( np.allclose( self.sim.rots[0, :, :], np.array([ [0.91675498, 0.2587233, 0.30433956], [0.39941773, -0.58404652, -0.70665065], [-0.00507853, 0.76938412, -0.63876622], ]), )) def testSimulationImages(self): images = self.sim.clean_images(0, 512).asnumpy() self.assertTrue( np.allclose( images, np.load(os.path.join(DATA_DIR, "sim_clean_images.npy")), rtol=1e-2, atol=utest_tolerance(self.sim.dtype), )) def testSimulationImagesNoisy(self): images = self.sim.images(0, 512).asnumpy() self.assertTrue( np.allclose( images, np.load(os.path.join(DATA_DIR, "sim_images_with_noise.npy")), rtol=1e-2, atol=utest_tolerance(self.sim.dtype), )) def testSimulationImagesDownsample(self): # The simulation already generates images of size 8 x 8; Downsampling to resolution 8 should thus have no effect self.sim.downsample(8) images = self.sim.clean_images(0, 512).asnumpy() self.assertTrue( np.allclose( images, np.load(os.path.join(DATA_DIR, "sim_clean_images.npy")), rtol=1e-2, atol=utest_tolerance(self.sim.dtype), )) def testSimulationImagesShape(self): # The 'images' method should be tolerant of bounds - here we ask for 1000 images starting at index 1000, # so we'll get back 25 images in return instead images = self.sim.images(1000, 1000) self.assertTrue(images.shape, (8, 8, 25)) def testSimulationImagesDownsampleShape(self): self.sim.downsample(6) first_image = self.sim.images(0, 1)[0] self.assertEqual(first_image.shape, (6, 6)) def testSimulationEigen(self): eigs_true, lambdas_true = self.sim.eigs() self.assertTrue( np.allclose( eigs_true[0, :, :, 2], np.array([ [ -1.67666201e-07, -7.95741380e-06, -1.49160041e-04, -1.10151654e-03, -3.11287888e-03, -3.09157884e-03, -9.91418026e-04, -1.31673165e-04, ], [ -1.15402077e-06, -2.49849709e-05, -3.51658906e-04, -2.21575261e-03, -7.83315487e-03, -9.44795180e-03, -4.07636259e-03, -9.02186439e-04, ], [ -1.88737249e-05, -1.91418396e-04, -1.09021540e-03, -1.02020288e-03, 1.39411855e-02, 8.58035963e-03, -5.54619730e-03, -3.86377703e-03, ], [ -1.21280536e-04, -9.51461843e-04, -3.22565017e-03, -1.05731178e-03, 2.61375736e-02, 3.11595201e-02, 6.40814053e-03, -2.31698658e-02, ], [ -2.44067283e-04, -1.40560151e-03, -6.73082832e-05, 1.44160679e-02, 2.99893934e-02, 5.92632964e-02, 7.75623545e-02, 3.06570008e-02, ], [ -1.53507499e-04, -7.21709803e-04, 8.54929152e-04, -1.27235036e-02, -5.34382043e-03, 2.18879692e-02, 6.22706190e-02, 4.51998860e-02, ], [ -3.00595184e-05, -1.43038429e-04, -2.15870258e-03, -9.99002904e-02, -7.79077187e-02, -1.53395887e-02, 1.88777559e-02, 1.68759506e-02, ], [ 3.22692649e-05, 4.07977635e-03, 1.63959339e-02, -8.68835449e-02, -7.86240026e-02, -1.75694861e-02, 3.24984640e-03, 1.95389288e-03, ], ]), )) def testSimulationMean(self): mean_vol = self.sim.mean_true() self.assertTrue( np.allclose( [ [ 0.00000930, 0.00033866, 0.00490734, 0.01998369, 0.03874487, 0.04617764, 0.02970645, 0.00967604, ], [ 0.00003904, 0.00247391, 0.03818476, 0.12325402, 0.22278425, 0.25246665, 0.14093882, 0.03683474, ], [ 0.00014177, 0.01191146, 0.14421064, 0.38428235, 0.78645319, 0.86522675, 0.44862473, 0.16382280, ], [ 0.00066036, 0.03137806, 0.29226971, 0.97105378, 2.39410496, 2.17099857, 1.23595858, 0.49233940, ], [ 0.00271748, 0.05491289, 0.49955708, 2.05356097, 3.70941424, 3.01578689, 1.51441932, 0.52054572, ], [ 0.00584845, 0.06962635, 0.50568032, 1.99643707, 3.77415895, 2.76039767, 1.04602003, 0.20633197, ], [ 0.00539583, 0.06068972, 0.47008955, 1.17128026, 1.82821035, 1.18743944, 0.30667788, 0.04851476, ], [ 0.00246362, 0.04867788, 0.65284950, 0.65238875, 0.65745538, 0.37955678, 0.08053055, 0.01210055, ], ], mean_vol[0, :, :, 4], )) def testSimulationVolCoords(self): coords, norms, inners = self.sim.vol_coords() self.assertTrue( np.allclose([4.72837704, -4.72837709], coords, atol=1e-4)) self.assertTrue( np.allclose([8.20515764e-07, 1.17550184e-06], norms, atol=1e-4)) self.assertTrue( np.allclose([3.78030562e-06, -4.20475816e-06], inners, atol=1e-4)) def testSimulationCovar(self): covar = self.sim.covar_true() result = [ [ -0.00000289, -0.00005839, -0.00018998, -0.00124722, -0.00003155, +0.00743356, +0.00798143, +0.00303416, ], [ -0.00000776, +0.00018371, +0.00448675, -0.00794970, -0.02988000, -0.00185446, +0.01786612, +0.00685990, ], [ +0.00001144, +0.00324029, +0.03364052, -0.00272520, -0.08976389, -0.05404807, +0.00268740, -0.03081760, ], [ +0.00003204, +0.00909853, +0.07859941, +0.07254293, -0.19365733, -0.09007251, -0.15731451, -0.15690306, ], [ -0.00040561, +0.00685139, +0.11074986, +0.35207557, +0.17264650, -0.16662873, -0.15010859, -0.14292650, ], [ -0.00107461, -0.00497393, +0.04630126, +0.38048555, +0.47915877, +0.05379957, -0.11833663, -0.03372971, ], [ -0.00029630, -0.00485664, -0.00640120, +0.22068169, +0.15419035, +0.08281200, +0.03373241, +0.00103902, ], [ +0.00044323, +0.00850533, +0.09683860, +0.16959519, +0.03629097, +0.03740599, +0.02212356, +0.00318127, ], ] self.assertTrue(np.allclose(result, covar[:, :, 4, 4, 4, 4], atol=1e-4)) def testSimulationEvalMean(self): mean_est = Volume(np.load(os.path.join(DATA_DIR, "mean_8_8_8.npy"))) result = self.sim.eval_mean(mean_est) self.assertTrue( np.allclose(result["err"], 2.664116055950763, atol=1e-4)) self.assertTrue( np.allclose(result["rel_err"], 0.1765943704851626, atol=1e-4)) self.assertTrue( np.allclose(result["corr"], 0.9849211540734224, atol=1e-4)) def testSimulationEvalCovar(self): covar_est = np.load(os.path.join(DATA_DIR, "covar_8_8_8_8_8_8.npy")) result = self.sim.eval_covar(covar_est) self.assertTrue( np.allclose(result["err"], 13.322721549011165, atol=1e-4)) self.assertTrue( np.allclose(result["rel_err"], 0.5958936073938558, atol=1e-4)) self.assertTrue( np.allclose(result["corr"], 0.8405347287741631, atol=1e-4)) def testSimulationEvalCoords(self): mean_est = Volume(np.load(os.path.join(DATA_DIR, "mean_8_8_8.npy"))) eigs_est = Volume( np.load(os.path.join(DATA_DIR, "eigs_est_8_8_8_1.npy"))[..., 0]) clustered_coords_est = np.load( os.path.join(DATA_DIR, "clustered_coords_est.npy")) result = self.sim.eval_coords(mean_est, eigs_est, clustered_coords_est) self.assertTrue( np.allclose( result["err"][:10], [ 1.58382394, 1.58382394, 3.72076112, 1.58382394, 1.58382394, 3.72076112, 3.72076112, 1.58382394, 1.58382394, 1.58382394, ], )) self.assertTrue( np.allclose( result["rel_err"][0, :10], [ 0.11048937, 0.11048937, 0.21684697, 0.11048937, 0.11048937, 0.21684697, 0.21684697, 0.11048937, 0.11048937, 0.11048937, ], )) self.assertTrue( np.allclose( result["corr"][0, :10], [ 0.99390133, 0.99390133, 0.97658719, 0.99390133, 0.99390133, 0.97658719, 0.97658719, 0.99390133, 0.99390133, 0.99390133, ], )) def testSimulationSaveFile(self): # Create a tmpdir in a context. It will be cleaned up on exit. with tempfile.TemporaryDirectory() as tmpdir: # Save the simulation object into STAR and MRCS files star_filepath = os.path.join(tmpdir, "save_test.star") # Save images into one single MRCS file self.sim.save(star_filepath, batch_size=512, save_mode="single", overwrite=False) imgs_org = self.sim.images(start=0, num=1024) # Input saved images into Relion object relion_src = RelionSource(star_filepath, tmpdir, max_rows=1024) imgs_sav = relion_src.images(start=0, num=1024) # Compare original images with saved images self.assertTrue(np.allclose(imgs_org.asnumpy(), imgs_sav.asnumpy())) # Save images into multiple MRCS files based on batch size self.sim.save(star_filepath, batch_size=512, overwrite=False) # Input saved images into Relion object relion_src = RelionSource(star_filepath, tmpdir, max_rows=1024) imgs_sav = relion_src.images(start=0, num=1024) # Compare original images with saved images self.assertTrue(np.allclose(imgs_org.asnumpy(), imgs_sav.asnumpy()))
logger = logging.getLogger(__name__) # %% # Create Simulation Object # ------------------------ # Specify parameters num_vols = 2 # number of volumes img_size = 8 # image size in square num_imgs = 1024 # number of images num_eigs = 16 # number of eigen-vectors to keep # Create a simulation object with specified filters sim = Simulation( L=img_size, n=num_imgs, C=num_vols, unique_filters=[RadialCTFFilter(defocus=d) for d in np.linspace(1.5e4, 2.5e4, 7)], ) # Specify the normal FB basis method for expending the 2D images basis = FBBasis3D((img_size, img_size, img_size)) # Estimate the noise variance. This is needed for the covariance estimation step below. noise_estimator = WhiteNoiseEstimator(sim, batchSize=500) noise_variance = noise_estimator.estimate() logger.info(f"Noise Variance = {noise_variance}") # %% # Estimate Mean Volume and Covariance # ----------------------------------- #
class SimTestCase(TestCase): def setUp(self): self.sim = Simulation(n=1024, L=8, filters=[ RadialCTFFilter(defocus=d) for d in np.linspace(1.5e4, 2.5e4, 7) ], seed=0, noise_filter=IdentityFilter(), dtype='single') def tearDown(self): pass def testGaussianBlob(self): blobs = self.sim.vols self.assertTrue( np.allclose(blobs, np.load(os.path.join(DATA_DIR, 'sim_blobs.npy')))) def testSimulationRots(self): self.assertTrue( np.allclose( self.sim.rots[0, :, :], np.array([[0.91675498, 0.2587233, 0.30433956], [0.39941773, -0.58404652, -0.70665065], [-0.00507853, 0.76938412, -0.63876622]]))) def testSimulationImages(self): images = self.sim.clean_images(0, 512).asnumpy() self.assertTrue( np.allclose(images, np.load(os.path.join(DATA_DIR, 'sim_clean_images.npy')), rtol=1e-2)) def testSimulationImagesNoisy(self): images = self.sim.images(0, 512).asnumpy() self.assertTrue( np.allclose(images, np.load( os.path.join(DATA_DIR, 'sim_images_with_noise.npy')), rtol=1e-2)) def testSimulationImagesDownsample(self): # The simulation already generates images of size 8 x 8; Downsampling to resolution 8 should thus have no effect self.sim.downsample(8) images = self.sim.clean_images(0, 512).asnumpy() self.assertTrue( np.allclose(images, np.load(os.path.join(DATA_DIR, 'sim_clean_images.npy')), rtol=1e-2)) def testSimulationImagesShape(self): # The 'images' method should be tolerant of bounds - here we ask for 1000 images starting at index 1000, # so we'll get back 25 images in return instead images = self.sim.images(1000, 1000) self.assertTrue(images.shape, (8, 8, 25)) def testSimulationEigen(self): eigs_true, lambdas_true = self.sim.eigs() self.assertTrue( np.allclose( eigs_true[:, :, 2, 0], np.array([[ -1.67666201e-07, -7.95741380e-06, -1.49160041e-04, -1.10151654e-03, -3.11287888e-03, -3.09157884e-03, -9.91418026e-04, -1.31673165e-04 ], [ -1.15402077e-06, -2.49849709e-05, -3.51658906e-04, -2.21575261e-03, -7.83315487e-03, -9.44795180e-03, -4.07636259e-03, -9.02186439e-04 ], [ -1.88737249e-05, -1.91418396e-04, -1.09021540e-03, -1.02020288e-03, 1.39411855e-02, 8.58035963e-03, -5.54619730e-03, -3.86377703e-03 ], [ -1.21280536e-04, -9.51461843e-04, -3.22565017e-03, -1.05731178e-03, 2.61375736e-02, 3.11595201e-02, 6.40814053e-03, -2.31698658e-02 ], [ -2.44067283e-04, -1.40560151e-03, -6.73082832e-05, 1.44160679e-02, 2.99893934e-02, 5.92632964e-02, 7.75623545e-02, 3.06570008e-02 ], [ -1.53507499e-04, -7.21709803e-04, 8.54929152e-04, -1.27235036e-02, -5.34382043e-03, 2.18879692e-02, 6.22706190e-02, 4.51998860e-02 ], [ -3.00595184e-05, -1.43038429e-04, -2.15870258e-03, -9.99002904e-02, -7.79077187e-02, -1.53395887e-02, 1.88777559e-02, 1.68759506e-02 ], [ 3.22692649e-05, 4.07977635e-03, 1.63959339e-02, -8.68835449e-02, -7.86240026e-02, -1.75694861e-02, 3.24984640e-03, 1.95389288e-03 ]]))) def testSimulationMean(self): mean_vol = self.sim.mean_true() self.assertTrue( np.allclose([ [ 0.00000930, 0.00033866, 0.00490734, 0.01998369, 0.03874487, 0.04617764, 0.02970645, 0.00967604 ], [ 0.00003904, 0.00247391, 0.03818476, 0.12325402, 0.22278425, 0.25246665, 0.14093882, 0.03683474 ], [ 0.00014177, 0.01191146, 0.14421064, 0.38428235, 0.78645319, 0.86522675, 0.44862473, 0.16382280 ], [ 0.00066036, 0.03137806, 0.29226971, 0.97105378, 2.39410496, 2.17099857, 1.23595858, 0.49233940 ], [ 0.00271748, 0.05491289, 0.49955708, 2.05356097, 3.70941424, 3.01578689, 1.51441932, 0.52054572 ], [ 0.00584845, 0.06962635, 0.50568032, 1.99643707, 3.77415895, 2.76039767, 1.04602003, 0.20633197 ], [ 0.00539583, 0.06068972, 0.47008955, 1.17128026, 1.82821035, 1.18743944, 0.30667788, 0.04851476 ], [ 0.00246362, 0.04867788, 0.65284950, 0.65238875, 0.65745538, 0.37955678, 0.08053055, 0.01210055 ], ], mean_vol[:, :, 4])) def testSimulationVolCoords(self): coords, norms, inners = self.sim.vol_coords() self.assertTrue( np.allclose([4.72837704, -4.72837709], coords, atol=1e-4)) self.assertTrue( np.allclose([8.20515764e-07, 1.17550184e-06], norms, atol=1e-4)) self.assertTrue( np.allclose([3.78030562e-06, -4.20475816e-06], inners, atol=1e-4)) def testSimulationCovar(self): covar = self.sim.covar_true() result = [ [ -0.00000289, -0.00005839, -0.00018998, -0.00124722, -0.00003155, +0.00743356, +0.00798143, +0.00303416 ], [ -0.00000776, +0.00018371, +0.00448675, -0.00794970, -0.02988000, -0.00185446, +0.01786612, +0.00685990 ], [ +0.00001144, +0.00324029, +0.03364052, -0.00272520, -0.08976389, -0.05404807, +0.00268740, -0.03081760 ], [ +0.00003204, +0.00909853, +0.07859941, +0.07254293, -0.19365733, -0.09007251, -0.15731451, -0.15690306 ], [ -0.00040561, +0.00685139, +0.11074986, +0.35207557, +0.17264650, -0.16662873, -0.15010859, -0.14292650 ], [ -0.00107461, -0.00497393, +0.04630126, +0.38048555, +0.47915877, +0.05379957, -0.11833663, -0.03372971 ], [ -0.00029630, -0.00485664, -0.00640120, +0.22068169, +0.15419035, +0.08281200, +0.03373241, +0.00103902 ], [ +0.00044323, +0.00850533, +0.09683860, +0.16959519, +0.03629097, +0.03740599, +0.02212356, +0.00318127 ], ] self.assertTrue(np.allclose(result, covar[:, :, 4, 4, 4, 4], atol=1e-4)) def testSimulationEvalMean(self): mean_est = np.load(os.path.join(DATA_DIR, 'mean_8_8_8.npy')) result = self.sim.eval_mean(mean_est) self.assertTrue( np.allclose(result['err'], 2.664116055950763, atol=1e-4)) self.assertTrue( np.allclose(result['rel_err'], 0.1765943704851626, atol=1e-4)) self.assertTrue( np.allclose(result['corr'], 0.9849211540734224, atol=1e-4)) def testSimulationEvalCovar(self): covar_est = np.load(os.path.join(DATA_DIR, 'covar_8_8_8_8_8_8.npy')) result = self.sim.eval_covar(covar_est) self.assertTrue( np.allclose(result['err'], 13.322721549011165, atol=1e-4)) self.assertTrue( np.allclose(result['rel_err'], 0.5958936073938558, atol=1e-4)) self.assertTrue( np.allclose(result['corr'], 0.8405347287741631, atol=1e-4)) def testSimulationEvalCoords(self): mean_est = np.load(os.path.join(DATA_DIR, 'mean_8_8_8.npy')) eigs_est = np.load(os.path.join(DATA_DIR, 'eigs_est_8_8_8_1.npy')) clustered_coords_est = np.load( os.path.join(DATA_DIR, 'clustered_coords_est.npy')) result = self.sim.eval_coords(mean_est, eigs_est, clustered_coords_est) self.assertTrue( np.allclose(result['err'][:10], [ 1.58382394, 1.58382394, 3.72076112, 1.58382394, 1.58382394, 3.72076112, 3.72076112, 1.58382394, 1.58382394, 1.58382394 ])) self.assertTrue( np.allclose(result['rel_err'][:10], [ 0.11048937, 0.11048937, 0.21684697, 0.11048937, 0.11048937, 0.21684697, 0.21684697, 0.11048937, 0.11048937, 0.11048937 ])) self.assertTrue( np.allclose(result['corr'][:10], [ 0.99390133, 0.99390133, 0.97658719, 0.99390133, 0.99390133, 0.97658719, 0.97658719, 0.99390133, 0.99390133, 0.99390133 ]))
class PreprocessPLTestCase(TestCase): def setUp(self): self.L = 64 self.n = 128 self.dtype = np.float32 self.noise_filter = FunctionFilter( lambda x, y: np.exp(-(x**2 + y**2) / 2)) self.sim = Simulation( L=self.L, n=self.n, unique_filters=[ RadialCTFFilter(defocus=d) for d in np.linspace(1.5e4, 2.5e4, 7) ], noise_filter=self.noise_filter, dtype=self.dtype, ) self.imgs_org = self.sim.images(start=0, num=self.n) def testPhaseFlip(self): self.sim.phase_flip() imgs_pf = self.sim.images(start=0, num=self.n) # check energy conservation self.assertTrue( np.allclose( anorm(self.imgs_org.asnumpy(), axes=(1, 2)), anorm(imgs_pf.asnumpy(), axes=(1, 2)), )) def testDownsample(self): # generate a 3D map with density decays as Gaussian function g3d = grid_3d(self.L, dtype=self.dtype) coords = np.array( [g3d["x"].flatten(), g3d["y"].flatten(), g3d["z"].flatten()]) sigma = 0.2 vol = np.exp(-0.5 * np.sum(np.abs(coords / sigma)**2, axis=0)).astype( self.dtype) vol = np.reshape(vol, g3d["x"].shape) vols = Volume(vol) # set noise to zero and CFT filters to unity for simulation object noise_var = 0 noise_filter = ScalarFilter(dim=2, value=noise_var) sim = Simulation( L=self.L, n=self.n, vols=vols, offsets=0.0, amplitudes=1.0, unique_filters=[ ScalarFilter(dim=2, value=1) for d in np.linspace(1.5e4, 2.5e4, 7) ], noise_filter=noise_filter, dtype=self.dtype, ) # get images before downsample imgs_org = sim.images(start=0, num=self.n) # get images after downsample max_resolution = 32 sim.downsample(max_resolution) imgs_ds = sim.images(start=0, num=self.n) # Check individual grid points self.assertTrue( np.allclose( imgs_org[:, 32, 32], imgs_ds[:, 16, 16], atol=utest_tolerance(self.dtype), )) # check resolution self.assertTrue(np.allclose(max_resolution, imgs_ds.shape[1])) # check energy conservation after downsample self.assertTrue( np.allclose( anorm(imgs_org.asnumpy(), axes=(1, 2)) / self.L, anorm(imgs_ds.asnumpy(), axes=(1, 2)) / max_resolution, atol=utest_tolerance(self.dtype), )) def testNormBackground(self): bg_radius = 1.0 grid = grid_2d(self.L) mask = grid["r"] > bg_radius self.sim.normalize_background() imgs_nb = self.sim.images(start=0, num=self.n).asnumpy() new_mean = np.mean(imgs_nb[:, mask]) new_variance = np.var(imgs_nb[:, mask]) # new mean of noise should be close to zero and variance should be close to 1 self.assertTrue(new_mean < 1e-7 and abs(new_variance - 1) < 1e-7) def testWhiten(self): noise_estimator = AnisotropicNoiseEstimator(self.sim) self.sim.whiten(noise_estimator.filter) imgs_wt = self.sim.images(start=0, num=self.n).asnumpy() # calculate correlation between two neighboring pixels from background corr_coef = np.corrcoef(imgs_wt[:, self.L - 1, self.L - 1], imgs_wt[:, self.L - 2, self.L - 1]) # correlation matrix should be close to identity self.assertTrue(np.allclose(np.eye(2), corr_coef, atol=1e-1)) def testInvertContrast(self): sim1 = self.sim imgs1 = sim1.images(start=0, num=128) sim1.invert_contrast() imgs1_rc = sim1.images(start=0, num=128) # need to set the negative images to the second simulation object sim2 = ArrayImageSource(-imgs1) sim2.invert_contrast() imgs2_rc = sim2.images(start=0, num=128) # all images should be the same after inverting contrast self.assertTrue(np.allclose(imgs1_rc.asnumpy(), imgs2_rc.asnumpy()))
infile = mrcfile.open(os.path.join(DATA_DIR, "clean70SRibosome_vol_65p.mrc")) logger.info(f"Load 3D map from mrc file, {infile}") vols = Volume(infile.data) # Downsample the volume to a desired resolution and increase density # by 1.0e5 time for a better graph view logger.info( f"Downsample map to a resolution of {img_size} x {img_size} x {img_size}") vols = vols.downsample((img_size, ) * 3) * 1.0e5 # Create a simulation object with specified filters and the downsampled 3D map logger.info("Use downsampled map to create simulation object.") source = Simulation( L=img_size, n=num_imgs, vols=vols, unique_filters=ctf_filters, noise_filter=noise_filter, ) logger.info("Obtain original images.") imgs_od = source.images(start=0, num=1).asnumpy() logger.info("Perform phase flip to input images.") source.phase_flip() imgs_pf = source.images(start=0, num=1).asnumpy() max_resolution = 15 logger.info(f"Downsample resolution to {max_resolution} X {max_resolution}") if max_resolution < source.L: source.downsample(max_resolution)
logger.info( f"Load 3D map and downsample 3D map to desired grids " f"of {img_size} x {img_size} x {img_size}." ) infile = mrcfile.open(os.path.join(DATA_DIR, "clean70SRibosome_vol_65p.mrc")) # We prefer that our various arrays have consistent dtype. vols = Volume(infile.data.astype(dtype) / np.max(infile.data)) vols = vols.downsample(img_size) # Create a simulation object with specified filters and the downsampled 3D map logger.info("Use downsampled map to creat simulation object.") sim = Simulation( L=img_size, n=num_imgs, vols=vols, unique_filters=ctf_filters, offsets=0.0, amplitudes=1.0, dtype=dtype, noise_filter=noise_filter, ) # Specify the fast FB basis method for expending the 2D images ffbbasis = FFBBasis2D((img_size, img_size), dtype=dtype) # Assign the CTF information and index for each image h_idx = sim.filter_indices # Evaluate CTF in the 8X8 FB basis h_ctf_fb = [filt.fb_mat(ffbbasis) for filt in ctf_filters] # Get clean images from projections of 3D map.
parser = ConfigArgumentParser(description='Generate a Simulation and run Covariance estimation.') parser.add_argument('--num_volumes', default=2, type=int) parser.add_argument('--image_size', default=8, type=int) parser.add_argument('--num_images', default=1024, type=int) parser.add_argument('--num_eigs', default=16, type=int) with parser.parse_args() as args: C = args.num_volumes L = args.image_size n = args.num_images sim = Simulation( n=n, C=C, filters=SourceFilter( [RadialCTFFilter(defocus=d) for d in np.linspace(1.5e4, 2.5e4, 7)], n=n ) ) basis = FBBasis3D((L, L, L)) noise_estimator = WhiteNoiseEstimator(sim, batchSize=500) # Estimate the noise variance. This is needed for the covariance estimation step below. noise_variance = noise_estimator.estimate() print(f'Noise Variance = {noise_variance}') """ Estimate the mean. This uses conjugate gradient on the normal equations for the least-squares estimator of the mean volume. The mean volume is represented internally using the basis object, but the output is in the form of an L-by-L-by-L array. """
class BatchedRotCov2DTestCase(TestCase): def setUp(self): n = 32 L = 8 filters = [ RadialCTFFilter(5, 200, defocus=d, Cs=2.0, alpha=0.1) for d in np.linspace(1.5e4, 2.5e4, 7) ] self.dtype = np.float32 self.noise_var = 0.1848 # Initial noise filter to generate noise images. # Noise variance is set to a value far away that is used to calculate # covariance matrix and CWF coefficients in order to check the function # for rebuilding positive definite covariance matrix. noise_filter = ScalarFilter(dim=2, value=self.noise_var * 0.001) self.src = Simulation(L, n, unique_filters=filters, dtype=self.dtype, noise_filter=noise_filter) self.basis = FFBBasis2D((L, L), dtype=self.dtype) self.coeff = self.basis.evaluate_t(self.src.images(0, self.src.n)) self.ctf_idx = self.src.filter_indices self.ctf_fb = [f.fb_mat(self.basis) for f in self.src.unique_filters] self.cov2d = RotCov2D(self.basis) self.bcov2d = BatchedRotCov2D(self.src, self.basis, batch_size=7) def tearDown(self): pass def blk_diag_allclose(self, blk_diag_a, blk_diag_b, atol=1e-8): close = True for blk_a, blk_b in zip(blk_diag_a, blk_diag_b): close = close and np.allclose(blk_a, blk_b, atol=atol) return close def testMeanCovar(self): # Test basic functionality against RotCov2D. mean_cov2d = self.cov2d.get_mean(self.coeff, ctf_fb=self.ctf_fb, ctf_idx=self.ctf_idx) covar_cov2d = self.cov2d.get_covar( self.coeff, mean_coeff=mean_cov2d, ctf_fb=self.ctf_fb, ctf_idx=self.ctf_idx, noise_var=self.noise_var, ) mean_bcov2d = self.bcov2d.get_mean() covar_bcov2d = self.bcov2d.get_covar(noise_var=self.noise_var) self.assertTrue( np.allclose(mean_cov2d, mean_bcov2d, atol=utest_tolerance(self.dtype))) self.assertTrue( self.blk_diag_allclose(covar_cov2d, covar_bcov2d, atol=utest_tolerance(self.dtype))) def testZeroMean(self): # Make sure it works with zero mean (pure second moment). zero_coeff = np.zeros((self.basis.count, ), dtype=self.dtype) covar_cov2d = self.cov2d.get_covar(self.coeff, mean_coeff=zero_coeff, ctf_fb=self.ctf_fb, ctf_idx=self.ctf_idx) covar_bcov2d = self.bcov2d.get_covar(mean_coeff=zero_coeff) self.assertTrue( self.blk_diag_allclose(covar_cov2d, covar_bcov2d, atol=utest_tolerance(self.dtype))) def testAutoMean(self): # Make sure it automatically calls get_mean if needed. covar_cov2d = self.cov2d.get_covar(self.coeff, ctf_fb=self.ctf_fb, ctf_idx=self.ctf_idx) covar_bcov2d = self.bcov2d.get_covar() self.assertTrue( self.blk_diag_allclose(covar_cov2d, covar_bcov2d, atol=utest_tolerance(self.dtype))) def testShrink(self): # Make sure it properly shrinks the right-hand side if specified. covar_est_opt = { "shrinker": "frobenius_norm", "verbose": 0, "max_iter": 250, "iter_callback": [], "store_iterates": False, "rel_tolerance": 1e-12, "precision": self.dtype, } covar_cov2d = self.cov2d.get_covar( self.coeff, ctf_fb=self.ctf_fb, ctf_idx=self.ctf_idx, covar_est_opt=covar_est_opt, ) covar_bcov2d = self.bcov2d.get_covar(covar_est_opt=covar_est_opt) self.assertTrue(self.blk_diag_allclose(covar_cov2d, covar_bcov2d)) def testAutoBasis(self): # Make sure basis is automatically created if not specified. nbcov2d = BatchedRotCov2D(self.src) covar_bcov2d = self.bcov2d.get_covar() covar_nbcov2d = nbcov2d.get_covar() self.assertTrue( self.blk_diag_allclose(covar_bcov2d, covar_nbcov2d, atol=utest_tolerance(self.dtype))) def testCWFCoeff(self): # Calculate CWF coefficients using Cov2D base class mean_cov2d = self.cov2d.get_mean(self.coeff, ctf_fb=self.ctf_fb, ctf_idx=self.ctf_idx) covar_cov2d = self.cov2d.get_covar( self.coeff, ctf_fb=self.ctf_fb, ctf_idx=self.ctf_idx, noise_var=self.noise_var, make_psd=True, ) coeff_cov2d = self.cov2d.get_cwf_coeffs( self.coeff, self.ctf_fb, self.ctf_idx, mean_coeff=mean_cov2d, covar_coeff=covar_cov2d, noise_var=self.noise_var, ) # Calculate CWF coefficients using Batched Cov2D class mean_bcov2d = self.bcov2d.get_mean() covar_bcov2d = self.bcov2d.get_covar(noise_var=self.noise_var, make_psd=True) coeff_bcov2d = self.bcov2d.get_cwf_coeffs( self.coeff, self.ctf_fb, self.ctf_idx, mean_bcov2d, covar_bcov2d, noise_var=self.noise_var, ) self.assertTrue( self.blk_diag_allclose( coeff_cov2d, coeff_bcov2d, atol=utest_tolerance(self.dtype), ))
# Load the map file of a 70S Ribosome and downsample the 3D map to desired resolution. # The downsampling should be done by the internal function of sim object in future. # Below we use alternative implementation to obtain the exact result with Matlab version. logger.info(f'Load 3D map and downsample 3D map to desired grids ' f'of {img_size} x {img_size} x {img_size}.') infile = mrcfile.open(os.path.join(DATA_DIR, 'clean70SRibosome_vol_65p.mrc')) vols = infile.data vols = vols[..., np.newaxis] vols = downsample(vols, (img_size * np.ones(3, dtype=int))) # Create a simulation object with specified filters and the downsampled 3D map logger.info('Use downsampled map to creat simulation object.') sim = Simulation(L=img_size, n=num_imgs, vols=vols, C=num_maps, filters=filters) # Specify the fast FB basis method for expending the 2D images ffbbasis = FFBBasis2D((img_size, img_size)) # Generate 2D clean images from input 3D map. The following statement can be used from the sim object: # imgs_clean = sim.clean_images(start=0, num=num_imgs) # To be consistent with the Matlab version in the numbers, we need to use the statements as below: logger.info( 'Generate random distributed rotation angles and obtain corresponding 2D clean images.' ) rots = qrand_rots(num_imgs, seed=0) imgs_clean = vol2img(sim.vols[..., 0], rots)
# The downsampling should be done by the internal function of Volume object in future. logger.info(f"Load 3D map and downsample 3D map to desired grids " f"of {img_size} x {img_size} x {img_size}.") infile = mrcfile.open(os.path.join(DATA_DIR, "clean70SRibosome_vol_65p.mrc")) vols = Volume(infile.data.astype(dtype)) vols = vols.downsample((img_size, ) * 3) # %% # Create Simulation Object and Obtain True Rotation Angles # -------------------------------------------------------- # Create a simulation object with specified filters and the downsampled 3D map logger.info("Use downsampled map to creat simulation object.") sim = Simulation(L=img_size, n=num_imgs, vols=vols, unique_filters=filters, dtype=dtype) logger.info( "Get true rotation angles generated randomly by the simulation object.") rots_true = sim.rots # %% # Estimate Orientation and Rotation Angles # ---------------------------------------- # Initialize an orientation estimation object and perform view angle estimation logger.info( "Estimate rotation angles using synchronization matrix and voting method.") orient_est = CLSyncVoting(sim, n_theta=36)
# %% # Setup Simulation Source # ----------------------- # Simulation will randomly shift and amplify images by default. # Instead we define the following parameters. shifts = np.zeros((n_img, 2)) amplitudes = np.ones(n_img) # Create a Simulation Source object src = Simulation( vols=v, # our Volume L=v.resolution, # resolution, should match Volume n=n_img, # number of projection images C=len(v), # Number of volumes in vols. 1 in this case angles=rots.angles, # pass our rotations as Euler angles offsets=shifts, # translations (wrt to origin) amplitudes=amplitudes, # amplification ( 1 is identity) seed=12345, # RNG seed for reproducibility dtype=v.dtype, # match our datatype to the Volume. noise_filter=white_noise_filter, # optionally prescribe noise ) # %% # Yield projection images from the Simulation Source # -------------------------------------------------- # Consume images from the source by providing # a starting index and number of images. # Here we generate the first 3 and peek at them. src.images(0, 3).show() src.projections(0, 3).show()