diff --git a/gallery/experiments/experimental_abinitio_pipeline_10028.py b/gallery/experiments/experimental_abinitio_pipeline_10028.py index 472c21a638..1608aad8b8 100644 --- a/gallery/experiments/experimental_abinitio_pipeline_10028.py +++ b/gallery/experiments/experimental_abinitio_pipeline_10028.py @@ -29,7 +29,6 @@ import numpy as np from aspire.abinitio import CLSyncVoting -from aspire.basis import FFBBasis3D from aspire.denoising import DefaultClassAvgSource, DenoisedSource, DenoiserCov2D from aspire.noise import AnisotropicNoiseEstimator from aspire.reconstruction import MeanEstimator @@ -179,11 +178,8 @@ logger.info("Begin Volume reconstruction") -# Create a reasonable Basis for the 3d Volume -basis = FFBBasis3D((img_size,) * 3, dtype=src.dtype) - # Setup an estimator to perform the back projection. -estimator = MeanEstimator(oriented_src, basis) +estimator = MeanEstimator(oriented_src) # Perform the estimation and save the volume. estimated_volume = estimator.estimate() diff --git a/gallery/experiments/experimental_abinitio_pipeline_10073.py b/gallery/experiments/experimental_abinitio_pipeline_10073.py index d47168ec13..4d641ed3b8 100644 --- a/gallery/experiments/experimental_abinitio_pipeline_10073.py +++ b/gallery/experiments/experimental_abinitio_pipeline_10073.py @@ -31,7 +31,7 @@ import numpy as np from aspire.abinitio import CLSyncVoting -from aspire.basis import FFBBasis2D, FFBBasis3D +from aspire.basis import FFBBasis2D from aspire.classification import ( BandedSNRImageQualityFunction, BFRAverager2D, @@ -149,11 +149,8 @@ logger.info("Begin Volume reconstruction") -# Create a reasonable Basis for the 3d Volume -basis = FFBBasis3D((img_size,) * 3, dtype=src.dtype) - # Setup an estimator to perform the back projection. -estimator = MeanEstimator(oriented_src, basis) +estimator = MeanEstimator(oriented_src) # Perform the estimation and save the volume. estimated_volume = estimator.estimate() diff --git a/gallery/experiments/experimental_abinitio_pipeline_10081.py b/gallery/experiments/experimental_abinitio_pipeline_10081.py index 45a6558a89..ff6139db37 100644 --- a/gallery/experiments/experimental_abinitio_pipeline_10081.py +++ b/gallery/experiments/experimental_abinitio_pipeline_10081.py @@ -26,7 +26,6 @@ from pathlib import Path from aspire.abinitio import CLSymmetryC3C4 -from aspire.basis import FFBBasis3D from aspire.denoising import DefaultClassAvgSource from aspire.noise import AnisotropicNoiseEstimator from aspire.reconstruction import MeanEstimator @@ -120,11 +119,9 @@ logger.info("Begin Volume reconstruction") -# Create a reasonable Basis for the 3d Volume -basis = FFBBasis3D((img_size,) * 3, dtype=src.dtype) # Setup an estimator to perform the back projection. -estimator = MeanEstimator(oriented_src, basis) +estimator = MeanEstimator(oriented_src) # Perform the estimation and save the volume. estimated_volume = estimator.estimate() diff --git a/gallery/experiments/simulated_abinitio_pipeline.py b/gallery/experiments/simulated_abinitio_pipeline.py index 63dce7e62a..ca0e4fbac5 100644 --- a/gallery/experiments/simulated_abinitio_pipeline.py +++ b/gallery/experiments/simulated_abinitio_pipeline.py @@ -21,7 +21,6 @@ import numpy as np from aspire.abinitio import CLSyncVoting -from aspire.basis import FFBBasis3D from aspire.denoising import DefaultClassAvgSource, DenoisedSource, DenoiserCov2D from aspire.downloader import emdb_2660 from aspire.noise import AnisotropicNoiseEstimator, CustomNoiseAdder @@ -209,11 +208,8 @@ def noise_function(x, y): logger.info("Begin Volume reconstruction") -# Create a reasonable Basis for the 3d Volume -basis = FFBBasis3D((v.resolution,) * 3, dtype=v.dtype) - # Setup an estimator to perform the back projection. -estimator = MeanEstimator(oriented_src, basis) +estimator = MeanEstimator(oriented_src) # Perform the estimation and save the volume. estimated_volume = estimator.estimate() diff --git a/gallery/tutorials/pipeline_demo.py b/gallery/tutorials/pipeline_demo.py index f2547d2c00..3f38ae8597 100644 --- a/gallery/tutorials/pipeline_demo.py +++ b/gallery/tutorials/pipeline_demo.py @@ -233,14 +233,10 @@ # estimate the mean volume by supplying the class averages and basis # for back projection. -from aspire.basis import FFBBasis3D from aspire.reconstruction import MeanEstimator -# Create a reasonable Basis for the 3d Volume -basis = FFBBasis3D(res, dtype=vol.dtype) - # Setup an estimator to perform the back projection. -estimator = MeanEstimator(oriented_src, basis) +estimator = MeanEstimator(oriented_src) # Perform the estimation and save the volume. estimated_volume = estimator.estimate() diff --git a/gallery/tutorials/tutorials/cov3d_simulation.py b/gallery/tutorials/tutorials/cov3d_simulation.py index 85b90a8d24..5fced70fbb 100644 --- a/gallery/tutorials/tutorials/cov3d_simulation.py +++ b/gallery/tutorials/tutorials/cov3d_simulation.py @@ -55,7 +55,7 @@ num_vols = sim.C # Specify the normal FB basis method for expending the 2D images -basis = FBBasis3D((img_size, img_size, img_size)) +basis = FBBasis3D(img_size) # Estimate the noise variance. This is needed for the covariance estimation step below. noise_estimator = WhiteNoiseEstimator(sim, batchSize=500) @@ -71,11 +71,13 @@ # using the basis object, but the output is in the form of an # L-by-L-by-L array. -mean_estimator = MeanEstimator(sim, basis) +mean_estimator = MeanEstimator(sim, basis=basis) mean_est = mean_estimator.estimate() # Passing in a mean_kernel argument to the following constructor speeds up some calculations -covar_estimator = CovarianceEstimator(sim, basis, mean_kernel=mean_estimator.kernel) +covar_estimator = CovarianceEstimator( + sim, basis=basis, mean_kernel=mean_estimator.kernel +) covar_est = covar_estimator.estimate(mean_est, noise_variance) # %% diff --git a/gallery/tutorials/tutorials/starfile.py b/gallery/tutorials/tutorials/starfile.py index c4c6ce7d50..f3827c923a 100644 --- a/gallery/tutorials/tutorials/starfile.py +++ b/gallery/tutorials/tutorials/starfile.py @@ -9,7 +9,6 @@ import matplotlib.pyplot as plt import numpy as np -from aspire.basis import FBBasis3D from aspire.noise import AnisotropicNoiseEstimator from aspire.reconstruction import MeanEstimator from aspire.source import RelionSource @@ -53,10 +52,8 @@ # Estimate Mean Volume # -------------------- -# We'll create a 3D Fourier Bessel Basis corresponding to volume resolution L. -basis = FBBasis3D((L, L, L)) # Estimate mean Volume -mean_estimator = MeanEstimator(source, basis, batch_size=8192) +mean_estimator = MeanEstimator(source, batch_size=8192) mean_est = mean_estimator.estimate() # %% diff --git a/src/aspire/reconstruction/estimator.py b/src/aspire/reconstruction/estimator.py index b0832be174..e002b7c3da 100644 --- a/src/aspire/reconstruction/estimator.py +++ b/src/aspire/reconstruction/estimator.py @@ -2,7 +2,7 @@ import os from pathlib import Path -from aspire.basis import Coef +from aspire.basis import Coef, FFBBasis3D from aspire.reconstruction.kernel import FourierKernel logger = logging.getLogger(__name__) @@ -12,7 +12,7 @@ class Estimator: def __init__( self, src, - basis, + basis=None, batch_size=512, preconditioner="circulant", checkpoint_iterations=10, @@ -50,6 +50,9 @@ def __init__( """ self.src = src + if basis is None: + logger.info("{self.__class__.__name__} instantiating default basis.") + basis = FFBBasis3D(src.L, dtype=src.dtype) self.basis = basis self.dtype = self.src.dtype self.batch_size = batch_size diff --git a/src/aspire/reconstruction/mean.py b/src/aspire/reconstruction/mean.py index 760bca3e3d..f7f5bf1a2c 100644 --- a/src/aspire/reconstruction/mean.py +++ b/src/aspire/reconstruction/mean.py @@ -296,12 +296,12 @@ class MeanEstimator(WeightedVolumesEstimator): for a single volume. """ - def __init__(self, src, basis, **kwargs): + def __init__(self, src, **kwargs): # Note, Handle boosting by adjusting weights based on symmetric order. weights = np.ones((src.n, 1)) / np.sqrt( src.n * len(src.symmetry_group.matrices) ) - super().__init__(weights, src, basis, **kwargs) + super().__init__(weights, src, **kwargs) def __getattr__(self, name): """ diff --git a/tests/test_array_image_source.py b/tests/test_array_image_source.py index c580ef3204..a43db1ee8e 100644 --- a/tests/test_array_image_source.py +++ b/tests/test_array_image_source.py @@ -100,7 +100,7 @@ def testArrayImageSourceAngGetterError(self): # We also test that a source consumer generates same error, # by instantiating a volume estimator. - estimator = MeanEstimator(src, self.basis, preconditioner="none") + estimator = MeanEstimator(src, basis=self.basis, preconditioner="none") # Test we raise with expected message with raises(RuntimeError, match=r"Consumer of ArrayImageSource.*"): @@ -127,7 +127,7 @@ def testArrayImageSourceMeanVol(self): """ # Run estimator with a Simulation source as a reference. - sim_estimator = MeanEstimator(self.sim, self.basis, preconditioner="none") + sim_estimator = MeanEstimator(self.sim, basis=self.basis, preconditioner="none") sim_est = sim_estimator.estimate() logger.info("Simulation source checkpoint") @@ -135,7 +135,7 @@ def testArrayImageSourceMeanVol(self): src = ArrayImageSource(self.im, angles=self.sim.angles) # Instantiate a volume estimator using ArrayImageSource - estimator = MeanEstimator(src, self.basis, preconditioner="none") + estimator = MeanEstimator(src, basis=self.basis, preconditioner="none") # Get estimate consuming ArrayImageSource est = estimator.estimate() diff --git a/tests/test_covar3d.py b/tests/test_covar3d.py index 9bf04b4594..8b8df4421a 100644 --- a/tests/test_covar3d.py +++ b/tests/test_covar3d.py @@ -36,18 +36,21 @@ def setUpClass(cls): basis = FBBasis3D((8, 8, 8), dtype=cls.dtype) cls.noise_variance = 0.0030762743633643615 - cls.mean_estimator = MeanEstimator(cls.sim, basis) + cls.mean_estimator = MeanEstimator(cls.sim, basis=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.sim, + basis=basis, + mean_kernel=cls.mean_estimator.kernel, + preconditioner="none", ) cls.covar_estimator_with_preconditioner = CovarianceEstimator( cls.sim, - basis, + basis=basis, mean_kernel=cls.mean_estimator.kernel, preconditioner="circulant", ) diff --git a/tests/test_mean_estimator.py b/tests/test_mean_estimator.py index 616c258fcc..4bbd64e6f3 100644 --- a/tests/test_mean_estimator.py +++ b/tests/test_mean_estimator.py @@ -29,10 +29,12 @@ def setUp(self): ) self.basis = FBBasis3D((self.resolution,) * 3, dtype=self.dtype) - self.estimator = MeanEstimator(self.sim, self.basis, preconditioner="none") + self.estimator = MeanEstimator( + self.sim, basis=self.basis, preconditioner="none" + ) self.estimator_with_preconditioner = MeanEstimator( - self.sim, self.basis, preconditioner="circulant" + self.sim, basis=self.basis, preconditioner="circulant" ) def tearDown(self): @@ -47,7 +49,7 @@ def testEstimateResolutionError(self): # This basis is intentionally the wrong resolution. incorrect_basis = FBBasis3D((2 * self.resolution,) * 3, dtype=self.dtype) - _ = MeanEstimator(self.sim, incorrect_basis, preconditioner="none") + _ = MeanEstimator(self.sim, basis=incorrect_basis, preconditioner="none") def testEstimate(self): estimate = self.estimator.estimate() @@ -352,7 +354,8 @@ def testOptimize1(self): -9.82705453e-04, 6.46337066e-05, ] - ] + ], + dtype=self.dtype, ) x = self.estimator.conj_grad(mean_b_coef) @@ -684,7 +687,7 @@ def testCheckpoint(self): prefix = os.path.join(tmp_input_dir, "new", "dirs", "chk") estimator = MeanEstimator( self.sim, - self.basis, + basis=self.basis, preconditioner="none", checkpoint_iterations=test_iter, maxiter=test_iter + 1, @@ -712,7 +715,7 @@ def testCheckpointArgs(self): ): _ = MeanEstimator( self.sim, - self.basis, + basis=self.basis, preconditioner="none", checkpoint_iterations=junk, checkpoint_prefix=prefix, @@ -723,7 +726,7 @@ def testCheckpointArgs(self): ): _ = MeanEstimator( self.sim, - self.basis, + basis=self.basis, preconditioner="none", maxiter=junk, checkpoint_prefix=prefix, diff --git a/tests/test_mean_estimator_boosting.py b/tests/test_mean_estimator_boosting.py index 005c0597af..9251dee09e 100644 --- a/tests/test_mean_estimator_boosting.py +++ b/tests/test_mean_estimator_boosting.py @@ -1,7 +1,6 @@ import numpy as np import pytest -from aspire.basis import FFBBasis3D from aspire.reconstruction import MeanEstimator, WeightedVolumesEstimator from aspire.source import ArrayImageSource, Simulation from aspire.utils import Rotation, utest_tolerance @@ -81,8 +80,7 @@ def source(volume): @pytest.fixture(scope="module") def estimated_volume(source): - basis = FFBBasis3D(source.L, dtype=source.dtype) - estimator = MeanEstimator(source, basis) + estimator = MeanEstimator(source) estimated_volume = estimator.estimate() return estimated_volume @@ -164,8 +162,7 @@ def test_boost_flag(source, estimated_volume): boosted_source = ArrayImageSource(ims_boosted, angles=rots_boosted.angles) # Estimate volume with boosting OFF. - basis = FFBBasis3D(boosted_source.L, dtype=boosted_source.dtype) - estimator = MeanEstimator(boosted_source, basis, boost=False) + estimator = MeanEstimator(boosted_source, boost=False) est_vol = estimator.estimate() # Check reconstructions are equal. @@ -192,8 +189,7 @@ def test_weighted_volumes(weighted_source): weights[:, 1] = weights[:, 1] / weights[:, 1].sum() * np.sqrt(n1) # Initialize estimator. - basis = FFBBasis3D(src.L, dtype=src.dtype) - estimator = WeightedVolumesEstimator(src=src, basis=basis, weights=weights) + estimator = WeightedVolumesEstimator(src=src, weights=weights) est_vols = estimator.estimate() # Check FSC (scaling may not be close enough to match mse) diff --git a/tests/test_weighted_mean_estimator.py b/tests/test_weighted_mean_estimator.py index 4d185ce142..23496428c1 100644 --- a/tests/test_weighted_mean_estimator.py +++ b/tests/test_weighted_mean_estimator.py @@ -32,10 +32,10 @@ def setUp(self): self.basis = FBBasis3D((L, L, L), dtype=self.dtype) self.weights = np.ones((self.n, self.r)) / np.sqrt(self.n) self.estimator = WeightedVolumesEstimator( - self.weights, self.sim, self.basis, preconditioner="none" + self.weights, self.sim, basis=self.basis, preconditioner="none" ) self.estimator_with_preconditioner = WeightedVolumesEstimator( - self.weights, self.sim, self.basis, preconditioner="circulant" + self.weights, self.sim, basis=self.basis, preconditioner="circulant" ) def tearDown(self): @@ -347,7 +347,8 @@ def testOptimize1(self): 6.46337066e-05, ] ] - * self.r + * self.r, + dtype=self.dtype, ) # Given equal weighting we should get the same result for all self.r volumes. @@ -687,7 +688,7 @@ def testNegativeWeightedEstimates(self): weights[:, 1] *= -1 # negate second set of weights estimator = WeightedVolumesEstimator( - weights, self.sim, self.basis, preconditioner="none" + weights, self.sim, basis=self.basis, preconditioner="none" ) estimate = estimator.estimate()