New logical (F)FBBasis3D tests

tests/_basis_util.py

@@ -4,12 +4,49 @@
 
 from aspire.image import Image
 from aspire.utils import gaussian_2d, utest_tolerance
-from aspire.utils.coor_trans import grid_2d
+from aspire.utils.coor_trans import grid_2d, grid_3d
+from aspire.utils.misc import gaussian_3d
 from aspire.utils.random import randn
 from aspire.volume import Volume
 
 
-class Steerable2DMixin:
+class SteerableMixin:
+    def testEvaluateExpand(self):
+        coef1 = randn(self.basis.count, seed=self.seed)
+        coef1 = coef1.astype(self.dtype)
+
+        x = self.basis.evaluate(coef1)
+        if isinstance(x, Image) or isinstance(x, Volume):
+            x = x.asnumpy()
+
+        coef2 = self.basis.expand(x)
+        if coef2.ndim == 2:
+            coef2 = coef2[0]
+
+        self.assertTrue(coef1.shape == coef2.shape)
+        self.assertTrue(np.allclose(coef1, coef2, atol=utest_tolerance(self.dtype)))
+
+    def testAdjoint(self):
+        u = randn(self.basis.count, seed=self.seed)
+        u = u.astype(self.dtype)
+
+        Au = self.basis.evaluate(u)
+        if isinstance(Au, Image) or isinstance(Au, Volume):
+            Au = Au.asnumpy()
+
+        x = randn(*self.basis.sz, seed=self.seed)
+        x = x.astype(self.dtype)
+
+        ATx = self.basis.evaluate_t(x)
+
+        Au_dot_x = np.sum(Au * x)
+        u_dot_ATx = np.sum(u * ATx)
+
+        self.assertTrue(Au_dot_x.shape == u_dot_ATx.shape)
+        self.assertTrue(np.isclose(Au_dot_x, u_dot_ATx))
+
+
+class Steerable2DMixin(SteerableMixin):
     def testIndices(self):
         ell_max = self.basis.ell_max
         k_max = self.basis.k_max
@@ -99,36 +136,97 @@ def testModulated(self):
 
             self.assertTrue(energy_ratio < 0.10)
 
-    def testEvaluateExpand(self):
-        coef1 = randn(self.basis.count, seed=self.seed)
-        coef1 = coef1.astype(self.dtype)
 
-        im = self.basis.evaluate(coef1)
-        if isinstance(im, Image):
-            im = im.asnumpy()
-        coef2 = self.basis.expand(im)[0]
+class Steerable3DMixin(SteerableMixin):
+    def testIndices(self):
+        ell_max = self.basis.ell_max
+        k_max = self.basis.k_max
 
-        self.assertTrue(coef1.shape == coef2.shape)
-        self.assertTrue(np.allclose(coef1, coef2, atol=utest_tolerance(self.dtype)))
+        indices = self.basis.indices()
 
-    def testAdjoint(self):
-        u = randn(self.basis.count, seed=self.seed)
-        u = u.astype(self.dtype)
+        i = 0
 
-        Au = self.basis.evaluate(u)
-        if isinstance(Au, Image):
-            Au = Au.asnumpy()
+        for ell in range(ell_max + 1):
+            for m in range(-ell, ell + 1):
+                for k in range(k_max[ell]):
+                    self.assertTrue(indices["ells"][i] == ell)
+                    self.assertTrue(indices["ms"][i] == m)
+                    self.assertTrue(indices["ks"][i] == k)
 
-        x = randn(*self.basis.sz, seed=self.seed)
-        x = x.astype(self.dtype)
+                    i += 1
 
-        ATx = self.basis.evaluate_t(x)
+    def testGaussianExpand(self):
+        # Offset slightly
+        x0 = 0.50
+        y0 = 0.75
+        z0 = 0.25
 
-        Au_dot_x = np.sum(Au * x)
-        u_dot_ATx = np.sum(u * ATx)
+        # Want sigma to be as large as possible without the Gaussian
+        # spilling too much outside the central disk.
+        sigma = self.L / 8
+        vol1 = gaussian_3d(
+            self.L, mu=(x0, y0, z0), sigma=(sigma, sigma, sigma), dtype=self.dtype
+        )
 
-        self.assertTrue(Au_dot_x.shape == u_dot_ATx.shape)
-        self.assertTrue(np.isclose(Au_dot_x, u_dot_ATx))
+        vol1 = vol1[np.newaxis]
+
+        coef = self.basis.expand(vol1)
+        vol2 = self.basis.evaluate(coef)
+
+        if isinstance(vol2, Volume):
+            vol2 = vol2.asnumpy()
+
+        # For small L there's too much clipping at high freqs to get 1e-3
+        # accuracy.
+        if self.L < 16:
+            atol = 1e-1
+        elif self.L < 32:
+            atol = 1e-2
+        else:
+            atol = 1e-3
+
+        self.assertEqual(vol1.shape, vol2.shape)
+        self.assertTrue(np.allclose(vol1, vol2, atol=atol))
+
+    def testIsotropic(self):
+        sigma = self.L / 8
+        im = gaussian_3d(self.L, sigma=(sigma, sigma, sigma), dtype=self.dtype)
+
+        coef = self.basis.expand(im)
+
+        ells = self.basis.indices()["ells"]
+
+        energy_outside = np.sum(np.abs(coef[ells != 0]) ** 2)
+        energy_total = np.sum(np.abs(coef) ** 2)
+
+        energy_ratio = energy_outside / energy_total
+
+        self.assertTrue(energy_ratio < 0.01)
+
+    def testModulated(self):
+        if self.L < 32:
+            raise SkipTest
+
+        ell = 1
+
+        sigma = self.L / 8
+        vol = gaussian_3d(self.L, sigma=(sigma, sigma, sigma), dtype=self.dtype)
+
+        g3d = grid_3d(self.L)
+
+        for trig_fun in (np.sin, np.cos):
+            vol1 = vol * trig_fun(ell * g3d["phi"])
+
+            coef = self.basis.expand(vol1)
+
+            ells = self.basis.indices()["ells"]
+
+            energy_outside = np.sum(np.abs(coef[ells != ell]) ** 2)
+            energy_total = np.sum(np.abs(coef) ** 2)
+
+            energy_ratio = energy_outside / energy_total
+
+            self.assertTrue(energy_ratio < 0.10)
 
 
 class UniversalBasisMixin: