Merge pull request #686 from ComputationalCryoEM/batched_cov2d_noctf

Batched cov2d noctf patch

Author: garrettwrong

src/aspire/covariance/covar2d.py

@@ -502,7 +502,7 @@ def _build(self):
 
             self.basis = FFBBasis2D((src.L, src.L), dtype=self.dtype)
 
-        if src.unique_filters is None:
+        if not src.unique_filters:
             logger.info("CTF filters are not included in Cov2D denoising")
             # set all CTF filters to an identity filter
             self.ctf_idx = np.zeros(src.n, dtype=int)

tests/test_batched_covar2d.py

@@ -10,13 +10,17 @@
 
 
 class BatchedRotCov2DTestCase(TestCase):
+    """
+    Tests batched cov2d without providing any CTF filters.
+    """
+
+    filters = None
+    ctf_idx = None
+    ctf_fb = None
+
     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
 
@@ -27,14 +31,15 @@ def setUp(self):
         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
+            L,
+            n,
+            unique_filters=self.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)
 
@@ -234,3 +239,24 @@ def testCWFCoeffCleanCTF(self):
                 atol=utest_tolerance(self.dtype),
             )
         )
+
+
+class BatchedRotCov2DTestCaseCTF(BatchedRotCov2DTestCase):
+    """
+    Tests batched cov2d with CTF information.
+    """
+
+    @property
+    def filters(self):
+        return [
+            RadialCTFFilter(5, 200, defocus=d, Cs=2.0, alpha=0.1)
+            for d in np.linspace(1.5e4, 2.5e4, 7)
+        ]
+
+    @property
+    def ctf_idx(self):
+        return self.src.filter_indices
+
+    @property
+    def ctf_fb(self):
+        return [f.fb_mat(self.basis) for f in self.src.unique_filters]

tests/test_covar2d.py

@@ -17,41 +17,39 @@
 
 
 class Cov2DTestCase(TestCase):
+    """
+    Cov2D Test without CTFFilters populated.
+    """
+
+    unique_filters = None
+    h_idx = None
+    h_ctf_fb = None
+
     # These class variables support parameterized arg checking in `testShrinkers`
     shrinkers = [(None,), "frobenius_norm", "operator_norm", "soft_threshold"]
     bad_shrinker_inputs = ["None", "notashrinker", ""]
 
     def setUp(self):
         self.dtype = np.float32
 
