Add coefficient_dtype property to Basis to enforce input and output types for evaluate and evaluate_t (#805)

* rewrite simplified coefficient_dtype

* add in for tests

* set coefficient_dtype in Basis __init__

* spacing

Author: chris-langfield

src/aspire/basis/basis.py

@@ -48,6 +48,9 @@ def __init__(self, size, ell_max=None, dtype=np.float32):
             raise NotImplementedError(
                 "Currently only implemented for float32 and float64 types"
             )
+        # dtype of coefficients is the same as self.dtype for real bases
+        # subclasses with complex coefficients override this attribute
+        self.coefficient_dtype = self.dtype
 
         self._build()
 
@@ -86,10 +89,10 @@ def evaluate(self, v):
             This is an Image or a Volume object containing one image/volume for each
             coefficient vector, and of size `self.sz`.
         """
-        if v.dtype != self.dtype:
+        if v.dtype != self.coefficient_dtype:
             logger.warning(
                 f"{self.__class__.__name__}::evaluate"
-                f" Inconsistent dtypes v: {v.dtype} self: {self.dtype}"
+                f" Inconsistent dtypes v: {v.dtype} self coefficient dtype: {self.coefficient_dtype}"
             )
 
         # Flatten stack, ndim is wrt Basis (2 or 3)
@@ -190,6 +193,12 @@ def expand(self, x):
         if isinstance(x, Image) or isinstance(x, Volume):
             x = x.asnumpy()
 
+        if x.dtype != self.dtype:
+            logger.warning(
+                f"{self.__class__.__name__}::expand"
+                f" Inconsistent dtypes x: {x.dtype} self: {self.dtype}"
+            )
+
         # check that last ndim values of input shape match
         # the shape of this basis
         assert (
@@ -212,7 +221,7 @@ def expand(self, x):
 
         # number of image samples
         n_data = x.shape[0]
-        v = np.zeros((n_data, self.count), dtype=x.dtype)
+        v = np.zeros((n_data, self.count), dtype=self.coefficient_dtype)
 
         for isample in range(0, n_data):
             b = self.evaluate_t(self._cls(x[isample])).T

src/aspire/basis/fpswf_2d.py

@@ -7,7 +7,6 @@
 
 from aspire.basis.basis_utils import lgwt, t_x_mat, t_x_mat_dot
 from aspire.basis.pswf_2d import PSWFBasis2D
-from aspire.image import Image
 from aspire.nufft import nufft
 from aspire.numeric import fft, xp
 from aspire.utils import complex_type
@@ -107,21 +106,13 @@ def _precomp(self):
         self.n_max = n_max
         self.size_x = len(self._disk_mask)
 
-    def evaluate_t(self, images):
+    def _evaluate_t(self, images):
         """
         Evaluate coefficient vectors in PSWF basis using the fast method.
 
         :param images: Image stack in the standard 2D coordinate basis.
         :return: Coefficient array in the PSWF basis.
         """
-
-        if not isinstance(images, Image):
-            logger.warning(
-                "FPSWFBasis2D.evaluate_t expects Image instance,"
-                " attempting conversion."
-            )
-            images = Image(images)
-
         # Construct array with zeros outside mask
         images_disk = np.zeros(images.shape, dtype=images.dtype)
         images_disk[:, self._disk_mask] = images[:, self._disk_mask]

src/aspire/basis/polar_2d.py

@@ -4,6 +4,7 @@
 
 from aspire.basis import Basis
 from aspire.nufft import anufft, nufft
+from aspire.utils import complex_type
 
 logger = logging.getLogger(__name__)
 
@@ -34,6 +35,9 @@ def __init__(self, size, nrad=None, ntheta=None, dtype=np.float32):
 
         super().__init__(size, dtype=dtype)
 
+        # this basis has complex coefficients
+        self.coefficient_dtype = complex_type(self.dtype)
+
     def _build(self):
         """
         Build the internal data structure to 2D polar Fourier grid

src/aspire/basis/pswf_2d.py

@@ -12,7 +12,6 @@
     t_x_mat,
 )
 from aspire.basis.pswf_utils import BNMatrix
-from aspire.image import Image
 from aspire.utils import complex_type
 
 logger = logging.getLogger(__name__)
@@ -59,6 +58,9 @@ def __init__(self, size, gamma_trunc=1.0, beta=1.0, dtype=np.float32):
         self.beta = beta
         super().__init__(size, dtype=dtype)
 
+        # this basis has complex coefficients
+        self.coefficient_dtype = complex_type(self.dtype)
+
     def _build(self):
         """
         Build internal data structures for the direct 2D PSWF method
@@ -149,27 +151,19 @@ def _generate_samples(self):
         # the column dimension of samples_conj_transpose is the number of basis coefficients
         self.count = self.samples_conj_transpose.shape[1]
 
-    def evaluate_t(self, images):
+    def _evaluate_t(self, images):
         """
         Evaluate coefficient vectors in PSWF basis using the direct method
 
         :param images: coefficient array in the standard 2D coordinate basis
             to be evaluated.
         :return: The evaluation of the coefficient array in the PSWF basis.
         """
-
-        if not isinstance(images, Image):
-            logger.warning(
-                "FPSWFBasis2D.evaluate_t expects Image instance,"
-                " attempting conversion."
-            )
-            images = Image(images)
-
         flattened_images = images[:, self._disk_mask]
 
         return flattened_images @ self.samples_conj_transpose
 
-    def evaluate(self, coefficients):
+    def _evaluate(self, coefficients):
         """
         Evaluate coefficients in standard 2D coordinate basis from those in PSWF basis
 
@@ -196,7 +190,7 @@ def evaluate(self, coefficients):
         )
         images[:, self._disk_mask] = np.real(flatten_images)
 
-        return Image(images)
+        return images
 
     def _init_pswf_func2d(self, c, eps):
         """

tests/_basis_util.py

@@ -138,26 +138,33 @@ def getClass(self):
             return Volume
 
     def testEvaluate(self):
-        # evaluate should take a NumPy array and return an Image/Volume
+        # evaluate should take a NumPy array of type basis.coefficient_dtype
+        # and return an Image/Volume
         _class = self.getClass()
-        result = self.basis.evaluate(np.zeros((self.basis.count), dtype=self.dtype))
+        result = self.basis.evaluate(
+            np.zeros((self.basis.count), dtype=self.basis.coefficient_dtype)
+        )
         self.assertTrue(isinstance(result, _class))
 
     def testEvaluate_t(self):
-        # evaluate_t should take an Image/Volume and return a NumPy array
+        # evaluate_t should take an Image/Volume and return a NumPy array of type
+        # basis.coefficient_dtype
         _class = self.getClass()
         result = self.basis.evaluate_t(
             _class(np.zeros((self.L,) * self.basis.ndim, dtype=self.dtype))
         )
         self.assertTrue(isinstance(result, np.ndarray))
+        self.assertEqual(result.dtype, self.basis.coefficient_dtype)
 
     def testExpand(self):
         _class = self.getClass()
-        # expand should take an Image/Volume and return a NumPy array
+        # expand should take an Image/Volume and return a NumPy array of type
+        # basis.coefficient_dtype
         result = self.basis.expand(
             _class(np.zeros((self.L,) * self.basis.ndim, dtype=self.dtype))
         )
         self.assertTrue(isinstance(result, np.ndarray))
+        self.assertEqual(result.dtype, self.basis.coefficient_dtype)
 
     def testInitWithIntSize(self):
         # make sure we can instantiate with just an int as a shortcut