Default Estimator 3D basis

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()

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()

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()

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()

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()

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)
 
 # %%

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()
 
 # %%

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

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):
         """

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,15 +127,15 @@ 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")
 
         # Construct the source for testing
         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()

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",
         )

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,

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)

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()