-        L = 8
+        self.L = 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(
+        self.sim = Simulation(
             n=n,
             L=L,
             vols=vols,
-            unique_filters=unique_filters,
+            unique_filters=self.unique_filters,
             offsets=0.0,
             amplitudes=1.0,
             dtype=self.dtype,
@@ -60,12 +58,9 @@ def setUp(self):
 
         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.imgs_clean = self.sim.projections()
+        self.imgs_ctf_clean = self.sim.clean_images()
+        self.imgs_ctf_noise = self.sim.images(start=0, num=n)
 
         self.cov2d = RotCov2D(self.basis)
         self.coeff_clean = self.basis.evaluate_t(self.imgs_clean)
@@ -76,89 +71,34 @@ def tearDown(self):
 
     def testGetMean(self):
         results = np.load(os.path.join(DATA_DIR, "clean70SRibosome_cov2d_mean.npy"))
-        self.mean_coeff = self.cov2d._get_mean(self.coeff_clean)
-        self.assertTrue(np.allclose(results, self.mean_coeff))
+        mean_coeff = self.cov2d._get_mean(self.coeff_clean)
+        self.assertTrue(np.allclose(results, mean_coeff))
 
     def testGetCovar(self):
         results = np.load(
             os.path.join(DATA_DIR, "clean70SRibosome_cov2d_covar.npy"),
             allow_pickle=True,
         )
-        self.covar_coeff = self.cov2d._get_covar(self.coeff_clean)
+        covar_coeff = self.cov2d._get_covar(self.coeff_clean)
 
         for im, mat in enumerate(results.tolist()):
-            self.assertTrue(np.allclose(mat, self.covar_coeff[im]))
+            self.assertTrue(np.allclose(mat, covar_coeff[im]))
 
     def testGetMeanCTF(self):
-        results = np.load(os.path.join(DATA_DIR, "clean70SRibosome_cov2d_meanctf.npy"))
-        self.mean_coeff_ctf = self.cov2d.get_mean(self.coeff, self.h_ctf_fb, self.h_idx)
-        self.assertTrue(np.allclose(results, self.mean_coeff_ctf))
-
-    def testGetCovarCTF(self):
-        results = np.load(
-            os.path.join(DATA_DIR, "clean70SRibosome_cov2d_covarctf.npy"),
-            allow_pickle=True,
-        )
-        self.covar_coeff_ctf = self.cov2d.get_covar(
-            self.coeff, self.h_ctf_fb, self.h_idx, noise_var=self.noise_var
-        )
-        for im, mat in enumerate(results.tolist()):
-            self.assertTrue(np.allclose(mat, self.covar_coeff_ctf[im]))
-
-    def testGetCovarCTFShrink(self):
-        results = np.load(
-            os.path.join(DATA_DIR, "clean70SRibosome_cov2d_covarctf_shrink.npy"),
-            allow_pickle=True,
-        )
-        covar_opt = {
-            "shrinker": "frobenius_norm",
-            "verbose": 0,
-            "max_iter": 250,
-            "iter_callback": [],
-            "store_iterates": False,
-            "rel_tolerance": 1e-12,
-            "precision": self.dtype,
-        }
-        self.covar_coeff_ctf_shrink = self.cov2d.get_covar(
-            self.coeff,
-            self.h_ctf_fb,
-            self.h_idx,
-            noise_var=self.noise_var,
-            covar_est_opt=covar_opt,
-        )
-
-        for im, mat in enumerate(results.tolist()):
-            self.assertTrue(np.allclose(mat, self.covar_coeff_ctf_shrink[im]))
-
-    def testGetCWFCoeffs(self):
-        results = np.load(
-            os.path.join(DATA_DIR, "clean70SRibosome_cov2d_cwf_coeff.npy")
-        )
-        self.coeff_cwf = self.cov2d.get_cwf_coeffs(
-            self.coeff, self.h_ctf_fb, self.h_idx, noise_var=self.noise_var
-        )
-        self.assertTrue(
-            np.allclose(results, self.coeff_cwf, atol=utest_tolerance(self.dtype))
-        )
-
-    def testGetCWFCoeffsIdentityCTF(self):
-        results = np.load(
-            os.path.join(DATA_DIR, "clean70SRibosome_cov2d_cwf_coeff_noCTF.npy")
-        )
-        self.coeff_cwf_noCTF = self.cov2d.get_cwf_coeffs(
-            self.coeff, noise_var=self.noise_var
-        )
-        self.assertTrue(
-            np.allclose(results, self.coeff_cwf_noCTF, atol=utest_tolerance(self.dtype))
-        )
+        """
+        Compare `get_mean` (no CTF args) with `_get_mean` (no CTF model).
+        """
+        mean_coeff_ctf = self.cov2d.get_mean(self.coeff, self.h_ctf_fb, self.h_idx)
+        mean_coeff = self.cov2d._get_mean(self.coeff_clean)
+        self.assertTrue(np.allclose(mean_coeff_ctf, mean_coeff, atol=0.002))
 
     def testGetCWFCoeffsClean(self):
         results = np.load(
             os.path.join(DATA_DIR, "clean70SRibosome_cov2d_cwf_coeff_clean.npy")
         )
-        self.coeff_cwf_clean = self.cov2d.get_cwf_coeffs(self.coeff_clean, noise_var=0)
+        coeff_cwf_clean = self.cov2d.get_cwf_coeffs(self.coeff_clean, noise_var=0)
         self.assertTrue(
-            np.allclose(results, self.coeff_cwf_clean, atol=utest_tolerance(self.dtype))
+            np.allclose(results, coeff_cwf_clean, atol=utest_tolerance(self.dtype))
         )
 
     def testGetCWFCoeffsCleanCTF(self):
@@ -180,17 +120,6 @@ def testGetCWFCoeffsCleanCTF(self):
         delta = np.mean(np.square((self.imgs_clean - img_est).asnumpy()))
         self.assertTrue(delta < 0.02)
 
-    def testGetCWFCoeffsCTFargs(self):
-        """
-        Test we raise when user supplies incorrect CTF arguments,
-        and that the error message matches.
-        """
-
-        with raises(RuntimeError, match=r".*Given ctf_fb.*"):
-            _ = self.cov2d.get_cwf_coeffs(
-                self.coeff, self.h_ctf_fb, None, noise_var=self.noise_var
-            )
-
     # Note, parameterized module can be removed at a later date
     # and replaced with pytest if ASPIRE-Python moves away from
     # the TestCase class style tests.
@@ -220,3 +149,110 @@ def testShrinkers(self, shrinker):
             self.assertTrue(
                 np.allclose(mat, covar_coeff[im], atol=utest_tolerance(self.dtype))
             )
+
+
+class Cov2DTestCaseCTF(Cov2DTestCase):
+    """
+    Cov2D Test with CTFFilters populated.
+    """
+
+    @property
+    def unique_filters(self):
+        return [
+            RadialCTFFilter(5.0 * 65 / self.L, 200, defocus=d, Cs=2.0, alpha=0.1)
+            for d in np.linspace(1.5e4, 2.5e4, 7)
+        ]
+
+    @property
+    def h_idx(self):
+        return self.sim.filter_indices
+
+    @property
+    def h_ctf_fb(self):
+        return [filt.fb_mat(self.basis) for filt in self.unique_filters]
+
+    def testGetCWFCoeffsCTFargs(self):
+        """
+        Test we raise when user supplies incorrect CTF arguments,
+        and that the error message matches.
+        """
+
+        with raises(RuntimeError, match=r".*Given ctf_fb.*"):
+            _ = self.cov2d.get_cwf_coeffs(
+                self.coeff, self.h_ctf_fb, None, noise_var=self.noise_var
+            )
+
+    def testGetMeanCTF(self):
+        """
+        Compare `get_mean` with saved legacy cov2d results.
+        """
+        results = np.load(os.path.join(DATA_DIR, "clean70SRibosome_cov2d_meanctf.npy"))
+        mean_coeff_ctf = self.cov2d.get_mean(self.coeff, self.h_ctf_fb, self.h_idx)
+        self.assertTrue(np.allclose(results, mean_coeff_ctf))
+
+    def testGetCWFCoeffs(self):
+        """
+        Tests `get_cwf_coeffs` with poulated CTF.
+        """
+        results = np.load(
+            os.path.join(DATA_DIR, "clean70SRibosome_cov2d_cwf_coeff.npy")
+        )
+        coeff_cwf = self.cov2d.get_cwf_coeffs(
+            self.coeff, self.h_ctf_fb, self.h_idx, noise_var=self.noise_var
+        )
+        self.assertTrue(
+            np.allclose(results, coeff_cwf, atol=utest_tolerance(self.dtype))
+        )
+
+    # Note, I think this file is incorrectly named...
+    #   It appears to have come from operations on images with ctf applied.
+    def testGetCWFCoeffsNoCTF(self):
+        """
+        Tests `get_cwf_coeffs` without providing CTF.  (Internally uses IdentityCTF).
+        """
+        results = np.load(
+            os.path.join(DATA_DIR, "clean70SRibosome_cov2d_cwf_coeff_noCTF.npy")
+        )
+        coeff_cwf_noCTF = self.cov2d.get_cwf_coeffs(
+            self.coeff, noise_var=self.noise_var
+        )
+
+        self.assertTrue(
+            np.allclose(results, coeff_cwf_noCTF, atol=utest_tolerance(self.dtype))
+        )
+
+    def testGetCovarCTF(self):
+        results = np.load(
+            os.path.join(DATA_DIR, "clean70SRibosome_cov2d_covarctf.npy"),
+            allow_pickle=True,
+        )
+        covar_coeff_ctf = self.cov2d.get_covar(
+            self.coeff, self.h_ctf_fb, self.h_idx, noise_var=self.noise_var
+        )
+        for im, mat in enumerate(results.tolist()):
+            self.assertTrue(np.allclose(mat, covar_coeff_ctf[im]))
+
+    def testGetCovarCTFShrink(self):
+        results = np.load(
+            os.path.join(DATA_DIR, "clean70SRibosome_cov2d_covarctf_shrink.npy"),
+            allow_pickle=True,
+        )
+        covar_opt = {
+            "shrinker": "frobenius_norm",
+            "verbose": 0,
+            "max_iter": 250,
+            "iter_callback": [],
+            "store_iterates": False,
+            "rel_tolerance": 1e-12,
+            "precision": self.dtype,
+        }
+        covar_coeff_ctf_shrink = self.cov2d.get_covar(
+            self.coeff,
+            self.h_ctf_fb,
+            self.h_idx,
+            noise_var=self.noise_var,
+            covar_est_opt=covar_opt,
+        )
+
+        for im, mat in enumerate(results.tolist()):
+            self.assertTrue(np.allclose(mat, covar_coeff_ctf_shrink[im]))