Simplify Fourier-Bessel basis list of zeros (#776)

* no m_reshape

* implement ragged list of 1d np arrays

* missed one r0 usage

Author: chris-langfield

src/aspire/basis/fb.py

@@ -3,7 +3,6 @@
 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__)
 
@@ -47,10 +46,6 @@ def _calc_k_max(self):
         #  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)
+        # set the zeros for each ell
+        # this is a ragged list of 1d ndarrays, where the i'th element is of size self.k_max[i]
+        self.r0 = zeros

src/aspire/basis/fb_2d.py

@@ -143,7 +143,7 @@ def _precomp(self):
             for k in range(1, self.k_max[ell] + 1):
                 # Only normalized by the radial part of basis function
                 radial[:, ind_radial] = (
-                    jv(ell, self.r0[k - 1, ell] * r_unique)
+                    jv(ell, self.r0[ell][k - 1] * r_unique)
                     / self.radial_norms[ind_radial]
                 )
                 ind_radial += 1
@@ -177,7 +177,7 @@ def basis_norm_2d(self, ell, k):
         Calculate the normalized factors from radial and angular parts of a specified basis function
         """
         rad_norm = (
-            np.abs(jv(ell + 1, self.r0[k - 1, ell]))
+            np.abs(jv(ell + 1, self.r0[ell][k - 1]))
             * np.sqrt(1 / 2.0)
             * self.nres
             / 2.0

src/aspire/basis/fb_3d.py

@@ -105,7 +105,7 @@ def _precomp(self):
 
         for ell in range(0, self.ell_max + 1):
             for k in range(1, self.k_max[ell] + 1):
-                radial[:, ind_radial] = sph_bessel(ell, self.r0[k - 1, ell] * r_unique)
+                radial[:, ind_radial] = sph_bessel(ell, self.r0[ell][k - 1] * r_unique)
                 ind_radial += 1
 
             for m in range(-ell, ell + 1):
@@ -136,7 +136,7 @@ def basis_norm_3d(self, ell, k):
         Calculate the normalized factor of a specified basis function.
         """
         return (
-            np.abs(sph_bessel(ell + 1, self.r0[k - 1, ell]))
+            np.abs(sph_bessel(ell + 1, self.r0[ell][k - 1]))
             / np.sqrt(2)
             * np.sqrt((self.nres / 2) ** 3)
         )

src/aspire/basis/ffb_2d.py

@@ -74,7 +74,7 @@ def _precomp(self):
         ind_radial = 0
         for ell in range(0, self.ell_max + 1):
             for k in range(1, self.k_max[ell] + 1):
-                radial[ind_radial] = jv(ell, self.r0[k - 1, ell] * r / self.kcut)
+                radial[ind_radial] = jv(ell, self.r0[ell][k - 1] * r / self.kcut)
                 # NOTE: We need to remove the factor due to the discretization here
                 # since it is already included in our quadrature weights
                 # Only normalized by the radial part of basis function

src/aspire/basis/ffb_3d.py

@@ -75,11 +75,11 @@ def _precomp(self):
         )
         for ell in range(0, self.ell_max + 1):
             k_max_ell = self.k_max[ell]
-            rmat = r * self.r0[0:k_max_ell, ell].T / self.kcut
+            rmat = r * self.r0[ell][0:k_max_ell].T / self.kcut
             radial_ell = np.zeros_like(rmat)
             for ik in range(0, k_max_ell):
                 radial_ell[:, ik] = sph_bessel(ell, rmat[:, ik])
-            nrm = np.abs(sph_bessel(ell + 1, self.r0[0:k_max_ell, ell].T) / 4)
+            nrm = np.abs(sph_bessel(ell + 1, self.r0[ell][0:k_max_ell].T) / 4)
             radial_ell = radial_ell / nrm
             radial_ell_wtd = r**2 * wt_r * radial_ell
             radial_wtd[:, 0:k_max_ell, ell] = radial_ell_wtd

src/aspire/utils/filter_to_fb_mat.py

@@ -43,7 +43,7 @@ def filter_to_fb_mat(h_fun, fbasis):
     ind_ell = 0
     for ell in range(0, fbasis.ell_max + 1):
         k_max = fbasis.k_max[ell]
-        rmat = 2 * k_vals.reshape(n_k, 1) * fbasis.r0[0:k_max, ell].T
+        rmat = 2 * k_vals.reshape(n_k, 1) * fbasis.r0[ell][0:k_max].T
         fb_vals = np.zeros_like(rmat)
         ind_radial = np.sum(fbasis.k_max[0:ell])
         fb_vals[:, 0:k_max] = radial[ind_radial : ind_radial + k_max].T

tests/test_FBbasis2D.py

@@ -51,7 +51,7 @@ def _testElement(self, ell, k, sgn):
         g2d = grid_2d(self.L, dtype=self.dtype)
         mask = g2d["r"] < 1
 
-        r0 = self.basis.r0[k, ell]
+        r0 = self.basis.r0[ell][k]
 
         im = np.zeros((self.L, self.L), dtype=self.dtype)
         im[mask] = jv(ell, g2d["r"][mask] * r0)

tests/test_FFBbasis2D.py

@@ -49,7 +49,7 @@ def _testElement(self, ell, k, sgn):
         g2d = grid_2d(self.L, dtype=self.dtype)
         mask = g2d["r"] < 1
 
-        r0 = self.basis.r0[k, ell]
+        r0 = self.basis.r0[ell][k]
 
         # TODO: Figure out where these factors of 1 / 2 are coming from.
         # Intuitively, the grid should go from -L / 2 to L / 2, not -L / 2 to