diff --git a/src/aspire/abinitio/commonline_base.py b/src/aspire/abinitio/commonline_base.py index d1831c8177..c33a63cec1 100644 --- a/src/aspire/abinitio/commonline_base.py +++ b/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,6 +78,16 @@ 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: @@ -84,15 +95,15 @@ def _build(self): 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): """ diff --git a/tests/test_oriented_source.py b/tests/test_oriented_source.py index 54c34ea404..b57c173eb7 100644 --- a/tests/test_oriented_source.py +++ b/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"