Fourier-Bessel Mixin class (#599)

Fourier Bessel Mixin Class: fb.py

Author: chris-langfield

src/aspire/basis/__init__.py

@@ -3,6 +3,7 @@
 
 from .basis import Basis
 from .steerable import SteerableBasis2D
+from .fb import FBBasisMixin
 
 # isort: on
 

src/aspire/basis/basis.py

@@ -3,10 +3,8 @@
 import numpy as np
 from scipy.sparse.linalg import LinearOperator, cg
 
-from aspire.basis.basis_utils import num_besselj_zeros
 from aspire.image import Image
 from aspire.utils import mdim_mat_fun_conj
-from aspire.utils.matlab_compat import m_reshape
 from aspire.volume import Volume
 
 logger = logging.getLogger(__name__)
@@ -47,46 +45,6 @@ def __init__(self, size, ell_max=None, dtype=np.float32):
 
         self._build()
 
-    def _getfbzeros(self):
-        """
-        Generate zeros of Bessel functions
-        """
-        # get upper_bound of zeros of Bessel functions
-        upper_bound = min(self.ell_max + 1, 2 * self.nres + 1)
-
-        # List of number of zeros
-        n = []
-        # List of zero values (each entry is an ndarray; all of possibly different lengths)
-        zeros = []
-
-        # generate zeros of Bessel functions for each ell
-        for ell in range(upper_bound):
-            # for each ell, num_besselj_zeros returns the zeros of the
-            # order ell Bessel function which are less than 2*pi*c*R = nres*pi/2,
-            # the truncation rule for the Fourier-Bessel expansion
-            _n, _zeros = num_besselj_zeros(
-                ell + (self.ndim - 2) / 2, self.nres * np.pi / 2
-            )
-            if _n == 0:
-                break
-            else:
-                n.append(_n)
-                zeros.append(_zeros)
-
-        #  get maximum number of ell
-        self.ell_max = len(n) - 1
-
-        #  set the maximum of k for each ell
-        self.k_max = np.array(n, dtype=int)
-
-        max_num_zeros = max(len(z) for z in zeros)
-        for i, z in enumerate(zeros):
-            zeros[i] = np.hstack(
-                (z, np.zeros(max_num_zeros - len(z), dtype=self.dtype))
-            )
-
-        self.r0 = m_reshape(np.hstack(zeros), (-1, self.ell_max + 1)).astype(self.dtype)
-
     def _build(self):
         """
         Build the internal data structure to represent basis

src/aspire/basis/basis_utils.py

@@ -247,7 +247,7 @@ def besselj_zeros(nu, k):
     return z
 
 
-def num_besselj_zeros(ell, r):
+def all_besselj_zeros(ell, r):
     """
     Compute the zeros of the order `ell` Bessel function which are less than `r`.
 

src/aspire/basis/fb.py

@@ -0,0 +1,56 @@
+import logging
+
+import numpy as np
+
+from aspire.basis.basis_utils import all_besselj_zeros
+from aspire.utils.matlab_compat import m_reshape
+
+logger = logging.getLogger(__name__)
+
+
+class FBBasisMixin(object):
+    """
+    FBBasisMixin is a mixin implementing methods specific to Fourier-Bessel expansions,
+    to be inherited by Fourier-Bessel subclasses of Basis.
+    """
+
+    def _calc_k_max(self):
+        """
+        Generate zeros of Bessel functions
+        """
+        # get upper_bound of zeros of Bessel functions
+        upper_bound = min(self.ell_max + 1, 2 * self.nres + 1)
+
+        # List of number of zeros
+        n = []
+        # List of zero values (each entry is an ndarray; all of possibly different lengths)
+        zeros = []
+
+        for ell in range(upper_bound):
+            # for each ell, num_besselj_zeros returns the zeros of the
+            # order ell Bessel function which are less than 2*pi*c*R = nres*pi/2,
+            # the truncation rule for the Fourier-Bessel expansion
+            if self.ndim == 2:
+                bessel_order = ell
+            elif self.ndim == 3:
+                bessel_order = ell + 1 / 2
+            _n, _zeros = all_besselj_zeros(bessel_order, self.nres * np.pi / 2)
+            if _n == 0:
+                break
+            else:
+                n.append(_n)
+                zeros.append(_zeros)
+
+        #  get maximum number of ell
+        self.ell_max = len(n) - 1
+
+        #  set the maximum of k for each ell
+        self.k_max = np.array(n, dtype=int)
+
+        max_num_zeros = max(len(z) for z in zeros)
+        for i, z in enumerate(zeros):
+            zeros[i] = np.hstack(
+                (z, np.zeros(max_num_zeros - len(z), dtype=self.dtype))
+            )
+
+        self.r0 = m_reshape(np.hstack(zeros), (-1, self.ell_max + 1)).astype(self.dtype)

src/aspire/basis/fb_2d.py

@@ -3,7 +3,7 @@
 import numpy as np
 from scipy.special import jv
 
-from aspire.basis import SteerableBasis2D
+from aspire.basis import FBBasisMixin, SteerableBasis2D
 from aspire.basis.basis_utils import unique_coords_nd
 from aspire.image import Image
 from aspire.utils import complex_type, real_type, roll_dim, unroll_dim
@@ -12,7 +12,7 @@
 logger = logging.getLogger(__name__)
 
 
-class FBBasis2D(SteerableBasis2D):
+class FBBasis2D(SteerableBasis2D, FBBasisMixin):
     """
     Define a derived class using the Fourier-Bessel basis for mapping 2D images
 
@@ -54,7 +54,7 @@ def _build(self):
         )
 
         # get upper bound of zeros, ells, and ks  of Bessel functions
-        self._getfbzeros()
+        self._calc_k_max()
 
         # calculate total number of basis functions
         self.count = self.k_max[0] + sum(2 * self.k_max[1:])

src/aspire/basis/fb_3d.py

@@ -2,15 +2,15 @@
 
 import numpy as np
 
-from aspire.basis import Basis
+from aspire.basis import Basis, FBBasisMixin
 from aspire.basis.basis_utils import real_sph_harmonic, sph_bessel, unique_coords_nd
 from aspire.utils import roll_dim, unroll_dim
 from aspire.utils.matlab_compat import m_flatten, m_reshape
 
 logger = logging.getLogger(__name__)
 
 
-class FBBasis3D(Basis):
+class FBBasis3D(Basis, FBBasisMixin):
     """
     Define a derived class for direct spherical Harmonics Bessel basis expanding 3D volumes
 
@@ -49,7 +49,7 @@ def _build(self):
         )
 
         # get upper bound of zeros, ells, and ks  of Bessel functions
