Extend general tests to handle FB3D

Author: janden

tests/_basis_util.py

@@ -3,19 +3,24 @@
 import numpy as np
 
 from aspire.image import Image
+from aspire.volume import Volume
 from aspire.utils import gaussian_2d, utest_tolerance
 from aspire.utils.coor_trans import grid_2d
 from aspire.utils.random import randn
 
+
 class SteerableMixin:
     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]
+        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)))
@@ -25,7 +30,7 @@ def testAdjoint(self):
         u = u.astype(self.dtype)
 
         Au = self.basis.evaluate(u)
-        if isinstance(Au, Image):
+        if isinstance(Au, Image) or isinstance(Au, Volume):
             Au = Au.asnumpy()
 
         x = randn(*self.basis.sz, seed=self.seed)
@@ -131,3 +136,7 @@ def testModulated(self):
             energy_ratio = energy_outside / energy_total
 
             self.assertTrue(energy_ratio < 0.10)
+
+
+class Steerable3DMixin(SteerableMixin):
+    pass

tests/test_FBbasis3D.py

@@ -6,13 +6,16 @@
 from aspire.basis import FBBasis3D
 from aspire.utils import utest_tolerance
 
+from ._basis_util import Steerable3DMixin
+
 DATA_DIR = os.path.join(os.path.dirname(__file__), "saved_test_data")
 
 
-class FBBasis3DTestCase(TestCase):
+class FBBasis3DTestCase(TestCase, Steerable3DMixin):
     def setUp(self):
         self.dtype = np.float32
         self.basis = FBBasis3D((8, 8, 8), dtype=self.dtype)
+        self.seed = 9161341
 
     def tearDown(self):
         pass