OrientedSource Lazy Evaluation

src/aspire/abinitio/commonline_base.py

@@ -57,6 +57,7 @@ def __init__(
         self.shift_step = shift_step
         self.mask = mask
         self.rotations = None
+        self._pf = None
 
         self._build()
 
@@ -77,22 +78,32 @@ def _build(self):
             logger.error(msg)
             raise NotImplementedError(msg)
 
+    @property
+    def pf(self):
+        if self._pf is None:
+            self._prepare_pf()
+        return self._pf
+
+    def _prepare_pf(self):
+        """
+        Prepare the polar Fourier transform used for correlations.
+        """
         imgs = self.src.images[:]
 
         if self.mask:
             fuzz_mask = fuzzy_mask((self.n_res, self.n_res), self.dtype)
             imgs = imgs * fuzz_mask
 
         # Obtain coefficients of polar Fourier transform for input 2D images
-        self.pft = PolarFT(
+        pft = PolarFT(
             (self.n_res, self.n_res), self.n_rad, self.n_theta, dtype=self.dtype
         )
-        self.pf = self.pft.transform(imgs)
+        pf = pft.transform(imgs)
 
         # We remove the DC the component. pf has size (n_img) x (n_theta/2) x (n_rad-1),
         # with pf[:, :, 0] containing low frequency content and pf[:, :, -1] containing
         # high frequency content.
-        self.pf = self.pf[:, :, 1:]
+        self._pf = pf[:, :, 1:]
 
     def estimate_rotations(self):
         """

tests/test_oriented_source.py

@@ -93,17 +93,27 @@ def test_estimator_error(src_fixture):
         _ = OrientedSource(og_src, junk_estimator)
 
 
-def test_lazy_orientation(src_fixture, caplog):
+def test_lazy_evaluation(src_fixture, caplog):
+    """
+    Test that both images and rotations are evaluated in a lazy fashion.
+    """
     _, oriented_src = src_fixture
 
-    # Check that instantiated oriented sources don't have rotation metadata.
+    # Check that instantiated oriented source does not have rotation metadata.
     rotation_metadata = ["_rlnAngleRot", "_rlnAngleTilt", "_rlnAnglePsi"]
     assert not oriented_src.has_metadata(rotation_metadata)
 
+    # Check that the oriented source's `orientation_estimator` does not have the attribute
+    # `_pf`, an indicator that images have not yet been requested by the estimator.
+    assert oriented_src.orientation_estimator._pf is None
+
     # Request rotations and check that metadata is populated.
     _ = oriented_src.rotations
     assert oriented_src.has_metadata(rotation_metadata)
 
+    # Chec that `_pf` is not None.
+    assert oriented_src.orientation_estimator._pf is not None
+
     # Check that requesting rotations again logs a debug message about skippin orientation.
     caplog.clear()
     msg = f"{oriented_src.__class__.__name__} already oriented, skipping"