continue filter_basis_mat cleanup

Author: garrettwrong

src/aspire/basis/fle_2d.py

@@ -780,16 +780,15 @@ def _radial_convolve_weights(self, b):
 
         return a
 
-    # def filter_to_basis_mat(self, f, **kwargs):
-    #     """
-    #     See `SteerableBasis2D.filter_stack_to_basis_mats`.
-    #     """
-    #     if len(f) != 1:
-    #         raise RuntimeError("Unexpected filter length.")
-    #     return self._filter_stack_to_basis_mats(f, **kwargs)[0]
-
-    # XXX TODO, convert to _filter_stack_to_basis_mats via broadcasting.
     def filter_to_basis_mat(self, f, **kwargs):
+        """
+        See `SteerableBasis2D.filter_stack_to_basis_mats`.
+        """
+        if len(f) != 1:
+            raise RuntimeError("Unexpected filter length.")
+        return self._filter_stack_to_basis_mats(f, **kwargs)
+
+    def _filter_stack_to_basis_mats(self, f, **kwargs):
         """
         See `SteerableBasis2D.filter_to_basis_mat`.
 
@@ -840,20 +839,28 @@ def filter_to_basis_mat(self, f, **kwargs):
 
         h_vals2d = (
             xp.asarray(h_fun(omega, pixel_size=pixel_size))
-            .reshape(n_k, n_theta)
+            .reshape(len(f), n_k, n_theta)
             .astype(self.dtype, copy=False)
         )
-        h_vals = xp.sum(h_vals2d, axis=1) / n_theta
+        h_vals = xp.sum(h_vals2d, axis=-1) / n_theta
 
-        h_basis = xp.zeros(self.count, dtype=self.dtype)
-        # For now we just need to handle 1D (stack of one ctf)
+        h_basis = xp.zeros((len(f), self.count), dtype=self.dtype)
+        # shape gymnastics to get a broadcast with csr A3
+        h_vals = h_vals.T
         for j in range(self.ell_p_max + 1):
-            h_basis[self.idx_list[j]] = self.A3[j] @ h_vals
+            h_basis[:, self.idx_list[j]] = (self.A3[j] @ h_vals).T
 
         # Convert from internal FLE ordering to FB convention
-        h_basis = h_basis[self._fle_to_fb_indices]
+        h_basis = h_basis[:, self._fle_to_fb_indices]
+        # who needs this as a list?
+
+        coefs = xp.asnumpy(h_basis)
+        if len(coefs) > 1:
+            coefs = [DiagMatrix(c) for c in coefs]
+        else:
+            coefs = DiagMatrix(coefs.flatten())
 
-        return DiagMatrix(xp.asnumpy(h_basis))
+        return coefs
 
     def expand_radial_vec(self, radial_vec, **kwargs):
         """

src/aspire/covariance/covar2d.py

@@ -567,40 +567,22 @@ def filters_to_basis_mats(self):
             getattr(self.basis.__class__, "expand_radial_vec", None)
         )
 
+        # Are the filters radial?
+        if self.src.filter_stack.radial:
+            logger.info("Found radial filter stack.")
+        else:
+            logger.info("Found non-radial filter stack.")
+
         if optimized_expand and self.src.filter_stack.radial:
-            logger.info(
-                "Found radial filter stack and `basis.expand_radial_vec` available."
-                "  Using bulk basis mat eval."
-            )
+            logger.info("Using optimized `basis.expand_radial_vec`.")
             return self._radial_filter_stack_to_basis_mats()
         else:
-            # Note, can come back and optmize the filter eval to bulk, just not radial
-            # For now use legacy path.
-            logger.info("Using sequential basis mat eval")
-            return self._filters_to_basis_mats()
-
-    def _filters_to_basis_mats(self):
-        """
-        old code, should work with all basis and filters. slow.
-        """
-        basis_mats = self.basis.filter_stack_to_basis_mats(
-            self.src.filter_stack,
-            pixel_size=self.src.pixel_size,
-            expand_method=self.expand_method,
-        )
-
-        # from tqdm import trange
-        # diff = 0
-        # for i in trange(len(basis_mats)):
-        #     a = basis_mats[i]
-        #     ref= old_basis_mats[i]
-        #     for j in range(len(ref)):
-        #         diff += np.sum(a[j]-ref[j])
-
-        # print("sum of diff across all filters", diff)
-        # breakpoint()
-
-        return basis_mats
+            logger.info("Using basis.filter_stack_to_basis_mats.")
+            return self.basis.filter_stack_to_basis_mats(
+                self.src.filter_stack,
+                pixel_size=self.src.pixel_size,
+                expand_method=self.expand_method,
+            )
 
     def _radial_filter_stack_to_basis_mats(self):
         logger.info("Generating filter eval points")