-        self._getfbzeros()
+        self._calc_k_max()
 
         # calculate total number of basis functions
         self.count = sum(self.k_max * (2 * np.arange(0, self.ell_max + 1) + 1))

src/aspire/basis/ffb_2d.py

@@ -44,7 +44,7 @@ def _build(self):
         self.n_theta = int((n_theta + np.mod(n_theta, 2)) / 2)
 
         # get upper bound of zeros, ells, and ks  of Bessel functions
-        self._getfbzeros()
+        self._calc_k_max()
 
         # calculate total number of basis functions
         self.count = self.k_max[0] + sum(2 * self.k_max[1:])

src/aspire/basis/ffb_3d.py

@@ -33,7 +33,7 @@ def _build(self):
         self.kcut = 0.5
 
         # get upper bound of zeros, ells, and ks  of Bessel functions
-        self._getfbzeros()
+        self._calc_k_max()
 
         # calculate total number of basis functions
         self.count = sum(self.k_max * (2 * np.arange(0, self.ell_max + 1) + 1))

tests/test_basis_utils.py

@@ -3,10 +3,10 @@
 import numpy as np
 
 from aspire.basis.basis_utils import (
+    all_besselj_zeros,
     besselj_zeros,
     lgwt,
     norm_assoc_legendre,
-    num_besselj_zeros,
     real_sph_harmonic,
     sph_bessel,
     unique_coords_nd,
@@ -41,7 +41,7 @@ def testBesselJZeros(self):
         )
 
     def testNumBesselJZeros(self):
-        n, zeros = num_besselj_zeros(10, 20)
+        n, zeros = all_besselj_zeros(10, 20)
         self.assertEqual(2, n)
         self.assertTrue(np.allclose(zeros, [14.47550069, 18.43346367]))