From adb02d0e4ac156429cb2715adb4b16d41da9e923 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Wed, 9 Oct 2024 10:09:39 -0400
Subject: [PATCH 001/184] Remove Numpy 1 limitation

---
 pyproject.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pyproject.toml b/pyproject.toml
index 712eaa8f11..c4eabef52f 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -36,7 +36,7 @@ dependencies = [
     "joblib",
     "matplotlib >= 3.2.0",
     "mrcfile",
-    "numpy>=1.21.5, <2.0.0",
+    "numpy>=1.21.5",
     "packaging",
     "pooch>=1.7.0",
     "pillow",

From c60d8f2fdc331e7e1a60a7bb9d868c0684770af1 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Wed, 9 Oct 2024 10:14:07 -0400
Subject: [PATCH 002/184] latest np2, Inf to inf

---
 src/aspire/basis/basis.py  | 4 ++--
 src/aspire/basis/fb_2d.py  | 4 ++--
 src/aspire/basis/fb_3d.py  | 4 ++--
 src/aspire/basis/fle_2d.py | 2 +-
 4 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/src/aspire/basis/basis.py b/src/aspire/basis/basis.py
index 23c770674a..ee6aee2d02 100644
--- a/src/aspire/basis/basis.py
+++ b/src/aspire/basis/basis.py
@@ -373,9 +373,9 @@ def __init__(self, size, ell_max=None, dtype=np.float32):
         :param size: The size of the vectors for which to define the basis.
             Currently only square images and cubic volumes are supported.
         :param ell_max: The maximum order ell of the basis elements. If no input
-            (= None), it will be set to np.Inf and the basis includes all
+            (= None), it will be set to np.inf and the basis includes all
             ell such that the resulting basis vectors are concentrated
-            below the Nyquist frequency (default Inf).
+            below the Nyquist frequency (default inf).
         """
         if ell_max is None:
             ell_max = np.inf
diff --git a/src/aspire/basis/fb_2d.py b/src/aspire/basis/fb_2d.py
index 2698e41ed6..18f7b83478 100644
--- a/src/aspire/basis/fb_2d.py
+++ b/src/aspire/basis/fb_2d.py
@@ -31,9 +31,9 @@ def __init__(self, size, ell_max=None, dtype=np.float32):
             May be a 2-tuple or an integer, in which case a square basis is assumed.
             Currently only square images are supported.
         :ell_max: The maximum order ell of the basis elements. If no input
-            (= None), it will be set to np.Inf and the basis includes all
+            (= None), it will be set to np.inf and the basis includes all
             ell such that the resulting basis vectors are concentrated
-            below the Nyquist frequency (default Inf).
+            below the Nyquist frequency (default inf).
         """
 
         if isinstance(size, int):
diff --git a/src/aspire/basis/fb_3d.py b/src/aspire/basis/fb_3d.py
index b787ff618f..dedd52025c 100644
--- a/src/aspire/basis/fb_3d.py
+++ b/src/aspire/basis/fb_3d.py
@@ -26,9 +26,9 @@ def __init__(self, size, ell_max=None, dtype=np.float32):
             May be a 3-tuple or an integer, in which case a cubic basis is assumed.
             Currently only cubic images are supported.
         :param ell_max: The maximum order ell of the basis elements. If no input
-            (= None), it will be set to np.Inf and the basis includes all
+            (= None), it will be set to np.inf and the basis includes all
             ell such that the resulting basis vectors are concentrated
-            below the Nyquist frequency (default Inf).
+            below the Nyquist frequency (default inf).
         """
         if isinstance(size, int):
             size = (size, size, size)
diff --git a/src/aspire/basis/fle_2d.py b/src/aspire/basis/fle_2d.py
index 76330e6fba..ffce0ea284 100644
--- a/src/aspire/basis/fle_2d.py
+++ b/src/aspire/basis/fle_2d.py
@@ -358,7 +358,7 @@ def _lap_eig_disk(self):
         num_ells = 1 + 2 * max_ell
         self._ells = np.zeros((num_ells, max_k), dtype=int)
         self._ks = np.zeros((num_ells, max_k), dtype=int)
-        self.bessel_zeros = np.ones((num_ells, max_k), dtype=np.float64) * np.Inf
+        self.bessel_zeros = np.ones((num_ells, max_k), dtype=np.float64) * np.inf
 
         # keep track of which order Bessel function we're on
         self._ells[0, :] = 0

From 23b88c31675bc6b57b0e24d77b0672c7cba16d91 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Wed, 9 Oct 2024 10:48:23 -0400
Subject: [PATCH 003/184] row_stack to vstack

---
 src/aspire/basis/fpswf_2d.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/aspire/basis/fpswf_2d.py b/src/aspire/basis/fpswf_2d.py
index 21be9270d2..6eac1a91db 100644
--- a/src/aspire/basis/fpswf_2d.py
+++ b/src/aspire/basis/fpswf_2d.py
@@ -87,7 +87,7 @@ def _precomp(self):
         self.quad_rule_radial_wts = e
         self.num_angular_pts = f
 
-        us_fft_pts = np.row_stack((self.quad_rule_pts_x, self.quad_rule_pts_y))
+        us_fft_pts = np.vstack((self.quad_rule_pts_x, self.quad_rule_pts_y))
         us_fft_pts = self.bandlimit / self.rcut * us_fft_pts
         (
             blk_r,

From 6235411d06d545e35ffa370aa3e4d4cd1cc5f049 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Wed, 9 Oct 2024 10:58:51 -0400
Subject: [PATCH 004/184] np2 copy False -> None

---
 src/aspire/abinitio/commonline_base.py       |  2 +-
 src/aspire/abinitio/commonline_c3_c4.py      |  2 +-
 src/aspire/abinitio/commonline_d2.py         |  4 ++--
 src/aspire/abinitio/commonline_sync3n.py     | 18 +++++++++---------
 src/aspire/apple/picking.py                  |  4 ++--
 src/aspire/basis/basis.py                    |  2 +-
 src/aspire/basis/fle_2d.py                   |  8 +++-----
 src/aspire/basis/fpswf_2d.py                 |  2 +-
 src/aspire/classification/averager2d.py      |  2 +-
 src/aspire/classification/rir_class2d.py     |  2 +-
 src/aspire/ctf/ctf_estimator.py              |  4 ++--
 src/aspire/image/image.py                    |  6 +++---
 src/aspire/nufft/cufinufft.py                |  2 +-
 src/aspire/nufft/finufft.py                  |  4 ++--
 src/aspire/numeric/complex_pca/validation.py | 10 +++++-----
 src/aspire/numeric/pyfftw_fft.py             |  2 +-
 src/aspire/operators/blk_diag_matrix.py      |  4 ++--
 src/aspire/operators/diag_matrix.py          |  2 +-
 src/aspire/reconstruction/mean.py            |  2 +-
 src/aspire/sinogram/sinogram.py              |  2 +-
 src/aspire/source/micrograph.py              |  2 +-
 src/aspire/utils/bot_align.py                |  8 ++++----
 src/aspire/utils/matrix.py                   |  2 +-
 src/aspire/utils/misc.py                     |  6 +++---
 src/aspire/utils/rotation.py                 |  6 +++---
 src/aspire/volume/volume.py                  |  2 +-
 tests/test_FFBbasis3D.py                     |  4 ++--
 tests/test_coef.py                           |  8 ++++----
 tests/test_complexPCA.py                     |  2 +-
 tests/test_diag_matrix.py                    |  2 +-
 tests/test_downsample.py                     |  4 ++--
 tests/test_matrix.py                         |  4 ++--
 tests/test_mean_estimator.py                 |  4 ++--
 tests/test_micrograph_source.py              |  2 +-
 tests/test_orient_sdp.py                     |  2 +-
 tests/test_polar_ft.py                       |  2 +-
 tests/test_simulation.py                     |  2 +-
 tests/test_volume.py                         |  6 +++---
 38 files changed, 75 insertions(+), 77 deletions(-)

diff --git a/src/aspire/abinitio/commonline_base.py b/src/aspire/abinitio/commonline_base.py
index fe456c645e..2cd80b61fa 100644
--- a/src/aspire/abinitio/commonline_base.py
+++ b/src/aspire/abinitio/commonline_base.py
@@ -149,7 +149,7 @@ def _prepare_pf(self):
             #   Apply mask in doubles (allow imgs to upcast as needed)
             imgs = imgs * fuzz_mask
             #   Cast to desired type
-            imgs = Image(imgs.asnumpy().astype(self.dtype, copy=False))
+            imgs = Image(imgs.asnumpy().astype(self.dtype, copy=None))
 
         # Obtain coefficients of polar Fourier transform for input 2D images
         pft = PolarFT(
diff --git a/src/aspire/abinitio/commonline_c3_c4.py b/src/aspire/abinitio/commonline_c3_c4.py
index 242094b745..bc2956ec0d 100644
--- a/src/aspire/abinitio/commonline_c3_c4.py
+++ b/src/aspire/abinitio/commonline_c3_c4.py
@@ -693,7 +693,7 @@ def _J_sync_power_method(self, vijs):
         while itr < max_iters and residual > epsilon:
             itr += 1
             # Note, this appears to need double precision for accuracy in the following division.
-            vec_new = self._signs_times_v(vijs, vec).astype(np.float64, copy=False)
+            vec_new = self._signs_times_v(vijs, vec).astype(np.float64, copy=None)
             vec_new = vec_new / norm(vec_new)
             residual = norm(vec_new - vec)
             vec = vec_new
diff --git a/src/aspire/abinitio/commonline_d2.py b/src/aspire/abinitio/commonline_d2.py
index a8e951c642..93b2e9bc0a 100644
--- a/src/aspire/abinitio/commonline_d2.py
+++ b/src/aspire/abinitio/commonline_d2.py
@@ -1149,12 +1149,12 @@ def _sync_colors(self, Rijs):
 
         # Seed eigs initial vector for iterative method.
         # scipy LinearOperator needs doubles for some architectures (arm).
-        v0 = randn(3 * n_pairs, seed=self.seed).astype(np.float64, copy=False)
+        v0 = randn(3 * n_pairs, seed=self.seed).astype(np.float64, copy=None)
 
         v0 = v0 / norm(v0)
         vals, colors = la.eigs(color_mat, k=3, which="LR", v0=v0)
         vals = np.real(vals)
-        colors = np.real(colors).astype(self.dtype, copy=False)
+        colors = np.real(colors).astype(self.dtype, copy=None)
         colors = np.sign(colors[0]) * colors  # Stable eigs
         cp, _ = self._unmix_colors(colors[:, :2])
 
diff --git a/src/aspire/abinitio/commonline_sync3n.py b/src/aspire/abinitio/commonline_sync3n.py
index 5afa771c69..aa11a2348a 100644
--- a/src/aspire/abinitio/commonline_sync3n.py
+++ b/src/aspire/abinitio/commonline_sync3n.py
@@ -197,7 +197,7 @@ def _sync3n_S_to_rot(self, S, W=None, n_eigs=4):
         """
 
         # Critical this occurs in double precision
-        S = S.astype(np.float64, copy=False)
+        S = S.astype(np.float64, copy=None)
 
         if n_eigs < 3:
             raise ValueError(
@@ -213,7 +213,7 @@ def _sync3n_S_to_rot(self, S, W=None, n_eigs=4):
                 )
             # Initialize D
             # Critical this occurs in double precision
-            W = W.astype(np.float64, copy=False)
+            W = W.astype(np.float64, copy=None)
             D = np.mean(W, axis=1)
 
             Dhalf = D
@@ -664,8 +664,8 @@ def _pairs_probabilities_cupy(self, Rijs, P2, A, a, B, b, x0):
         )
 
         # accumulate over thread results
-        ln_f_arb = ln_f_arb_dev.get().astype(self.dtype, copy=False)
-        ln_f_ind = ln_f_ind_dev.get().astype(self.dtype, copy=False)
+        ln_f_arb = ln_f_arb_dev.get().astype(self.dtype, copy=None)
+        ln_f_ind = ln_f_ind_dev.get().astype(self.dtype, copy=None)
 
         return ln_f_ind, ln_f_arb
 
@@ -748,8 +748,8 @@ def fun(x, B, P, b, x0, A=A, a=a):
 
         popt, pcov = curve_fit(
             fun,
-            hist_x.astype(np.float64, copy=False),
-            scores_hist.astype(np.float64, copy=False),
+            hist_x.astype(np.float64, copy=None),
+            scores_hist.astype(np.float64, copy=None),
             p0=start_values,
             bounds=(lower_bounds, upper_bounds),
             method="trf",  # MATLAB used method "LAR" with algo "Trust-Region"
@@ -1043,7 +1043,7 @@ def _J_sync_power_method(self, Rijs):
         while itr < max_iters and residual > epsilon:
             itr += 1
             # Todo, this code code actually needs double precision for accuracy... forcing.
-            vec_new = self._signs_times_v(Rijs, vec).astype(np.float64, copy=False)
+            vec_new = self._signs_times_v(Rijs, vec).astype(np.float64, copy=None)
             vec_new = vec_new / norm(vec_new)
             residual = norm(vec_new - vec)
             vec = vec_new
@@ -1073,7 +1073,7 @@ def _signs_times_v(self, Rijs, vec):
         else:
             new_vec = self._signs_times_v_host(Rijs, vec)
 
-        return new_vec.astype(vec.dtype, copy=False)
+        return new_vec.astype(vec.dtype, copy=None)
 
     def _signs_times_v_host(self, Rijs, vec):
         """
@@ -1177,7 +1177,7 @@ def _signs_times_v_cupy(self, Rijs, vec):
         )
 
         # dtoh
-        new_vec = new_vec_dev.get().astype(vec.dtype, copy=False)
+        new_vec = new_vec_dev.get().astype(vec.dtype, copy=None)
 
         return new_vec
 
diff --git a/src/aspire/apple/picking.py b/src/aspire/apple/picking.py
index d64d16c773..d54cf61f1d 100644
--- a/src/aspire/apple/picking.py
+++ b/src/aspire/apple/picking.py
@@ -149,7 +149,7 @@ def read_micrograph(self):
         self.original_im = (
             DiskMicrographSource(self.filename)
             .asnumpy()[0]
-            .astype(np.float32, copy=False)
+            .astype(np.float32, copy=None)
         )
 
         # Discard outer pixels
@@ -167,7 +167,7 @@ def read_micrograph(self):
         # Note, float64 required for signal.correlate call accuracy.
         im = np.asarray(
             PILImage.fromarray(im).resize(size, PILImage.Resampling.BICUBIC)
-        ).astype(np.float64, copy=False)
+        ).astype(np.float64, copy=None)
 
         im = signal.correlate(
             im,
diff --git a/src/aspire/basis/basis.py b/src/aspire/basis/basis.py
index ee6aee2d02..6baf16f769 100644
--- a/src/aspire/basis/basis.py
+++ b/src/aspire/basis/basis.py
@@ -62,7 +62,7 @@ def __init__(self, basis, data, dtype=None):
             )
         self.basis = basis
 
-        self._data = data.astype(self.dtype, copy=False)
+        self._data = data.astype(self.dtype, copy=None)
         self.ndim = self._data.ndim
         self.shape = self._data.shape
         self.stack_ndim = self._data.ndim - 1
diff --git a/src/aspire/basis/fle_2d.py b/src/aspire/basis/fle_2d.py
index ffce0ea284..6499919110 100644
--- a/src/aspire/basis/fle_2d.py
+++ b/src/aspire/basis/fle_2d.py
@@ -314,7 +314,7 @@ def _compute_nufft_points(self):
         )
         grid_xy[0] = xp.cos(phi)  # x
         grid_xy[1] = xp.sin(phi)  # y
-        grid_xy = grid_xy * nodes * h
+        grid_xy[:] = grid_xy * nodes * h
         self.grid_xy = grid_xy.reshape(2, -1)
 
     def _build_interpolation_matrix(self):
@@ -643,7 +643,7 @@ def _step1(self, z):
         num_img = z.shape[0]
         z = z[:, :, : self.num_angular_nodes // 2].reshape(num_img, -1)
         im = anufft(
-            z.astype(complex_type(self.dtype), copy=False),
+            z.astype(complex_type(self.dtype), copy=None),
             self.grid_xy,
             (self.nres, self.nres),
             epsilon=self.epsilon,
@@ -798,9 +798,7 @@ def filter_to_basis_mat(self, f, **kwargs):
         omega = 2 * xp.pi * xp.vstack((omegax.flatten("C"), omegay.flatten("C")))
 
         h_vals2d = (
-            xp.asarray(h_fun(omega))
-            .reshape(n_k, n_theta)
-            .astype(self.dtype, copy=False)
+            xp.asarray(h_fun(omega)).reshape(n_k, n_theta).astype(self.dtype, copy=None)
         )
         h_vals = xp.sum(h_vals2d, axis=1) / n_theta
 
diff --git a/src/aspire/basis/fpswf_2d.py b/src/aspire/basis/fpswf_2d.py
index 6eac1a91db..aa7b9913d9 100644
--- a/src/aspire/basis/fpswf_2d.py
+++ b/src/aspire/basis/fpswf_2d.py
@@ -99,7 +99,7 @@ def _precomp(self):
         ) = self._pswf_integration_sub_routine()
 
         self.us_fft_pts = us_fft_pts.astype(
-            self.dtype, copy=False
+            self.dtype, copy=None
         )  # TODO, debug where this is incorrect dtype
         self.blk_r = blk_r
         self.num_angular_pts = num_angular_pts
diff --git a/src/aspire/classification/averager2d.py b/src/aspire/classification/averager2d.py
index b447a6d91c..644359c9e7 100644
--- a/src/aspire/classification/averager2d.py
+++ b/src/aspire/classification/averager2d.py
@@ -82,7 +82,7 @@ def _cls_images(self, cls, src=None):
             source for a certain operation (ie alignment).
         """
         src = src or self.src
-        return src.images[cls].asnumpy().astype(self.dtype, copy=False)
+        return src.images[cls].asnumpy().astype(self.dtype, copy=None)
 
 
 class AligningAverager2D(Averager2D):
diff --git a/src/aspire/classification/rir_class2d.py b/src/aspire/classification/rir_class2d.py
index 19a3479ffe..dde35ef050 100644
--- a/src/aspire/classification/rir_class2d.py
+++ b/src/aspire/classification/rir_class2d.py
@@ -390,7 +390,7 @@ def _sk_pca(self, M):
         #   We use an extension of SK that is hacked to admit complex.
         pca = ComplexPCA(
             self.bispectrum_components,
-            copy=False,  # careful, overwrites data matrix... we'll handle the copies.
+            copy=None,  # careful, overwrites data matrix... we'll handle the copies.
             svd_solver="auto",  # use randomized (Halko) for larger problems
             random_state=self.seed,
         )
diff --git a/src/aspire/ctf/ctf_estimator.py b/src/aspire/ctf/ctf_estimator.py
index b6983cb071..39f2998a21 100644
--- a/src/aspire/ctf/ctf_estimator.py
+++ b/src/aspire/ctf/ctf_estimator.py
@@ -241,7 +241,7 @@ def estimate_psd(self, blocks, tapers_1d):
         """
 
         num_1d_tapers = tapers_1d.shape[-1]
-        tapers_1d = tapers_1d.astype(complex_type(self.dtype), copy=False)
+        tapers_1d = tapers_1d.astype(complex_type(self.dtype), copy=None)
 
         blocks_mt = np.zeros(blocks[0, :, :].shape, dtype=self.dtype)
 
@@ -768,7 +768,7 @@ def estimate_ctf(
             micrograph = mrc.data
 
         # Try to match dtype used in Basis instance
-        micrograph = micrograph.astype(dtype, copy=False)
+        micrograph = micrograph.astype(dtype, copy=None)
 
         micrograph_blocks = ctf_object.preprocess_micrograph(micrograph, psd_size)
 
diff --git a/src/aspire/image/image.py b/src/aspire/image/image.py
index 757362870d..6d0f3b8960 100644
--- a/src/aspire/image/image.py
+++ b/src/aspire/image/image.py
@@ -195,7 +195,7 @@ def __init__(self, data, pixel_size=None, dtype=None):
         if not data.shape[-1] == data.shape[-2]:
             raise ValueError("Only square ndarrays are supported.")
 
-        self._data = data.astype(self.dtype, copy=False)
+        self._data = data.astype(self.dtype, copy=None)
         self.ndim = self._data.ndim
         self.shape = self._data.shape
         self.stack_ndim = self._data.ndim - 2
@@ -543,7 +543,7 @@ def load(filepath, dtype=None):
 
         # Attempt casting when user provides dtype
         if dtype is not None:
-            im = im.astype(dtype, copy=False)
+            im = im.astype(dtype, copy=None)
 
         # Return as Image instance
         return Image(im, pixel_size=pixel_size)
@@ -658,7 +658,7 @@ def backproject(self, rot_matrices, symmetry_group=None, zero_nyquist=True):
 
             # TODO: rotated_grids might as well give us correctly shaped array in the first place
             pts_rot = aspire.volume.rotated_grids(L, rotations).astype(
-                self.dtype, copy=False
+                self.dtype, copy=None
             )
             pts_rot = pts_rot.reshape((3, -1))
 
diff --git a/src/aspire/nufft/cufinufft.py b/src/aspire/nufft/cufinufft.py
index fd869aacfd..2660f4ea46 100644
--- a/src/aspire/nufft/cufinufft.py
+++ b/src/aspire/nufft/cufinufft.py
@@ -30,7 +30,7 @@ def __init__(self, sz, fourier_pts, epsilon=1e-8, ntransforms=1, **kwargs):
         # ASPIRE-Python/703
         # Cast to doubles.
         self._original_dtype = fourier_pts.dtype
-        fourier_pts = fourier_pts.astype(np.float64, copy=False)
+        fourier_pts = fourier_pts.astype(np.float64, copy=None)
 
         # Basic dtype passthough.
         dtype = fourier_pts.dtype
diff --git a/src/aspire/nufft/finufft.py b/src/aspire/nufft/finufft.py
index efc1a062f2..e396d76124 100644
--- a/src/aspire/nufft/finufft.py
+++ b/src/aspire/nufft/finufft.py
@@ -99,7 +99,7 @@ def transform(self, signal):
         ), f"Signal frame to be transformed must have shape {self.sz}"
 
         # FINUFFT was designed for a complex input array
-        signal = np.array(signal, copy=False, dtype=self.complex_dtype, order="C")
+        signal = np.array(signal, copy=None, dtype=self.complex_dtype, order="C")
 
         result = self._transform_plan.execute(signal)
 
@@ -130,7 +130,7 @@ def adjoint(self, signal):
                 signal = signal.reshape(self.num_pts)
 
         # FINUFFT was designed for a complex input array
-        signal = np.array(signal, copy=False, dtype=self.complex_dtype, order="C")
+        signal = np.array(signal, copy=None, dtype=self.complex_dtype, order="C")
 
         result = self._adjoint_plan.execute(signal)
 
diff --git a/src/aspire/numeric/complex_pca/validation.py b/src/aspire/numeric/complex_pca/validation.py
index 7fafedfe18..11d7aaba9f 100644
--- a/src/aspire/numeric/complex_pca/validation.py
+++ b/src/aspire/numeric/complex_pca/validation.py
@@ -81,7 +81,7 @@ def _ensure_sparse_format(
         Data type of result. If None, the dtype of the input is preserved.
 
     copy : boolean
-        Whether a forced copy will be triggered. If copy=False, a copy might
+        Whether a forced copy will be triggered. If copy=None, a copy might
         be triggered by a conversion.
 
     force_all_finite : boolean or 'allow-nan', (default=True)
@@ -173,7 +173,7 @@ def check_array(
     accept_large_sparse=True,
     dtype="numeric",
     order=None,
-    copy=False,
+    copy=None,
     force_all_finite=True,
     ensure_2d=True,
     allow_nd=False,
@@ -216,13 +216,13 @@ def check_array(
 
     order : 'F', 'C' or None (default=None)
         Whether an array will be forced to be fortran or c-style.
-        When order is None (default), then if copy=False, nothing is ensured
+        When order is None (default), then if copy=None, nothing is ensured
         about the memory layout of the output array; otherwise (copy=True)
         the memory layout of the returned array is kept as close as possible
         to the original array.
 
     copy : boolean (default=False)
-        Whether a forced copy will be triggered. If copy=False, a copy might
+        Whether a forced copy will be triggered. If copy=None, a copy might
         be triggered by a conversion.
 
     force_all_finite : boolean or 'allow-nan', (default=True)
@@ -365,7 +365,7 @@ def check_array(
                     array = np.asarray(array, order=order)
                     if array.dtype.kind == "f":
                         _assert_all_finite(array, allow_nan=False, msg_dtype=dtype)
-                    array = array.astype(dtype, casting="unsafe", copy=False)
+                    array = array.astype(dtype, casting="unsafe", copy=None)
                 else:
                     array = np.asarray(array, order=order, dtype=dtype)
             except ComplexWarning:
diff --git a/src/aspire/numeric/pyfftw_fft.py b/src/aspire/numeric/pyfftw_fft.py
index 95a8ea80f7..4503b42b38 100644
--- a/src/aspire/numeric/pyfftw_fft.py
+++ b/src/aspire/numeric/pyfftw_fft.py
@@ -139,7 +139,7 @@ def ifftn(self, a, axes=None, workers=-1):
         # FFTW_BACKWARD requires complex input array, cast as needed.
         # See https://pyfftw.readthedocs.io/en/latest/source/pyfftw/pyfftw.html#scheme-table
         comp_type = complex_type(a.dtype)
-        a = a.astype(comp_type, copy=False)
+        a = a.astype(comp_type, copy=None)
 
         mutex.acquire()
         try:
diff --git a/src/aspire/operators/blk_diag_matrix.py b/src/aspire/operators/blk_diag_matrix.py
index 3e493f9525..af882ecc68 100644
--- a/src/aspire/operators/blk_diag_matrix.py
+++ b/src/aspire/operators/blk_diag_matrix.py
@@ -91,7 +91,7 @@ def append(self, blk):
         :param blk: Block to append (ndarray).
         """
 
-        self.data.append(blk.astype(self.dtype, copy=False))
+        self.data.append(blk.astype(self.dtype, copy=None))
         self.nblocks += 1
         self.reset_cache()
 
@@ -130,7 +130,7 @@ def __setitem__(self, key, value):
         Convenience wrapper, setter on self.data.
         """
 
-        self.data[key] = value.astype(self.dtype, copy=False)
+        self.data[key] = value.astype(self.dtype, copy=None)
         self.reset_cache()
 
     def __len__(self):
diff --git a/src/aspire/operators/diag_matrix.py b/src/aspire/operators/diag_matrix.py
index 4c94ab83d1..68d6a728ae 100644
--- a/src/aspire/operators/diag_matrix.py
+++ b/src/aspire/operators/diag_matrix.py
@@ -45,7 +45,7 @@ def __init__(self, data, dtype=None):
         self.dtype = np.dtype(dtype)
 
         # Assign the `data`
-        self._data = data.astype(self.dtype, copy=False)
+        self._data = data.astype(self.dtype, copy=None)
 
         # Assign shapes from `data`
         self.count = self._data.shape[-1]
diff --git a/src/aspire/reconstruction/mean.py b/src/aspire/reconstruction/mean.py
index d0cade9754..56f0624e8d 100644
--- a/src/aspire/reconstruction/mean.py
+++ b/src/aspire/reconstruction/mean.py
@@ -102,7 +102,7 @@ def _compute_kernel(self):
             _range = np.arange(i, min(self.src.n, i + self.batch_size), dtype=int)
             sq_filters_f = evaluate_src_filters_on_grid(self.src, _range) ** 2
             amplitudes_sq = (self.src.amplitudes[_range] ** 2).astype(
-                self.dtype, copy=False
+                self.dtype, copy=None
             )
 
             for k in range(self.r):
diff --git a/src/aspire/sinogram/sinogram.py b/src/aspire/sinogram/sinogram.py
index 7c7bb43662..0f82b0c479 100644
--- a/src/aspire/sinogram/sinogram.py
+++ b/src/aspire/sinogram/sinogram.py
@@ -37,7 +37,7 @@ def __init__(self, data, dtype=None):
                 f"Invalid data shape: {data.shape}. Expected shape: (..., angles, radial_points), where '...' is the stack number."
             )
 
-        self._data = data.astype(self.dtype, copy=False)
+        self._data = data.astype(self.dtype, copy=None)
         self.ndim = self._data.ndim
         self.shape = self._data.shape
         self.stack_shape = self._data.shape[:-2]
diff --git a/src/aspire/source/micrograph.py b/src/aspire/source/micrograph.py
index 5aa35e389c..86b161d822 100644
--- a/src/aspire/source/micrograph.py
+++ b/src/aspire/source/micrograph.py
@@ -155,7 +155,7 @@ def __init__(self, micrographs, dtype=None, pixel_size=None):
         )
 
         # We're already backed by an array, access it directly.
-        self._data = micrographs.astype(self.dtype, copy=False)
+        self._data = micrographs.astype(self.dtype, copy=None)
 
     def _images(self, indices):
         """
diff --git a/src/aspire/utils/bot_align.py b/src/aspire/utils/bot_align.py
index fc8fbe4fdc..e94888d92e 100644
--- a/src/aspire/utils/bot_align.py
+++ b/src/aspire/utils/bot_align.py
@@ -139,7 +139,7 @@ def cost(new, t=t, q=q):
             Loss function for surrogate problems.
             """
             kx = np.array(
-                [cov_fun(new.astype(dtype, copy=False), R[j]) for j in range(t)]
+                [cov_fun(new.astype(dtype, copy=None), R[j]) for j in range(t)]
             )
             mu = kx @ q
             return mu
@@ -151,7 +151,7 @@ def euclidean_grad(new, t=t, q=q):
             """
             # Corresponds to equation 11 in the paper.
             kx_grad = np.array(
-                [cov_fun_grad(new.astype(dtype, copy=False), R[j]) for j in range(t)]
+                [cov_fun_grad(new.astype(dtype, copy=None), R[j]) for j in range(t)]
             )
 
             kx_grad = kx_grad.reshape((t, 9))
@@ -168,7 +168,7 @@ def euclidean_grad(new, t=t, q=q):
             verbosity=verbosity,
         )
         result = optimizer.run(problem)
-        R_new = result.point.astype(dtype, copy=False)
+        R_new = result.point.astype(dtype, copy=None)
 
         loss[t] = loss_fun(R_new)
         R[t] = R_new
@@ -190,7 +190,7 @@ def loss_u(u):
             """
             Convert the alignment loss function over SO(3) to be over R^3.
             """
-            u = u.astype(dtype, copy=False)
+            u = u.astype(dtype, copy=None)
             R = sign * Rotation.from_rotvec(u).matrices[0]
             v_rot = vol_given_ds.rotate(Rotation(R)).asnumpy()[0]
             return norm(vol_ref_ds - v_rot)
diff --git a/src/aspire/utils/matrix.py b/src/aspire/utils/matrix.py
index 71c709608b..bc802f87f8 100644
--- a/src/aspire/utils/matrix.py
+++ b/src/aspire/utils/matrix.py
@@ -462,7 +462,7 @@ def nearest_rotations(A, allow_reflection=False):
         # If det(U)*det(V) = -1, we negate the third singular value to
         # ensure we have a rotation.
         neg_det_idx = np.linalg.det(U) * np.linalg.det(V) < 0
-        U[neg_det_idx] = U[neg_det_idx] @ np.diag((1, 1, -1)).astype(dtype, copy=False)
+        U[neg_det_idx] = U[neg_det_idx] @ np.diag((1, 1, -1)).astype(dtype, copy=None)
 
     rots = U @ V
 
diff --git a/src/aspire/utils/misc.py b/src/aspire/utils/misc.py
index a3f9917024..087454de5e 100644
--- a/src/aspire/utils/misc.py
+++ b/src/aspire/utils/misc.py
@@ -135,7 +135,7 @@ def gaussian_1d(size, mu=0, sigma=1, dtype=np.float64):
 
     p = (g["x"] - mu) ** 2 / (2 * sigma**2)
 
-    return np.exp(-p).astype(dtype, copy=False)
+    return np.exp(-p).astype(dtype, copy=None)
 
 
 def gaussian_2d(size, mu=(0, 0), sigma=(1, 1), indexing="yx", dtype=np.float64):
@@ -180,7 +180,7 @@ def gaussian_2d(size, mu=(0, 0), sigma=(1, 1), indexing="yx", dtype=np.float64):
         2 * sigma[1] ** 2
     )
 
-    return np.exp(-p).astype(dtype, copy=False)
+    return np.exp(-p).astype(dtype, copy=None)
 
 
 def gaussian_3d(size, mu=(0, 0, 0), sigma=(1, 1, 1), indexing="zyx", dtype=np.float64):
@@ -228,7 +228,7 @@ def gaussian_3d(size, mu=(0, 0, 0), sigma=(1, 1, 1), indexing="zyx", dtype=np.fl
         + (g["z"] - mu[2]) ** 2 / (2 * sigma[2] ** 2)
     )
 
-    return np.exp(-p).astype(dtype, copy=False)
+    return np.exp(-p).astype(dtype, copy=None)
 
 
 def bump_3d(size, spread=1, dtype=np.float64):
diff --git a/src/aspire/utils/rotation.py b/src/aspire/utils/rotation.py
index 07a31df9df..fb643914bf 100644
--- a/src/aspire/utils/rotation.py
+++ b/src/aspire/utils/rotation.py
@@ -256,7 +256,7 @@ def from_euler(values, dtype=None):
         """
         dtype = dtype or getattr(values, "dtype", np.float32)
         rotations = sp_rot.from_euler("ZYZ", values, degrees=False)
-        matrices = rotations.as_matrix().astype(dtype, copy=False)
+        matrices = rotations.as_matrix().astype(dtype, copy=None)
         return Rotation(matrices)
 
     @staticmethod
@@ -324,7 +324,7 @@ def from_rotvec(vec, dtype=None):
         """
         dtype = dtype or vec.dtype
         rots = sp_rot.from_rotvec(vec)
-        matrices = rots.as_matrix().astype(dtype, copy=False)
+        matrices = rots.as_matrix().astype(dtype, copy=None)
         return Rotation(matrices)
 
     @staticmethod
@@ -337,7 +337,7 @@ def from_matrix(values, dtype=None):
         :return: new Rotation object
         """
         dtype = dtype or values.dtype
-        return Rotation(values.astype(dtype, copy=False))
+        return Rotation(values.astype(dtype, copy=None))
 
     @staticmethod
     def generate_random_rotations(
diff --git a/src/aspire/volume/volume.py b/src/aspire/volume/volume.py
index eae985a794..1954a1cbd5 100644
--- a/src/aspire/volume/volume.py
+++ b/src/aspire/volume/volume.py
@@ -105,7 +105,7 @@ def __init__(self, data, dtype=None, pixel_size=None, symmetry_group=None):
         if not (data.shape[-1] == data.shape[-2] == data.shape[-3]):
             raise ValueError("Only cubed ndarrays are supported.")
 
-        self._data = data.astype(self.dtype, copy=False)
+        self._data = data.astype(self.dtype, copy=None)
         self.ndim = self._data.ndim
         self.shape = self._data.shape
         self.stack_ndim = self._data.ndim - 3
diff --git a/tests/test_FFBbasis3D.py b/tests/test_FFBbasis3D.py
index 03bdfdd21b..3992f5d374 100644
--- a/tests/test_FFBbasis3D.py
+++ b/tests/test_FFBbasis3D.py
@@ -470,7 +470,7 @@ def testFFBBasis3DEvaluate(self, basis):
 
     def testFFBBasis3DEvaluate_t(self, basis):
         x = np.load(os.path.join(DATA_DIR, "ffbbasis3d_xcoef_in_8_8_8.npy")).T  # RCOPT
-        x = x.astype(basis.dtype, copy=False)
+        x = x.astype(basis.dtype, copy=None)
         result = basis.evaluate_t(Volume(x))
 
         ref = np.load(os.path.join(DATA_DIR, "ffbbasis3d_vcoef_out_8_8_8.npy"))[..., 0]
@@ -479,7 +479,7 @@ def testFFBBasis3DEvaluate_t(self, basis):
 
     def testFFBBasis3DExpand(self, basis):
         x = np.load(os.path.join(DATA_DIR, "ffbbasis3d_xcoef_in_8_8_8.npy")).T  # RCOPT
-        x = x.astype(basis.dtype, copy=False)
+        x = x.astype(basis.dtype, copy=None)
         result = basis.expand(x)
 
         ref = np.load(os.path.join(DATA_DIR, "ffbbasis3d_vcoef_out_exp_8_8_8.npy"))[
diff --git a/tests/test_coef.py b/tests/test_coef.py
index 3ace0ddec5..2a283d16e4 100644
--- a/tests/test_coef.py
+++ b/tests/test_coef.py
@@ -90,7 +90,7 @@ def coef_fixture(basis, stack, dtype):
     # shape.
     size = stack + (basis.count,)
 
-    coef_np = np.random.random(size=size).astype(dtype, copy=False)
+    coef_np = np.random.random(size=size).astype(dtype, copy=None)
 
     return Coef(basis, coef_np, dtype=dtype)
 
@@ -212,7 +212,7 @@ def test_add(basis, coef_fixture):
     Tests addition operation against pure Numpy.
     """
     # Make array
-    x = np.random.random(size=coef_fixture.shape).astype(coef_fixture.dtype, copy=False)
+    x = np.random.random(size=coef_fixture.shape).astype(coef_fixture.dtype, copy=None)
     # Construct Coef
     c = Coef(basis, x)
 
@@ -231,7 +231,7 @@ def test_sub(basis, coef_fixture):
     Tests subtraction operation against pure Numpy.
     """
     # Make array
-    x = np.random.random(size=coef_fixture.shape).astype(coef_fixture.dtype, copy=False)
+    x = np.random.random(size=coef_fixture.shape).astype(coef_fixture.dtype, copy=None)
     # Construct Coef
     c = Coef(basis, x)
 
@@ -264,7 +264,7 @@ def test_mul(basis, coef_fixture):
     Tests multiplication operation against pure Numpy.
     """
     # Make array
-    x = np.random.random(size=coef_fixture.shape).astype(coef_fixture.dtype, copy=False)
+    x = np.random.random(size=coef_fixture.shape).astype(coef_fixture.dtype, copy=None)
     # Construct Coef
     c = Coef(basis, x)
 
diff --git a/tests/test_complexPCA.py b/tests/test_complexPCA.py
index 509b0b84a5..1c224324d7 100644
--- a/tests/test_complexPCA.py
+++ b/tests/test_complexPCA.py
@@ -53,7 +53,7 @@ def testLargeFitTransform(self):
         """
 
         pca = ComplexPCA(
-            n_components=self.components_large, copy=False, svd_solver="full"
+            n_components=self.components_large, copy=None, svd_solver="full"
         )
 
         # Input data matrix X should be (n_samples, m_features)
diff --git a/tests/test_diag_matrix.py b/tests/test_diag_matrix.py
index 05805912c9..30f2b76faa 100644
--- a/tests/test_diag_matrix.py
+++ b/tests/test_diag_matrix.py
@@ -65,7 +65,7 @@ def diag_matrix_fixture(stack, matrix_size, dtype):
     """
     shape = (2,) + stack + (matrix_size,)
     # Internally convert dtype.  Passthrough will be checked explicitly in `test_dtype_passthrough` and `test_dtype_cast`
-    d_np = np.random.random(shape).astype(dtype, copy=False)
+    d_np = np.random.random(shape).astype(dtype, copy=None)
     d1 = DiagMatrix(d_np[0])
     d2 = DiagMatrix(d_np[1])
 
diff --git a/tests/test_downsample.py b/tests/test_downsample.py
index 6c1cac82dd..1af37d7281 100644
--- a/tests/test_downsample.py
+++ b/tests/test_downsample.py
@@ -148,7 +148,7 @@ def emdb_vol():
 
 @pytest.fixture(scope="module")
 def volume(emdb_vol, res, dtype):
-    vol = emdb_vol.astype(dtype, copy=False)
+    vol = emdb_vol.astype(dtype, copy=None)
     vol = vol.downsample(res)
     return vol
 
@@ -266,7 +266,7 @@ def test_pixel_size():
     dsL = 5  # downsampled
 
     # Construct a small test Image
-    img = Image(np.random.random((1, L, L)).astype(DTYPE, copy=False), pixel_size=1.23)
+    img = Image(np.random.random((1, L, L)).astype(DTYPE, copy=None), pixel_size=1.23)
 
     # Downsample the image
     result = img.downsample(dsL)
diff --git a/tests/test_matrix.py b/tests/test_matrix.py
index 728200b9b3..be6829d3fa 100644
--- a/tests/test_matrix.py
+++ b/tests/test_matrix.py
@@ -356,7 +356,7 @@ def test_nearest_rotations(dtype):
     rots = Rotation.generate_random_rotations(n_rots, seed=0, dtype=dtype).matrices
 
     # Add some noise to the rotations.
-    noise = 1e-3 * randn(n_rots * 9, seed=0).astype(dtype, copy=False).reshape(
+    noise = 1e-3 * randn(n_rots * 9, seed=0).astype(dtype, copy=None).reshape(
         n_rots, 3, 3
     )
     noisy_rots = rots + noise
@@ -381,7 +381,7 @@ def test_nearest_rotations_reflection(dtype):
 
     # Add a reflection and some noise to the rotation.
     refl = rot @ np.diag((1, -1, 1)).astype(dtype)
-    noise = 1e-3 * randn(9, seed=0).astype(dtype, copy=False).reshape(3, 3)
+    noise = 1e-3 * randn(9, seed=0).astype(dtype, copy=None).reshape(3, 3)
     noisy_refl = refl + noise
 
     # Find nearest rotation.
diff --git a/tests/test_mean_estimator.py b/tests/test_mean_estimator.py
index e6b2a2f837..3e7f915193 100644
--- a/tests/test_mean_estimator.py
+++ b/tests/test_mean_estimator.py
@@ -115,8 +115,8 @@ def test_adjoint(sim, basis, estimator):
     L = sim.L
     n = sim.n
 
-    u = np.random.rand(n, L, L).astype(sim.dtype, copy=False)
-    v = np.random.rand(L, L, L).astype(sim.dtype, copy=False)
+    u = np.random.rand(n, L, L).astype(sim.dtype, copy=None)
+    v = np.random.rand(L, L, L).astype(sim.dtype, copy=None)
 
     proj = Volume(v).project(rots)
     backproj = Image(u).backproject(rots)
diff --git a/tests/test_micrograph_source.py b/tests/test_micrograph_source.py
index d4793cf61f..c44e703cbb 100644
--- a/tests/test_micrograph_source.py
+++ b/tests/test_micrograph_source.py
@@ -56,7 +56,7 @@ def image_data_fixture(micrograph_count, micrograph_size, dtype):
     This generates a Numpy array with prescribed shape and dtype.
     """
     img_np = np.random.rand(micrograph_count, micrograph_size, micrograph_size)
-    return img_np.astype(dtype, copy=False)
+    return img_np.astype(dtype, copy=None)
 
 
 # =====
diff --git a/tests/test_orient_sdp.py b/tests/test_orient_sdp.py
index 22658ee06a..aa215d1766 100644
--- a/tests/test_orient_sdp.py
+++ b/tests/test_orient_sdp.py
@@ -152,7 +152,7 @@ def test_ATA_solver():
     rots = Rotation.generate_random_rotations(n=n_rots, seed=seed, dtype=dtype).matrices
 
     # Create a simple reference linear transformation A that is rank-3.
-    A_ref = np.diag([1, 2, 3]).astype(dtype, copy=False)
+    A_ref = np.diag([1, 2, 3]).astype(dtype, copy=None)
 
     # Create v1 and v2 such that A_ref*v1=R1 and A_ref*v2=R2, R1 and R2 are the first
     # and second columns of all rotations.
diff --git a/tests/test_polar_ft.py b/tests/test_polar_ft.py
index 425d5e14bb..4ec69b67e4 100644
--- a/tests/test_polar_ft.py
+++ b/tests/test_polar_ft.py
@@ -203,7 +203,7 @@ def test_half_to_full_transform(stack_shape):
     """
     img_size = 32
     image = Image(
-        np.random.rand(*stack_shape, img_size, img_size).astype(np.float32, copy=False)
+        np.random.rand(*stack_shape, img_size, img_size).astype(np.float32, copy=None)
     )
     pft = PolarFT(size=img_size)
     pf = pft.transform(image)
diff --git a/tests/test_simulation.py b/tests/test_simulation.py
index 944e2e7c06..46b00d2ce9 100644
--- a/tests/test_simulation.py
+++ b/tests/test_simulation.py
@@ -82,7 +82,7 @@ def testPassthroughFromVol(self):
         without an explcit Simulation dtype.
         """
         for dtype in (np.float32, np.float64):
-            sim = Simulation(vols=self.vol.astype(dtype, copy=False))
+            sim = Simulation(vols=self.vol.astype(dtype, copy=None))
             # Did we assign the right type?
             self.assertTrue(sim.dtype == dtype)
 
diff --git a/tests/test_volume.py b/tests/test_volume.py
index d220860784..74c77c5cc5 100644
--- a/tests/test_volume.py
+++ b/tests/test_volume.py
@@ -163,10 +163,10 @@ def test_astype(vols_1, dtype):
 
 def test_astype_copy(vols_1):
     """
-    `astype(copy=False)` is an optimization partially mimicked from numpy.
+    `astype(copy=None)` is an optimization partially mimicked from numpy.
     """
-    # Same dtype, copy=False
-    v2 = vols_1.astype(vols_1.dtype, copy=False)
+    # Same dtype, copy=None
+    v2 = vols_1.astype(vols_1.dtype, copy=None)
     # Details should match,
     assert isinstance(v2, Volume)
     assert np.allclose(v2.asnumpy(), vols_1.asnumpy())

From eecb8d18ff18718a3c746fba63bb9019530e15a5 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Wed, 9 Oct 2024 14:19:48 -0400
Subject: [PATCH 005/184] np2 breaks np.sign for complex

---
 src/aspire/utils/matrix.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/aspire/utils/matrix.py b/src/aspire/utils/matrix.py
index bc802f87f8..534c027b30 100644
--- a/src/aspire/utils/matrix.py
+++ b/src/aspire/utils/matrix.py
@@ -493,7 +493,7 @@ def fix_signs(u):
     # Now we only care about the sign +1/-1.
     #  The following corrects for any numerical division noise,
     #  and also remaps 0 to +1.
-    signs = np.sign(signs * 2 + 1)
+    signs = np.sign(signs.real * 2 + 1)
 
     # Apply signs elementwise to matrix
     return u * signs

From 6ecda629590da5721ad8473559684b08d3673c4c Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Wed, 9 Oct 2024 15:06:13 -0400
Subject: [PATCH 006/184] np2 multiply xform dtype cast

---
 src/aspire/image/xform.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/aspire/image/xform.py b/src/aspire/image/xform.py
index 0759c819e8..38f65d9bee 100644
--- a/src/aspire/image/xform.py
+++ b/src/aspire/image/xform.py
@@ -144,8 +144,9 @@ def __init__(self, factor):
         self.multipliers = np.array(factor)
 
     def _forward(self, im, indices):
+
         if self.multipliers.size == 1:  # if we have a scalar multiplier
-            im_new = im * self.multipliers
+            im_new = im * self.multipliers.astype(im.dtype)
         else:
             im_new = im * self.multipliers[indices]
 

From f23d36e905b903afa6255a8f27d15dd943aba35f Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Wed, 9 Oct 2024 15:11:03 -0400
Subject: [PATCH 007/184] sk pca copy False

---
 src/aspire/classification/rir_class2d.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/aspire/classification/rir_class2d.py b/src/aspire/classification/rir_class2d.py
index dde35ef050..19a3479ffe 100644
--- a/src/aspire/classification/rir_class2d.py
+++ b/src/aspire/classification/rir_class2d.py
@@ -390,7 +390,7 @@ def _sk_pca(self, M):
         #   We use an extension of SK that is hacked to admit complex.
         pca = ComplexPCA(
             self.bispectrum_components,
-            copy=None,  # careful, overwrites data matrix... we'll handle the copies.
+            copy=False,  # careful, overwrites data matrix... we'll handle the copies.
             svd_solver="auto",  # use randomized (Halko) for larger problems
             random_state=self.seed,
         )

From ebaeb04d06bd460488970cf3d173efdd285417cb Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Wed, 9 Oct 2024 15:11:51 -0400
Subject: [PATCH 008/184] sk pca copy False

---
 tests/test_complexPCA.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/test_complexPCA.py b/tests/test_complexPCA.py
index 1c224324d7..509b0b84a5 100644
--- a/tests/test_complexPCA.py
+++ b/tests/test_complexPCA.py
@@ -53,7 +53,7 @@ def testLargeFitTransform(self):
         """
 
         pca = ComplexPCA(
-            n_components=self.components_large, copy=None, svd_solver="full"
+            n_components=self.components_large, copy=False, svd_solver="full"
         )
 
         # Input data matrix X should be (n_samples, m_features)

From d9ccdefdb1189a390eee58968bf75ade4d5963b6 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Wed, 9 Oct 2024 15:28:01 -0400
Subject: [PATCH 009/184] another great np2 improvement

---
 src/aspire/image/image.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/aspire/image/image.py b/src/aspire/image/image.py
index 6d0f3b8960..bd9dfa22a1 100644
--- a/src/aspire/image/image.py
+++ b/src/aspire/image/image.py
@@ -72,7 +72,7 @@ def normalize_bg(imgs, bg_radius=1.0, do_ramp=True):
 
     # Apply mask images and calculate mean and std values of background
     imgs_masked = imgs * mask
-    denominator = np.sum(mask)
+    denominator = np.count_nonzero(mask)  # scalar int
     first_moment = np.sum(imgs_masked, axis=(1, 2)) / denominator
     second_moment = np.sum(imgs_masked**2, axis=(1, 2)) / denominator
     mean = first_moment.reshape(-1, 1, 1)

From e975604a921059d9f45b68a08a871675f0bfe67c Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Wed, 9 Oct 2024 15:39:22 -0400
Subject: [PATCH 010/184] sqrt(weight) dtype not same as weight

---
 src/aspire/__init__.py                       | 15 +++++++++++++++
 src/aspire/abinitio/commonline_base.py       |  2 +-
 src/aspire/abinitio/commonline_c3_c4.py      |  2 +-
 src/aspire/abinitio/commonline_d2.py         |  4 ++--
 src/aspire/abinitio/commonline_sync3n.py     | 20 ++++++++++----------
 src/aspire/apple/picking.py                  |  4 ++--
 src/aspire/basis/basis.py                    |  2 +-
 src/aspire/basis/fle_2d.py                   |  6 ++++--
 src/aspire/basis/fpswf_2d.py                 |  2 +-
 src/aspire/classification/averager2d.py      |  2 +-
 src/aspire/covariance/covar2d.py             |  2 +-
 src/aspire/ctf/ctf_estimator.py              |  4 ++--
 src/aspire/image/image.py                    |  6 +++---
 src/aspire/nufft/cufinufft.py                |  2 +-
 src/aspire/nufft/finufft.py                  |  4 ++--
 src/aspire/numeric/complex_pca/validation.py | 10 +++++-----
 src/aspire/numeric/pyfftw_fft.py             |  2 +-
 src/aspire/operators/blk_diag_matrix.py      |  4 ++--
 src/aspire/operators/diag_matrix.py          |  2 +-
 src/aspire/reconstruction/estimator.py       |  1 -
 src/aspire/reconstruction/kernel.py          |  4 ++--
 src/aspire/reconstruction/mean.py            | 19 ++++++++++---------
 src/aspire/sinogram/sinogram.py              |  2 +-
 src/aspire/source/micrograph.py              |  2 +-
 src/aspire/utils/bot_align.py                |  8 ++++----
 src/aspire/utils/matrix.py                   |  4 ++--
 src/aspire/utils/misc.py                     |  8 ++++----
 src/aspire/utils/rotation.py                 |  6 +++---
 src/aspire/volume/volume.py                  |  2 +-
 tests/test_FFBbasis3D.py                     |  4 ++--
 tests/test_coef.py                           |  8 ++++----
 tests/test_diag_matrix.py                    |  2 +-
 tests/test_downsample.py                     |  4 ++--
 tests/test_matrix.py                         |  4 ++--
 tests/test_mean_estimator.py                 |  4 ++--
 tests/test_micrograph_source.py              |  2 +-
 tests/test_orient_sdp.py                     |  2 +-
 tests/test_polar_ft.py                       |  2 +-
 tests/test_simulation.py                     |  2 +-
 tests/test_volume.py                         |  6 +++---
 40 files changed, 104 insertions(+), 87 deletions(-)

diff --git a/src/aspire/__init__.py b/src/aspire/__init__.py
index 9c80750c9c..8273099c75 100644
--- a/src/aspire/__init__.py
+++ b/src/aspire/__init__.py
@@ -2,11 +2,14 @@
 import logging.config
 import os
 import pkgutil
+import warnings
 from datetime import datetime
 from pathlib import Path
 
 import confuse
+import numpy as np
 import pooch
+from packaging.version import parse as parse_version
 
 import aspire
 from aspire.exceptions import handle_exception
@@ -84,3 +87,15 @@ def __getattr__(attr):
         return importlib.import_module(f"aspire.{attr}")
     else:
         raise AttributeError(f"module `{__name__}` has no attribute `{attr}`.")
+
+
+if parse_version(np.version.version) >= parse_version("2.0.0"):
+    # ImportWarnings are generally filtered, but raise this one for now.
+    with warnings.catch_warnings():
+        warnings.simplefilter("default")
+        warnings.warn(
+            "ASPIRE's Numpy 2 support is in beta.  If you experience a runtime"
+            "crash relating to mismatched dtypes or a Numpy call please try"
+            ' `pip install "numpy<2"` and report to ASPIRE developers.',
+            ImportWarning,
+        )
diff --git a/src/aspire/abinitio/commonline_base.py b/src/aspire/abinitio/commonline_base.py
index 2cd80b61fa..fe456c645e 100644
--- a/src/aspire/abinitio/commonline_base.py
+++ b/src/aspire/abinitio/commonline_base.py
@@ -149,7 +149,7 @@ def _prepare_pf(self):
             #   Apply mask in doubles (allow imgs to upcast as needed)
             imgs = imgs * fuzz_mask
             #   Cast to desired type
-            imgs = Image(imgs.asnumpy().astype(self.dtype, copy=None))
+            imgs = Image(imgs.asnumpy().astype(self.dtype, copy=False))
 
         # Obtain coefficients of polar Fourier transform for input 2D images
         pft = PolarFT(
diff --git a/src/aspire/abinitio/commonline_c3_c4.py b/src/aspire/abinitio/commonline_c3_c4.py
index bc2956ec0d..242094b745 100644
--- a/src/aspire/abinitio/commonline_c3_c4.py
+++ b/src/aspire/abinitio/commonline_c3_c4.py
@@ -693,7 +693,7 @@ def _J_sync_power_method(self, vijs):
         while itr < max_iters and residual > epsilon:
             itr += 1
             # Note, this appears to need double precision for accuracy in the following division.
-            vec_new = self._signs_times_v(vijs, vec).astype(np.float64, copy=None)
+            vec_new = self._signs_times_v(vijs, vec).astype(np.float64, copy=False)
             vec_new = vec_new / norm(vec_new)
             residual = norm(vec_new - vec)
             vec = vec_new
diff --git a/src/aspire/abinitio/commonline_d2.py b/src/aspire/abinitio/commonline_d2.py
index 93b2e9bc0a..a8e951c642 100644
--- a/src/aspire/abinitio/commonline_d2.py
+++ b/src/aspire/abinitio/commonline_d2.py
@@ -1149,12 +1149,12 @@ def _sync_colors(self, Rijs):
 
         # Seed eigs initial vector for iterative method.
         # scipy LinearOperator needs doubles for some architectures (arm).
-        v0 = randn(3 * n_pairs, seed=self.seed).astype(np.float64, copy=None)
+        v0 = randn(3 * n_pairs, seed=self.seed).astype(np.float64, copy=False)
 
         v0 = v0 / norm(v0)
         vals, colors = la.eigs(color_mat, k=3, which="LR", v0=v0)
         vals = np.real(vals)
-        colors = np.real(colors).astype(self.dtype, copy=None)
+        colors = np.real(colors).astype(self.dtype, copy=False)
         colors = np.sign(colors[0]) * colors  # Stable eigs
         cp, _ = self._unmix_colors(colors[:, :2])
 
diff --git a/src/aspire/abinitio/commonline_sync3n.py b/src/aspire/abinitio/commonline_sync3n.py
index aa11a2348a..a0db2263ec 100644
--- a/src/aspire/abinitio/commonline_sync3n.py
+++ b/src/aspire/abinitio/commonline_sync3n.py
@@ -197,7 +197,7 @@ def _sync3n_S_to_rot(self, S, W=None, n_eigs=4):
         """
 
         # Critical this occurs in double precision
-        S = S.astype(np.float64, copy=None)
+        S = S.astype(np.float64, copy=False)
 
         if n_eigs < 3:
             raise ValueError(
@@ -213,7 +213,7 @@ def _sync3n_S_to_rot(self, S, W=None, n_eigs=4):
                 )
             # Initialize D
             # Critical this occurs in double precision
-            W = W.astype(np.float64, copy=None)
+            W = W.astype(np.float64, copy=False)
             D = np.mean(W, axis=1)
 
             Dhalf = D
@@ -266,7 +266,7 @@ def _sync3n_S_to_rot(self, S, W=None, n_eigs=4):
         # Enforce we are returning actual rotations
         rotations = nearest_rotations(rotations, allow_reflection=True)
 
-        return rotations.astype(self.dtype)
+        return rotations.astype(self.dtype, copy=False)
 
     def _construct_sync3n_matrix(self, Rij):
         """
@@ -664,8 +664,8 @@ def _pairs_probabilities_cupy(self, Rijs, P2, A, a, B, b, x0):
         )
 
         # accumulate over thread results
-        ln_f_arb = ln_f_arb_dev.get().astype(self.dtype, copy=None)
-        ln_f_ind = ln_f_ind_dev.get().astype(self.dtype, copy=None)
+        ln_f_arb = ln_f_arb_dev.get().astype(self.dtype, copy=False)
+        ln_f_ind = ln_f_ind_dev.get().astype(self.dtype, copy=False)
 
         return ln_f_ind, ln_f_arb
 
@@ -748,8 +748,8 @@ def fun(x, B, P, b, x0, A=A, a=a):
 
         popt, pcov = curve_fit(
             fun,
-            hist_x.astype(np.float64, copy=None),
-            scores_hist.astype(np.float64, copy=None),
+            hist_x.astype(np.float64, copy=False),
+            scores_hist.astype(np.float64, copy=False),
             p0=start_values,
             bounds=(lower_bounds, upper_bounds),
             method="trf",  # MATLAB used method "LAR" with algo "Trust-Region"
@@ -1043,7 +1043,7 @@ def _J_sync_power_method(self, Rijs):
         while itr < max_iters and residual > epsilon:
             itr += 1
             # Todo, this code code actually needs double precision for accuracy... forcing.
-            vec_new = self._signs_times_v(Rijs, vec).astype(np.float64, copy=None)
+            vec_new = self._signs_times_v(Rijs, vec).astype(np.float64, copy=False)
             vec_new = vec_new / norm(vec_new)
             residual = norm(vec_new - vec)
             vec = vec_new
@@ -1073,7 +1073,7 @@ def _signs_times_v(self, Rijs, vec):
         else:
             new_vec = self._signs_times_v_host(Rijs, vec)
 
-        return new_vec.astype(vec.dtype, copy=None)
+        return new_vec.astype(vec.dtype, copy=False)
 
     def _signs_times_v_host(self, Rijs, vec):
         """
@@ -1177,7 +1177,7 @@ def _signs_times_v_cupy(self, Rijs, vec):
         )
 
         # dtoh
-        new_vec = new_vec_dev.get().astype(vec.dtype, copy=None)
+        new_vec = new_vec_dev.get().astype(vec.dtype, copy=False)
 
         return new_vec
 
diff --git a/src/aspire/apple/picking.py b/src/aspire/apple/picking.py
index d54cf61f1d..d64d16c773 100644
--- a/src/aspire/apple/picking.py
+++ b/src/aspire/apple/picking.py
@@ -149,7 +149,7 @@ def read_micrograph(self):
         self.original_im = (
             DiskMicrographSource(self.filename)
             .asnumpy()[0]
-            .astype(np.float32, copy=None)
+            .astype(np.float32, copy=False)
         )
 
         # Discard outer pixels
@@ -167,7 +167,7 @@ def read_micrograph(self):
         # Note, float64 required for signal.correlate call accuracy.
         im = np.asarray(
             PILImage.fromarray(im).resize(size, PILImage.Resampling.BICUBIC)
-        ).astype(np.float64, copy=None)
+        ).astype(np.float64, copy=False)
 
         im = signal.correlate(
             im,
diff --git a/src/aspire/basis/basis.py b/src/aspire/basis/basis.py
index 6baf16f769..ee6aee2d02 100644
--- a/src/aspire/basis/basis.py
+++ b/src/aspire/basis/basis.py
@@ -62,7 +62,7 @@ def __init__(self, basis, data, dtype=None):
             )
         self.basis = basis
 
-        self._data = data.astype(self.dtype, copy=None)
+        self._data = data.astype(self.dtype, copy=False)
         self.ndim = self._data.ndim
         self.shape = self._data.shape
         self.stack_ndim = self._data.ndim - 1
diff --git a/src/aspire/basis/fle_2d.py b/src/aspire/basis/fle_2d.py
index 6499919110..fc22352c18 100644
--- a/src/aspire/basis/fle_2d.py
+++ b/src/aspire/basis/fle_2d.py
@@ -643,7 +643,7 @@ def _step1(self, z):
         num_img = z.shape[0]
         z = z[:, :, : self.num_angular_nodes // 2].reshape(num_img, -1)
         im = anufft(
-            z.astype(complex_type(self.dtype), copy=None),
+            z.astype(complex_type(self.dtype), copy=False),
             self.grid_xy,
             (self.nres, self.nres),
             epsilon=self.epsilon,
@@ -798,7 +798,9 @@ def filter_to_basis_mat(self, f, **kwargs):
         omega = 2 * xp.pi * xp.vstack((omegax.flatten("C"), omegay.flatten("C")))
 
         h_vals2d = (
-            xp.asarray(h_fun(omega)).reshape(n_k, n_theta).astype(self.dtype, copy=None)
+            xp.asarray(h_fun(omega))
+            .reshape(n_k, n_theta)
+            .astype(self.dtype, copy=False)
         )
         h_vals = xp.sum(h_vals2d, axis=1) / n_theta
 
diff --git a/src/aspire/basis/fpswf_2d.py b/src/aspire/basis/fpswf_2d.py
index aa7b9913d9..6eac1a91db 100644
--- a/src/aspire/basis/fpswf_2d.py
+++ b/src/aspire/basis/fpswf_2d.py
@@ -99,7 +99,7 @@ def _precomp(self):
         ) = self._pswf_integration_sub_routine()
 
         self.us_fft_pts = us_fft_pts.astype(
-            self.dtype, copy=None
+            self.dtype, copy=False
         )  # TODO, debug where this is incorrect dtype
         self.blk_r = blk_r
         self.num_angular_pts = num_angular_pts
diff --git a/src/aspire/classification/averager2d.py b/src/aspire/classification/averager2d.py
index 644359c9e7..b447a6d91c 100644
--- a/src/aspire/classification/averager2d.py
+++ b/src/aspire/classification/averager2d.py
@@ -82,7 +82,7 @@ def _cls_images(self, cls, src=None):
             source for a certain operation (ie alignment).
         """
         src = src or self.src
-        return src.images[cls].asnumpy().astype(self.dtype, copy=None)
+        return src.images[cls].asnumpy().astype(self.dtype, copy=False)
 
 
 class AligningAverager2D(Averager2D):
diff --git a/src/aspire/covariance/covar2d.py b/src/aspire/covariance/covar2d.py
index a4f1971b78..d2e7dafe43 100644
--- a/src/aspire/covariance/covar2d.py
+++ b/src/aspire/covariance/covar2d.py
@@ -607,7 +607,7 @@ def _calc_op(self):
             ctf_basis_k_sq = ctf_basis_k_t @ ctf_basis_k
             A_mean_k = weight * ctf_basis_k_sq
             A_mean += A_mean_k
-            A_covar_k = np.sqrt(weight) * ctf_basis_k_sq
+            A_covar_k = np.sqrt(weight).astype(self.dtype) * ctf_basis_k_sq
             A_covar[k] = A_covar_k
 
             M_covar += A_covar_k
diff --git a/src/aspire/ctf/ctf_estimator.py b/src/aspire/ctf/ctf_estimator.py
index 39f2998a21..b6983cb071 100644
--- a/src/aspire/ctf/ctf_estimator.py
+++ b/src/aspire/ctf/ctf_estimator.py
@@ -241,7 +241,7 @@ def estimate_psd(self, blocks, tapers_1d):
         """
 
         num_1d_tapers = tapers_1d.shape[-1]
-        tapers_1d = tapers_1d.astype(complex_type(self.dtype), copy=None)
+        tapers_1d = tapers_1d.astype(complex_type(self.dtype), copy=False)
 
         blocks_mt = np.zeros(blocks[0, :, :].shape, dtype=self.dtype)
 
@@ -768,7 +768,7 @@ def estimate_ctf(
             micrograph = mrc.data
 
         # Try to match dtype used in Basis instance
-        micrograph = micrograph.astype(dtype, copy=None)
+        micrograph = micrograph.astype(dtype, copy=False)
 
         micrograph_blocks = ctf_object.preprocess_micrograph(micrograph, psd_size)
 
diff --git a/src/aspire/image/image.py b/src/aspire/image/image.py
index bd9dfa22a1..1e980e88dd 100644
--- a/src/aspire/image/image.py
+++ b/src/aspire/image/image.py
@@ -195,7 +195,7 @@ def __init__(self, data, pixel_size=None, dtype=None):
         if not data.shape[-1] == data.shape[-2]:
             raise ValueError("Only square ndarrays are supported.")
 
-        self._data = data.astype(self.dtype, copy=None)
+        self._data = data.astype(self.dtype, copy=False)
         self.ndim = self._data.ndim
         self.shape = self._data.shape
         self.stack_ndim = self._data.ndim - 2
@@ -543,7 +543,7 @@ def load(filepath, dtype=None):
 
         # Attempt casting when user provides dtype
         if dtype is not None:
-            im = im.astype(dtype, copy=None)
+            im = im.astype(dtype, copy=False)
 
         # Return as Image instance
         return Image(im, pixel_size=pixel_size)
@@ -658,7 +658,7 @@ def backproject(self, rot_matrices, symmetry_group=None, zero_nyquist=True):
 
             # TODO: rotated_grids might as well give us correctly shaped array in the first place
             pts_rot = aspire.volume.rotated_grids(L, rotations).astype(
-                self.dtype, copy=None
+                self.dtype, copy=False
             )
             pts_rot = pts_rot.reshape((3, -1))
 
diff --git a/src/aspire/nufft/cufinufft.py b/src/aspire/nufft/cufinufft.py
index 2660f4ea46..fd869aacfd 100644
--- a/src/aspire/nufft/cufinufft.py
+++ b/src/aspire/nufft/cufinufft.py
@@ -30,7 +30,7 @@ def __init__(self, sz, fourier_pts, epsilon=1e-8, ntransforms=1, **kwargs):
         # ASPIRE-Python/703
         # Cast to doubles.
         self._original_dtype = fourier_pts.dtype
-        fourier_pts = fourier_pts.astype(np.float64, copy=None)
+        fourier_pts = fourier_pts.astype(np.float64, copy=False)
 
         # Basic dtype passthough.
         dtype = fourier_pts.dtype
diff --git a/src/aspire/nufft/finufft.py b/src/aspire/nufft/finufft.py
index e396d76124..a955ab20d0 100644
--- a/src/aspire/nufft/finufft.py
+++ b/src/aspire/nufft/finufft.py
@@ -99,7 +99,7 @@ def transform(self, signal):
         ), f"Signal frame to be transformed must have shape {self.sz}"
 
         # FINUFFT was designed for a complex input array
-        signal = np.array(signal, copy=None, dtype=self.complex_dtype, order="C")
+        signal = np.asarray(signal, dtype=self.complex_dtype, order="C")
 
         result = self._transform_plan.execute(signal)
 
@@ -130,7 +130,7 @@ def adjoint(self, signal):
                 signal = signal.reshape(self.num_pts)
 
         # FINUFFT was designed for a complex input array
-        signal = np.array(signal, copy=None, dtype=self.complex_dtype, order="C")
+        signal = np.asarray(signal, dtype=self.complex_dtype, order="C")
 
         result = self._adjoint_plan.execute(signal)
 
diff --git a/src/aspire/numeric/complex_pca/validation.py b/src/aspire/numeric/complex_pca/validation.py
index 11d7aaba9f..7fafedfe18 100644
--- a/src/aspire/numeric/complex_pca/validation.py
+++ b/src/aspire/numeric/complex_pca/validation.py
@@ -81,7 +81,7 @@ def _ensure_sparse_format(
         Data type of result. If None, the dtype of the input is preserved.
 
     copy : boolean
-        Whether a forced copy will be triggered. If copy=None, a copy might
+        Whether a forced copy will be triggered. If copy=False, a copy might
         be triggered by a conversion.
 
     force_all_finite : boolean or 'allow-nan', (default=True)
@@ -173,7 +173,7 @@ def check_array(
     accept_large_sparse=True,
     dtype="numeric",
     order=None,
-    copy=None,
+    copy=False,
     force_all_finite=True,
     ensure_2d=True,
     allow_nd=False,
@@ -216,13 +216,13 @@ def check_array(
 
     order : 'F', 'C' or None (default=None)
         Whether an array will be forced to be fortran or c-style.
-        When order is None (default), then if copy=None, nothing is ensured
+        When order is None (default), then if copy=False, nothing is ensured
         about the memory layout of the output array; otherwise (copy=True)
         the memory layout of the returned array is kept as close as possible
         to the original array.
 
     copy : boolean (default=False)
-        Whether a forced copy will be triggered. If copy=None, a copy might
+        Whether a forced copy will be triggered. If copy=False, a copy might
         be triggered by a conversion.
 
     force_all_finite : boolean or 'allow-nan', (default=True)
@@ -365,7 +365,7 @@ def check_array(
                     array = np.asarray(array, order=order)
                     if array.dtype.kind == "f":
                         _assert_all_finite(array, allow_nan=False, msg_dtype=dtype)
-                    array = array.astype(dtype, casting="unsafe", copy=None)
+                    array = array.astype(dtype, casting="unsafe", copy=False)
                 else:
                     array = np.asarray(array, order=order, dtype=dtype)
             except ComplexWarning:
diff --git a/src/aspire/numeric/pyfftw_fft.py b/src/aspire/numeric/pyfftw_fft.py
index 4503b42b38..95a8ea80f7 100644
--- a/src/aspire/numeric/pyfftw_fft.py
+++ b/src/aspire/numeric/pyfftw_fft.py
@@ -139,7 +139,7 @@ def ifftn(self, a, axes=None, workers=-1):
         # FFTW_BACKWARD requires complex input array, cast as needed.
         # See https://pyfftw.readthedocs.io/en/latest/source/pyfftw/pyfftw.html#scheme-table
         comp_type = complex_type(a.dtype)
-        a = a.astype(comp_type, copy=None)
+        a = a.astype(comp_type, copy=False)
 
         mutex.acquire()
         try:
diff --git a/src/aspire/operators/blk_diag_matrix.py b/src/aspire/operators/blk_diag_matrix.py
index af882ecc68..3e493f9525 100644
--- a/src/aspire/operators/blk_diag_matrix.py
+++ b/src/aspire/operators/blk_diag_matrix.py
@@ -91,7 +91,7 @@ def append(self, blk):
         :param blk: Block to append (ndarray).
         """
 
-        self.data.append(blk.astype(self.dtype, copy=None))
+        self.data.append(blk.astype(self.dtype, copy=False))
         self.nblocks += 1
         self.reset_cache()
 
@@ -130,7 +130,7 @@ def __setitem__(self, key, value):
         Convenience wrapper, setter on self.data.
         """
 
-        self.data[key] = value.astype(self.dtype, copy=None)
+        self.data[key] = value.astype(self.dtype, copy=False)
         self.reset_cache()
 
     def __len__(self):
diff --git a/src/aspire/operators/diag_matrix.py b/src/aspire/operators/diag_matrix.py
index 68d6a728ae..4c94ab83d1 100644
--- a/src/aspire/operators/diag_matrix.py
+++ b/src/aspire/operators/diag_matrix.py
@@ -45,7 +45,7 @@ def __init__(self, data, dtype=None):
         self.dtype = np.dtype(dtype)
 
         # Assign the `data`
-        self._data = data.astype(self.dtype, copy=None)
+        self._data = data.astype(self.dtype, copy=False)
 
         # Assign shapes from `data`
         self.count = self._data.shape[-1]
diff --git a/src/aspire/reconstruction/estimator.py b/src/aspire/reconstruction/estimator.py
index 2bf0ec3366..ac5e9caeca 100644
--- a/src/aspire/reconstruction/estimator.py
+++ b/src/aspire/reconstruction/estimator.py
@@ -159,7 +159,6 @@ def apply_kernel(self, vol_coef, kernel=None):
 
         if kernel is None:
             kernel = self.kernel
-
         vol = Coef(self.basis, vol_coef).evaluate()  # returns a Volume
         vol = kernel.convolve_volume(vol)  # returns a Volume
         vol_coef = self.basis.evaluate_t(vol)
diff --git a/src/aspire/reconstruction/kernel.py b/src/aspire/reconstruction/kernel.py
index 84d4e61409..9a86a244f2 100644
--- a/src/aspire/reconstruction/kernel.py
+++ b/src/aspire/reconstruction/kernel.py
@@ -32,7 +32,7 @@ def __add__(self, delta):
             to be able to use it within optimization loops. This operator allows one to use the FourierKernel object
             with the underlying 'kernel' attribute tweaked with a regularization parameter.
         """
-        new_kernel = self.kernel + delta
+        new_kernel = self.kernel + float(delta)
         return FourierKernel(new_kernel)
 
     def circularize(self):
@@ -191,7 +191,7 @@ def __add__(self, delta):
 
     def circularize(self):
         _L = self.M // 2
-        xx = np.empty((self.r, self.r, _L, _L, _L))
+        xx = np.empty((self.r, self.r, _L, _L, _L), self.dtype)
         for k in range(self.r):
             for j in range(self.r):
                 xx[k, j] = FourierKernel(self.kermat[k, j]).circularize().real
diff --git a/src/aspire/reconstruction/mean.py b/src/aspire/reconstruction/mean.py
index 56f0624e8d..b380c8788c 100644
--- a/src/aspire/reconstruction/mean.py
+++ b/src/aspire/reconstruction/mean.py
@@ -5,10 +5,9 @@
 from scipy.linalg import norm
 from scipy.sparse.linalg import LinearOperator
 
-from aspire import config
 from aspire.basis import Coef
 from aspire.nufft import anufft
-from aspire.numeric import fft
+from aspire.numeric import config, fft
 from aspire.numeric.scipy import cg
 from aspire.operators import evaluate_src_filters_on_grid
 from aspire.reconstruction import Estimator, FourierKernel, FourierKernelMatrix
@@ -102,7 +101,7 @@ def _compute_kernel(self):
             _range = np.arange(i, min(self.src.n, i + self.batch_size), dtype=int)
             sq_filters_f = evaluate_src_filters_on_grid(self.src, _range) ** 2
             amplitudes_sq = (self.src.amplitudes[_range] ** 2).astype(
-                self.dtype, copy=None
+                self.dtype, copy=False
             )
 
             for k in range(self.r):
@@ -139,7 +138,7 @@ def _compute_kernel(self):
                     if j != k:
                         kernel[j, k] += batch_kernel
 
-        kermat_f = np.zeros((self.r, self.r, _2L, _2L, _2L))
+        kermat_f = np.zeros((self.r, self.r, _2L, _2L, _2L), dtype=self.dtype)
         logger.info("Computing non-centered Fourier Transform Kernel Mat")
         for k in range(self.r):
             for j in range(self.r):
@@ -184,10 +183,12 @@ def src_backward(self):
                     i,
                     self.weights[:, k],
                     symmetry_group=symmetry_group,
-                ) / (self.src.n * sym_order)
+                ) / float(self.src.n * sym_order)
                 vol_rhs[k] += batch_vol_rhs.astype(self.dtype)
 
-        res = np.sqrt(self.src.n * sym_order) * self.basis.evaluate_t(vol_rhs)
+        res = np.sqrt(self.src.n * sym_order, dtype=self.dtype) * self.basis.evaluate_t(
+            vol_rhs
+        )
         logger.info(f"Determined weighted adjoint mappings. Shape = {res.shape}")
 
         return res
@@ -210,6 +211,7 @@ def conj_grad(self, b_coef, x0=None, tol=1e-5, regularizer=0):
             precond_kernel = self.precond_kernel
             if regularizer > 0:
                 precond_kernel += regularizer
+
             M = LinearOperator(
                 (self.r * count, self.r * count),
                 matvec=partial(self.apply_kernel, kernel=precond_kernel),
@@ -262,7 +264,7 @@ def cb(xk):
                 f"Conjugate gradient unable to converge after {info} iterations."
             )
 
-        return x.reshape(self.r, self.basis.count)
+        return x.reshape(self.r, self.basis.count).astype(self.dtype, copy=False)
 
     def apply_kernel(self, vol_coef, kernel=None):
         """
@@ -283,8 +285,7 @@ def apply_kernel(self, vol_coef, kernel=None):
         vols_out = Volume(
             np.zeros((self.r, self.src.L, self.src.L, self.src.L), dtype=self.dtype)
         )
-
-        vol = Coef(self.basis, vol_coef).evaluate()
+        vol = Coef(self.basis, vol_coef.astype(self.dtype, copy=False)).evaluate()
 
         for k in range(self.r):
             for j in range(self.r):
diff --git a/src/aspire/sinogram/sinogram.py b/src/aspire/sinogram/sinogram.py
index 0f82b0c479..7c7bb43662 100644
--- a/src/aspire/sinogram/sinogram.py
+++ b/src/aspire/sinogram/sinogram.py
@@ -37,7 +37,7 @@ def __init__(self, data, dtype=None):
                 f"Invalid data shape: {data.shape}. Expected shape: (..., angles, radial_points), where '...' is the stack number."
             )
 
-        self._data = data.astype(self.dtype, copy=None)
+        self._data = data.astype(self.dtype, copy=False)
         self.ndim = self._data.ndim
         self.shape = self._data.shape
         self.stack_shape = self._data.shape[:-2]
diff --git a/src/aspire/source/micrograph.py b/src/aspire/source/micrograph.py
index 86b161d822..5aa35e389c 100644
--- a/src/aspire/source/micrograph.py
+++ b/src/aspire/source/micrograph.py
@@ -155,7 +155,7 @@ def __init__(self, micrographs, dtype=None, pixel_size=None):
         )
 
         # We're already backed by an array, access it directly.
-        self._data = micrographs.astype(self.dtype, copy=None)
+        self._data = micrographs.astype(self.dtype, copy=False)
 
     def _images(self, indices):
         """
diff --git a/src/aspire/utils/bot_align.py b/src/aspire/utils/bot_align.py
index e94888d92e..fc8fbe4fdc 100644
--- a/src/aspire/utils/bot_align.py
+++ b/src/aspire/utils/bot_align.py
@@ -139,7 +139,7 @@ def cost(new, t=t, q=q):
             Loss function for surrogate problems.
             """
             kx = np.array(
-                [cov_fun(new.astype(dtype, copy=None), R[j]) for j in range(t)]
+                [cov_fun(new.astype(dtype, copy=False), R[j]) for j in range(t)]
             )
             mu = kx @ q
             return mu
@@ -151,7 +151,7 @@ def euclidean_grad(new, t=t, q=q):
             """
             # Corresponds to equation 11 in the paper.
             kx_grad = np.array(
-                [cov_fun_grad(new.astype(dtype, copy=None), R[j]) for j in range(t)]
+                [cov_fun_grad(new.astype(dtype, copy=False), R[j]) for j in range(t)]
             )
 
             kx_grad = kx_grad.reshape((t, 9))
@@ -168,7 +168,7 @@ def euclidean_grad(new, t=t, q=q):
             verbosity=verbosity,
         )
         result = optimizer.run(problem)
-        R_new = result.point.astype(dtype, copy=None)
+        R_new = result.point.astype(dtype, copy=False)
 
         loss[t] = loss_fun(R_new)
         R[t] = R_new
@@ -190,7 +190,7 @@ def loss_u(u):
             """
             Convert the alignment loss function over SO(3) to be over R^3.
             """
-            u = u.astype(dtype, copy=None)
+            u = u.astype(dtype, copy=False)
             R = sign * Rotation.from_rotvec(u).matrices[0]
             v_rot = vol_given_ds.rotate(Rotation(R)).asnumpy()[0]
             return norm(vol_ref_ds - v_rot)
diff --git a/src/aspire/utils/matrix.py b/src/aspire/utils/matrix.py
index 534c027b30..f33aa47b23 100644
--- a/src/aspire/utils/matrix.py
+++ b/src/aspire/utils/matrix.py
@@ -1,5 +1,5 @@
 """
-Utilties for arrays/n-dimensional matrices.
+Utilities for arrays/n-dimensional matrices.
 """
 
 import numpy as np
@@ -462,7 +462,7 @@ def nearest_rotations(A, allow_reflection=False):
         # If det(U)*det(V) = -1, we negate the third singular value to
         # ensure we have a rotation.
         neg_det_idx = np.linalg.det(U) * np.linalg.det(V) < 0
-        U[neg_det_idx] = U[neg_det_idx] @ np.diag((1, 1, -1)).astype(dtype, copy=None)
+        U[neg_det_idx] = U[neg_det_idx] @ np.diag((1, 1, -1)).astype(dtype, copy=False)
 
     rots = U @ V
 
diff --git a/src/aspire/utils/misc.py b/src/aspire/utils/misc.py
index 087454de5e..5485c42788 100644
--- a/src/aspire/utils/misc.py
+++ b/src/aspire/utils/misc.py
@@ -1,5 +1,5 @@
 """
-Miscellaneous Utilities that have no better place (yet).
+Miscellaneous utilities that have no better place (yet).
 """
 
 import hashlib
@@ -135,7 +135,7 @@ def gaussian_1d(size, mu=0, sigma=1, dtype=np.float64):
 
     p = (g["x"] - mu) ** 2 / (2 * sigma**2)
 
-    return np.exp(-p).astype(dtype, copy=None)
+    return np.exp(-p).astype(dtype, copy=False)
 
 
 def gaussian_2d(size, mu=(0, 0), sigma=(1, 1), indexing="yx", dtype=np.float64):
@@ -180,7 +180,7 @@ def gaussian_2d(size, mu=(0, 0), sigma=(1, 1), indexing="yx", dtype=np.float64):
         2 * sigma[1] ** 2
     )
 
-    return np.exp(-p).astype(dtype, copy=None)
+    return np.exp(-p).astype(dtype, copy=False)
 
 
 def gaussian_3d(size, mu=(0, 0, 0), sigma=(1, 1, 1), indexing="zyx", dtype=np.float64):
@@ -228,7 +228,7 @@ def gaussian_3d(size, mu=(0, 0, 0), sigma=(1, 1, 1), indexing="zyx", dtype=np.fl
         + (g["z"] - mu[2]) ** 2 / (2 * sigma[2] ** 2)
     )
 
-    return np.exp(-p).astype(dtype, copy=None)
+    return np.exp(-p).astype(dtype, copy=False)
 
 
 def bump_3d(size, spread=1, dtype=np.float64):
diff --git a/src/aspire/utils/rotation.py b/src/aspire/utils/rotation.py
index fb643914bf..07a31df9df 100644
--- a/src/aspire/utils/rotation.py
+++ b/src/aspire/utils/rotation.py
@@ -256,7 +256,7 @@ def from_euler(values, dtype=None):
         """
         dtype = dtype or getattr(values, "dtype", np.float32)
         rotations = sp_rot.from_euler("ZYZ", values, degrees=False)
-        matrices = rotations.as_matrix().astype(dtype, copy=None)
+        matrices = rotations.as_matrix().astype(dtype, copy=False)
         return Rotation(matrices)
 
     @staticmethod
@@ -324,7 +324,7 @@ def from_rotvec(vec, dtype=None):
         """
         dtype = dtype or vec.dtype
         rots = sp_rot.from_rotvec(vec)
-        matrices = rots.as_matrix().astype(dtype, copy=None)
+        matrices = rots.as_matrix().astype(dtype, copy=False)
         return Rotation(matrices)
 
     @staticmethod
@@ -337,7 +337,7 @@ def from_matrix(values, dtype=None):
         :return: new Rotation object
         """
         dtype = dtype or values.dtype
-        return Rotation(values.astype(dtype, copy=None))
+        return Rotation(values.astype(dtype, copy=False))
 
     @staticmethod
     def generate_random_rotations(
diff --git a/src/aspire/volume/volume.py b/src/aspire/volume/volume.py
index 1954a1cbd5..eae985a794 100644
--- a/src/aspire/volume/volume.py
+++ b/src/aspire/volume/volume.py
@@ -105,7 +105,7 @@ def __init__(self, data, dtype=None, pixel_size=None, symmetry_group=None):
         if not (data.shape[-1] == data.shape[-2] == data.shape[-3]):
             raise ValueError("Only cubed ndarrays are supported.")
 
-        self._data = data.astype(self.dtype, copy=None)
+        self._data = data.astype(self.dtype, copy=False)
         self.ndim = self._data.ndim
         self.shape = self._data.shape
         self.stack_ndim = self._data.ndim - 3
diff --git a/tests/test_FFBbasis3D.py b/tests/test_FFBbasis3D.py
index 3992f5d374..03bdfdd21b 100644
--- a/tests/test_FFBbasis3D.py
+++ b/tests/test_FFBbasis3D.py
@@ -470,7 +470,7 @@ def testFFBBasis3DEvaluate(self, basis):
 
     def testFFBBasis3DEvaluate_t(self, basis):
         x = np.load(os.path.join(DATA_DIR, "ffbbasis3d_xcoef_in_8_8_8.npy")).T  # RCOPT
-        x = x.astype(basis.dtype, copy=None)
+        x = x.astype(basis.dtype, copy=False)
         result = basis.evaluate_t(Volume(x))
 
         ref = np.load(os.path.join(DATA_DIR, "ffbbasis3d_vcoef_out_8_8_8.npy"))[..., 0]
@@ -479,7 +479,7 @@ def testFFBBasis3DEvaluate_t(self, basis):
 
     def testFFBBasis3DExpand(self, basis):
         x = np.load(os.path.join(DATA_DIR, "ffbbasis3d_xcoef_in_8_8_8.npy")).T  # RCOPT
-        x = x.astype(basis.dtype, copy=None)
+        x = x.astype(basis.dtype, copy=False)
         result = basis.expand(x)
 
         ref = np.load(os.path.join(DATA_DIR, "ffbbasis3d_vcoef_out_exp_8_8_8.npy"))[
diff --git a/tests/test_coef.py b/tests/test_coef.py
index 2a283d16e4..3ace0ddec5 100644
--- a/tests/test_coef.py
+++ b/tests/test_coef.py
@@ -90,7 +90,7 @@ def coef_fixture(basis, stack, dtype):
     # shape.
     size = stack + (basis.count,)
 
-    coef_np = np.random.random(size=size).astype(dtype, copy=None)
+    coef_np = np.random.random(size=size).astype(dtype, copy=False)
 
     return Coef(basis, coef_np, dtype=dtype)
 
@@ -212,7 +212,7 @@ def test_add(basis, coef_fixture):
     Tests addition operation against pure Numpy.
     """
     # Make array
-    x = np.random.random(size=coef_fixture.shape).astype(coef_fixture.dtype, copy=None)
+    x = np.random.random(size=coef_fixture.shape).astype(coef_fixture.dtype, copy=False)
     # Construct Coef
     c = Coef(basis, x)
 
@@ -231,7 +231,7 @@ def test_sub(basis, coef_fixture):
     Tests subtraction operation against pure Numpy.
     """
     # Make array
-    x = np.random.random(size=coef_fixture.shape).astype(coef_fixture.dtype, copy=None)
+    x = np.random.random(size=coef_fixture.shape).astype(coef_fixture.dtype, copy=False)
     # Construct Coef
     c = Coef(basis, x)
 
@@ -264,7 +264,7 @@ def test_mul(basis, coef_fixture):
     Tests multiplication operation against pure Numpy.
     """
     # Make array
-    x = np.random.random(size=coef_fixture.shape).astype(coef_fixture.dtype, copy=None)
+    x = np.random.random(size=coef_fixture.shape).astype(coef_fixture.dtype, copy=False)
     # Construct Coef
     c = Coef(basis, x)
 
diff --git a/tests/test_diag_matrix.py b/tests/test_diag_matrix.py
index 30f2b76faa..05805912c9 100644
--- a/tests/test_diag_matrix.py
+++ b/tests/test_diag_matrix.py
@@ -65,7 +65,7 @@ def diag_matrix_fixture(stack, matrix_size, dtype):
     """
     shape = (2,) + stack + (matrix_size,)
     # Internally convert dtype.  Passthrough will be checked explicitly in `test_dtype_passthrough` and `test_dtype_cast`
-    d_np = np.random.random(shape).astype(dtype, copy=None)
+    d_np = np.random.random(shape).astype(dtype, copy=False)
     d1 = DiagMatrix(d_np[0])
     d2 = DiagMatrix(d_np[1])
 
diff --git a/tests/test_downsample.py b/tests/test_downsample.py
index 1af37d7281..6c1cac82dd 100644
--- a/tests/test_downsample.py
+++ b/tests/test_downsample.py
@@ -148,7 +148,7 @@ def emdb_vol():
 
 @pytest.fixture(scope="module")
 def volume(emdb_vol, res, dtype):
-    vol = emdb_vol.astype(dtype, copy=None)
+    vol = emdb_vol.astype(dtype, copy=False)
     vol = vol.downsample(res)
     return vol
 
@@ -266,7 +266,7 @@ def test_pixel_size():
     dsL = 5  # downsampled
 
     # Construct a small test Image
-    img = Image(np.random.random((1, L, L)).astype(DTYPE, copy=None), pixel_size=1.23)
+    img = Image(np.random.random((1, L, L)).astype(DTYPE, copy=False), pixel_size=1.23)
 
     # Downsample the image
     result = img.downsample(dsL)
diff --git a/tests/test_matrix.py b/tests/test_matrix.py
index be6829d3fa..728200b9b3 100644
--- a/tests/test_matrix.py
+++ b/tests/test_matrix.py
@@ -356,7 +356,7 @@ def test_nearest_rotations(dtype):
     rots = Rotation.generate_random_rotations(n_rots, seed=0, dtype=dtype).matrices
 
     # Add some noise to the rotations.
-    noise = 1e-3 * randn(n_rots * 9, seed=0).astype(dtype, copy=None).reshape(
+    noise = 1e-3 * randn(n_rots * 9, seed=0).astype(dtype, copy=False).reshape(
         n_rots, 3, 3
     )
     noisy_rots = rots + noise
@@ -381,7 +381,7 @@ def test_nearest_rotations_reflection(dtype):
 
     # Add a reflection and some noise to the rotation.
     refl = rot @ np.diag((1, -1, 1)).astype(dtype)
-    noise = 1e-3 * randn(9, seed=0).astype(dtype, copy=None).reshape(3, 3)
+    noise = 1e-3 * randn(9, seed=0).astype(dtype, copy=False).reshape(3, 3)
     noisy_refl = refl + noise
 
     # Find nearest rotation.
diff --git a/tests/test_mean_estimator.py b/tests/test_mean_estimator.py
index 3e7f915193..e6b2a2f837 100644
--- a/tests/test_mean_estimator.py
+++ b/tests/test_mean_estimator.py
@@ -115,8 +115,8 @@ def test_adjoint(sim, basis, estimator):
     L = sim.L
     n = sim.n
 
-    u = np.random.rand(n, L, L).astype(sim.dtype, copy=None)
-    v = np.random.rand(L, L, L).astype(sim.dtype, copy=None)
+    u = np.random.rand(n, L, L).astype(sim.dtype, copy=False)
+    v = np.random.rand(L, L, L).astype(sim.dtype, copy=False)
 
     proj = Volume(v).project(rots)
     backproj = Image(u).backproject(rots)
diff --git a/tests/test_micrograph_source.py b/tests/test_micrograph_source.py
index c44e703cbb..d4793cf61f 100644
--- a/tests/test_micrograph_source.py
+++ b/tests/test_micrograph_source.py
@@ -56,7 +56,7 @@ def image_data_fixture(micrograph_count, micrograph_size, dtype):
     This generates a Numpy array with prescribed shape and dtype.
     """
     img_np = np.random.rand(micrograph_count, micrograph_size, micrograph_size)
-    return img_np.astype(dtype, copy=None)
+    return img_np.astype(dtype, copy=False)
 
 
 # =====
diff --git a/tests/test_orient_sdp.py b/tests/test_orient_sdp.py
index aa215d1766..22658ee06a 100644
--- a/tests/test_orient_sdp.py
+++ b/tests/test_orient_sdp.py
@@ -152,7 +152,7 @@ def test_ATA_solver():
     rots = Rotation.generate_random_rotations(n=n_rots, seed=seed, dtype=dtype).matrices
 
     # Create a simple reference linear transformation A that is rank-3.
-    A_ref = np.diag([1, 2, 3]).astype(dtype, copy=None)
+    A_ref = np.diag([1, 2, 3]).astype(dtype, copy=False)
 
     # Create v1 and v2 such that A_ref*v1=R1 and A_ref*v2=R2, R1 and R2 are the first
     # and second columns of all rotations.
diff --git a/tests/test_polar_ft.py b/tests/test_polar_ft.py
index 4ec69b67e4..425d5e14bb 100644
--- a/tests/test_polar_ft.py
+++ b/tests/test_polar_ft.py
@@ -203,7 +203,7 @@ def test_half_to_full_transform(stack_shape):
     """
     img_size = 32
     image = Image(
-        np.random.rand(*stack_shape, img_size, img_size).astype(np.float32, copy=None)
+        np.random.rand(*stack_shape, img_size, img_size).astype(np.float32, copy=False)
     )
     pft = PolarFT(size=img_size)
     pf = pft.transform(image)
diff --git a/tests/test_simulation.py b/tests/test_simulation.py
index 46b00d2ce9..944e2e7c06 100644
--- a/tests/test_simulation.py
+++ b/tests/test_simulation.py
@@ -82,7 +82,7 @@ def testPassthroughFromVol(self):
         without an explcit Simulation dtype.
         """
         for dtype in (np.float32, np.float64):
-            sim = Simulation(vols=self.vol.astype(dtype, copy=None))
+            sim = Simulation(vols=self.vol.astype(dtype, copy=False))
             # Did we assign the right type?
             self.assertTrue(sim.dtype == dtype)
 
diff --git a/tests/test_volume.py b/tests/test_volume.py
index 74c77c5cc5..d220860784 100644
--- a/tests/test_volume.py
+++ b/tests/test_volume.py
@@ -163,10 +163,10 @@ def test_astype(vols_1, dtype):
 
 def test_astype_copy(vols_1):
     """
-    `astype(copy=None)` is an optimization partially mimicked from numpy.
+    `astype(copy=False)` is an optimization partially mimicked from numpy.
     """
-    # Same dtype, copy=None
-    v2 = vols_1.astype(vols_1.dtype, copy=None)
+    # Same dtype, copy=False
+    v2 = vols_1.astype(vols_1.dtype, copy=False)
     # Details should match,
     assert isinstance(v2, Volume)
     assert np.allclose(v2.asnumpy(), vols_1.asnumpy())

From 7fe6b87eea5e1325c061878dc8cb2de30086f8d5 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Wed, 9 Oct 2024 16:19:09 -0400
Subject: [PATCH 011/184] looks like np2 impacted mrcfile as well

---
 tests/test_mrc.py | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/tests/test_mrc.py b/tests/test_mrc.py
index 363db73a6c..723a45464c 100644
--- a/tests/test_mrc.py
+++ b/tests/test_mrc.py
@@ -81,15 +81,13 @@ def testUpdate(self):
             with mrcfile.new_mmap(
                 files[1], shape=(self.n, self.n), mrc_mode=2, overwrite=True
             ) as mrc:
-                mrc.set_data(self.a.astype(np.float32))
+                mrc.set_data(self.a)
+                self.stats.update_header(mrc)
                 mrc.header.time = epoch
                 mrc.header.label[0] = label
 
             # Our homebrew and mrcfile files should now match to the bit.
             comparison = sha256sum(files[0]) == sha256sum(files[1])
-            # Expected hash:
-            # 71355fa0bcd5b989ff88166962ea5d2b78ea032933bd6fda41fbdcc1c6d1a009
             logging.debug(f"sha256(file0): {sha256sum(files[0])}")
             logging.debug(f"sha256(file1): {sha256sum(files[1])}")
-
             self.assertTrue(comparison)

From 202f76f8434bfd4cfcb3a055730c22fc345dc8b1 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Mon, 21 Oct 2024 12:46:55 -0400
Subject: [PATCH 012/184] Tox wants explicit stack level

---
 src/aspire/__init__.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/aspire/__init__.py b/src/aspire/__init__.py
index 8273099c75..a5145dcd5f 100644
--- a/src/aspire/__init__.py
+++ b/src/aspire/__init__.py
@@ -98,4 +98,5 @@ def __getattr__(attr):
             "crash relating to mismatched dtypes or a Numpy call please try"
             ' `pip install "numpy<2"` and report to ASPIRE developers.',
             ImportWarning,
+            stacklevel=1,
         )

From 7bcfd1c9dc6a6c8042dda1c76ebc01631b1b505f Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Tue, 22 Oct 2024 08:29:17 -0400
Subject: [PATCH 013/184] missing space in warning message

---
 src/aspire/__init__.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/aspire/__init__.py b/src/aspire/__init__.py
index a5145dcd5f..686280e8e9 100644
--- a/src/aspire/__init__.py
+++ b/src/aspire/__init__.py
@@ -95,7 +95,7 @@ def __getattr__(attr):
         warnings.simplefilter("default")
         warnings.warn(
             "ASPIRE's Numpy 2 support is in beta.  If you experience a runtime"
-            "crash relating to mismatched dtypes or a Numpy call please try"
+            " crash relating to mismatched dtypes or a Numpy call please try"
             ' `pip install "numpy<2"` and report to ASPIRE developers.',
             ImportWarning,
             stacklevel=1,

From f36f51bf6d31a947c099d8cc927c549c35610bbd Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Tue, 22 Oct 2024 08:31:25 -0400
Subject: [PATCH 014/184] misc cleanup, reducing pr delta from draft work

---
 src/aspire/image/xform.py              | 1 -
 src/aspire/reconstruction/estimator.py | 1 +
 src/aspire/reconstruction/mean.py      | 4 ++--
 3 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/aspire/image/xform.py b/src/aspire/image/xform.py
index 38f65d9bee..8003b4a1c5 100644
--- a/src/aspire/image/xform.py
+++ b/src/aspire/image/xform.py
@@ -144,7 +144,6 @@ def __init__(self, factor):
         self.multipliers = np.array(factor)
 
     def _forward(self, im, indices):
-
         if self.multipliers.size == 1:  # if we have a scalar multiplier
             im_new = im * self.multipliers.astype(im.dtype)
         else:
diff --git a/src/aspire/reconstruction/estimator.py b/src/aspire/reconstruction/estimator.py
index ac5e9caeca..2bf0ec3366 100644
--- a/src/aspire/reconstruction/estimator.py
+++ b/src/aspire/reconstruction/estimator.py
@@ -159,6 +159,7 @@ def apply_kernel(self, vol_coef, kernel=None):
 
         if kernel is None:
             kernel = self.kernel
+
         vol = Coef(self.basis, vol_coef).evaluate()  # returns a Volume
         vol = kernel.convolve_volume(vol)  # returns a Volume
         vol_coef = self.basis.evaluate_t(vol)
diff --git a/src/aspire/reconstruction/mean.py b/src/aspire/reconstruction/mean.py
index b380c8788c..e3888ee200 100644
--- a/src/aspire/reconstruction/mean.py
+++ b/src/aspire/reconstruction/mean.py
@@ -5,9 +5,10 @@
 from scipy.linalg import norm
 from scipy.sparse.linalg import LinearOperator
 
+from aspire import config
 from aspire.basis import Coef
 from aspire.nufft import anufft
-from aspire.numeric import config, fft
+from aspire.numeric import fft
 from aspire.numeric.scipy import cg
 from aspire.operators import evaluate_src_filters_on_grid
 from aspire.reconstruction import Estimator, FourierKernel, FourierKernelMatrix
@@ -211,7 +212,6 @@ def conj_grad(self, b_coef, x0=None, tol=1e-5, regularizer=0):
             precond_kernel = self.precond_kernel
             if regularizer > 0:
                 precond_kernel += regularizer
-
             M = LinearOperator(
                 (self.r * count, self.r * count),
                 matvec=partial(self.apply_kernel, kernel=precond_kernel),

From d24bca98ae2e48686a2926f3a628c082af1cf51f Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Tue, 19 Nov 2024 09:03:38 -0500
Subject: [PATCH 015/184] Remove sklearn workaround

---
 src/aspire/numeric/complex_pca/complex_pca.py | 17 ++---------------
 1 file changed, 2 insertions(+), 15 deletions(-)

diff --git a/src/aspire/numeric/complex_pca/complex_pca.py b/src/aspire/numeric/complex_pca/complex_pca.py
index bd270ef64b..e38820c702 100644
--- a/src/aspire/numeric/complex_pca/complex_pca.py
+++ b/src/aspire/numeric/complex_pca/complex_pca.py
@@ -13,8 +13,6 @@
 
 import numpy as np
 import scipy.sparse as sp
-import sklearn
-from packaging.version import Version
 from sklearn.decomposition import PCA
 from sklearn.utils._array_api import get_namespace
 
@@ -71,22 +69,11 @@ def _fit(self, X):
             else:
                 self._fit_svd_solver = "full"
 
-        # sci-kit changed `_fit_*()` API in latest release v1.5.0
-        # which supports Python 3.9 - 3.12. This can be removed after
-        # our minimal support is Python 3.9.
-        API_dep = Version(sklearn.__version__) < Version("1.5.0")
-
         # Call different fits for either full or truncated SVD
         if self._fit_svd_solver == "full":
-            if API_dep:
-                return self._fit_full(X, n_components)
-            else:
-                return self._fit_full(X, n_components, xp, is_array_api_compliant)
+            return self._fit_full(X, n_components, xp, is_array_api_compliant)
         elif self._fit_svd_solver in ["arpack", "randomized"]:
-            if API_dep:
-                return self._fit_truncated(X, n_components, self._fit_svd_solver)
-            else:
-                return self._fit_truncated(X, n_components, xp)
+            return self._fit_truncated(X, n_components, xp)
         else:
             raise ValueError(
                 "Unrecognized svd_solver='{0}'" "".format(self._fit_svd_solver)

From 71b300260424bd190a34cbc0ecaefa010b28023b Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Tue, 19 Nov 2024 10:10:21 -0500
Subject: [PATCH 016/184] Min version for sklearn.

---
 pyproject.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pyproject.toml b/pyproject.toml
index c4eabef52f..9db57652e0 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -44,7 +44,7 @@ dependencies = [
     "pymanopt",
     "PyWavelets",
     "scipy >= 1.10.0",
-    "scikit-learn",
+    "scikit-learn >= 1.5.0",
     "scikit-image",
     "setuptools >= 0.41",
     "tqdm",

From 74ce51b8691edd90e14ae407a7c9cb128eb8c7ab Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Mon, 18 Nov 2024 13:55:01 -0500
Subject: [PATCH 017/184] Convert gallery logging to print;

Sphinx Gallery logging is broken.
---
 .../experimental_abinitio_pipeline_10028.py   | 17 +++----
 .../experimental_abinitio_pipeline_10073.py   | 13 ++---
 .../experimental_abinitio_pipeline_10081.py   | 15 +++---
 gallery/tutorials/tutorials/apple_picker.py   |  7 +--
 .../tutorials/tutorials/class_averaging.py    | 15 +++---
 .../tutorials/tutorials/cov2d_simulation.py   | 29 +++++------
 .../tutorials/tutorials/cov3d_simulation.py   | 23 ++++-----
 .../generating_volume_projections.py          |  2 -
 .../tutorials/tutorials/image_expansion.py    | 45 ++++++++---------
 .../tutorials/orient3d_simulation.py          | 19 +++----
 .../tutorials/preprocess_imgs_sim.py          | 49 +++++++++----------
 11 files changed, 101 insertions(+), 133 deletions(-)

diff --git a/gallery/experiments/experimental_abinitio_pipeline_10028.py b/gallery/experiments/experimental_abinitio_pipeline_10028.py
index 7892302758..7cdfbe9d89 100644
--- a/gallery/experiments/experimental_abinitio_pipeline_10028.py
+++ b/gallery/experiments/experimental_abinitio_pipeline_10028.py
@@ -22,7 +22,6 @@
 # In addition, import some classes from
 # the ASPIRE package that will be used throughout this experiment.
 
-import logging
 from pathlib import Path
 
 import matplotlib.pyplot as plt
@@ -34,8 +33,6 @@
 from aspire.reconstruction import MeanEstimator
 from aspire.source import OrientedSource, RelionSource
 
-logger = logging.getLogger(__name__)
-
 
 # %%
 # Parameters
@@ -68,7 +65,7 @@
 )
 
 # Downsample the images
-logger.info(f"Set the resolution to {img_size} X {img_size}")
+print(f"Set the resolution to {img_size} X {img_size}")
 src = src.downsample(img_size)
 
 # Peek
@@ -76,7 +73,7 @@
     src.images[:10].show()
 
 # Use phase_flip to attempt correcting for CTF.
-logger.info("Perform phase flip to input images.")
+print("Perform phase flip to input images.")
 src = src.phase_flip()
 
 # Estimate the noise and `Whiten` based on the estimated noise
@@ -94,7 +91,7 @@
 
 # # Optionally invert image contrast, depends on data convention.
 # # This is not needed for 10028, but included anyway.
-# logger.info("Invert the global density contrast")
+# print("Invert the global density contrast")
 # src = src.invert_contrast()
 
 # Caching is used for speeding up large datasets on high memory machines.
@@ -132,7 +129,7 @@
 #
 # Now perform classification and averaging for each class.
 
-logger.info("Begin Class Averaging")
+print("Begin Class Averaging")
 
 # Now perform classification and averaging for each class.
 # This also demonstrates the potential to use a different source for classification and averaging.
@@ -159,7 +156,7 @@
 # Next create a CL instance for estimating orientation of projections
 # using the Common Line with Synchronization Voting method.
 
-logger.info("Begin Orientation Estimation")
+print("Begin Orientation Estimation")
 
 # Create a custom orientation estimation object for ``avgs``.
 # This is done to customize the ``n_theta`` value.
@@ -175,7 +172,7 @@
 #
 # Using the oriented source, attempt to reconstruct a volume.
 
-logger.info("Begin Volume reconstruction")
+print("Begin Volume reconstruction")
 
 # Setup an estimator to perform the back projection.
 estimator = MeanEstimator(oriented_src)
@@ -183,7 +180,7 @@
 # Perform the estimation and save the volume.
 estimated_volume = estimator.estimate()
 estimated_volume.save(volume_output_filename, overwrite=True)
-logger.info(f"Saved Volume to {str(Path(volume_output_filename).resolve())}")
+print(f"Saved Volume to {str(Path(volume_output_filename).resolve())}")
 
 # Peek at result
 if interactive:
diff --git a/gallery/experiments/experimental_abinitio_pipeline_10073.py b/gallery/experiments/experimental_abinitio_pipeline_10073.py
index d936a98d65..6e4ca8cfcb 100644
--- a/gallery/experiments/experimental_abinitio_pipeline_10073.py
+++ b/gallery/experiments/experimental_abinitio_pipeline_10073.py
@@ -25,7 +25,6 @@
 # In addition, import some classes from
 # the ASPIRE package that will be used throughout this experiment.
 
-import logging
 from pathlib import Path
 
 import numpy as np
@@ -42,8 +41,6 @@
 from aspire.reconstruction import MeanEstimator
 from aspire.source import OrientedSource, RelionSource
 
-logger = logging.getLogger(__name__)
-
 
 # %%
 # Parameters
@@ -74,7 +71,7 @@
 )
 
 # Downsample the images
-logger.info(f"Set the resolution to {img_size} X {img_size}")
+print(f"Set the resolution to {img_size} X {img_size}")
 src = src.downsample(img_size)
 
 src = src.cache()
@@ -85,7 +82,7 @@
 #
 # Now perform classification and averaging for each class.
 
-logger.info("Begin Class Averaging")
+print("Begin Class Averaging")
 
 # Now perform classification and averaging for each class.
 # This also demonstrates customizing a ClassAvgSource, by using global
@@ -132,7 +129,7 @@
 # Next create a CL instance for estimating orientation of projections
 # using the Common Line with Synchronization Voting method.
 
-logger.info("Begin Orientation Estimation")
+print("Begin Orientation Estimation")
 
 # Create a custom orientation estimation object for ``avgs``.
 orient_est = CLSync3N(avgs, n_theta=72)
@@ -147,7 +144,7 @@
 #
 # Using the oriented source, attempt to reconstruct a volume.
 
-logger.info("Begin Volume reconstruction")
+print("Begin Volume reconstruction")
 
 # Setup an estimator to perform the back projection.
 estimator = MeanEstimator(oriented_src)
@@ -155,4 +152,4 @@
 # Perform the estimation and save the volume.
 estimated_volume = estimator.estimate()
 estimated_volume.save(volume_output_filename, overwrite=True)
-logger.info(f"Saved Volume to {str(Path(volume_output_filename).resolve())}")
+print(f"Saved Volume to {str(Path(volume_output_filename).resolve())}")
diff --git a/gallery/experiments/experimental_abinitio_pipeline_10081.py b/gallery/experiments/experimental_abinitio_pipeline_10081.py
index 7d83842c24..7b753b6f2b 100644
--- a/gallery/experiments/experimental_abinitio_pipeline_10081.py
+++ b/gallery/experiments/experimental_abinitio_pipeline_10081.py
@@ -22,7 +22,6 @@
 # In addition, import some classes from
 # the ASPIRE package that will be used throughout this experiment.
 
-import logging
 from pathlib import Path
 
 from aspire.abinitio import CLSymmetryC3C4
@@ -31,8 +30,6 @@
 from aspire.reconstruction import MeanEstimator
 from aspire.source import OrientedSource, RelionSource
 
-logger = logging.getLogger(__name__)
-
 
 # %%
 # Parameters
@@ -67,11 +64,11 @@
 )
 
 # Downsample the images
-logger.info(f"Set the resolution to {img_size} X {img_size}")
+print(f"Set the resolution to {img_size} X {img_size}")
 src = src.downsample(img_size)
 
 # Use phase_flip to attempt correcting for CTF.
-logger.info("Perform phase flip to input images.")
+print("Perform phase flip to input images.")
 src = src.phase_flip()
 
 # Estimate the noise and `Whiten` based on the estimated noise
@@ -87,7 +84,7 @@
 #
 # Now perform classification and averaging for each class.
 
-logger.info("Begin Class Averaging")
+print("Begin Class Averaging")
 
 # Now perform classification and averaging for each class.
 # Automatically configure parallel processing
@@ -106,7 +103,7 @@
 # Next create a CL instance for estimating orientation of projections
 # using the Common Line with Synchronization Voting method.
 
-logger.info("Begin Orientation Estimation")
+print("Begin Orientation Estimation")
 
 # Create a custom orientation estimation object for ``avgs``.
 orient_est = CLSymmetryC3C4(avgs, symmetry="C4", n_theta=72, max_shift=0)
@@ -127,7 +124,7 @@
 # back-projection.  This boosts the effective number of images used in
 # the reconstruction from ``n_classes`` to ``4*n_classes``.
 
-logger.info("Begin Volume reconstruction")
+print("Begin Volume reconstruction")
 
 
 # Setup an estimator to perform the back projection.
@@ -136,4 +133,4 @@
 # Perform the estimation and save the volume.
 estimated_volume = estimator.estimate()
 estimated_volume.save(volume_output_filename, overwrite=True)
-logger.info(f"Saved Volume to {str(Path(volume_output_filename).resolve())}")
+print(f"Saved Volume to {str(Path(volume_output_filename).resolve())}")
diff --git a/gallery/tutorials/tutorials/apple_picker.py b/gallery/tutorials/tutorials/apple_picker.py
index 08ce582a65..3c3891dd2a 100644
--- a/gallery/tutorials/tutorials/apple_picker.py
+++ b/gallery/tutorials/tutorials/apple_picker.py
@@ -5,7 +5,6 @@
 We demonstrate ASPIRE's particle picking methods using the ``Apple`` class.
 """
 
-import logging
 import os
 
 import matplotlib.pyplot as plt
@@ -13,8 +12,6 @@
 
 from aspire.apple.apple import Apple
 
-logger = logging.getLogger(__name__)
-
 # %%
 # Read and Plot Micrograph
 # ------------------------
@@ -63,8 +60,8 @@
 
 img_dim = micro_img.shape
 particles = centers.shape[0]
-logger.info(f"Dimensions of the micrograph are {img_dim}")
-logger.info(f"{particles} particles were picked")
+print(f"Dimensions of the micrograph are {img_dim}")
+print(f"{particles} particles were picked")
 
 # sphinx_gallery_thumbnail_number = 2
 plt.imshow(particles_img, cmap="gray")
diff --git a/gallery/tutorials/tutorials/class_averaging.py b/gallery/tutorials/tutorials/class_averaging.py
index 0e4768f9be..407f86e263 100644
--- a/gallery/tutorials/tutorials/class_averaging.py
+++ b/gallery/tutorials/tutorials/class_averaging.py
@@ -6,8 +6,6 @@
 representation algorithm.
 """
 
-import logging
-
 import matplotlib.pyplot as plt
 import numpy as np
 from PIL import Image as PILImage
@@ -17,7 +15,6 @@
 from aspire.source import ArrayImageSource  # Helpful hint if you want to BYO array.
 from aspire.utils import gaussian_2d
 
-logger = logging.getLogger(__name__)
 
 # %%
 # Build Simulated Data
@@ -197,8 +194,8 @@
 # ``noise_src``.
 classes = avgs.class_indices[review_class]
 reflections = avgs.class_refl[review_class]
-logger.info(f"Class {review_class}'s neighors: {classes}")
-logger.info(f"Class {review_class}'s reflections: {reflections}")
+print(f"Class {review_class}'s neighors: {classes}")
+print(f"Class {review_class}'s reflections: {reflections}")
 
 # The original image is the initial image in the class array.
 original_image_idx = classes[0]
@@ -230,9 +227,9 @@
 est_shifts = avgs.averager.shifts
 est_correlations = avgs.averager.correlations
 
-logger.info(f"Estimated Rotations: {est_rotations}")
-logger.info(f"Estimated Shifts: {est_shifts}")
-logger.info(f"Estimated Correlations: {est_correlations}")
+print(f"Estimated Rotations: {est_rotations}")
+print(f"Estimated Shifts: {est_shifts}")
+print(f"Estimated Correlations: {est_correlations}")
 
 # Compare the original unaligned images with the estimated alignment.
 # Get the indices from the classification results.
@@ -247,7 +244,7 @@
 # Rotate using estimated rotations.
 angle = est_rotations[0, nbr] * 180 / np.pi
 if reflections[nbr]:
-    logger.info("Reflection reported.")
+    print("Reflection reported.")
     original_img_nbr = np.flipud(original_img_nbr)
 rotated_img_nbr = np.asarray(PILImage.fromarray(original_img_nbr).rotate(angle))
 
diff --git a/gallery/tutorials/tutorials/cov2d_simulation.py b/gallery/tutorials/tutorials/cov2d_simulation.py
index e0ed8328e0..7eff7675fa 100644
--- a/gallery/tutorials/tutorials/cov2d_simulation.py
+++ b/gallery/tutorials/tutorials/cov2d_simulation.py
@@ -8,7 +8,6 @@
 that covariance matrix.
 """
 
-import logging
 import os
 
 import matplotlib.pyplot as plt
@@ -22,14 +21,12 @@
 from aspire.utils import anorm
 from aspire.volume import Volume
 
-logger = logging.getLogger(__name__)
-
 file_path = os.path.join(
     os.path.dirname(os.getcwd()), "data", "clean70SRibosome_vol_65p.mrc"
 )
 
 
-logger.info(
+print(
     "This script illustrates 2D covariance Wiener filtering functionality in ASPIRE package."
 )
 
@@ -44,7 +41,7 @@
 num_imgs = 1024
 # Set dtype for this experiment
 dtype = np.float32
-logger.info(f"Simulation running in {dtype} precision.")
+print(f"Simulation running in {dtype} precision.")
 
 
 # %%
@@ -73,7 +70,7 @@
 # Initialize Simulation Object and CTF Filters
 # --------------------------------------------
 
-logger.info("Initialize simulation object and CTF filters.")
+print("Initialize simulation object and CTF filters.")
 # Create filters
 ctf_filters = [
     RadialCTFFilter(pixel_size, voltage, defocus=d, Cs=2.0, alpha=0.1)
@@ -81,7 +78,7 @@
 ]
 
 # Load the map file of a 70S Ribosome
-logger.info(
+print(
     f"Load 3D map and downsample 3D map to desired grids "
     f"of {img_size} x {img_size} x {img_size}."
 )
@@ -92,7 +89,7 @@
 vols = vols.downsample(img_size)
 
 # Create a simulation object with specified filters and the downsampled 3D map
-logger.info("Use downsampled map to creat simulation object.")
+print("Use downsampled map to creat simulation object.")
 sim = Simulation(
     L=img_size,
     n=num_imgs,
@@ -119,12 +116,12 @@
 h_ctf_fb = [ffbbasis.filter_to_basis_mat(filt) for filt in ctf_filters]
 
 # Get clean images from projections of 3D map.
-logger.info("Apply CTF filters to clean images.")
+print("Apply CTF filters to clean images.")
 imgs_clean = sim.projections[:]
 imgs_ctf_clean = sim.clean_images[:]
 power_clean = imgs_ctf_clean.norm() ** 2 / imgs_ctf_clean.size
 sn_ratio = power_clean / noise_var
-logger.info(f"Signal to noise ratio is {sn_ratio}.")
+print(f"Signal to noise ratio is {sn_ratio}.")
 
 # get noisy images after applying CTF and noise filters
 imgs_noise = sim.images[:num_imgs]
@@ -139,7 +136,7 @@
 # expansion by applying the adjoint of the evaluation mapping using
 # ``basis.evaluate_t``.
 
-logger.info("Get coefficients of clean and noisy images in FFB basis.")
+print("Get coefficients of clean and noisy images in FFB basis.")
 coef_clean = ffbbasis.evaluate_t(imgs_clean)
 coef_noise = ffbbasis.evaluate_t(imgs_noise)
 
@@ -157,7 +154,7 @@
 # mean and covariance estimates will allow us to evaluate the mean and
 # covariance estimates from the filtered, noisy data, later.
 
-logger.info(
+print(
     "Get 2D covariance matrices of clean and noisy images using FB coefficients."
 )
 cov2d = RotCov2D(ffbbasis)
@@ -202,7 +199,7 @@
 # covariance, and the variance of the noise. The resulting estimator has
 # the lowest expected mean square error out of all linear estimators.
 
-logger.info("Get the CWF coefficients of noising images.")
+print("Get the CWF coefficients of noising images.")
 coef_est = cov2d.get_cwf_coefs(
     coef_noise,
     h_ctf_fb,
@@ -230,9 +227,9 @@
 
 # Calculate the normalized RMSE of the estimated images.
 nrmse_ims = (imgs_est - imgs_clean).norm() / imgs_clean.norm()
-logger.info(f"Deviation of the noisy mean estimate: {diff_mean}")
-logger.info(f"Deviation of the noisy covariance estimate: {diff_covar}")
-logger.info(f"Estimated images normalized RMSE: {nrmse_ims}")
+print(f"Deviation of the noisy mean estimate: {diff_mean}")
+print(f"Deviation of the noisy covariance estimate: {diff_covar}")
+print(f"Estimated images normalized RMSE: {nrmse_ims}")
 
 # plot the first images at different stages
 idm = 0
diff --git a/gallery/tutorials/tutorials/cov3d_simulation.py b/gallery/tutorials/tutorials/cov3d_simulation.py
index 741a47de99..86d26116ea 100644
--- a/gallery/tutorials/tutorials/cov3d_simulation.py
+++ b/gallery/tutorials/tutorials/cov3d_simulation.py
@@ -6,8 +6,6 @@
 generated from Gaussian blob volumes.
 """
 
-import logging
-
 import numpy as np
 from scipy.cluster.vq import kmeans2
 
@@ -22,7 +20,6 @@
 from aspire.utils.random import Random
 from aspire.volume import LegacyVolume, Volume
 
-logger = logging.getLogger(__name__)
 
 # %%
 # Create Simulation Object
@@ -62,7 +59,7 @@
 # Estimate the noise variance. This is needed for the covariance estimation step below.
 noise_estimator = WhiteNoiseEstimator(sim, batchSize=500)
 noise_variance = noise_estimator.estimate()
-logger.info(f"Noise Variance = {noise_variance}")
+print(f"Noise Variance = {noise_variance}")
 
 # %%
 # Estimate Mean Volume and Covariance
@@ -150,19 +147,19 @@
 
 # Output estimated covariance spectrum.
 
-logger.info(f"Population Covariance Spectrum = {np.diag(lambdas_est)}")
+print(f"Population Covariance Spectrum = {np.diag(lambdas_est)}")
 
 
 # Output performance results.
 
-logger.info(f'Mean (rel. error) = {mean_perf["rel_err"]}')
-logger.info(f'Mean (correlation) = {mean_perf["corr"]}')
-logger.info(f'Covariance (rel. error) = {covar_perf["rel_err"]}')
-logger.info(f'Covariance (correlation) = {covar_perf["corr"]}')
-logger.info(f'Eigendecomposition (rel. error) = {eigs_perf["rel_err"]}')
-logger.info(f"Clustering (accuracy) = {clustering_accuracy}")
-logger.info(f'Coordinates (mean rel. error) = {coords_perf["rel_err"]}')
-logger.info(f'Coordinates (mean correlation) = {np.mean(coords_perf["corr"])}')
+print(f'Mean (rel. error) = {mean_perf["rel_err"]}')
+print(f'Mean (correlation) = {mean_perf["corr"]}')
+print(f'Covariance (rel. error) = {covar_perf["rel_err"]}')
+print(f'Covariance (correlation) = {covar_perf["corr"]}')
+print(f'Eigendecomposition (rel. error) = {eigs_perf["rel_err"]}')
+print(f"Clustering (accuracy) = {clustering_accuracy}")
+print(f'Coordinates (mean rel. error) = {coords_perf["rel_err"]}')
+print(f'Coordinates (mean correlation) = {np.mean(coords_perf["corr"])}')
 
 # Basic Check
 assert covar_perf["rel_err"] <= 0.80
diff --git a/gallery/tutorials/tutorials/generating_volume_projections.py b/gallery/tutorials/tutorials/generating_volume_projections.py
index 48d82ce231..f9e342dab5 100644
--- a/gallery/tutorials/tutorials/generating_volume_projections.py
+++ b/gallery/tutorials/tutorials/generating_volume_projections.py
@@ -7,7 +7,6 @@
 
 """
 
-import logging
 import os
 
 import numpy as np
@@ -17,7 +16,6 @@
 from aspire.utils import Rotation
 from aspire.volume import Volume
 
-logger = logging.getLogger(__name__)
 
 # %%
 # Configure how many images we'd like to project
diff --git a/gallery/tutorials/tutorials/image_expansion.py b/gallery/tutorials/tutorials/image_expansion.py
index 4c3f36aac4..55cb6c5db4 100644
--- a/gallery/tutorials/tutorials/image_expansion.py
+++ b/gallery/tutorials/tutorials/image_expansion.py
@@ -6,7 +6,6 @@
 based on the basis functions of  Fourier-Bessel (FB) and prolate spheroidal wave function (PSWF).
 """
 
-import logging
 import os
 import timeit
 
@@ -16,10 +15,8 @@
 from aspire.basis import FBBasis2D, FFBBasis2D, FPSWFBasis2D, PSWFBasis2D
 from aspire.utils import anorm
 
-logger = logging.getLogger(__name__)
 
-
-logger.info(
+print(
     "This script illustrates different image expansion methods in ASPIRE package."
 )
 
@@ -49,26 +46,26 @@
 fb_basis = FBBasis2D((img_size, img_size), dtype=org_images.dtype)
 
 # Get the expansion coefficients based on FB basis
-logger.info("Start normal FB expansion of original images.")
+print("Start normal FB expansion of original images.")
 tstart = timeit.default_timer()
 fb_coefs = fb_basis.evaluate_t(org_images)
 tstop = timeit.default_timer()
 dtime = tstop - tstart
-logger.info(f"Finish normal FB expansion of original images in {dtime:.4f} seconds.")
+print(f"Finish normal FB expansion of original images in {dtime:.4f} seconds.")
 
 # Reconstruct images from the expansion coefficients based on FB basis
 fb_images = fb_basis.evaluate(fb_coefs).asnumpy()
-logger.info("Finish reconstruction of images from normal FB expansion coefficients.")
+print("Finish reconstruction of images from normal FB expansion coefficients.")
 
 # Calculate the mean value of maximum differences between the FB estimated images and the original images
 fb_meanmax = np.mean(np.max(abs(fb_images - org_images), axis=(1, 2)))
-logger.info(
+print(
     f"Mean value of maximum differences between FB estimated images and original images: {fb_meanmax}"
 )
 
 # Calculate the normalized RMSE of the FB estimated images
 fb_nrmse_ims = anorm(fb_images - org_images) / anorm(org_images)
-logger.info(f"FB estimated images normalized RMSE: {fb_nrmse_ims}")
+print(f"FB estimated images normalized RMSE: {fb_nrmse_ims}")
 
 # plot the first images using the normal FB method
 plt.subplot(1, 3, 1)
@@ -92,27 +89,27 @@
 ffb_basis = FFBBasis2D((img_size, img_size), dtype=org_images.dtype)
 
 # Get the expansion coefficients based on fast FB basis
-logger.info("start fast FB expansion of original images.")
+print("start fast FB expansion of original images.")
 tstart = timeit.default_timer()
 ffb_coefs = ffb_basis.evaluate_t(org_images)
 tstop = timeit.default_timer()
 dtime = tstop - tstart
-logger.info(f"Finish fast FB expansion of original images in {dtime:.4f} seconds.")
+print(f"Finish fast FB expansion of original images in {dtime:.4f} seconds.")
 
 # Reconstruct images from the expansion coefficients based on fast FB basis
 ffb_images = ffb_basis.evaluate(ffb_coefs).asnumpy()
-logger.info("Finish reconstruction of images from fast FB expansion coefficients.")
+print("Finish reconstruction of images from fast FB expansion coefficients.")
 
 # Calculate the mean value of maximum differences between the fast FB estimated images to the original images
 diff = ffb_images - org_images
 ffb_meanmax = np.mean(np.max(abs(diff), axis=(1, 2)))
-logger.info(
+print(
     f"Mean value of maximum differences between FFB estimated images and original images: {ffb_meanmax}"
 )
 
 # Calculate the normalized RMSE of the estimated images
 ffb_nrmse_ims = anorm(diff) / anorm(org_images)
-logger.info(f"FFB Estimated images normalized RMSE: {ffb_nrmse_ims}")
+print(f"FFB Estimated images normalized RMSE: {ffb_nrmse_ims}")
 
 # plot the first images using the fast FB method
 plt.subplot(1, 3, 1)
@@ -136,27 +133,27 @@
 pswf_basis = PSWFBasis2D((img_size, img_size), dtype=org_images.dtype)
 
 # Get the expansion coefficients based on direct PSWF basis
-logger.info("Start direct PSWF expansion of original images.")
+print("Start direct PSWF expansion of original images.")
 tstart = timeit.default_timer()
 pswf_coefs = pswf_basis.evaluate_t(org_images)
 tstop = timeit.default_timer()
 dtime = tstop - tstart
-logger.info(f"Finish direct PSWF expansion of original images in {dtime:.4f} seconds.")
+print(f"Finish direct PSWF expansion of original images in {dtime:.4f} seconds.")
 
 # Reconstruct images from the expansion coefficients based on direct PSWF basis
 pswf_images = pswf_basis.evaluate(pswf_coefs).asnumpy()
-logger.info("Finish reconstruction of images from direct PSWF expansion coefficients.")
+print("Finish reconstruction of images from direct PSWF expansion coefficients.")
 
 # Calculate the mean value of maximum differences between direct PSWF estimated images and original images
 diff = pswf_images - org_images
 pswf_meanmax = np.mean(np.max(abs(diff), axis=(1, 2)))
-logger.info(
+print(
     f"Mean value of maximum differences between PSWF estimated images and original images: {pswf_meanmax}"
 )
 
 # Calculate the normalized RMSE of the estimated images
 pswf_nrmse_ims = anorm(diff) / anorm(org_images)
-logger.info(f"PSWF Estimated images normalized RMSE: {pswf_nrmse_ims}")
+print(f"PSWF Estimated images normalized RMSE: {pswf_nrmse_ims}")
 
 # plot the first images using the direct PSWF method
 plt.subplot(1, 3, 1)
@@ -180,27 +177,27 @@
 fpswf_basis = FPSWFBasis2D((img_size, img_size), dtype=org_images.dtype)
 
 # Get the expansion coefficients based on fast PSWF basis
-logger.info("Start fast PSWF expansion of original images.")
+print("Start fast PSWF expansion of original images.")
 tstart = timeit.default_timer()
 fpswf_coefs = fpswf_basis.evaluate_t(org_images)
 tstop = timeit.default_timer()
 dtime = tstop - tstart
-logger.info(f"Finish fast PSWF expansion of original images in {dtime:.4f} seconds.")
+print(f"Finish fast PSWF expansion of original images in {dtime:.4f} seconds.")
 
 # Reconstruct images from the expansion coefficients based on direct PSWF basis
 fpswf_images = fpswf_basis.evaluate(fpswf_coefs).asnumpy()
-logger.info("Finish reconstruction of images from fast PSWF expansion coefficients.")
+print("Finish reconstruction of images from fast PSWF expansion coefficients.")
 
 # Calculate mean value of maximum differences between the fast PSWF estimated images and the original images
 diff = fpswf_images - org_images
 fpswf_meanmax = np.mean(np.max(abs(diff), axis=(1, 2)))
-logger.info(
+print(
     f"Mean value of maximum differences between FPSWF estimated images and original images: {fpswf_meanmax}"
 )
 
 # Calculate the normalized RMSE of the estimated images
 fpswf_nrmse_ims = anorm(diff) / anorm(org_images)
-logger.info(f"FPSWF Estimated images normalized RMSE: {fpswf_nrmse_ims}")
+print(f"FPSWF Estimated images normalized RMSE: {fpswf_nrmse_ims}")
 
 # plot the first images using the fast PSWF method
 plt.subplot(1, 3, 1)
diff --git a/gallery/tutorials/tutorials/orient3d_simulation.py b/gallery/tutorials/tutorials/orient3d_simulation.py
index 83cbc6d3b3..56752a711f 100644
--- a/gallery/tutorials/tutorials/orient3d_simulation.py
+++ b/gallery/tutorials/tutorials/orient3d_simulation.py
@@ -6,7 +6,6 @@
 matrix and the voting method, based on simulated data projected from a 3D cryo-EM map.
 """
 
-import logging
 import os
 
 import numpy as np
@@ -17,13 +16,11 @@
 from aspire.utils import mean_aligned_angular_distance
 from aspire.volume import Volume
 
-logger = logging.getLogger(__name__)
-
 file_path = os.path.join(
     os.path.dirname(os.getcwd()), "data", "clean70SRibosome_vol_65p.mrc"
 )
 
-logger.info(
+print(
     "This script illustrates orientation estimation using "
     "synchronization matrix and voting method"
 )
@@ -51,7 +48,7 @@
 Cs = 2.0  # Spherical aberration
 alpha = 0.1  # Amplitude contrast
 
-logger.info("Initialize simulation object and CTF filters.")
+print("Initialize simulation object and CTF filters.")
 # Create CTF filters
 filters = [
     RadialCTFFilter(pixel_size, voltage, defocus=d, Cs=2.0, alpha=0.1)
@@ -64,7 +61,7 @@
 
 # Load the map file of a 70S Ribosome and downsample the 3D map to desired resolution.
 # The downsampling can be done by the internal function of Volume object.
-logger.info(
+print(
     f"Load 3D map and downsample 3D map to desired grids "
     f"of {img_size} x {img_size} x {img_size}."
 )
@@ -76,10 +73,10 @@
 # --------------------------------------------------------
 
 # Create a simulation object with specified filters and the downsampled 3D map
-logger.info("Use downsampled map to creat simulation object.")
+print("Use downsampled map to creat simulation object.")
 sim = Simulation(L=img_size, n=num_imgs, vols=vols, unique_filters=filters, dtype=dtype)
 
-logger.info("Get true rotation angles generated randomly by the simulation object.")
+print("Get true rotation angles generated randomly by the simulation object.")
 rots_true = sim.rotations
 
 # %%
@@ -94,7 +91,7 @@
 # images. Additionally, since we are processing images with no noise,
 # we opt not to use a ``fuzzy_mask``, an option that improves
 # common-line detection in higher noise regimes.
-logger.info(
+print(
     "Estimate rotation angles and offsets using synchronization matrix and voting method."
 )
 orient_est = CLSyncVoting(sim, n_theta=36, mask=False)
@@ -109,7 +106,7 @@
 # ``mean_aligned_angular_distance`` will perform global alignment of the estimated rotations
 # to the ground truth and find the mean angular distance between them (in degrees).
 mean_ang_dist = mean_aligned_angular_distance(rots_est, rots_true)
-logger.info(
+print(
     f"Mean angular distance between estimates and ground truth: {mean_ang_dist} degrees"
 )
 
@@ -128,6 +125,6 @@
 
 # Calculate the mean error in pixels across all images.
 offs_err = offs_diff.mean()
-logger.info(
+print(
     f"Mean offset error in pixels {offs_err}, approx {offs_err/img_size*100:.1f}%"
 )
diff --git a/gallery/tutorials/tutorials/preprocess_imgs_sim.py b/gallery/tutorials/tutorials/preprocess_imgs_sim.py
index a9dd9d78b4..af6c7a4314 100644
--- a/gallery/tutorials/tutorials/preprocess_imgs_sim.py
+++ b/gallery/tutorials/tutorials/preprocess_imgs_sim.py
@@ -6,7 +6,6 @@
 reconstructing a 3D map using simulated 2D images.
 """
 
-import logging
 import os
 
 import matplotlib.pyplot as plt
@@ -17,8 +16,6 @@
 from aspire.source.simulation import Simulation
 from aspire.volume import Volume
 
-logger = logging.getLogger(__name__)
-
 file_path = os.path.join(
     os.path.dirname(os.getcwd()), "data", "clean70SRibosome_vol_65p.mrc"
 )
@@ -54,7 +51,7 @@
 # Build Simulation Object and Apply Noise
 # ---------------------------------------
 
-logger.info("Initialize simulation object and CTF filters.")
+print("Initialize simulation object and CTF filters.")
 # Create CTF filters
 ctf_filters = [
     RadialCTFFilter(pixel_size, voltage, defocus=d, Cs=2.0, alpha=0.1)
@@ -62,16 +59,16 @@
 ]
 
 # Load the map file of a 70S ribosome and downsample the 3D map to desired image size.
-logger.info("Load 3D map from mrc file")
+print("Load 3D map from mrc file")
 vols = Volume.load(file_path)
 
 # Downsample the volume to a desired image size and increase density
 # by 1.0e5 time for a better graph view
-logger.info(f"Downsample map to a image size of {img_size} x {img_size} x {img_size}")
+print(f"Downsample map to a image size of {img_size} x {img_size} x {img_size}")
 vols = vols.downsample(img_size) * 1.0e5
 
 # Create a simulation object with specified filters and the downsampled 3D map
-logger.info("Use downsampled map to create simulation object.")
+print("Use downsampled map to create simulation object.")
 source = Simulation(
     L=img_size,
     n=num_imgs,
@@ -89,22 +86,22 @@
 # variable ``source_*``.  That leaves the original ``source`` object
 # untouched.
 
-logger.info("Obtain original images.")
+print("Obtain original images.")
 imgs_od = source.images[0].asnumpy()
 
-logger.info("Perform phase flip to input images.")
+print("Perform phase flip to input images.")
 source_pf = source.phase_flip()
 
-logger.info(f"Downsample image size to {ds_img_size} X {ds_img_size}")
+print(f"Downsample image size to {ds_img_size} X {ds_img_size}")
 source_ds = source.downsample(ds_img_size)
 
-logger.info("Normalize images to background noise.")
+print("Normalize images to background noise.")
 source_nb = source.normalize_background()
 
-logger.info("Whiten noise of images")
+print("Whiten noise of images")
 source_wt = source.whiten()
 
-logger.info("Invert the global density contrast if need")
+print("Invert the global density contrast if need")
 source_rc = source.invert_contrast()
 
 
@@ -113,7 +110,7 @@
 # ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 # Plot the first images.
 
-logger.info("Plot first image from each preprocess steps")
+print("Plot first image from each preprocess steps")
 idm = 0
 plt.subplot(2, 3, 1)
 plt.imshow(imgs_od[idm], cmap="gray")
@@ -160,26 +157,26 @@
 # we can copy a source just by copying the object.
 source_copy = source
 
-logger.info("Obtain original images.")
+print("Obtain original images.")
 imgs_seq_od = source.images[0].asnumpy()
 
-logger.info("Perform phase flip to input images.")
+print("Perform phase flip to input images.")
 source = source.phase_flip()
 imgs_seq_pf = source.images[0].asnumpy()
 
-logger.info(f"Downsample image size to {ds_img_size} X {ds_img_size}")
+print(f"Downsample image size to {ds_img_size} X {ds_img_size}")
 source = source.downsample(ds_img_size)
 imgs_seq_ds = source.images[0].asnumpy()
 
-logger.info("Normalize images to background noise.")
+print("Normalize images to background noise.")
 source = source.normalize_background()
 imgs_seq_nb = source.images[0].asnumpy()
 
-logger.info("Whiten noise of images")
+print("Whiten noise of images")
 source = source.whiten()
 imgs_seq_wt = source.images[0].asnumpy()
 
-logger.info("Invert the global density contrast if need")
+print("Invert the global density contrast if need")
 source = source.invert_contrast()
 imgs_seq_rc = source.images[0].asnumpy()
 
@@ -189,7 +186,7 @@
 # ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 # Plot the first images.
 
-logger.info("Plot first image from each preprocess steps")
+print("Plot first image from each preprocess steps")
 idm = 0
 plt.subplot(2, 3, 1)
 plt.imshow(imgs_od[idm], cmap="gray")
@@ -231,11 +228,11 @@
 # We'll reset our ``source`` to the reference copy we started with.
 source = source_copy
 
-logger.info("Perform phase flip to input images.")
-logger.info(f"Downsample image size to {ds_img_size} X {ds_img_size}")
-logger.info("Normalize images to background noise.")
-logger.info("Whiten noise of images")
-logger.info("Invert the global density contrast if need")
+print("Perform phase flip to input images.")
+print(f"Downsample image size to {ds_img_size} X {ds_img_size}")
+print("Normalize images to background noise.")
+print("Whiten noise of images")
+print("Invert the global density contrast if need")
 source = (
     source.phase_flip()
     .downsample(ds_img_size)

From e341699cf9a8f9de6387ccb43516e40e50769b7c Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Mon, 18 Nov 2024 14:26:36 -0500
Subject: [PATCH 018/184] typo

---
 .../experiments/simulated_abinitio_pipeline.py | 18 +++++++-----------
 gallery/tutorials/tutorials/class_averaging.py |  2 +-
 2 files changed, 8 insertions(+), 12 deletions(-)

diff --git a/gallery/experiments/simulated_abinitio_pipeline.py b/gallery/experiments/simulated_abinitio_pipeline.py
index e9b81065de..60a3593879 100644
--- a/gallery/experiments/simulated_abinitio_pipeline.py
+++ b/gallery/experiments/simulated_abinitio_pipeline.py
@@ -15,8 +15,6 @@
 # In addition, import some classes from
 # the ASPIRE package that will be used throughout this experiment.
 
-import logging
-
 import matplotlib.pyplot as plt
 import numpy as np
 
@@ -29,8 +27,6 @@
 from aspire.source import OrientedSource, Simulation
 from aspire.utils import mean_aligned_angular_distance
 
-logger = logging.getLogger(__name__)
-
 
 # %%
 # Parameters
@@ -55,7 +51,7 @@
 # Start with the hi-res volume map EMDB-2660 sourced from EMDB,
 # https://www.ebi.ac.uk/emdb/EMD-2660, and dowloaded via ASPIRE's downloader utility.
 og_v = emdb_2660()
-logger.info("Original volume map data" f" shape: {og_v.shape} dtype:{og_v.dtype}")
+print("Original volume map data" f" shape: {og_v.shape} dtype:{og_v.dtype}")
 
 
 # Then create a filter based on that variance
@@ -72,7 +68,7 @@ def noise_function(x, y):
 
 custom_noise = CustomNoiseAdder(noise_filter=FunctionFilter(noise_function))
 
-logger.info("Initialize CTF filters.")
+print("Initialize CTF filters.")
 # Create some CTF effects
 pixel_size = og_v.pixel_size  # Pixel size (in angstroms)
 voltage = 200  # Voltage (in KV)
@@ -104,7 +100,7 @@ def noise_function(x, y):
     src.images[:10].show()
 
 # Use phase_flip to attempt correcting for CTF.
-logger.info("Perform phase flip to input images.")
+print("Perform phase flip to input images.")
 src = src.phase_flip()
 
 # Estimate the noise and `Whiten` based on the estimated noise
@@ -175,7 +171,7 @@ def noise_function(x, y):
 # Next create a CL instance for estimating orientation of projections
 # using the Common Line with Synchronization Voting method.
 
-logger.info("Begin Orientation Estimation")
+print("Begin Orientation Estimation")
 
 # Stash true rotations for later comparison.
 # Note class selection re-ordered our images, so we remap the indices back to the original source.
@@ -189,11 +185,11 @@ def noise_function(x, y):
 # estimation in a lazy fashion upon request of images or rotations.
 oriented_src = OrientedSource(avgs, orient_est)
 
-logger.info("Compare with known rotations")
+print("Compare with known rotations")
 # Compare with known true rotations. ``mean_aligned_angular_distance`` globally aligns the estimated
 # rotations to the ground truth and finds the mean angular distance between them.
 mean_ang_dist = mean_aligned_angular_distance(oriented_src.rotations, true_rotations)
-logger.info(
+print(
     f"Mean angular distance between globally aligned estimates and ground truth rotations: {mean_ang_dist}\n"
 )
 
@@ -203,7 +199,7 @@ def noise_function(x, y):
 #
 # Using the estimated rotations, attempt to reconstruct a volume.
 
-logger.info("Begin Volume reconstruction")
+print("Begin Volume reconstruction")
 
 # Setup an estimator to perform the back projection.
 estimator = MeanEstimator(oriented_src)
diff --git a/gallery/tutorials/tutorials/class_averaging.py b/gallery/tutorials/tutorials/class_averaging.py
index 407f86e263..0fcda8616b 100644
--- a/gallery/tutorials/tutorials/class_averaging.py
+++ b/gallery/tutorials/tutorials/class_averaging.py
@@ -194,7 +194,7 @@
 # ``noise_src``.
 classes = avgs.class_indices[review_class]
 reflections = avgs.class_refl[review_class]
-print(f"Class {review_class}'s neighors: {classes}")
+print(f"Class {review_class}'s neighbors: {classes}")
 print(f"Class {review_class}'s reflections: {reflections}")
 
 # The original image is the initial image in the class array.

From d830a3d7c122fdf6455e2fd20be2c162ebb11e10 Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Mon, 18 Nov 2024 14:59:28 -0500
Subject: [PATCH 019/184] tox.ini and formatting

---
 .../experimental_abinitio_pipeline_10028.py           |  1 -
 .../experimental_abinitio_pipeline_10073.py           |  1 -
 .../experimental_abinitio_pipeline_10081.py           |  1 -
 gallery/experiments/simulated_abinitio_pipeline.py    |  1 -
 gallery/tutorials/tutorials/class_averaging.py        |  1 -
 gallery/tutorials/tutorials/cov2d_simulation.py       | 11 ++---------
 gallery/tutorials/tutorials/cov3d_simulation.py       |  1 -
 .../tutorials/generating_volume_projections.py        |  1 -
 gallery/tutorials/tutorials/image_expansion.py        |  5 +----
 gallery/tutorials/tutorials/orient3d_simulation.py    |  4 +---
 tox.ini                                               |  2 ++
 11 files changed, 6 insertions(+), 23 deletions(-)

diff --git a/gallery/experiments/experimental_abinitio_pipeline_10028.py b/gallery/experiments/experimental_abinitio_pipeline_10028.py
index 7cdfbe9d89..1b4d309fb5 100644
--- a/gallery/experiments/experimental_abinitio_pipeline_10028.py
+++ b/gallery/experiments/experimental_abinitio_pipeline_10028.py
@@ -33,7 +33,6 @@
 from aspire.reconstruction import MeanEstimator
 from aspire.source import OrientedSource, RelionSource
 
-
 # %%
 # Parameters
 # ---------------
diff --git a/gallery/experiments/experimental_abinitio_pipeline_10073.py b/gallery/experiments/experimental_abinitio_pipeline_10073.py
index 6e4ca8cfcb..f31c22ce34 100644
--- a/gallery/experiments/experimental_abinitio_pipeline_10073.py
+++ b/gallery/experiments/experimental_abinitio_pipeline_10073.py
@@ -41,7 +41,6 @@
 from aspire.reconstruction import MeanEstimator
 from aspire.source import OrientedSource, RelionSource
 
-
 # %%
 # Parameters
 # ---------------
diff --git a/gallery/experiments/experimental_abinitio_pipeline_10081.py b/gallery/experiments/experimental_abinitio_pipeline_10081.py
index 7b753b6f2b..9a5ce82ae2 100644
--- a/gallery/experiments/experimental_abinitio_pipeline_10081.py
+++ b/gallery/experiments/experimental_abinitio_pipeline_10081.py
@@ -30,7 +30,6 @@
 from aspire.reconstruction import MeanEstimator
 from aspire.source import OrientedSource, RelionSource
 
-
 # %%
 # Parameters
 # ---------------
diff --git a/gallery/experiments/simulated_abinitio_pipeline.py b/gallery/experiments/simulated_abinitio_pipeline.py
index 60a3593879..8370d96370 100644
--- a/gallery/experiments/simulated_abinitio_pipeline.py
+++ b/gallery/experiments/simulated_abinitio_pipeline.py
@@ -27,7 +27,6 @@
 from aspire.source import OrientedSource, Simulation
 from aspire.utils import mean_aligned_angular_distance
 
-
 # %%
 # Parameters
 # ---------------
diff --git a/gallery/tutorials/tutorials/class_averaging.py b/gallery/tutorials/tutorials/class_averaging.py
index 0fcda8616b..b609c59985 100644
--- a/gallery/tutorials/tutorials/class_averaging.py
+++ b/gallery/tutorials/tutorials/class_averaging.py
@@ -15,7 +15,6 @@
 from aspire.source import ArrayImageSource  # Helpful hint if you want to BYO array.
 from aspire.utils import gaussian_2d
 
-
 # %%
 # Build Simulated Data
 # --------------------
diff --git a/gallery/tutorials/tutorials/cov2d_simulation.py b/gallery/tutorials/tutorials/cov2d_simulation.py
index 7eff7675fa..bf64684db4 100644
--- a/gallery/tutorials/tutorials/cov2d_simulation.py
+++ b/gallery/tutorials/tutorials/cov2d_simulation.py
@@ -26,11 +26,6 @@
 )
 
 
-print(
-    "This script illustrates 2D covariance Wiener filtering functionality in ASPIRE package."
-)
-
-
 # %%
 # Image Formatting
 # ----------------
@@ -89,7 +84,7 @@
 vols = vols.downsample(img_size)
 
 # Create a simulation object with specified filters and the downsampled 3D map
-print("Use downsampled map to creat simulation object.")
+print("Use downsampled map to create simulation object.")
 sim = Simulation(
     L=img_size,
     n=num_imgs,
@@ -154,9 +149,7 @@
 # mean and covariance estimates will allow us to evaluate the mean and
 # covariance estimates from the filtered, noisy data, later.
 
-print(
-    "Get 2D covariance matrices of clean and noisy images using FB coefficients."
-)
+print("Get 2D covariance matrices of clean and noisy images using FB coefficients.")
 cov2d = RotCov2D(ffbbasis)
 mean_coef = cov2d.get_mean(coef_clean)
 covar_coef = cov2d.get_covar(coef_clean, mean_coef, noise_var=0)
diff --git a/gallery/tutorials/tutorials/cov3d_simulation.py b/gallery/tutorials/tutorials/cov3d_simulation.py
index 86d26116ea..8bb72906b7 100644
--- a/gallery/tutorials/tutorials/cov3d_simulation.py
+++ b/gallery/tutorials/tutorials/cov3d_simulation.py
@@ -20,7 +20,6 @@
 from aspire.utils.random import Random
 from aspire.volume import LegacyVolume, Volume
 
-
 # %%
 # Create Simulation Object
 # ------------------------
diff --git a/gallery/tutorials/tutorials/generating_volume_projections.py b/gallery/tutorials/tutorials/generating_volume_projections.py
index f9e342dab5..4ed0b60258 100644
--- a/gallery/tutorials/tutorials/generating_volume_projections.py
+++ b/gallery/tutorials/tutorials/generating_volume_projections.py
@@ -16,7 +16,6 @@
 from aspire.utils import Rotation
 from aspire.volume import Volume
 
-
 # %%
 # Configure how many images we'd like to project
 # ----------------------------------------------
diff --git a/gallery/tutorials/tutorials/image_expansion.py b/gallery/tutorials/tutorials/image_expansion.py
index 55cb6c5db4..9dc5346087 100644
--- a/gallery/tutorials/tutorials/image_expansion.py
+++ b/gallery/tutorials/tutorials/image_expansion.py
@@ -15,10 +15,7 @@
 from aspire.basis import FBBasis2D, FFBBasis2D, FPSWFBasis2D, PSWFBasis2D
 from aspire.utils import anorm
 
-
-print(
-    "This script illustrates different image expansion methods in ASPIRE package."
-)
+print("This script illustrates different image expansion methods in ASPIRE package.")
 
 # %%
 # Load Initial Images
diff --git a/gallery/tutorials/tutorials/orient3d_simulation.py b/gallery/tutorials/tutorials/orient3d_simulation.py
index 56752a711f..085eb9ac23 100644
--- a/gallery/tutorials/tutorials/orient3d_simulation.py
+++ b/gallery/tutorials/tutorials/orient3d_simulation.py
@@ -125,6 +125,4 @@
 
 # Calculate the mean error in pixels across all images.
 offs_err = offs_diff.mean()
-print(
-    f"Mean offset error in pixels {offs_err}, approx {offs_err/img_size*100:.1f}%"
-)
+print(f"Mean offset error in pixels {offs_err}, approx {offs_err/img_size*100:.1f}%")
diff --git a/tox.ini b/tox.ini
index f05f9896d4..f72dde7939 100644
--- a/tox.ini
+++ b/tox.ini
@@ -67,6 +67,8 @@ max-line-length = 88
 extend-ignore = E203, E501
 per-file-ignores =
   __init__.py: F401
+ gallery/experiments/*.py: T201
+ gallery/tutorials/tutorials/*.py: T201
  gallery/tutorials/aspire_introduction.py: T201, F401, E402
  gallery/tutorials/configuration.py: T201, E402
  gallery/tutorials/pipeline_demo.py: T201

From bf2348e1f51042fcc0a196a82b170e89233a8657 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Tue, 19 Nov 2024 13:08:59 -0500
Subject: [PATCH 020/184] revert experiment prints back to logging

---
 .../experimental_abinitio_pipeline_10028.py   | 18 +++++++++++-------
 .../experimental_abinitio_pipeline_10073.py   | 14 +++++++++-----
 .../experimental_abinitio_pipeline_10081.py   | 16 ++++++++++------
 .../simulated_abinitio_pipeline.py            | 19 ++++++++++++-------
 4 files changed, 42 insertions(+), 25 deletions(-)

diff --git a/gallery/experiments/experimental_abinitio_pipeline_10028.py b/gallery/experiments/experimental_abinitio_pipeline_10028.py
index 1b4d309fb5..7892302758 100644
--- a/gallery/experiments/experimental_abinitio_pipeline_10028.py
+++ b/gallery/experiments/experimental_abinitio_pipeline_10028.py
@@ -22,6 +22,7 @@
 # In addition, import some classes from
 # the ASPIRE package that will be used throughout this experiment.
 
+import logging
 from pathlib import Path
 
 import matplotlib.pyplot as plt
@@ -33,6 +34,9 @@
 from aspire.reconstruction import MeanEstimator
 from aspire.source import OrientedSource, RelionSource
 
+logger = logging.getLogger(__name__)
+
+
 # %%
 # Parameters
 # ---------------
@@ -64,7 +68,7 @@
 )
 
 # Downsample the images
-print(f"Set the resolution to {img_size} X {img_size}")
+logger.info(f"Set the resolution to {img_size} X {img_size}")
 src = src.downsample(img_size)
 
 # Peek
@@ -72,7 +76,7 @@
     src.images[:10].show()
 
 # Use phase_flip to attempt correcting for CTF.
-print("Perform phase flip to input images.")
+logger.info("Perform phase flip to input images.")
 src = src.phase_flip()
 
 # Estimate the noise and `Whiten` based on the estimated noise
@@ -90,7 +94,7 @@
 
 # # Optionally invert image contrast, depends on data convention.
 # # This is not needed for 10028, but included anyway.
-# print("Invert the global density contrast")
+# logger.info("Invert the global density contrast")
 # src = src.invert_contrast()
 
 # Caching is used for speeding up large datasets on high memory machines.
@@ -128,7 +132,7 @@
 #
 # Now perform classification and averaging for each class.
 
-print("Begin Class Averaging")
+logger.info("Begin Class Averaging")
 
 # Now perform classification and averaging for each class.
 # This also demonstrates the potential to use a different source for classification and averaging.
@@ -155,7 +159,7 @@
 # Next create a CL instance for estimating orientation of projections
 # using the Common Line with Synchronization Voting method.
 
-print("Begin Orientation Estimation")
+logger.info("Begin Orientation Estimation")
 
 # Create a custom orientation estimation object for ``avgs``.
 # This is done to customize the ``n_theta`` value.
@@ -171,7 +175,7 @@
 #
 # Using the oriented source, attempt to reconstruct a volume.
 
-print("Begin Volume reconstruction")
+logger.info("Begin Volume reconstruction")
 
 # Setup an estimator to perform the back projection.
 estimator = MeanEstimator(oriented_src)
@@ -179,7 +183,7 @@
 # Perform the estimation and save the volume.
 estimated_volume = estimator.estimate()
 estimated_volume.save(volume_output_filename, overwrite=True)
-print(f"Saved Volume to {str(Path(volume_output_filename).resolve())}")
+logger.info(f"Saved Volume to {str(Path(volume_output_filename).resolve())}")
 
 # Peek at result
 if interactive:
diff --git a/gallery/experiments/experimental_abinitio_pipeline_10073.py b/gallery/experiments/experimental_abinitio_pipeline_10073.py
index f31c22ce34..d936a98d65 100644
--- a/gallery/experiments/experimental_abinitio_pipeline_10073.py
+++ b/gallery/experiments/experimental_abinitio_pipeline_10073.py
@@ -25,6 +25,7 @@
 # In addition, import some classes from
 # the ASPIRE package that will be used throughout this experiment.
 
+import logging
 from pathlib import Path
 
 import numpy as np
@@ -41,6 +42,9 @@
 from aspire.reconstruction import MeanEstimator
 from aspire.source import OrientedSource, RelionSource
 
+logger = logging.getLogger(__name__)
+
+
 # %%
 # Parameters
 # ---------------
@@ -70,7 +74,7 @@
 )
 
 # Downsample the images
-print(f"Set the resolution to {img_size} X {img_size}")
+logger.info(f"Set the resolution to {img_size} X {img_size}")
 src = src.downsample(img_size)
 
 src = src.cache()
@@ -81,7 +85,7 @@
 #
 # Now perform classification and averaging for each class.
 
-print("Begin Class Averaging")
+logger.info("Begin Class Averaging")
 
 # Now perform classification and averaging for each class.
 # This also demonstrates customizing a ClassAvgSource, by using global
@@ -128,7 +132,7 @@
 # Next create a CL instance for estimating orientation of projections
 # using the Common Line with Synchronization Voting method.
 
-print("Begin Orientation Estimation")
+logger.info("Begin Orientation Estimation")
 
 # Create a custom orientation estimation object for ``avgs``.
 orient_est = CLSync3N(avgs, n_theta=72)
@@ -143,7 +147,7 @@
 #
 # Using the oriented source, attempt to reconstruct a volume.
 
-print("Begin Volume reconstruction")
+logger.info("Begin Volume reconstruction")
 
 # Setup an estimator to perform the back projection.
 estimator = MeanEstimator(oriented_src)
@@ -151,4 +155,4 @@
 # Perform the estimation and save the volume.
 estimated_volume = estimator.estimate()
 estimated_volume.save(volume_output_filename, overwrite=True)
-print(f"Saved Volume to {str(Path(volume_output_filename).resolve())}")
+logger.info(f"Saved Volume to {str(Path(volume_output_filename).resolve())}")
diff --git a/gallery/experiments/experimental_abinitio_pipeline_10081.py b/gallery/experiments/experimental_abinitio_pipeline_10081.py
index 9a5ce82ae2..7d83842c24 100644
--- a/gallery/experiments/experimental_abinitio_pipeline_10081.py
+++ b/gallery/experiments/experimental_abinitio_pipeline_10081.py
@@ -22,6 +22,7 @@
 # In addition, import some classes from
 # the ASPIRE package that will be used throughout this experiment.
 
+import logging
 from pathlib import Path
 
 from aspire.abinitio import CLSymmetryC3C4
@@ -30,6 +31,9 @@
 from aspire.reconstruction import MeanEstimator
 from aspire.source import OrientedSource, RelionSource
 
+logger = logging.getLogger(__name__)
+
+
 # %%
 # Parameters
 # ---------------
@@ -63,11 +67,11 @@
 )
 
 # Downsample the images
-print(f"Set the resolution to {img_size} X {img_size}")
+logger.info(f"Set the resolution to {img_size} X {img_size}")
 src = src.downsample(img_size)
 
 # Use phase_flip to attempt correcting for CTF.
-print("Perform phase flip to input images.")
+logger.info("Perform phase flip to input images.")
 src = src.phase_flip()
 
 # Estimate the noise and `Whiten` based on the estimated noise
@@ -83,7 +87,7 @@
 #
 # Now perform classification and averaging for each class.
 
-print("Begin Class Averaging")
+logger.info("Begin Class Averaging")
 
 # Now perform classification and averaging for each class.
 # Automatically configure parallel processing
@@ -102,7 +106,7 @@
 # Next create a CL instance for estimating orientation of projections
 # using the Common Line with Synchronization Voting method.
 
-print("Begin Orientation Estimation")
+logger.info("Begin Orientation Estimation")
 
 # Create a custom orientation estimation object for ``avgs``.
 orient_est = CLSymmetryC3C4(avgs, symmetry="C4", n_theta=72, max_shift=0)
@@ -123,7 +127,7 @@
 # back-projection.  This boosts the effective number of images used in
 # the reconstruction from ``n_classes`` to ``4*n_classes``.
 
-print("Begin Volume reconstruction")
+logger.info("Begin Volume reconstruction")
 
 
 # Setup an estimator to perform the back projection.
@@ -132,4 +136,4 @@
 # Perform the estimation and save the volume.
 estimated_volume = estimator.estimate()
 estimated_volume.save(volume_output_filename, overwrite=True)
-print(f"Saved Volume to {str(Path(volume_output_filename).resolve())}")
+logger.info(f"Saved Volume to {str(Path(volume_output_filename).resolve())}")
diff --git a/gallery/experiments/simulated_abinitio_pipeline.py b/gallery/experiments/simulated_abinitio_pipeline.py
index 8370d96370..e9b81065de 100644
--- a/gallery/experiments/simulated_abinitio_pipeline.py
+++ b/gallery/experiments/simulated_abinitio_pipeline.py
@@ -15,6 +15,8 @@
 # In addition, import some classes from
 # the ASPIRE package that will be used throughout this experiment.
 
+import logging
+
 import matplotlib.pyplot as plt
 import numpy as np
 
@@ -27,6 +29,9 @@
 from aspire.source import OrientedSource, Simulation
 from aspire.utils import mean_aligned_angular_distance
 
+logger = logging.getLogger(__name__)
+
+
 # %%
 # Parameters
 # ---------------
@@ -50,7 +55,7 @@
 # Start with the hi-res volume map EMDB-2660 sourced from EMDB,
 # https://www.ebi.ac.uk/emdb/EMD-2660, and dowloaded via ASPIRE's downloader utility.
 og_v = emdb_2660()
-print("Original volume map data" f" shape: {og_v.shape} dtype:{og_v.dtype}")
+logger.info("Original volume map data" f" shape: {og_v.shape} dtype:{og_v.dtype}")
 
 
 # Then create a filter based on that variance
@@ -67,7 +72,7 @@ def noise_function(x, y):
 
 custom_noise = CustomNoiseAdder(noise_filter=FunctionFilter(noise_function))
 
-print("Initialize CTF filters.")
+logger.info("Initialize CTF filters.")
 # Create some CTF effects
 pixel_size = og_v.pixel_size  # Pixel size (in angstroms)
 voltage = 200  # Voltage (in KV)
@@ -99,7 +104,7 @@ def noise_function(x, y):
     src.images[:10].show()
 
 # Use phase_flip to attempt correcting for CTF.
-print("Perform phase flip to input images.")
+logger.info("Perform phase flip to input images.")
 src = src.phase_flip()
 
 # Estimate the noise and `Whiten` based on the estimated noise
@@ -170,7 +175,7 @@ def noise_function(x, y):
 # Next create a CL instance for estimating orientation of projections
 # using the Common Line with Synchronization Voting method.
 
-print("Begin Orientation Estimation")
+logger.info("Begin Orientation Estimation")
 
 # Stash true rotations for later comparison.
 # Note class selection re-ordered our images, so we remap the indices back to the original source.
@@ -184,11 +189,11 @@ def noise_function(x, y):
 # estimation in a lazy fashion upon request of images or rotations.
 oriented_src = OrientedSource(avgs, orient_est)
 
-print("Compare with known rotations")
+logger.info("Compare with known rotations")
 # Compare with known true rotations. ``mean_aligned_angular_distance`` globally aligns the estimated
 # rotations to the ground truth and finds the mean angular distance between them.
 mean_ang_dist = mean_aligned_angular_distance(oriented_src.rotations, true_rotations)
-print(
+logger.info(
     f"Mean angular distance between globally aligned estimates and ground truth rotations: {mean_ang_dist}\n"
 )
 
@@ -198,7 +203,7 @@ def noise_function(x, y):
 #
 # Using the estimated rotations, attempt to reconstruct a volume.
 
-print("Begin Volume reconstruction")
+logger.info("Begin Volume reconstruction")
 
 # Setup an estimator to perform the back projection.
 estimator = MeanEstimator(oriented_src)

From 24c150dfe4165eb86fd6173e03388a5b56634e2c Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Tue, 19 Nov 2024 15:48:03 -0500
Subject: [PATCH 021/184] remove redundant ignores

---
 tox.ini | 9 +++------
 1 file changed, 3 insertions(+), 6 deletions(-)

diff --git a/tox.ini b/tox.ini
index f72dde7939..9795975041 100644
--- a/tox.ini
+++ b/tox.ini
@@ -67,17 +67,14 @@ max-line-length = 88
 extend-ignore = E203, E501
 per-file-ignores =
   __init__.py: F401
- gallery/experiments/*.py: T201
  gallery/tutorials/tutorials/*.py: T201
  gallery/tutorials/aspire_introduction.py: T201, F401, E402
  gallery/tutorials/configuration.py: T201, E402
  gallery/tutorials/pipeline_demo.py: T201
  gallery/tutorials/turorials/data_downloader.py: E402
- gallery/tutorials/tutorials/ctf.py: T201, E402
- gallery/tutorials/tutorials/image_class.py: T201
- gallery/tutorials/tutorials/micrograph_source.py: T201, E402
- gallery/tutorials/tutorials/weighted_volume_estimation.py: T201, E402
- gallery/tutorials/tutorials/relion_projection_interop.py: T201
+ gallery/tutorials/tutorials/ctf.py: E402
+ gallery/tutorials/tutorials/micrograph_source.py: E402
+ gallery/tutorials/tutorials/weighted_volume_estimation.py: E402
  # Ignore Sphinx gallery builds
  docs/build/html/_downloads/*/*.py: T201, E402, F401, E265
  docs/source/auto*/*.py: T201, E402, F401, E265

From 6c59064f7c9b5b64678a9eb759c28140f8b428c6 Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Tue, 19 Nov 2024 16:01:07 -0500
Subject: [PATCH 022/184] remove conflicting ignores

---
 tox.ini | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/tox.ini b/tox.ini
index 9795975041..afa6003cc5 100644
--- a/tox.ini
+++ b/tox.ini
@@ -67,14 +67,10 @@ max-line-length = 88
 extend-ignore = E203, E501
 per-file-ignores =
   __init__.py: F401
- gallery/tutorials/tutorials/*.py: T201
+ gallery/tutorials/tutorials/*.py: T201, E402
  gallery/tutorials/aspire_introduction.py: T201, F401, E402
  gallery/tutorials/configuration.py: T201, E402
  gallery/tutorials/pipeline_demo.py: T201
- gallery/tutorials/turorials/data_downloader.py: E402
- gallery/tutorials/tutorials/ctf.py: E402
- gallery/tutorials/tutorials/micrograph_source.py: E402
- gallery/tutorials/tutorials/weighted_volume_estimation.py: E402
  # Ignore Sphinx gallery builds
  docs/build/html/_downloads/*/*.py: T201, E402, F401, E265
  docs/source/auto*/*.py: T201, E402, F401, E265

From 05dccc60ec0478766678f678ec7aae39fbfcbe5b Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Wed, 16 Oct 2024 14:55:21 -0400
Subject: [PATCH 023/184] Remove m_reshape from ffb2d.

---
 src/aspire/basis/ffb_2d.py | 32 +++++++++++++-------------------
 1 file changed, 13 insertions(+), 19 deletions(-)

diff --git a/src/aspire/basis/ffb_2d.py b/src/aspire/basis/ffb_2d.py
index 8d46e8419c..776c03cac8 100644
--- a/src/aspire/basis/ffb_2d.py
+++ b/src/aspire/basis/ffb_2d.py
@@ -10,7 +10,6 @@
 from aspire.numeric import fft, xp
 from aspire.operators import BlkDiagMatrix
 from aspire.utils import complex_type
-from aspire.utils.matlab_compat import m_reshape
 
 logger = logging.getLogger(__name__)
 
@@ -93,14 +92,9 @@ def _precomp(self):
                 ind_radial += 1
 
         # Only calculate "positive" frequencies in one half-plane.
-        freqs_x = m_reshape(r, (n_r, 1)) @ m_reshape(
-            np.cos(np.arange(n_theta, dtype=self.dtype) * 2 * pi / (2 * n_theta)),
-            (1, n_theta),
-        )
-        freqs_y = m_reshape(r, (n_r, 1)) @ m_reshape(
-            np.sin(np.arange(n_theta, dtype=self.dtype) * 2 * pi / (2 * n_theta)),
-            (1, n_theta),
-        )
+        theta_grid = np.arange(n_theta, dtype=self.dtype) * 2 * pi / (2 * n_theta)
+        freqs_x = r[:, None] @ np.cos(theta_grid)[None, :]
+        freqs_y = r[:, None] @ np.sin(theta_grid)[None, :]
         freqs = np.vstack((freqs_y[np.newaxis, ...], freqs_x[np.newaxis, ...]))
 
         return {"gl_nodes": r, "gl_weights": w, "radial": radial, "freqs": freqs}
@@ -127,13 +121,13 @@ def _evaluate(self, v):
         n_r = self._precomp["freqs"].shape[1]
 
         # go through  each basis function and find corresponding coefficient
-        pf = xp.zeros((n_data, 2 * n_theta, n_r), dtype=complex_type(self.dtype))
+        pf = xp.zeros((n_data, n_r, 2 * n_theta), dtype=complex_type(self.dtype))
 
         ind = 0
 
         idx = ind + np.arange(self.k_max[0], dtype=int)
 
-        pf[:, 0, :] = v[:, self._zero_angular_inds] @ self.radial_norm[idx]
+        pf[:, :, 0] = v[:, self._zero_angular_inds] @ self.radial_norm[idx]
         ind = ind + idx.size
 
         ind_pos = ind
@@ -149,27 +143,27 @@ def _evaluate(self, v):
                 v_ell = 1j * v_ell
 
             pf_ell = v_ell @ self.radial_norm[idx]
-            pf[:, ell, :] = pf_ell
+            pf[:, :, ell] = pf_ell
 
             if np.mod(ell, 2) == 0:
-                pf[:, 2 * n_theta - ell, :] = pf_ell.conjugate()
+                pf[:, :, 2 * n_theta - ell] = pf_ell.conjugate()
             else:
-                pf[:, 2 * n_theta - ell, :] = -pf_ell.conjugate()
+                pf[:, :, 2 * n_theta - ell] = -pf_ell.conjugate()
 
             ind = ind + idx.size
             ind_pos = ind_pos + 2 * self.k_max[ell]
 
         # 1D inverse FFT in the degree of polar angle
-        pf = 2 * xp.pi * fft.ifft(pf, axis=1)
+        pf = 2 * xp.pi * fft.ifft(pf, axis=2)
 
         # Only need "positive" frequencies.
-        hsize = int(pf.shape[1] / 2)
-        pf = pf[:, 0:hsize, :]
-        pf *= self.gl_weighted_nodes[None, None, :]
+        hsize = int(pf.shape[2] / 2)
+        pf = pf[:, :, 0:hsize]
+        pf *= self.gl_weighted_nodes[None, :, None]
         pf = pf.reshape(n_data, n_r * n_theta)
 
         # perform inverse non-uniformly FFT transform back to 2D coordinate basis
-        freqs = m_reshape(self._precomp["freqs"], (2, n_r * n_theta))
+        freqs = self._precomp["freqs"].reshape(2, n_r * n_theta)
 
         x = 2 * anufft(pf, 2 * pi * freqs, self.sz, real=True)
 

From ea7c11465045ddb383571a6ca1b86d4446d1bb39 Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Fri, 18 Oct 2024 15:00:13 -0400
Subject: [PATCH 024/184] remove m_funcs from fb2d evaluate_t

---
 src/aspire/basis/fb_2d.py | 19 +++++++------------
 1 file changed, 7 insertions(+), 12 deletions(-)

diff --git a/src/aspire/basis/fb_2d.py b/src/aspire/basis/fb_2d.py
index 18f7b83478..b581b3a946 100644
--- a/src/aspire/basis/fb_2d.py
+++ b/src/aspire/basis/fb_2d.py
@@ -220,32 +220,28 @@ def _evaluate(self, v):
 
         return x
 
-    def _evaluate_t(self, v):
+    def _evaluate_t(self, x):
         """
         Evaluate coefficient in FB basis from those in standard 2D coordinate basis
 
-        :param v: The coefficient array to be evaluated. The last dimensions
+        :param x: The coefficient array to be evaluated. The last dimensions
             must equal `self.sz`.
         :return: The evaluation of the coefficient array `v` in the dual basis
             of `basis`. This is an array of vectors whose last dimension equals
             `self.count` and whose first dimensions correspond to
             first dimensions of `v`.
         """
-        v = v.T
-        x, sz_roll = unroll_dim(v, self.ndim + 1)
-        x = m_reshape(
-            x, new_shape=tuple([np.prod(self.sz)] + list(x.shape[self.ndim :]))
-        )
+        x = x.reshape(x.shape[0], -1)
 
         r_idx = self.basis_coords["r_idx"]
         ang_idx = self.basis_coords["ang_idx"]
-        mask = m_flatten(self.basis_coords["mask"])
+        mask = self.basis_coords["mask"].flatten()
 
         ind = 0
         ind_radial = 0
         ind_ang = 0
 
-        v = np.zeros(shape=tuple([self.count] + list(x.shape[1:])), dtype=v.dtype)
+        v = np.zeros((x.shape[0], self.count), dtype=x.dtype)
         for ell in range(0, self.ell_max + 1):
             k_max = self.k_max[ell]
             idx_radial = ind_radial + np.arange(0, k_max)
@@ -259,14 +255,13 @@ def _evaluate_t(self, v):
                 ang = self._precomp["ang"][:, ind_ang]
                 ang_radial = np.expand_dims(ang[ang_idx], axis=1) * radial[r_idx]
                 idx = ind + np.arange(0, k_max)
-                v[idx] = ang_radial.T @ x[mask]
+                v[:, idx] = x[:, mask] @ ang_radial
                 ind += len(idx)
                 ind_ang += 1
 
             ind_radial += len(idx_radial)
 
-        v = roll_dim(v, sz_roll)
-        return v.T  # RCOPT
+        return v
 
     def calculate_bispectrum(
         self, coef, flatten=False, filter_nonzero_freqs=False, freq_cutoff=None

From 89ae2844b2e249267498e988779ce713f5e93e30 Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Fri, 18 Oct 2024 15:35:43 -0400
Subject: [PATCH 025/184] remove m_funcs imports from fb2d

---
 src/aspire/basis/fb_2d.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/src/aspire/basis/fb_2d.py b/src/aspire/basis/fb_2d.py
index b581b3a946..39b2207b62 100644
--- a/src/aspire/basis/fb_2d.py
+++ b/src/aspire/basis/fb_2d.py
@@ -5,8 +5,6 @@
 
 from aspire.basis import FBBasisMixin, SteerableBasis2D
 from aspire.basis.basis_utils import 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__)
 
@@ -189,7 +187,7 @@ def _evaluate(self, v):
 
         r_idx = self.basis_coords["r_idx"]
         ang_idx = self.basis_coords["ang_idx"]
-        mask = m_flatten(self.basis_coords["mask"])
+        mask = self.basis_coords["mask"].flatten()
 
         ind = 0
         ind_radial = 0

From 25fc5bb0bda709fcf02f896215475424d60b19f4 Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Mon, 21 Oct 2024 11:04:27 -0400
Subject: [PATCH 026/184] remove m_reshape from covar2d.py

---
 src/aspire/covariance/covar2d.py | 30 +++++++++++++-----------------
 1 file changed, 13 insertions(+), 17 deletions(-)

diff --git a/src/aspire/covariance/covar2d.py b/src/aspire/covariance/covar2d.py
index d2e7dafe43..0162915c72 100644
--- a/src/aspire/covariance/covar2d.py
+++ b/src/aspire/covariance/covar2d.py
@@ -8,7 +8,6 @@
 from aspire.operators import BlkDiagMatrix, DiagMatrix
 from aspire.optimization import conj_grad, fill_struct
 from aspire.utils import make_symmat
-from aspire.utils.matlab_compat import m_reshape
 
 logger = logging.getLogger(__name__)
 
@@ -347,36 +346,35 @@ def identity(x):
         cg_opt = covar_est_opt
 
         covar_coef = BlkDiagMatrix.zeros(self.basis.blk_diag_cov_shape)
+        covar_coef2 = BlkDiagMatrix.zeros(self.basis.blk_diag_cov_shape)
 
         def precond_fun(S, x):
             p = np.size(S, 0)
             assert np.size(x) == p * p, "The sizes of S and x are not consistent."
-            x = m_reshape(x, (p, p))
+            x = x.reshape(p, p)
             y = S @ x @ S
-            y = m_reshape(y, (p**2,))
-            return y
+            return y.flatten()
 
         def apply(A, x):
             p = np.size(A[0], 0)
-            x = m_reshape(x, (p, p))
+            x = x.reshape(p, p)
             y = np.zeros_like(x)
             for k in range(0, len(A)):
                 y = y + A[k] @ x @ A[k].T
-            y = m_reshape(y, (p**2,))
-            return y
+            return y.flatten()
 
         for ell in range(0, len(b)):
             A_ell = []
             for k in range(0, len(A)):
                 A_ell.append(A[k][ell])
             p = np.size(A_ell[0], 0)
-            b_ell = m_reshape(b[ell], (p**2,))
+            b_ell = b[ell].flatten()
             S = inv(M[ell])
             cg_opt["preconditioner"] = lambda x, S=S: precond_fun(S, x)
             covar_coef_ell, _, _ = conj_grad(
                 lambda x, A_ell=A_ell: apply(A_ell, x), b_ell, cg_opt
             )
-            covar_coef[ell] = m_reshape(covar_coef_ell, (p, p))
+            covar_coef[ell] = covar_coef_ell.reshape(p, p)
 
         if not covar_coef.check_psd():
             logger.warning("Covariance matrix in Cov2D is not positive semidefinite.")
@@ -685,19 +683,17 @@ def _solve_covar_cg(self, A_covar, b_covar, M, covar_est_opt):
         def precond_fun(S, x):
             p = np.size(S, 0)
             assert np.size(x) == p * p, "The sizes of S and x are not consistent."
-            x = m_reshape(x, (p, p))
+            x = x.reshape(p, p)
             y = S @ x @ S
-            y = m_reshape(y, (p**2,))
-            return y
+            return y.flatten()
 
         def apply(A, x):
             p = np.size(A[0], 0)
-            x = m_reshape(x, (p, p))
+            x = x.reshape(p, p)
             y = np.zeros_like(x)
             for k in range(0, len(A)):
                 y = y + A[k] @ x @ A[k].T
-            y = m_reshape(y, (p**2,))
-            return y
+            return y.flatten()
 
         cg_opt = covar_est_opt
         covar_coef = BlkDiagMatrix.zeros(
@@ -709,13 +705,13 @@ def apply(A, x):
             for k in range(0, len(A_covar)):
                 A_ell.append(A_covar[k][ell])
             p = np.size(A_ell[0], 0)
-            b_ell = m_reshape(b_covar[ell], (p**2,))
+            b_ell = b_covar[ell].flatten()
             S = inv(M[ell])
             cg_opt["preconditioner"] = lambda x, S=S: precond_fun(S, x)
             covar_coef_ell, _, _ = conj_grad(
                 lambda x, A_ell=A_ell: apply(A_ell, x), b_ell, cg_opt
             )
-            covar_coef[ell] = m_reshape(covar_coef_ell, (p, p))
+            covar_coef[ell] = covar_coef_ell.reshape(p, p)
 
         return covar_coef
 

From ac25c3c29cd0610d452b7edaed0336528aa83b85 Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Mon, 21 Oct 2024 11:06:01 -0400
Subject: [PATCH 027/184] remove temp test variable

---
 src/aspire/covariance/covar2d.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/src/aspire/covariance/covar2d.py b/src/aspire/covariance/covar2d.py
index 0162915c72..eebfd5a2c7 100644
--- a/src/aspire/covariance/covar2d.py
+++ b/src/aspire/covariance/covar2d.py
@@ -346,7 +346,6 @@ def identity(x):
         cg_opt = covar_est_opt
 
         covar_coef = BlkDiagMatrix.zeros(self.basis.blk_diag_cov_shape)
-        covar_coef2 = BlkDiagMatrix.zeros(self.basis.blk_diag_cov_shape)
 
         def precond_fun(S, x):
             p = np.size(S, 0)

From 30fabb8d679c2ffd8edf809595376224eaff9715 Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Fri, 25 Oct 2024 09:00:10 -0400
Subject: [PATCH 028/184] remove m_reshape from fb_3d

---
 src/aspire/basis/fb_3d.py | 39 ++++++++++++---------------------------
 1 file changed, 12 insertions(+), 27 deletions(-)

diff --git a/src/aspire/basis/fb_3d.py b/src/aspire/basis/fb_3d.py
index dedd52025c..66e1ae7082 100644
--- a/src/aspire/basis/fb_3d.py
+++ b/src/aspire/basis/fb_3d.py
@@ -4,8 +4,6 @@
 
 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__)
 
@@ -150,21 +148,17 @@ def _evaluate(self, v):
             This is an array whose first dimensions equal `self.z` and the
             remaining dimensions correspond to dimensions two and higher of `v`.
         """
-
-        v = v.T
-        v, sz_roll = unroll_dim(v, 2)
+        v = v.reshape(v.shape[0], -1)
 
         r_idx = self.basis_coords["r_idx"]
         ang_idx = self.basis_coords["ang_idx"]
-        mask = m_flatten(self.basis_coords["mask"])
+        mask = self.basis_coords["mask"].flatten()
 
         ind = 0
         ind_radial = 0
         ind_ang = 0
 
-        x = np.zeros(
-            shape=tuple([np.prod(self.sz)] + list(v.shape[1:])), dtype=self.dtype
-        )
+        x = np.zeros((v.shape[0], np.prod(self.sz)), dtype=self.dtype)
         for ell in range(0, self.ell_max + 1):
             k_max = self.k_max[ell]
             idx_radial = ind_radial + np.arange(0, k_max)
@@ -176,43 +170,35 @@ def _evaluate(self, v):
                 ang = self._precomp["ang"][:, ind_ang]
                 ang_radial = np.expand_dims(ang[ang_idx], axis=1) * radial[r_idx]
                 idx = ind + np.arange(0, len(idx_radial))
-                x[mask] += ang_radial @ v[idx]
+                x[:, mask] += v[:, idx] @ ang_radial.T
                 ind += len(idx)
                 ind_ang += 1
 
             ind_radial += len(idx_radial)
 
-        x = m_reshape(x, self.sz + x.shape[1:])
-        x = roll_dim(x, sz_roll)
-
-        return x.T
+        return x.reshape(v.shape[0], *self.sz)
 
-    def _evaluate_t(self, v):
+    def _evaluate_t(self, x):
         """
         Evaluate coefficient in FB basis from those in standard 3D coordinate basis
 
-        :param v: The coefficient array to be evaluated. The first dimensions
+        :param x: The coefficient array to be evaluated. The first dimensions
             must equal `self.sz`.
         :return: The evaluation of the coefficient array `v` in the dual
             basis of `basis`. This is an array of vectors whose first dimension
             equals `self.count` and whose remaining dimensions correspond
             to higher dimensions of `v`.
         """
-        v = v.T
-        x, sz_roll = unroll_dim(v, self.ndim + 1)
-        x = m_reshape(
-            x, new_shape=tuple([np.prod(self.sz)] + list(x.shape[self.ndim :]))
-        )
-
+        x = x.reshape(x.shape[0], -1)
         r_idx = self.basis_coords["r_idx"]
         ang_idx = self.basis_coords["ang_idx"]
-        mask = m_flatten(self.basis_coords["mask"])
+        mask = self.basis_coords["mask"].flatten()
 
         ind = 0
         ind_radial = 0
         ind_ang = 0
 
-        v = np.zeros(shape=tuple([self.count] + list(x.shape[1:])), dtype=self.dtype)
+        v = np.zeros((x.shape[0], self.count), dtype=self.dtype)
         for ell in range(0, self.ell_max + 1):
             k_max = self.k_max[ell]
             idx_radial = ind_radial + np.arange(0, k_max)
@@ -224,11 +210,10 @@ def _evaluate_t(self, v):
                 ang = self._precomp["ang"][:, ind_ang]
                 ang_radial = np.expand_dims(ang[ang_idx], axis=1) * radial[r_idx]
                 idx = ind + np.arange(0, len(idx_radial))
-                v[idx] = np.real(ang_radial.T @ x[mask])
+                v[:, idx] = x[:, mask] @ ang_radial
                 ind += len(idx)
                 ind_ang += 1
 
             ind_radial += len(idx_radial)
 
-        v = roll_dim(v, sz_roll)
-        return v.T
+        return v

From 6aa36a0949a30b7484dcaeed9d927f6fc7ff06d7 Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Mon, 28 Oct 2024 13:41:38 -0400
Subject: [PATCH 029/184] Reorder pf in ffb_2d to favor speed over images

---
 src/aspire/basis/ffb_2d.py | 19 ++++++++++---------
 1 file changed, 10 insertions(+), 9 deletions(-)

diff --git a/src/aspire/basis/ffb_2d.py b/src/aspire/basis/ffb_2d.py
index 776c03cac8..8806bad78a 100644
--- a/src/aspire/basis/ffb_2d.py
+++ b/src/aspire/basis/ffb_2d.py
@@ -121,13 +121,13 @@ def _evaluate(self, v):
         n_r = self._precomp["freqs"].shape[1]
 
         # go through  each basis function and find corresponding coefficient
-        pf = xp.zeros((n_data, n_r, 2 * n_theta), dtype=complex_type(self.dtype))
+        pf = xp.zeros((2 * n_theta, n_data, n_r), dtype=complex_type(self.dtype))
 
         ind = 0
 
         idx = ind + np.arange(self.k_max[0], dtype=int)
 
-        pf[:, :, 0] = v[:, self._zero_angular_inds] @ self.radial_norm[idx]
+        pf[0] = v[:, self._zero_angular_inds] @ self.radial_norm[idx]
         ind = ind + idx.size
 
         ind_pos = ind
@@ -143,23 +143,24 @@ def _evaluate(self, v):
                 v_ell = 1j * v_ell
 
             pf_ell = v_ell @ self.radial_norm[idx]
-            pf[:, :, ell] = pf_ell
+            pf[ell] = pf_ell
 
             if np.mod(ell, 2) == 0:
-                pf[:, :, 2 * n_theta - ell] = pf_ell.conjugate()
+                pf[2 * n_theta - ell] = pf_ell.conjugate()
             else:
-                pf[:, :, 2 * n_theta - ell] = -pf_ell.conjugate()
+                pf[2 * n_theta - ell] = -pf_ell.conjugate()
 
             ind = ind + idx.size
             ind_pos = ind_pos + 2 * self.k_max[ell]
 
         # 1D inverse FFT in the degree of polar angle
-        pf = 2 * xp.pi * fft.ifft(pf, axis=2)
+        pf = 2 * xp.pi * fft.ifft(pf, axis=0)
 
         # Only need "positive" frequencies.
-        hsize = int(pf.shape[2] / 2)
-        pf = pf[:, :, 0:hsize]
-        pf *= self.gl_weighted_nodes[None, :, None]
+        hsize = int(pf.shape[0] / 2)
+        pf = pf[0:hsize]
+        pf *= self.gl_weighted_nodes[None, None, :]
+        pf = pf.transpose(1, 2, 0)
         pf = pf.reshape(n_data, n_r * n_theta)
 
         # perform inverse non-uniformly FFT transform back to 2D coordinate basis

From 0eede1769c96f6f2003163bda420c1bc4b40c7d8 Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Tue, 29 Oct 2024 15:04:45 -0400
Subject: [PATCH 030/184] fix shapes in FB_3D

---
 src/aspire/basis/fb_3d.py | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/src/aspire/basis/fb_3d.py b/src/aspire/basis/fb_3d.py
index 66e1ae7082..b2910c4262 100644
--- a/src/aspire/basis/fb_3d.py
+++ b/src/aspire/basis/fb_3d.py
@@ -148,8 +148,8 @@ def _evaluate(self, v):
             This is an array whose first dimensions equal `self.z` and the
             remaining dimensions correspond to dimensions two and higher of `v`.
         """
-        v = v.reshape(v.shape[0], -1)
-
+        stack_shape = v.shape[:-1]
+        v = v.reshape(-1, v.shape[-1])
         r_idx = self.basis_coords["r_idx"]
         ang_idx = self.basis_coords["ang_idx"]
         mask = self.basis_coords["mask"].flatten()
@@ -176,7 +176,7 @@ def _evaluate(self, v):
 
             ind_radial += len(idx_radial)
 
-        return x.reshape(v.shape[0], *self.sz)
+        return x.reshape(*stack_shape, *self.sz)
 
     def _evaluate_t(self, x):
         """
@@ -189,7 +189,8 @@ def _evaluate_t(self, x):
             equals `self.count` and whose remaining dimensions correspond
             to higher dimensions of `v`.
         """
-        x = x.reshape(x.shape[0], -1)
+        stack_shape = x.shape[: -self.ndim]
+        x = x.reshape(-1, np.prod(self.sz))
         r_idx = self.basis_coords["r_idx"]
         ang_idx = self.basis_coords["ang_idx"]
         mask = self.basis_coords["mask"].flatten()
@@ -216,4 +217,4 @@ def _evaluate_t(self, x):
 
             ind_radial += len(idx_radial)
 
-        return v
+        return v.reshape(*stack_shape, self.count)

From 9a8fcb34cfd65b4a8bf56d2aeda2f76987f8e333 Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Thu, 31 Oct 2024 11:10:37 -0400
Subject: [PATCH 031/184] remove unnecesary m_reshapes

---
 src/aspire/basis/ffb_3d.py | 20 +++++++++-----------
 1 file changed, 9 insertions(+), 11 deletions(-)

diff --git a/src/aspire/basis/ffb_3d.py b/src/aspire/basis/ffb_3d.py
index 7f0821b99a..47e8b2d088 100644
--- a/src/aspire/basis/ffb_3d.py
+++ b/src/aspire/basis/ffb_3d.py
@@ -60,15 +60,15 @@ def _precomp(self):
 
         r, wt_r = lgwt(n_r, 0.0, self.kcut, dtype=self.dtype)
         z, wt_z = lgwt(n_phi, -1, 1, dtype=self.dtype)
-        r = m_reshape(xp.asarray(r), (n_r, 1))
+        r = r.reshape(n_r, 1)
         rh = xp.asnumpy(r)
-        wt_r = m_reshape(xp.asarray(wt_r), (n_r, 1))
-        z = m_reshape(xp.asarray(z), (n_phi, 1))
-        wt_z = m_reshape(xp.asarray(wt_z), (n_phi, 1))
+        wt_r = xp.asarray(wt_r.reshape(n_r, 1))
+        z = xp.asarray(z.reshape(n_phi, 1))
+        wt_z = xp.asarray(wt_z.reshape(n_phi, 1))
         phi = xp.arccos(z)
         wt_phi = wt_z
         theta = 2 * xp.pi * xp.arange(n_theta, dtype=self.dtype).T / (2 * n_theta)
-        theta = m_reshape(theta, (n_theta, 1))
+        theta = theta.reshape(n_theta, 1)
 
         # evaluate basis function in the radial dimension
         radial_wtd = xp.zeros(
@@ -119,13 +119,12 @@ def _precomp(self):
 
         # evaluate basis function in the theta dimension
         ang_theta = xp.zeros((n_theta, 2 * self.ell_max + 1), dtype=theta.dtype)
-
         ang_theta[:, 0 : self.ell_max] = np.sqrt(2) * xp.sin(
-            theta @ m_reshape(xp.arange(self.ell_max, 0, -1), (1, self.ell_max))
+            theta @ xp.arange(self.ell_max, 0, -1).reshape(1, self.ell_max)
         )
         ang_theta[:, self.ell_max] = xp.ones(n_theta, dtype=theta.dtype)
         ang_theta[:, self.ell_max + 1 : 2 * self.ell_max + 1] = np.sqrt(2) * xp.cos(
-            theta @ m_reshape(xp.arange(1, self.ell_max + 1), (1, self.ell_max))
+            theta @ xp.arange(1, self.ell_max + 1).reshape(1, self.ell_max)
         )
 
         ang_theta_wtd = (2 * np.pi / n_theta) * ang_theta
@@ -133,6 +132,7 @@ def _precomp(self):
         theta_grid, phi_grid, r_grid = xp.meshgrid(
             theta.flatten(), phi.flatten(), r.flatten(), sparse=False, indexing="ij"
         )
+
         fourier_x = m_flatten(r_grid * xp.cos(theta_grid) * xp.sin(phi_grid))
         fourier_y = m_flatten(r_grid * xp.sin(theta_grid) * xp.sin(phi_grid))
         fourier_z = m_flatten(r_grid * xp.cos(phi_grid))
@@ -200,7 +200,6 @@ def _evaluate(self, v):
             radial_wtd = self._precomp["radial_wtd"][:, 0:k_max_ell, ell]
 
             ind = self._indices["ells"] == ell
-
             v_ell = m_reshape(v[:, ind].T, (k_max_ell, (2 * ell + 1) * n_data))
             v_ell = radial_wtd @ v_ell
             v_ell = m_reshape(v_ell, (n_r, 2 * ell + 1, n_data))
@@ -273,8 +272,7 @@ def _evaluate(self, v):
         pf = xp.moveaxis(pf, 0, -1)
 
         # perform inverse non-uniformly FFT transformation back to 3D rectangular coordinates
-        freqs = m_reshape(self._precomp["fourier_pts"], (3, n_r * n_theta * n_phi))
-        x = anufft(pf, freqs, self.sz, real=True)
+        x = anufft(pf, self._precomp["fourier_pts"], self.sz, real=True)
 
         # Roll, return the x with the last three dimensions as self.sz
         # Higher dimensions should be like v.

From 28eb9fb362a7ec6c62ad92f88c0f882d0c84a571 Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Thu, 7 Nov 2024 15:19:04 -0500
Subject: [PATCH 032/184] Refactor FFBBasis3D evaluate_t to remove m_reshape
 and speed up.

---
 src/aspire/basis/ffb_3d.py | 62 +++++++++++++++++---------------------
 1 file changed, 28 insertions(+), 34 deletions(-)

diff --git a/src/aspire/basis/ffb_3d.py b/src/aspire/basis/ffb_3d.py
index 47e8b2d088..cc119e8ef3 100644
--- a/src/aspire/basis/ffb_3d.py
+++ b/src/aspire/basis/ffb_3d.py
@@ -300,32 +300,30 @@ def _evaluate_t(self, x):
         n_phi = np.size(self._precomp["ang_phi_wtd_even"][0], 0)
         n_theta = np.size(self._precomp["ang_theta_wtd"], 0)
 
-        # resamping x in a polar Fourier gird using nonuniform discrete Fourier transform
+        # resamping x in a polar Fourier grid using nonuniform discrete Fourier transform
         pf = nufft(x, self._precomp["fourier_pts"])
-
-        pf = m_reshape(pf.T, (n_theta, n_phi * n_r * n_data))
+        pf = pf.reshape(n_data * n_r * n_phi, n_theta)
 
         # evaluate the theta parts
-        ang_theta_wtd_trans = self._precomp["ang_theta_wtd"].T
-        u_even = ang_theta_wtd_trans @ pf.real
-        u_odd = ang_theta_wtd_trans @ pf.imag
-
-        u_even = m_reshape(u_even, (2 * self.ell_max + 1, n_phi, n_r, n_data))
-        u_odd = m_reshape(u_odd, (2 * self.ell_max + 1, n_phi, n_r, n_data))
+        u_even = pf.real @ self._precomp["ang_theta_wtd"]
+        u_odd = pf.imag @ self._precomp["ang_theta_wtd"]
 
-        u_even = u_even.transpose((1, 2, 3, 0))
-        u_odd = u_odd.transpose((1, 2, 3, 0))
+        u_even = u_even.reshape(n_data, n_r, n_phi, 2 * self.ell_max + 1)
+        u_odd = u_odd.reshape(n_data, n_r, n_phi, 2 * self.ell_max + 1)
 
         w_even = xp.zeros(
-            (int(np.floor(self.ell_max / 2) + 1), n_r, 2 * self.ell_max + 1, n_data),
+            (n_data, 2 * self.ell_max + 1, n_r, int(np.floor(self.ell_max / 2) + 1)),
             dtype=x.dtype,
         )
         w_odd = xp.zeros(
-            (int(np.ceil(self.ell_max / 2)), n_r, 2 * self.ell_max + 1, n_data),
+            (n_data, 2 * self.ell_max + 1, n_r, int(np.ceil(self.ell_max / 2))),
             dtype=x.dtype,
         )
 
-        # evaluate the phi parts
+        # Transpose and copy as contiguous for faster slicing and matrix multiplication.
+        u_even = np.ascontiguousarray(u_even.transpose(3, 0, 1, 2))
+        u_odd = np.ascontiguousarray(u_odd.transpose(3, 0, 1, 2))
+
         for m in range(0, self.ell_max + 1):
             ang_phi_wtd_m_even = self._precomp["ang_phi_wtd_even"][m]
             ang_phi_wtd_m_odd = self._precomp["ang_phi_wtd_odd"][m]
@@ -339,24 +337,22 @@ def _evaluate_t(self, x):
                 sgns = (1, -1)
 
             for sgn in sgns:
-                u_m_even = u_even[:, :, :, self.ell_max + sgn * m]
-                u_m_odd = u_odd[:, :, :, self.ell_max + sgn * m]
+                u_m_even = u_even[self.ell_max + sgn * m]
+                u_m_odd = u_odd[self.ell_max + sgn * m]
 
-                u_m_even = m_reshape(u_m_even, (n_phi, n_r * n_data))
-                u_m_odd = m_reshape(u_m_odd, (n_phi, n_r * n_data))
+                u_m_even = u_m_even.reshape(n_data * n_r, n_phi)
+                u_m_odd = u_m_odd.reshape(n_data * n_r, n_phi)
 
-                w_m_even = ang_phi_wtd_m_even.T @ u_m_even
-                w_m_odd = ang_phi_wtd_m_odd.T @ u_m_odd
+                w_m_even = u_m_even @ ang_phi_wtd_m_even
+                w_m_odd = u_m_odd @ ang_phi_wtd_m_odd
 
-                w_m_even = m_reshape(w_m_even, (n_even_ell, n_r, n_data))
-                w_m_odd = m_reshape(w_m_odd, (n_odd_ell, n_r, n_data))
-                end = np.size(w_even, 0)
-                w_even[end - n_even_ell : end, :, self.ell_max + sgn * m, :] = w_m_even
-                end = np.size(w_odd, 0)
-                w_odd[end - n_odd_ell : end, :, self.ell_max + sgn * m, :] = w_m_odd
+                w_m_even = w_m_even.reshape(n_data, n_r, n_even_ell)
+                w_m_odd = w_m_odd.reshape(n_data, n_r, n_odd_ell)
 
-        w_even = w_even.transpose((1, 2, 3, 0))
-        w_odd = w_odd.transpose((1, 2, 3, 0))
+                end = np.size(w_even, -1)
+                w_even[:, self.ell_max + sgn * m, :, end - n_even_ell : end] = w_m_even
+                end = np.size(w_odd, -1)
+                w_odd[:, self.ell_max + sgn * m, :, end - n_odd_ell : end] = w_m_odd
 
         # evaluate the radial parts
         v = xp.zeros((n_data, self.count), dtype=x.dtype)
@@ -379,15 +375,13 @@ def _evaluate_t(self, x):
                     int((ell - 1) / 2),
                 ]
 
-            v_ell = m_reshape(v_ell, (n_r, (2 * ell + 1) * n_data))
-
-            v_ell = radial_wtd.T @ v_ell
-
-            v_ell = m_reshape(v_ell, (k_max_ell * (2 * ell + 1), n_data))
+            v_ell = v_ell.reshape(n_data * (2 * ell + 1), n_r)
+            v_ell = v_ell @ radial_wtd
+            v_ell = v_ell.reshape(n_data, (2 * ell + 1) * k_max_ell)
 
             # TODO: Fix this to avoid lookup each time.
             ind = self._indices["ells"] == ell
-            v[:, ind] = v_ell.T
+            v[:, ind] = v_ell
 
         # Roll dimensions, last dimension should be self.count,
         # Higher dimensions like x.

From 581497e7f92da219dc6a70866e70579eed4d2d18 Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Thu, 7 Nov 2024 15:29:46 -0500
Subject: [PATCH 033/184] missing xp array

---
 src/aspire/basis/ffb_3d.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/aspire/basis/ffb_3d.py b/src/aspire/basis/ffb_3d.py
index cc119e8ef3..51b80acf8c 100644
--- a/src/aspire/basis/ffb_3d.py
+++ b/src/aspire/basis/ffb_3d.py
@@ -60,7 +60,7 @@ def _precomp(self):
 
         r, wt_r = lgwt(n_r, 0.0, self.kcut, dtype=self.dtype)
         z, wt_z = lgwt(n_phi, -1, 1, dtype=self.dtype)
-        r = r.reshape(n_r, 1)
+        r = xp.asarray(r.reshape(n_r, 1))
         rh = xp.asnumpy(r)
         wt_r = xp.asarray(wt_r.reshape(n_r, 1))
         z = xp.asarray(z.reshape(n_phi, 1))

From 7325ebb3677b63914b64e566b9de667476d3ba57 Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Thu, 7 Nov 2024 16:03:16 -0500
Subject: [PATCH 034/184] xp.contiguousarray

---
 src/aspire/basis/ffb_3d.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/aspire/basis/ffb_3d.py b/src/aspire/basis/ffb_3d.py
index 51b80acf8c..68122aceb3 100644
--- a/src/aspire/basis/ffb_3d.py
+++ b/src/aspire/basis/ffb_3d.py
@@ -321,8 +321,8 @@ def _evaluate_t(self, x):
         )
 
         # Transpose and copy as contiguous for faster slicing and matrix multiplication.
-        u_even = np.ascontiguousarray(u_even.transpose(3, 0, 1, 2))
-        u_odd = np.ascontiguousarray(u_odd.transpose(3, 0, 1, 2))
+        u_even = xp.ascontiguousarray(u_even.transpose(3, 0, 1, 2))
+        u_odd = xp.ascontiguousarray(u_odd.transpose(3, 0, 1, 2))
 
         for m in range(0, self.ell_max + 1):
             ang_phi_wtd_m_even = self._precomp["ang_phi_wtd_even"][m]

From 72f0675172fe8363eb4e8666c489c91d8106f1cc Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Mon, 11 Nov 2024 10:54:56 -0500
Subject: [PATCH 035/184] Remove m_funcs from ffb_3d.evaluate.

---
 src/aspire/basis/ffb_3d.py | 52 ++++++++++++++++++--------------------
 1 file changed, 24 insertions(+), 28 deletions(-)

diff --git a/src/aspire/basis/ffb_3d.py b/src/aspire/basis/ffb_3d.py
index 68122aceb3..80e887e7e3 100644
--- a/src/aspire/basis/ffb_3d.py
+++ b/src/aspire/basis/ffb_3d.py
@@ -6,7 +6,6 @@
 from aspire.basis.basis_utils import lgwt, norm_assoc_legendre, sph_bessel
 from aspire.nufft import anufft, nufft
 from aspire.numeric import xp
-from aspire.utils.matlab_compat import m_flatten, m_reshape
 
 logger = logging.getLogger(__name__)
 
@@ -181,15 +180,15 @@ def _evaluate(self, v):
 
         u_even = xp.zeros(
             (
-                n_r,
-                int(2 * self.ell_max + 1),
-                n_data,
                 int(np.floor(self.ell_max / 2) + 1),
+                n_data,
+                int(2 * self.ell_max + 1),
+                n_r,
             ),
             dtype=v.dtype,
         )
         u_odd = xp.zeros(
-            (n_r, int(2 * self.ell_max + 1), n_data, int(np.ceil(self.ell_max / 2))),
+            (int(np.ceil(self.ell_max / 2)), n_data, int(2 * self.ell_max + 1), n_r),
             dtype=v.dtype,
         )
 
@@ -200,29 +199,27 @@ def _evaluate(self, v):
             radial_wtd = self._precomp["radial_wtd"][:, 0:k_max_ell, ell]
 
             ind = self._indices["ells"] == ell
-            v_ell = m_reshape(v[:, ind].T, (k_max_ell, (2 * ell + 1) * n_data))
-            v_ell = radial_wtd @ v_ell
-            v_ell = m_reshape(v_ell, (n_r, 2 * ell + 1, n_data))
+            v_ell = v[:, ind].reshape((2 * ell + 1) * n_data, k_max_ell)
+            v_ell = v_ell @ radial_wtd.T
+            v_ell = v_ell.reshape(n_data, 2 * ell + 1, n_r)
 
             if np.mod(ell, 2) == 0:
                 u_even[
+                    int(ell / 2),
                     :,
                     int(self.ell_max - ell) : int(self.ell_max + ell + 1),
                     :,
-                    int(ell / 2),
                 ] = v_ell
             else:
                 u_odd[
+                    int((ell - 1) / 2),
                     :,
                     int(self.ell_max - ell) : int(self.ell_max + ell + 1),
                     :,
-                    int((ell - 1) / 2),
                 ] = v_ell
 
-        u_even = u_even.transpose((3, 0, 1, 2))
-        u_odd = u_odd.transpose((3, 0, 1, 2))
-        w_even = xp.zeros((n_phi, n_r, n_data, 2 * self.ell_max + 1), dtype=v.dtype)
-        w_odd = xp.zeros((n_phi, n_r, n_data, 2 * self.ell_max + 1), dtype=v.dtype)
+        w_even = xp.zeros((2 * self.ell_max + 1, n_phi, n_data, n_r), dtype=v.dtype)
+        w_odd = xp.zeros((2 * self.ell_max + 1, n_phi, n_data, n_r), dtype=v.dtype)
 
         # evaluate the phi parts
         for m in range(0, self.ell_max + 1):
@@ -243,33 +240,32 @@ def _evaluate(self, v):
                 end = np.size(u_odd, 0)
                 u_m_odd = u_odd[end - n_odd_ell : end, :, self.ell_max + sgn * m, :]
 
-                u_m_even = m_reshape(u_m_even, (n_even_ell, n_r * n_data))
-                u_m_odd = m_reshape(u_m_odd, (n_odd_ell, n_r * n_data))
+                u_m_even = u_m_even.reshape(n_even_ell, n_data * n_r)
+                u_m_odd = u_m_odd.reshape(n_odd_ell, n_data * n_r)
 
                 w_m_even = ang_phi_wtd_m_even @ u_m_even
                 w_m_odd = ang_phi_wtd_m_odd @ u_m_odd
 
-                w_m_even = m_reshape(w_m_even, (n_phi, n_r, n_data))
-                w_m_odd = m_reshape(w_m_odd, (n_phi, n_r, n_data))
+                w_m_even = w_m_even.reshape(n_phi, n_data, n_r)
+                w_m_odd = w_m_odd.reshape(n_phi, n_data, n_r)
 
-                w_even[:, :, :, self.ell_max + sgn * m] = w_m_even
-                w_odd[:, :, :, self.ell_max + sgn * m] = w_m_odd
+                w_even[self.ell_max + sgn * m] = w_m_even
+                w_odd[self.ell_max + sgn * m] = w_m_odd
 
-        w_even = w_even.transpose((3, 0, 1, 2))
-        w_odd = w_odd.transpose((3, 0, 1, 2))
+        w_even = w_even.transpose((2, 3, 1, 0))
+        w_odd = w_odd.transpose((2, 3, 1, 0))
         u_even = w_even
         u_odd = w_odd
 
-        u_even = m_reshape(u_even, (2 * self.ell_max + 1, n_phi * n_r * n_data))
-        u_odd = m_reshape(u_odd, (2 * self.ell_max + 1, n_phi * n_r * n_data))
+        u_even = u_even.reshape(n_data * n_r * n_phi, 2 * self.ell_max + 1)
+        u_odd = u_odd.reshape(n_data * n_r * n_phi, 2 * self.ell_max + 1)
 
         # evaluate the theta parts
-        w_even = self._precomp["ang_theta_wtd"] @ u_even
-        w_odd = self._precomp["ang_theta_wtd"] @ u_odd
+        w_even = u_even @ self._precomp["ang_theta_wtd"].T
+        w_odd = u_odd @ self._precomp["ang_theta_wtd"].T
 
         pf = w_even + 1j * w_odd
-        pf = m_reshape(pf, (n_theta * n_phi * n_r, n_data))
-        pf = xp.moveaxis(pf, 0, -1)
+        pf = pf.reshape(n_data, n_r * n_phi * n_theta)
 
         # perform inverse non-uniformly FFT transformation back to 3D rectangular coordinates
         x = anufft(pf, self._precomp["fourier_pts"], self.sz, real=True)

From 1c89ff8f4d3daadf914e5c48087577e358de9444 Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Mon, 11 Nov 2024 11:07:15 -0500
Subject: [PATCH 036/184] remove m_funcs from fourier_pts

---
 src/aspire/basis/ffb_3d.py | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/src/aspire/basis/ffb_3d.py b/src/aspire/basis/ffb_3d.py
index 80e887e7e3..53155253d7 100644
--- a/src/aspire/basis/ffb_3d.py
+++ b/src/aspire/basis/ffb_3d.py
@@ -128,13 +128,13 @@ def _precomp(self):
 
         ang_theta_wtd = (2 * np.pi / n_theta) * ang_theta
 
-        theta_grid, phi_grid, r_grid = xp.meshgrid(
-            theta.flatten(), phi.flatten(), r.flatten(), sparse=False, indexing="ij"
+        r_grid, phi_grid, theta_grid = xp.meshgrid(
+            r.flatten(), phi.flatten(), theta.flatten(), sparse=False, indexing="ij"
         )
 
-        fourier_x = m_flatten(r_grid * xp.cos(theta_grid) * xp.sin(phi_grid))
-        fourier_y = m_flatten(r_grid * xp.sin(theta_grid) * xp.sin(phi_grid))
-        fourier_z = m_flatten(r_grid * xp.cos(phi_grid))
+        fourier_x = (r_grid * xp.cos(theta_grid) * xp.sin(phi_grid)).flatten()
+        fourier_y = (r_grid * xp.sin(theta_grid) * xp.sin(phi_grid)).flatten()
+        fourier_z = (r_grid * xp.cos(phi_grid)).flatten()
         fourier_pts = (
             2
             * xp.pi

From fd7ba2ab5d30385ce9f191e5100ad26aa9ba24bb Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Mon, 11 Nov 2024 11:36:02 -0500
Subject: [PATCH 037/184] remove im_to_vec vec_to_im

---
 src/aspire/utils/__init__.py |  2 --
 src/aspire/utils/matrix.py   | 28 ----------------------------
 tests/test_matrix.py         | 26 --------------------------
 3 files changed, 56 deletions(-)

diff --git a/src/aspire/utils/__init__.py b/src/aspire/utils/__init__.py
index c4b46bcdec..b1fc612ff8 100644
--- a/src/aspire/utils/__init__.py
+++ b/src/aspire/utils/__init__.py
@@ -42,7 +42,6 @@
     best_rank1_approximation,
     eigs,
     fix_signs,
-    im_to_vec,
     make_psd,
     make_symmat,
     mat_to_vec,
@@ -52,7 +51,6 @@
     symmat_to_vec,
     symmat_to_vec_iso,
     unroll_dim,
-    vec_to_im,
     vec_to_mat,
     vec_to_symmat,
     vec_to_symmat_iso,
diff --git a/src/aspire/utils/matrix.py b/src/aspire/utils/matrix.py
index f33aa47b23..49335f3968 100644
--- a/src/aspire/utils/matrix.py
+++ b/src/aspire/utils/matrix.py
@@ -35,20 +35,6 @@ def roll_dim(X, dim):
     return Y
 
 
-def im_to_vec(im):
-    """
-    Roll up images into vectors
-
-    :param im: An N-by-N-by-... array.
-    :return: An N^2-by-... array.
-    """
-    shape = im.shape
-    assert im.ndim >= 2, "Array should have at least 2 dimensions"
-    assert shape[0] == shape[1], "Array should have first 2 dimensions identical"
-
-    return m_reshape(im, (shape[0] ** 2,) + (shape[2:]))
-
-
 def vol_to_vec(X):
     """
     Roll up volumes into vectors
@@ -65,20 +51,6 @@ def vol_to_vec(X):
     return m_reshape(X, (shape[0] ** 3,) + (shape[3:]))
 
 
-def vec_to_im(X):
-    """
-    Unroll vectors to images
-
-    :param X: N^2-by-... array.
-    :return: An N-by-N-by-... array.
-    """
-    shape = X.shape
-    N = round(shape[0] ** (1 / 2))
-    assert N**2 == shape[0], "First dimension of X must be square"
-
-    return m_reshape(X, (N, N) + (shape[1:]))
-
-
 def vec_to_vol(X):
     """
     Unroll vectors to volumes
diff --git a/tests/test_matrix.py b/tests/test_matrix.py
index 728200b9b3..fe1a713a49 100644
--- a/tests/test_matrix.py
+++ b/tests/test_matrix.py
@@ -8,7 +8,6 @@
     Rotation,
     best_rank1_approximation,
     fix_signs,
-    im_to_vec,
     mat_to_vec,
     mean_aligned_angular_distance,
     nearest_rotations,
@@ -17,7 +16,6 @@
     symmat_to_vec_iso,
     unroll_dim,
     utest_tolerance,
-    vec_to_im,
     vec_to_symmat,
     vec_to_symmat_iso,
     vec_to_vol,
@@ -59,30 +57,6 @@ def testRollDims(self):
         # The values should still be filled in with the first axis values changing fastest
         self.assertTrue(np.allclose(m2[:, 0, 0, 0, 0], np.array([1, 2, 3, 4, 5])))
 
-    def testImToVec1(self):
-        m = np.empty((3, 3, 10))
-        m2 = im_to_vec(m)
-
-        self.assertEqual(m2.shape, (9, 10))
-
-    def testImToVec2(self):
-        m = np.empty((3, 3))
-        m2 = im_to_vec(m)
-
-        self.assertEqual(m2.shape, (9,))
-
-    def testVecToIm1(self):
-        m = np.empty((25, 10))
-        m2 = vec_to_im(m)
-
-        self.assertEqual(m2.shape, (5, 5, 10))
-
-    def testVecToIm2(self):
-        m = np.empty((16,))
-        m2 = vec_to_im(m)
-
-        self.assertEqual(m2.shape, (4, 4))
-
     def testVolToVec1(self):
         m = np.empty((3, 3, 3, 10))
         m2 = vol_to_vec(m)

From e1894d54416c5c2bc96272fc4da00d322685fcd9 Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Tue, 12 Nov 2024 10:50:08 -0500
Subject: [PATCH 038/184] remove m_reshape from reconstruction/kernel

---
 src/aspire/reconstruction/kernel.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/src/aspire/reconstruction/kernel.py b/src/aspire/reconstruction/kernel.py
index 9a86a244f2..4a68ed49c3 100644
--- a/src/aspire/reconstruction/kernel.py
+++ b/src/aspire/reconstruction/kernel.py
@@ -3,7 +3,6 @@
 import numpy as np
 
 from aspire.numeric import fft
-from aspire.utils.matlab_compat import m_reshape
 from aspire.volume import Volume
 
 logger = logging.getLogger(__name__)
@@ -59,10 +58,10 @@ def circularize_1d(self, kernel, dim):
         mult_shape[dim] = N
         mult_shape = tuple(mult_shape)
 
-        mult = m_reshape((np.arange(N, dtype=self.dtype) / N), mult_shape)
+        mult = (np.arange(N, dtype=self.dtype) / N).reshape(mult_shape)
         kernel_circ = mult * top
 
-        mult = m_reshape((np.arange(N, 0, -1, dtype=self.dtype) / N), mult_shape)
+        mult = (np.arange(N, 0, -1, dtype=self.dtype) / N).reshape(mult_shape)
         kernel_circ += mult * bottom
 
         return fft.fftshift(kernel_circ, dim)

From 7cca4276411a6a500488766f92ce5c91fc40ac0f Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Tue, 12 Nov 2024 14:51:53 -0500
Subject: [PATCH 039/184] remove vec_to_mat from qr_vols_forward.

---
 src/aspire/volume/volume.py | 11 ++++-------
 1 file changed, 4 insertions(+), 7 deletions(-)

diff --git a/src/aspire/volume/volume.py b/src/aspire/volume/volume.py
index eae985a794..60c8bef033 100644
--- a/src/aspire/volume/volume.py
+++ b/src/aspire/volume/volume.py
@@ -16,9 +16,7 @@
     crop_pad_3d,
     grid_2d,
     grid_3d,
-    mat_to_vec,
     rename_with_timestamp,
-    vec_to_mat,
 )
 from aspire.volume import IdentitySymmetryGroup, SymmetryGroup
 
@@ -40,17 +38,16 @@ def qr_vols_forward(sim, s, n, vols, k):
     ims = np.zeros((k, n, sim.L, sim.L), dtype=vols.dtype)
     for ell in range(k):
         ims[ell] = sim.vol_forward(vols[ell], s, n).asnumpy()
-
-    ims = np.swapaxes(ims, 1, 3)
-    ims = np.swapaxes(ims, 0, 2)
+    ims = ims.transpose((2, 3, 0, 1))
 
     Q_vecs = np.zeros((sim.L**2, k, n), dtype=vols.dtype)
     Rs = np.zeros((k, k, n), dtype=vols.dtype)
 
-    im_vecs = mat_to_vec(ims)
+    im_vecs = ims.reshape(sim.L**2, k, n)
+
     for i in range(n):
         Q_vecs[:, :, i], Rs[:, :, i] = qr(im_vecs[:, :, i])
-    Qs = vec_to_mat(Q_vecs)
+    Qs = Q_vecs.reshape(sim.L, sim.L, k, n)
 
     return Qs, Rs
 

From 580348524cc5c17726769bbc9847e3e0923c3fda Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Tue, 12 Nov 2024 15:36:45 -0500
Subject: [PATCH 040/184] remove vec_to_vol and vol_to_vec

---
 src/aspire/utils/__init__.py |  2 --
 src/aspire/utils/matrix.py   | 30 ------------------------------
 tests/test_matrix.py         | 26 --------------------------
 3 files changed, 58 deletions(-)

diff --git a/src/aspire/utils/__init__.py b/src/aspire/utils/__init__.py
index b1fc612ff8..83a681373d 100644
--- a/src/aspire/utils/__init__.py
+++ b/src/aspire/utils/__init__.py
@@ -54,9 +54,7 @@
     vec_to_mat,
     vec_to_symmat,
     vec_to_symmat_iso,
-    vec_to_vol,
     vecmat_to_volmat,
-    vol_to_vec,
     volmat_to_vecmat,
 )
 from .multiprocessing import (
diff --git a/src/aspire/utils/matrix.py b/src/aspire/utils/matrix.py
index 49335f3968..a8415929c0 100644
--- a/src/aspire/utils/matrix.py
+++ b/src/aspire/utils/matrix.py
@@ -35,36 +35,6 @@ def roll_dim(X, dim):
     return Y
 
 
-def vol_to_vec(X):
-    """
-    Roll up volumes into vectors
-
-    :param X: N-by-N-by-N-by-... array.
-    :return: An N^3-by-... array.
-    """
-    shape = X.shape
-    assert X.ndim >= 3, "Array should have at least 3 dimensions"
-    assert (
-        shape[0] == shape[1] == shape[2]
-    ), "Array should have first 3 dimensions identical"
-
-    return m_reshape(X, (shape[0] ** 3,) + (shape[3:]))
-
-
-def vec_to_vol(X):
-    """
-    Unroll vectors to volumes
-
-    :param X: N^3-by-... array.
-    :return: An N-by-N-by-N-by-... array.
-    """
-    shape = X.shape
-    N = round(shape[0] ** (1 / 3))
-    assert N**3 == shape[0], "First dimension of X must be cubic"
-
-    return m_reshape(X, (N, N, N) + (shape[1:]))
-
-
 def vecmat_to_volmat(X):
     """
     Roll up vector matrices into volume matrices
diff --git a/tests/test_matrix.py b/tests/test_matrix.py
index fe1a713a49..5274125ac0 100644
--- a/tests/test_matrix.py
+++ b/tests/test_matrix.py
@@ -18,9 +18,7 @@
     utest_tolerance,
     vec_to_symmat,
     vec_to_symmat_iso,
-    vec_to_vol,
     vecmat_to_volmat,
-    vol_to_vec,
     volmat_to_vecmat,
 )
 
@@ -57,30 +55,6 @@ def testRollDims(self):
         # The values should still be filled in with the first axis values changing fastest
         self.assertTrue(np.allclose(m2[:, 0, 0, 0, 0], np.array([1, 2, 3, 4, 5])))
 
-    def testVolToVec1(self):
-        m = np.empty((3, 3, 3, 10))
-        m2 = vol_to_vec(m)
-
-        self.assertEqual(m2.shape, (27, 10))
-
-    def testVolToVec2(self):
-        m = np.empty((3, 3, 3))
-        m2 = vol_to_vec(m)
-
-        self.assertEqual(m2.shape, (27,))
-
-    def testVecToVol1(self):
-        m = np.empty((27, 10))
-        m2 = vec_to_vol(m)
-
-        self.assertEqual(m2.shape, (3, 3, 3, 10))
-
-    def testVecToVol2(self):
-        m = np.empty((27,))
-        m2 = vec_to_vol(m)
-
-        self.assertEqual(m2.shape, (3, 3, 3))
-
     def testVecmatToVolmat(self):
         m = np.empty((8, 27, 10))
         m2 = vecmat_to_volmat(m)

From a928c356532bbd46944ec3f0f61b02f9ee56da0b Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Thu, 14 Nov 2024 10:28:58 -0500
Subject: [PATCH 041/184] No need to keep track of stack_shape.

---
 src/aspire/basis/fb_3d.py | 7 ++-----
 1 file changed, 2 insertions(+), 5 deletions(-)

diff --git a/src/aspire/basis/fb_3d.py b/src/aspire/basis/fb_3d.py
index b2910c4262..684935ada1 100644
--- a/src/aspire/basis/fb_3d.py
+++ b/src/aspire/basis/fb_3d.py
@@ -148,8 +148,6 @@ def _evaluate(self, v):
             This is an array whose first dimensions equal `self.z` and the
             remaining dimensions correspond to dimensions two and higher of `v`.
         """
-        stack_shape = v.shape[:-1]
-        v = v.reshape(-1, v.shape[-1])
         r_idx = self.basis_coords["r_idx"]
         ang_idx = self.basis_coords["ang_idx"]
         mask = self.basis_coords["mask"].flatten()
@@ -176,7 +174,7 @@ def _evaluate(self, v):
 
             ind_radial += len(idx_radial)
 
-        return x.reshape(*stack_shape, *self.sz)
+        return x.reshape(-1, *self.sz)
 
     def _evaluate_t(self, x):
         """
@@ -189,7 +187,6 @@ def _evaluate_t(self, x):
             equals `self.count` and whose remaining dimensions correspond
             to higher dimensions of `v`.
         """
-        stack_shape = x.shape[: -self.ndim]
         x = x.reshape(-1, np.prod(self.sz))
         r_idx = self.basis_coords["r_idx"]
         ang_idx = self.basis_coords["ang_idx"]
@@ -217,4 +214,4 @@ def _evaluate_t(self, x):
 
             ind_radial += len(idx_radial)
 
-        return v.reshape(*stack_shape, self.count)
+        return v

From 8f2b459e84bf06ff760f2a7735e4fe62ac5df7dd Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Thu, 14 Nov 2024 10:31:44 -0500
Subject: [PATCH 042/184] reshape on gpu

---
 src/aspire/basis/ffb_3d.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/aspire/basis/ffb_3d.py b/src/aspire/basis/ffb_3d.py
index 53155253d7..350a146e06 100644
--- a/src/aspire/basis/ffb_3d.py
+++ b/src/aspire/basis/ffb_3d.py
@@ -59,11 +59,11 @@ def _precomp(self):
 
         r, wt_r = lgwt(n_r, 0.0, self.kcut, dtype=self.dtype)
         z, wt_z = lgwt(n_phi, -1, 1, dtype=self.dtype)
-        r = xp.asarray(r.reshape(n_r, 1))
+        r = xp.asarray(r).reshape(n_r, 1)
         rh = xp.asnumpy(r)
-        wt_r = xp.asarray(wt_r.reshape(n_r, 1))
-        z = xp.asarray(z.reshape(n_phi, 1))
-        wt_z = xp.asarray(wt_z.reshape(n_phi, 1))
+        wt_r = xp.asarray(wt_r).reshape(n_r, 1)
+        z = xp.asarray(z).reshape(n_phi, 1)
+        wt_z = xp.asarray(wt_z).reshape(n_phi, 1)
         phi = xp.arccos(z)
         wt_phi = wt_z
         theta = 2 * xp.pi * xp.arange(n_theta, dtype=self.dtype).T / (2 * n_theta)

From e01619f267bbd0cb57dce73c8cda5519e492a09e Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Thu, 14 Nov 2024 14:25:12 -0500
Subject: [PATCH 043/184] Bring stack_shape back for FB_3D _evaluate(_t). Used
 in Cov3D. Also fix dtype bug in mean kernel.

---
 src/aspire/basis/fb_3d.py      | 7 +++++--
 src/aspire/covariance/covar.py | 1 +
 2 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/src/aspire/basis/fb_3d.py b/src/aspire/basis/fb_3d.py
index 684935ada1..b2910c4262 100644
--- a/src/aspire/basis/fb_3d.py
+++ b/src/aspire/basis/fb_3d.py
@@ -148,6 +148,8 @@ def _evaluate(self, v):
             This is an array whose first dimensions equal `self.z` and the
             remaining dimensions correspond to dimensions two and higher of `v`.
         """
+        stack_shape = v.shape[:-1]
+        v = v.reshape(-1, v.shape[-1])
         r_idx = self.basis_coords["r_idx"]
         ang_idx = self.basis_coords["ang_idx"]
         mask = self.basis_coords["mask"].flatten()
@@ -174,7 +176,7 @@ def _evaluate(self, v):
 
             ind_radial += len(idx_radial)
 
-        return x.reshape(-1, *self.sz)
+        return x.reshape(*stack_shape, *self.sz)
 
     def _evaluate_t(self, x):
         """
@@ -187,6 +189,7 @@ def _evaluate_t(self, x):
             equals `self.count` and whose remaining dimensions correspond
             to higher dimensions of `v`.
         """
+        stack_shape = x.shape[: -self.ndim]
         x = x.reshape(-1, np.prod(self.sz))
         r_idx = self.basis_coords["r_idx"]
         ang_idx = self.basis_coords["ang_idx"]
@@ -214,4 +217,4 @@ def _evaluate_t(self, x):
 
             ind_radial += len(idx_radial)
 
-        return v
+        return v.reshape(*stack_shape, self.count)
diff --git a/src/aspire/covariance/covar.py b/src/aspire/covariance/covar.py
index 89977052fe..9a399231f3 100644
--- a/src/aspire/covariance/covar.py
+++ b/src/aspire/covariance/covar.py
@@ -154,6 +154,7 @@ def apply_kernel(self, coef, kernel=None, packed=False):
         result = self.basis.mat_evaluate_t(
             kernel.convolve_volume_matrix(self.basis.mat_evaluate(coef))
         )
+
         return symmat_to_vec_iso(result) if packed else result
 
     def src_backward(self, mean_vol, noise_variance, shrink_method=None):

From e2b6039d88f6cc6cf8efc9561aac22494e5ccada Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Thu, 14 Nov 2024 15:12:26 -0500
Subject: [PATCH 044/184] blank line

---
 src/aspire/covariance/covar.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/src/aspire/covariance/covar.py b/src/aspire/covariance/covar.py
index 9a399231f3..89977052fe 100644
--- a/src/aspire/covariance/covar.py
+++ b/src/aspire/covariance/covar.py
@@ -154,7 +154,6 @@ def apply_kernel(self, coef, kernel=None, packed=False):
         result = self.basis.mat_evaluate_t(
             kernel.convolve_volume_matrix(self.basis.mat_evaluate(coef))
         )
-
         return symmat_to_vec_iso(result) if packed else result
 
     def src_backward(self, mean_vol, noise_variance, shrink_method=None):

From fcc03275cfa583344aaf49c80f0a1ef26313904c Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Wed, 20 Nov 2024 14:01:13 -0500
Subject: [PATCH 045/184] remove redundant reshape

---
 src/aspire/basis/fb_2d.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/aspire/basis/fb_2d.py b/src/aspire/basis/fb_2d.py
index 39b2207b62..e737f0fa8b 100644
--- a/src/aspire/basis/fb_2d.py
+++ b/src/aspire/basis/fb_2d.py
@@ -183,7 +183,7 @@ def _evaluate(self, v):
             dimensions correspond to first dimensions of `v`.
         """
         # Transpose here once, instead of several times below  #RCOPT
-        v = v.reshape(-1, self.count).T
+        v = v.T
 
         r_idx = self.basis_coords["r_idx"]
         ang_idx = self.basis_coords["ang_idx"]

From 77b76ba3957b3b1f49f74c575b02f2c480417cfb Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Mon, 4 Nov 2024 14:50:04 -0500
Subject: [PATCH 046/184] Initial grid cache implementation

---
 src/aspire/utils/coor_trans.py |  4 ++++
 src/aspire/utils/misc.py       | 18 ++++++++++--------
 tests/test_utils.py            |  4 ++--
 3 files changed, 16 insertions(+), 10 deletions(-)

diff --git a/src/aspire/utils/coor_trans.py b/src/aspire/utils/coor_trans.py
index b33098d3b4..8d5cb0ce1b 100644
--- a/src/aspire/utils/coor_trans.py
+++ b/src/aspire/utils/coor_trans.py
@@ -4,6 +4,7 @@
 
 import logging
 import math
+from functools import cache
 
 import numpy as np
 from numpy.linalg import norm
@@ -83,6 +84,7 @@ def _mgrid_slice(n, shifted, normalized):
     return slice(start, end, num_points)
 
 
+@cache
 def grid_1d(n, shifted=False, normalized=True, dtype=np.float32):
     """
     Generate one dimensional grid.
@@ -98,6 +100,7 @@ def grid_1d(n, shifted=False, normalized=True, dtype=np.float32):
     return {"x": x, "r": r}
 
 
+@cache
 def grid_2d(n, shifted=False, normalized=True, indexing="yx", dtype=np.float32):
     """
     Generate two dimensional grid.
@@ -124,6 +127,7 @@ def grid_2d(n, shifted=False, normalized=True, indexing="yx", dtype=np.float32):
     return {"x": x, "y": y, "phi": phi, "r": r}
 
 
+@cache
 def grid_3d(n, shifted=False, normalized=True, indexing="zyx", dtype=np.float32):
     """
     Generate three dimensional grid.
diff --git a/src/aspire/utils/misc.py b/src/aspire/utils/misc.py
index 5485c42788..7ce6602f86 100644
--- a/src/aspire/utils/misc.py
+++ b/src/aspire/utils/misc.py
@@ -170,15 +170,15 @@ def gaussian_2d(size, mu=(0, 0), sigma=(1, 1), indexing="yx", dtype=np.float64):
     # Construct centered mesh
     g = grid_2d(size, shifted=False, normalized=False, indexing=indexing, dtype=dtype)
 
+    X, Y = g["x"], g["y"]
+
     if indexing == "yx":
         mu, sigma = mu[::-1], sigma[::-1]
-        g["x"], g["y"] = g["y"], g["x"]
+        X, Y = Y, X
     elif indexing != "xy":
         raise ValueError("Indexing must be `yx` or `xy`.")
 
-    p = (g["x"] - mu[0]) ** 2 / (2 * sigma[0] ** 2) + (g["y"] - mu[1]) ** 2 / (
-        2 * sigma[1] ** 2
-    )
+    p = (X - mu[0]) ** 2 / (2 * sigma[0] ** 2) + (Y - mu[1]) ** 2 / (2 * sigma[1] ** 2)
 
     return np.exp(-p).astype(dtype, copy=False)
 
@@ -216,16 +216,18 @@ def gaussian_3d(size, mu=(0, 0, 0), sigma=(1, 1, 1), indexing="zyx", dtype=np.fl
     # Construct centered mesh
     g = grid_3d(size, shifted=False, normalized=False, indexing=indexing, dtype=dtype)
 
+    X, Y, Z = g["x"], g["y"], g["z"]
+
     if indexing == "zyx":
         mu, sigma = mu[::-1], sigma[::-1]
-        g["x"], g["y"], g["z"] = g["z"], g["y"], g["x"]
+        X, Y, Z = Z, Y, X
     elif indexing != "xyz":
         raise ValueError("Indexing must be `zyx` or `xyz`.")
 
     p = (
-        (g["x"] - mu[0]) ** 2 / (2 * sigma[0] ** 2)
-        + (g["y"] - mu[1]) ** 2 / (2 * sigma[1] ** 2)
-        + (g["z"] - mu[2]) ** 2 / (2 * sigma[2] ** 2)
+        (X - mu[0]) ** 2 / (2 * sigma[0] ** 2)
+        + (Y - mu[1]) ** 2 / (2 * sigma[1] ** 2)
+        + (Z - mu[2]) ** 2 / (2 * sigma[2] ** 2)
     )
 
     return np.exp(-p).astype(dtype, copy=False)
diff --git a/tests/test_utils.py b/tests/test_utils.py
index a81ccaf11a..67a8caf4ea 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -287,8 +287,8 @@ def test_gaussian_scalar_param():
     g_3d = gaussian_3d(L, mu_3d, sigma_3d)
     g_3d_scalar = gaussian_3d(L, mu_3d, sigma)
 
-    assert np.allclose(g_2d, g_2d_scalar)
-    assert np.allclose(g_3d, g_3d_scalar)
+    np.testing.assert_allclose(g_2d, g_2d_scalar)
+    np.testing.assert_allclose(g_3d, g_3d_scalar)
 
 
 @pytest.mark.parametrize("L", [29, 30])

From 88a53d56e41224d4d214e42a21db454e15adfb33 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Wed, 6 Nov 2024 09:56:20 -0500
Subject: [PATCH 047/184] add cache config section

migrates cache_dir to config block
---
 .github/workflows/workflow.yml                 | 5 +++--
 gallery/tutorials/tutorials/data_downloader.py | 2 +-
 src/aspire/__init__.py                         | 4 ++--
 src/aspire/config_default.yaml                 | 8 ++++++++
 src/aspire/downloader/data_fetcher.py          | 2 +-
 src/aspire/operators/filters.py                | 3 +++
 src/aspire/utils/coor_trans.py                 | 9 +++++----
 7 files changed, 23 insertions(+), 10 deletions(-)

diff --git a/.github/workflows/workflow.yml b/.github/workflows/workflow.yml
index 5c013beacf..fc31e18968 100644
--- a/.github/workflows/workflow.yml
+++ b/.github/workflows/workflow.yml
@@ -149,15 +149,16 @@ jobs:
         echo "Stash the WORK_DIR to GitHub env so we can clean it up later."
         echo "WORK_DIR=${WORK_DIR}" >> $GITHUB_ENV
         echo -e "common:" >> ${WORK_DIR}/config.yaml
-        echo -e "    cache_dir: ${CI_CACHE_DIR}" >> ${WORK_DIR}/config.yaml
         echo -e "    numeric: cupy" >> ${WORK_DIR}/config.yaml
         echo -e "    fft: cupy\n" >> ${WORK_DIR}/config.yaml
+        echo -e "cache:" >> ${WORK_DIR}/config.yaml
+        echo -e "    cache_dir: ${CI_CACHE_DIR}" >> ${WORK_DIR}/config.yaml
         echo "Log the config: ${WORK_DIR}/config.yaml"
         cat ${WORK_DIR}/config.yaml
     - name: Cache Data
       run: |
         ASPIREDIR=${{ env.WORK_DIR }} python -c \
-        "import aspire; print(aspire.config['common']['cache_dir']); import aspire.downloader; aspire.downloader.emdb_2660()"
+        "import aspire; print(aspire.config['cache']['cache_dir']); import aspire.downloader; aspire.downloader.emdb_2660()"
     - name: Run
       run: |
         ASPIREDIR=${{ env.WORK_DIR }} PYTHONWARNINGS=error python -m pytest --durations=50 --cov=aspire --cov-report=xml
diff --git a/gallery/tutorials/tutorials/data_downloader.py b/gallery/tutorials/tutorials/data_downloader.py
index 7333570320..d3f7e11493 100644
--- a/gallery/tutorials/tutorials/data_downloader.py
+++ b/gallery/tutorials/tutorials/data_downloader.py
@@ -46,7 +46,7 @@
 #
 #    .. code-block:: yaml
 #
-#      common:
+#      cache:
 #        cache_dir: /tmp/ASPIRE-data
 #
 # See the `ASPIRE Conguration
diff --git a/src/aspire/__init__.py b/src/aspire/__init__.py
index 686280e8e9..960892ef7f 100644
--- a/src/aspire/__init__.py
+++ b/src/aspire/__init__.py
@@ -63,8 +63,8 @@
 logging.debug(f"Resolved config.yaml:\n{aspire.config.dump()}\n")
 
 # Set cache location for ASPIRE example data.
-if not config["common"]["cache_dir"].get():
-    config["common"]["cache_dir"] = pooch.os_cache("ASPIRE-data").as_posix()
+if not config["cache"]["cache_dir"].get():
+    config["cache"]["cache_dir"] = pooch.os_cache("ASPIRE-data").as_posix()
 
 # Implements some code that writes out exceptions to 'aspire.err.log'.
 if config["logging"]["log_exceptions"].get(int):
diff --git a/src/aspire/config_default.yaml b/src/aspire/config_default.yaml
index 0270e0ec24..b614b8fcf0 100644
--- a/src/aspire/config_default.yaml
+++ b/src/aspire/config_default.yaml
@@ -5,6 +5,7 @@ common:
     # fft backend to use - one of pyfftw/scipy/cupy/mkl
     fft: scipy
 
+cache:
     # Set cache directory for ASPIRE example data.
     # By default the cache location will be set by pooch.os_cache(),
     # which sets cache based on operating system as follows:
@@ -13,6 +14,13 @@ common:
     # Windows: C:\Users\<user>\AppData\Local\<AppAuthor>\<AppName>\Cache
     cache_dir: ""
 
+    # The following control runtime cache sizes for various components,
+    # where `size` is the number of cached function calls.
+    # In YAML `null` translates to a limit of `None` in Python,
+    # which corresponds to unlimited calls.
+    grid_cache_size: null
+    filter_cache_size: null
+
 logging:
     # Set log_dir to a relative or absolute directory
     # Default is a subfolder `logs` in your current working directory.
diff --git a/src/aspire/downloader/data_fetcher.py b/src/aspire/downloader/data_fetcher.py
index 3fe163521b..f655d06ab3 100644
--- a/src/aspire/downloader/data_fetcher.py
+++ b/src/aspire/downloader/data_fetcher.py
@@ -16,7 +16,7 @@
     # folder operating system dependent, set by `pooch.os_cache`.
     # Pooch uses appdirs (https://github.com/ActiveState/appdirs) to
     # select an appropriate directory for the cache on each platform.
-    path=config["common"]["cache_dir"].as_filename(),
+    path=config["cache"]["cache_dir"].as_filename(),
     # The remote data is on Zenodo, `base_url` is a required param,
     # even though we override using individual urls in the registry.
     base_url="https://zenodo.org/communities/computationalcryoem/",
diff --git a/src/aspire/operators/filters.py b/src/aspire/operators/filters.py
index c58331d069..1a9f0d12b8 100644
--- a/src/aspire/operators/filters.py
+++ b/src/aspire/operators/filters.py
@@ -1,9 +1,11 @@
 import inspect
 import logging
+from functools import lru_cache
 
 import numpy as np
 from scipy.interpolate import RegularGridInterpolator
 
+from aspire import config
 from aspire.utils import grid_2d, voltage_to_wavelength
 
 logger = logging.getLogger(__name__)
@@ -110,6 +112,7 @@ def scale(self, c=1):
         """
         return ScaledFilter(self, c)
 
+    @lru_cache(maxsize=config["cache"]["filter_cache_size"].get())
     def evaluate_grid(self, L, *args, dtype=np.float32, **kwargs):
         """
         Generates a two dimensional grid with prescribed dtype,
diff --git a/src/aspire/utils/coor_trans.py b/src/aspire/utils/coor_trans.py
index 8d5cb0ce1b..ce5e5e8a49 100644
--- a/src/aspire/utils/coor_trans.py
+++ b/src/aspire/utils/coor_trans.py
@@ -4,12 +4,13 @@
 
 import logging
 import math
-from functools import cache
+from functools import lru_cache
 
 import numpy as np
 from numpy.linalg import norm
 from scipy.linalg import svd
 
+from aspire import config
 from aspire.numeric import xp
 from aspire.utils.random import Random
 from aspire.utils.rotation import Rotation
@@ -84,7 +85,7 @@ def _mgrid_slice(n, shifted, normalized):
     return slice(start, end, num_points)
 
 
-@cache
+@lru_cache(maxsize=config["cache"]["grid_cache_size"].get())
 def grid_1d(n, shifted=False, normalized=True, dtype=np.float32):
     """
     Generate one dimensional grid.
@@ -100,7 +101,7 @@ def grid_1d(n, shifted=False, normalized=True, dtype=np.float32):
     return {"x": x, "r": r}
 
 
-@cache
+@lru_cache(maxsize=config["cache"]["grid_cache_size"].get())
 def grid_2d(n, shifted=False, normalized=True, indexing="yx", dtype=np.float32):
     """
     Generate two dimensional grid.
@@ -127,7 +128,7 @@ def grid_2d(n, shifted=False, normalized=True, indexing="yx", dtype=np.float32):
     return {"x": x, "y": y, "phi": phi, "r": r}
 
 
-@cache
+@lru_cache(maxsize=config["cache"]["grid_cache_size"].get())
 def grid_3d(n, shifted=False, normalized=True, indexing="zyx", dtype=np.float32):
     """
     Generate three dimensional grid.

From 5cc8682c441995b7ba6ff5f55786c8f708c49878 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Wed, 6 Nov 2024 10:14:19 -0500
Subject: [PATCH 048/184] add another evaluate_grid cache call

---
 src/aspire/operators/filters.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/src/aspire/operators/filters.py b/src/aspire/operators/filters.py
index 1a9f0d12b8..539b053602 100644
--- a/src/aspire/operators/filters.py
+++ b/src/aspire/operators/filters.py
@@ -205,6 +205,7 @@ def __init__(self, filter, power=1, epsilon=None):
     def _evaluate(self, omega):
         return self._filter.evaluate(omega) ** self._power
 
+    @lru_cache(maxsize=config["cache"]["filter_cache_size"].get())
     def evaluate_grid(self, L, *args, dtype=np.float32, **kwargs):
         """
         Calls the provided filter's evaluate_grid method in case there is an optimization.
@@ -355,6 +356,7 @@ def _evaluate(self, omega):
 
         return result
 
+    # No need to cache ArrayFilter
     def evaluate_grid(self, L, *args, dtype=np.float32, **kwargs):
         """
         Optimized evaluate_grid method for ArrayFilter.

From 4f0071abea3d0bd0e2e788ebedc37515f3803ce7 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Wed, 6 Nov 2024 11:14:43 -0500
Subject: [PATCH 049/184] limit filter cache, qa B019

---
 src/aspire/config_default.yaml  | 3 ++-
 src/aspire/operators/filters.py | 4 ++--
 2 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/src/aspire/config_default.yaml b/src/aspire/config_default.yaml
index b614b8fcf0..591d8f2dc5 100644
--- a/src/aspire/config_default.yaml
+++ b/src/aspire/config_default.yaml
@@ -19,7 +19,8 @@ cache:
     # In YAML `null` translates to a limit of `None` in Python,
     # which corresponds to unlimited calls.
     grid_cache_size: null
-    filter_cache_size: null
+    # Using unlimited `filter_cache_size` may cause excessive memory use.
+    filter_cache_size: 2
 
 logging:
     # Set log_dir to a relative or absolute directory
diff --git a/src/aspire/operators/filters.py b/src/aspire/operators/filters.py
index 539b053602..f290d88eb1 100644
--- a/src/aspire/operators/filters.py
+++ b/src/aspire/operators/filters.py
@@ -112,7 +112,7 @@ def scale(self, c=1):
         """
         return ScaledFilter(self, c)
 
-    @lru_cache(maxsize=config["cache"]["filter_cache_size"].get())
+    @lru_cache(maxsize=config["cache"]["filter_cache_size"].get())  # noqa: B019
     def evaluate_grid(self, L, *args, dtype=np.float32, **kwargs):
         """
         Generates a two dimensional grid with prescribed dtype,
@@ -205,7 +205,7 @@ def __init__(self, filter, power=1, epsilon=None):
     def _evaluate(self, omega):
         return self._filter.evaluate(omega) ** self._power
 
-    @lru_cache(maxsize=config["cache"]["filter_cache_size"].get())
+    @lru_cache(maxsize=config["cache"]["filter_cache_size"].get())  # noqa: B019
     def evaluate_grid(self, L, *args, dtype=np.float32, **kwargs):
         """
         Calls the provided filter's evaluate_grid method in case there is an optimization.

From f04a7bec43546dfe716ba02b1c86d6fbf74b9d0e Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Thu, 14 Nov 2024 14:11:11 -0500
Subject: [PATCH 050/184] rm vecmat, volmat

---
 src/aspire/covariance/covar.py  | 20 ++++++++------------
 src/aspire/source/simulation.py |  4 ++--
 2 files changed, 10 insertions(+), 14 deletions(-)

diff --git a/src/aspire/covariance/covar.py b/src/aspire/covariance/covar.py
index 89977052fe..380b082b71 100644
--- a/src/aspire/covariance/covar.py
+++ b/src/aspire/covariance/covar.py
@@ -10,14 +10,7 @@
 from aspire.numeric.scipy import cg
 from aspire.operators import evaluate_src_filters_on_grid
 from aspire.reconstruction import Estimator, FourierKernel, MeanEstimator
-from aspire.utils import (
-    make_symmat,
-    symmat_to_vec_iso,
-    trange,
-    vec_to_symmat_iso,
-    vecmat_to_volmat,
-    volmat_to_vecmat,
-)
+from aspire.utils import make_symmat, symmat_to_vec_iso, trange, vec_to_symmat_iso
 from aspire.volume import Volume, rotated_grids
 
 logger = logging.getLogger(__name__)
@@ -90,12 +83,15 @@ def estimate(self, mean_vol, noise_variance, tol=1e-5, regularizer=0):
         b_coef = self.src_backward(mean_vol, noise_variance)
         est_coef = self.conj_grad(b_coef, tol=tol, regularizer=regularizer)
         covar_est = self.basis.mat_evaluate(est_coef)
-        # Note, notice these cancel out, but requires a lot of changes elsewhere in this file,
-        # basically totally removing all the `utils/matrix` hacks ... todo.
-        covar_est = vecmat_to_volmat(make_symmat(volmat_to_vecmat(covar_est)))
+
+        L = covar_est.shape[-1]
+        covar_est = covar_est.reshape((L**3, L**3))
+        covar_est = make_symmat(covar_est)
+        covar_est = covar_est.reshape((L,) * 6)
+
         return covar_est
 
-    def conj_grad(self, b_coef, tol=1e-5, regularizer=0):
+    def conj_grad(self, b_coef, tol=1e-5, regularizer=0):        
         b_coef = symmat_to_vec_iso(b_coef)
         N = b_coef.shape[0]
         kernel = self.kernel
diff --git a/src/aspire/source/simulation.py b/src/aspire/source/simulation.py
index 2ee06cc5ee..15dd1827e4 100644
--- a/src/aspire/source/simulation.py
+++ b/src/aspire/source/simulation.py
@@ -16,7 +16,6 @@
     anorm,
     make_symmat,
     uniform_random_angles,
-    vecmat_to_volmat,
 )
 from aspire.utils.random import rand, randi, randn
 from aspire.volume import AsymmetricVolume, Volume
@@ -399,7 +398,8 @@ def covar_true(self):
         eigs_true = eigs_true.T.to_vec()
 
         covar_true = eigs_true.T @ lamdbas_true @ eigs_true
-        covar_true = vecmat_to_volmat(covar_true)
+        # Hrmm
+        covar_true = covar_true.reshape((self.L,) * 6).transpose(2, 1, 0, 5, 4, 3)
 
         return covar_true
 

From 4ae9ae6d525e5c58aba438d4c8003b878f287ca6 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Wed, 20 Nov 2024 15:29:42 -0500
Subject: [PATCH 051/184] remove m_ from cov3d

need to rebase on other m_ work
---
 src/aspire/covariance/covar.py         |   2 +-
 src/aspire/denoising/denoiser_cov2d.py |  11 +-
 src/aspire/utils/matrix.py             |  88 +++++++--------
 src/aspire/volume/volume.py            |  11 +-
 tests/test_matrix.py                   | 143 +++++++++----------------
 5 files changed, 100 insertions(+), 155 deletions(-)

diff --git a/src/aspire/covariance/covar.py b/src/aspire/covariance/covar.py
index 380b082b71..cf48d43aa3 100644
--- a/src/aspire/covariance/covar.py
+++ b/src/aspire/covariance/covar.py
@@ -91,7 +91,7 @@ def estimate(self, mean_vol, noise_variance, tol=1e-5, regularizer=0):
 
         return covar_est
 
-    def conj_grad(self, b_coef, tol=1e-5, regularizer=0):        
+    def conj_grad(self, b_coef, tol=1e-5, regularizer=0):
         b_coef = symmat_to_vec_iso(b_coef)
         N = b_coef.shape[0]
         kernel = self.kernel
diff --git a/src/aspire/denoising/denoiser_cov2d.py b/src/aspire/denoising/denoiser_cov2d.py
index 545f6ce642..89a850e559 100644
--- a/src/aspire/denoising/denoiser_cov2d.py
+++ b/src/aspire/denoising/denoiser_cov2d.py
@@ -84,16 +84,13 @@ def src_wiener_coords(
         Qs, Rs = qr_vols_forward(sim, i, batch_n, eig_vols, k)
 
         Q_vecs = mat_to_vec(Qs)
-
-        # RCOPT
-        ims = np.moveaxis(ims.asnumpy(), 0, 2)
-        im_vecs = mat_to_vec(ims)
+        im_vecs = mat_to_vec(ims.asnumpy())
 
         for j in range(batch_n):
-            im_coords = Q_vecs[:, :, j].T @ im_vecs[:, j]
-            covar_im = (Rs[:, :, j] @ lambdas @ Rs[:, :, j].T) + covar_noise
+            im_coords = Q_vecs[j] @ im_vecs[j]
+            covar_im = (Rs[j] @ lambdas @ Rs[j].T) + covar_noise
             xx = solve(covar_im, im_coords)
-            im_coords = lambdas @ Rs[:, :, j].T @ xx
+            im_coords = lambdas @ Rs[j].T @ xx
             coords[:, i + j] = im_coords
 
     return coords
diff --git a/src/aspire/utils/matrix.py b/src/aspire/utils/matrix.py
index a8415929c0..38375c745f 100644
--- a/src/aspire/utils/matrix.py
+++ b/src/aspire/utils/matrix.py
@@ -121,18 +121,17 @@ def symmat_to_vec_iso(mat):
     """
     Isometrically maps a symmetric matrix to a packed vector
 
-    :param mat: An array of size N-by-N-by-... where the first two dimensions constitute symmetric or Hermitian
+    :param mat: An array of size ...-by-N-by-N where the last two dimensions constitute symmetric or Hermitian
         matrices.
-    :return: A vector of size N*(N+1)/2-by-... consisting of the lower triangular part of each matrix, reweighted so
+    :return: A vector of size ...-by-N*(N+1)/2 consisting of the lower triangular part of each matrix, reweighted so
         that the Frobenius inner product is mapped to the Euclidean inner product.
     """
-    mat, sz_roll = unroll_dim(mat, 3)
-    N = mat.shape[0]
-    mat = mat_to_vec(mat)
-    mat[np.arange(0, N**2, N + 1)] *= SQRT2_R
-    mat *= SQRT2
-    mat = vec_to_mat(mat)
-    mat = roll_dim(mat, sz_roll)
+    sz = mat.shape[:-2]
+    N = mat.shape[-1]
+    vec = mat.reshape(*sz, -1)
+    vec[..., np.arange(0, N**2, N + 1)] *= SQRT2_R
+    vec *= SQRT2
+    mat = vec.reshape(*sz, N, N)
     vec = symmat_to_vec(mat)
 
     return vec
@@ -142,65 +141,59 @@ def vec_to_symmat_iso(vec):
     """
     Isometrically map packed vector to symmetric matrix
 
-    :param vec: A vector of size N*(N+1)/2-by-... describing a symmetric (or Hermitian) matrix.
-    :return: An array of size N-by-N-by-... which indexes symmetric/Hermitian matrices that occupy the first two
+    :param vec: A vector of size ...-by-N*(N+1)/2 describing a symmetric (or Hermitian) matrix.
+    :return: An array of size ...-by-N-by-N which indexes symmetric/Hermitian matrices that occupy the first two
         dimensions. The lower triangular parts of these matrices consists of the corresponding vectors in vec,
         reweighted so that the Euclidean inner product maps to the Frobenius inner product.
     """
+
     mat = vec_to_symmat(vec)
-    mat, sz_roll = unroll_dim(mat, 3)
-    N = mat.shape[0]
+    N = mat.shape[-1]
     mat = mat_to_vec(mat)
-    mat[np.arange(0, N**2, N + 1)] *= SQRT2
+    mat[..., np.arange(0, N**2, N + 1)] *= SQRT2
     mat *= SQRT2_R
     mat = vec_to_mat(mat)
-    mat = roll_dim(mat, sz_roll)
     return mat
 
 
 def symmat_to_vec(mat):
     """
-    Packs a symmetric matrix into a lower triangular vector
+    Packs a symmetric matrix into a upper triangular vector
 
-    :param mat: An array of size N-by-N-by-... where the first two dimensions constitute symmetric or
+    :param mat: An array of size ...-by-N-by-N where the first two dimensions constitute symmetric or
         Hermitian matrices.
-    :return: A vector of size N*(N+1)/2-by-... consisting of the lower triangular part of each matrix.
-
-    Note that a lot of acrobatics happening here (swapaxes/triu instead of tril etc.) are so that we can get
-    column-major ordering of elements (to get behavior consistent with MATLAB), since masking in numpy only returns
-    data in row-major order.
+    :return: A vector of size ...-by-N*(N+1)/2 consisting of the upper triangular part of each matrix.
     """
-    N = mat.shape[0]
-    assert mat.shape[1] == N, "Matrix must be square"
 
-    mat, sz_roll = unroll_dim(mat, 3)
-    triu_indices = np.triu_indices(N)
-    vec = mat.swapaxes(0, 1)[triu_indices]
-    vec = roll_dim(vec, sz_roll)
+    N = mat.shape[-1]
+    assert mat.shape[-2] == N, "Matrix must be square"
+
+    sz = mat.shape[:-2]
+    tri_indices = np.triu_indices(N)
+    vec = mat[..., *tri_indices].reshape(*sz, N * (N + 1) // 2)
 
     return vec
 
 
 def vec_to_symmat(vec):
     """
-    Convert packed lower triangular vector to symmetric matrix
+    Convert packed upper triangular vector to symmetric matrix
 
-    :param vec: A vector of size N*(N+1)/2-by-... describing a symmetric (or Hermitian) matrix.
-    :return: An array of size N-by-N-by-... which indexes symmetric/Hermitian matrices that occupy the first two
-        dimensions. The lower triangular parts of these matrices consists of the corresponding vectors in vec.
+    :param vec: A vector of size ...-by-N*(N+1)/2 describing a symmetric (or Hermitian) matrix.
+    :return: An array of size ...-by-N-by-N which indexes symmetric/Hermitian matrices that occupy the first two
+        dimensions. The upper triangular parts of these matrices consists of the corresponding vectors in vec.
     """
     # TODO: Handle complex values in vec
     if np.iscomplex(vec).any():
         raise NotImplementedError("Coming soon")
 
-    # M represents N(N+1)/2
-    M = vec.shape[0]
+    M = vec.shape[-1]
     N = int(round(np.sqrt(2 * M + 0.25) - 0.5))
     assert (
         M == 0.5 * N * (N + 1)
     ) and N != 0, "Vector must be of size N*(N+1)/2 for some N>0."
 
-    vec, sz_roll = unroll_dim(vec, 2)
+    sz = vec.shape[:-1]
     index_matrix = np.empty((N, N))
     i_upper = np.triu_indices_from(index_matrix)
     index_matrix[i_upper] = np.arange(
@@ -208,9 +201,8 @@ def vec_to_symmat(vec):
     )  # Incrementally populate upper triangle in row major order
     index_matrix.T[i_upper] = index_matrix[i_upper]  # Copy to lower triangle
 
-    mat = vec[index_matrix.flatten("F").astype("int")]
-    mat = m_reshape(mat, (N, N) + mat.shape[1:])
-    mat = roll_dim(mat, sz_roll)
+    mat = vec[..., index_matrix.astype("int")]
+    mat = mat.reshape(*sz, N, N)
 
     return mat
 
@@ -219,17 +211,17 @@ def mat_to_vec(mat, is_symmat=False):
     """
     Converts a matrix into vectorized form
 
-    :param mat: An array of size N-by-N-by-... containing the matrices to be vectorized.
+    :param mat: An array of size ...-by-N-by-N-by containing the matrices to be vectorized.
     :param is_symmat: Specifies whether the matrices are symmetric/Hermitian, in which case they are stored in packed
         form using symmat_to_vec (default False).
-    :return: The vectorized form of the matrices, with dimension N^2-by-... or N*(N+1)/2-by-... depending on the value
+    :return: The vectorized form of the matrices, with dimension ...-by-N^2 or ...-by-N*(N+1)/2 depending on the value
         of is_symmat.
     """
     if not is_symmat:
         sz = mat.shape
-        N = sz[0]
-        assert sz[1] == N, "Matrix must be square"
-        return m_reshape(mat, (N**2,) + sz[2:])
+        N = sz[-1]
+        assert sz[-2] == N, "Matrix must be square"
+        return mat.reshape(*sz[:-2], N**2)
     else:
         return symmat_to_vec(mat)
 
@@ -239,15 +231,15 @@ def vec_to_mat(vec, is_symmat=False):
     Converts a vectorized matrix into a matrix
 
     :param vec: The vectorized representations. If the matrix is non-symmetric, this array has the dimensions
-        N^2-by-..., but if the matrix is symmetric, the dimensions are N*(N+1)/2-by-... .
+        ...-by-N^2, but if the matrix is symmetric, the dimensions are ...-by-N*(N+1)/2 .
     :param is_symmat: True if the vectors represent symmetric matrices (default False)
-    :return: The array of size N-by-N-by-... representing the matrices.
+    :return: The array of size ...-by-N-by-N representing the matrices.
     """
     if not is_symmat:
         sz = vec.shape
-        N = int(round(np.sqrt(sz[0])))
-        assert sz[0] == N**2, "Vector must represent square matrix."
-        return m_reshape(vec, (N, N) + sz[1:])
+        N = int(round(np.sqrt(sz[-1])))
+        assert sz[-1] == N**2, "Vector must represent square matrix."
+        return vec.reshape(*sz[:-1], N, N)
     else:
         return vec_to_symmat(vec)
 
diff --git a/src/aspire/volume/volume.py b/src/aspire/volume/volume.py
index 60c8bef033..8af6caa1b5 100644
--- a/src/aspire/volume/volume.py
+++ b/src/aspire/volume/volume.py
@@ -38,16 +38,17 @@ def qr_vols_forward(sim, s, n, vols, k):
     ims = np.zeros((k, n, sim.L, sim.L), dtype=vols.dtype)
     for ell in range(k):
         ims[ell] = sim.vol_forward(vols[ell], s, n).asnumpy()
-    ims = ims.transpose((2, 3, 0, 1))
 
-    Q_vecs = np.zeros((sim.L**2, k, n), dtype=vols.dtype)
-    Rs = np.zeros((k, k, n), dtype=vols.dtype)
+    ims = ims.transpose((1, 0, 2, 3))  # n, k, L, L
+
+    Q_vecs = np.zeros((n, sim.L**2, k), dtype=vols.dtype)
+    Rs = np.zeros((n, k, k), dtype=vols.dtype)
 
     im_vecs = ims.reshape(sim.L**2, k, n)
 
     for i in range(n):
-        Q_vecs[:, :, i], Rs[:, :, i] = qr(im_vecs[:, :, i])
-    Qs = Q_vecs.reshape(sim.L, sim.L, k, n)
+        Q_vecs[i], Rs[i] = qr(im_vecs[i].T)  # column vectors
+    Qs = vec_to_mat(Q_vecs.transpose(0, 2, 1))  # n, k, L, L
 
     return Qs, Rs
 
diff --git a/tests/test_matrix.py b/tests/test_matrix.py
index 5274125ac0..cd101a3654 100644
--- a/tests/test_matrix.py
+++ b/tests/test_matrix.py
@@ -71,37 +71,20 @@ def testMatToVec1(self):
         m = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
 
         v = mat_to_vec(m)
-        self.assertTrue(np.allclose(v, np.array([1, 4, 7, 2, 5, 8, 3, 6, 9])))
+        self.assertTrue(np.allclose(v, np.array([1, 2, 3, 4, 5, 6, 7, 8, 9])))
 
     def testMatToVec2(self):
-        m = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
+        _m = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
         # Make 2 copies depthwise
-        m = np.dstack((m, m))
+        m = np.stack((_m, _m))
 
         v = mat_to_vec(m)
-        self.assertTrue(
-            np.allclose(
-                v,
-                np.array(
-                    [
-                        [1, 1],
-                        [4, 4],
-                        [7, 7],
-                        [2, 2],
-                        [5, 5],
-                        [8, 8],
-                        [3, 3],
-                        [6, 6],
-                        [9, 9],
-                    ]
-                ),
-            )
-        )
+        self.assertTrue(np.allclose(v, np.stack((_m.flatten(),) * 2)))
 
     def testMatToVecSymm1(self):
         # We create an unsymmetric matrix and pass it to the functions as a symmetric matrix,
         # just so we can closely inspect the returned values without confusion
-        m = np.array([[0, 4, 8, 12], [1, 5, 9, 13], [2, 6, 10, 14], [3, 7, 11, 15]])
+        m = np.arange(16).reshape(4, 4)
         v = mat_to_vec(m, is_symmat=True)
         # Notice the order of the elements in symmetric matrix - axis 0 first, then axis 1
         self.assertTrue(np.allclose(v, np.array([0, 1, 2, 3, 5, 6, 7, 10, 11, 15])))
@@ -109,29 +92,18 @@ def testMatToVecSymm1(self):
     def testMatToVecSymm2(self):
         # We create an unsymmetric matrix and pass it to the functions as a symmetric matrix,
         # just so we can closely inspect the returned values without confusion
-        m = np.array([[0, 4, 8, 12], [1, 5, 9, 13], [2, 6, 10, 14], [3, 7, 11, 15]])
+        m = np.arange(16).reshape(4, 4)
+
         # Make 2 copies depthwise
-        m = np.dstack((m, m))
+        m = np.stack((m, m))
 
         v = mat_to_vec(m, is_symmat=True)
         # Notice the order of the elements in symmetric matrix - axis 0 first, then axis 1
+
         self.assertTrue(
             np.allclose(
                 v,
-                np.array(
-                    [
-                        [0, 0],
-                        [1, 1],
-                        [2, 2],
-                        [3, 3],
-                        [5, 5],
-                        [6, 6],
-                        [7, 7],
-                        [10, 10],
-                        [11, 11],
-                        [15, 15],
-                    ]
-                ),
+                np.stack((np.array([0, 1, 2, 3, 5, 6, 7, 10, 11, 15]),) * 2),
             )
         )
 
@@ -140,32 +112,35 @@ def testMatToVecSymmIso(self):
 
         # We create an unsymmetric matrix and pass it to the functions as a symmetric matrix,
         # just so we can closely inspect the returned values without confusion
-        m = np.array(
-            [[0, 4, 8, 12], [1, 5, 9, 13], [2, 6, 10, 14], [3, 7, 11, 15]],
-            dtype=np.float32,
-        )
+        m = np.arange(16).reshape(4, 4).astype(dtype=np.float64)
 
         # Make 2 copies depthwise
-        m = np.dstack((m, m))
+        m = np.stack((m, m))
 
         v = symmat_to_vec_iso(m)
         # Notice the order of the elements in symmetric matrix - axis 0 first, then axis 1
+
         self.assertTrue(
             np.allclose(
                 v,
-                np.array(
-                    [
-                        [0, 0],
-                        [1.4142, 1.4142],
-                        [2.8284, 2.8284],
-                        [4.2426, 4.2426],
-                        [5, 5],
-                        [8.4853, 8.4853],
-                        [9.8995, 9.8995],
-                        [10, 10],
-                        [15.5563, 15.5563],
-                        [15, 15],
-                    ]
+                np.stack(
+                    (
+                        np.array(
+                            [
+                                0.0,
+                                1.41421356,
+                                2.82842712,
+                                4.24264069,
+                                5.0,
+                                8.48528137,
+                                9.89949494,
+                                10.0,
+                                15.55634919,
+                                15.0,
+                            ]
+                        ),
+                    )
+                    * 2
                 ),
             )
         )
@@ -184,18 +159,18 @@ def testVecToMatSymm1(self):
                 [11, 11],
                 [15, 15],
             ]
-        )
+        ).transpose(1, 0)
 
         m = vec_to_symmat(v)
         self.assertTrue(
             np.allclose(
-                m[:, :, 0],
+                m[0],
                 np.array([[0, 1, 2, 3], [1, 5, 6, 7], [2, 6, 10, 11], [3, 7, 11, 15]]),
             )
         )
         self.assertTrue(
             np.allclose(
-                m[:, :, 1],
+                m[1],
                 np.array([[0, 1, 2, 3], [1, 5, 6, 7], [2, 6, 10, 11], [3, 7, 11, 15]]),
             )
         )
@@ -213,46 +188,26 @@ def testVecToMatSymm2(self):
 
     def testVecToMatSymmIso(self):
         # Very similar to the case above, except that the resulting matrix is reweighted.
-        v = np.array(
-            [
-                [0, 0],
-                [1, 1],
-                [2, 2],
-                [3, 3],
-                [5, 5],
-                [6, 6],
-                [7, 7],
-                [10, 10],
-                [11, 11],
-                [15, 15],
-            ],
-            dtype=np.float32,
+        v = np.stack((np.array([0, 1, 2, 3, 5, 6, 7, 10, 11, 15]),) * 2).astype(
+            np.float32
         )
 
         m = vec_to_symmat_iso(v)
         self.assertTrue(
             np.allclose(
-                m[:, :, 0],
-                np.array(
-                    [
-                        [0, 0.70710678, 1.41421356, 2.12132034],
-                        [0.70710678, 5, 4.24264069, 4.94974747],
-                        [1.41421356, 4.24264069, 10, 7.77817459],
-                        [2.12132034, 4.94974747, 7.77817459, 15],
-                    ]
-                ),
-            )
-        )
-        self.assertTrue(
-            np.allclose(
-                m[:, :, 1],
-                np.array(
-                    [
-                        [0, 0.70710678, 1.41421356, 2.12132034],
-                        [0.70710678, 5, 4.24264069, 4.94974747],
-                        [1.41421356, 4.24264069, 10, 7.77817459],
-                        [2.12132034, 4.94974747, 7.77817459, 15],
-                    ]
+                m,
+                np.stack(
+                    (
+                        np.array(
+                            [
+                                [0, 0.70710678, 1.41421356, 2.12132034],
+                                [0.70710678, 5, 4.24264069, 4.94974747],
+                                [1.41421356, 4.24264069, 10, 7.77817459],
+                                [2.12132034, 4.94974747, 7.77817459, 15],
+                            ]
+                        ),
+                    )
+                    * 2
                 ),
             )
         )

From a5db7ff096a47395cf99e3b120f4652f9c48675a Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Thu, 21 Nov 2024 09:27:14 -0500
Subject: [PATCH 052/184] resolve volume conflict

---
 src/aspire/volume/volume.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/src/aspire/volume/volume.py b/src/aspire/volume/volume.py
index 8af6caa1b5..04ae8429be 100644
--- a/src/aspire/volume/volume.py
+++ b/src/aspire/volume/volume.py
@@ -16,7 +16,9 @@
     crop_pad_3d,
     grid_2d,
     grid_3d,
+    mat_to_vec,
     rename_with_timestamp,
+    vec_to_mat,
 )
 from aspire.volume import IdentitySymmetryGroup, SymmetryGroup
 
@@ -44,8 +46,7 @@ def qr_vols_forward(sim, s, n, vols, k):
     Q_vecs = np.zeros((n, sim.L**2, k), dtype=vols.dtype)
     Rs = np.zeros((n, k, k), dtype=vols.dtype)
 
-    im_vecs = ims.reshape(sim.L**2, k, n)
-
+    im_vecs = mat_to_vec(ims)
     for i in range(n):
         Q_vecs[i], Rs[i] = qr(im_vecs[i].T)  # column vectors
     Qs = vec_to_mat(Q_vecs.transpose(0, 2, 1))  # n, k, L, L

From 3900576637b4320a6712cafb0a5279647ae82fe7 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Thu, 21 Nov 2024 11:42:37 -0500
Subject: [PATCH 053/184] resolve cov3d mdim regression/conflicts

---
 src/aspire/basis/basis.py  | 14 ++++++++++++--
 src/aspire/basis/ffb_3d.py | 12 ------------
 src/aspire/utils/matrix.py |  9 ---------
 3 files changed, 12 insertions(+), 23 deletions(-)

diff --git a/src/aspire/basis/basis.py b/src/aspire/basis/basis.py
index ee6aee2d02..f099fde570 100644
--- a/src/aspire/basis/basis.py
+++ b/src/aspire/basis/basis.py
@@ -506,7 +506,12 @@ def mat_evaluate(self, V):
             -`self.sz` corresponding to the evaluation of `V` in
             this basis.
         """
-        return mdim_mat_fun_conj(V, 1, len(self.sz), self._evaluate)
+
+        def f(V):
+            """Wrapper to handle dimension rolling."""
+            return self.evaluate(Coef(self, V)).asnumpy()
+
+        return mdim_mat_fun_conj(V, 1, len(self.sz), f)
 
     def mat_evaluate_t(self, X):
         """
@@ -522,7 +527,12 @@ def mat_evaluate_t(self, X):
             function calculates V = B' * X * B, where the rows of `B`, rows
             of 'X', and columns of `X` are read as vectorized arrays.
         """
-        return mdim_mat_fun_conj(X, len(self.sz), 1, self._evaluate_t)
+
+        def f(X):
+            """Wrapper to handle dimension rolling."""
+            return self.evaluate_t(Volume(X)).asnumpy()
+
+        return mdim_mat_fun_conj(X, len(self.sz), 1, f)
 
     def expand(self, x, tol=None, atol=0):
         """
diff --git a/src/aspire/basis/ffb_3d.py b/src/aspire/basis/ffb_3d.py
index 350a146e06..7b50f3af61 100644
--- a/src/aspire/basis/ffb_3d.py
+++ b/src/aspire/basis/ffb_3d.py
@@ -166,9 +166,6 @@ def _evaluate(self, v):
             `self.sz` and the remaining dimensions correspond to `v`.
         """
         v = xp.asarray(v)
-        # roll dimensions of v
-        sz_roll = v.shape[:-1]
-        v = v.reshape((-1, self.count))
 
         # get information on polar grids from precomputed data
         n_theta = np.size(self._precomp["ang_theta_wtd"], 0)
@@ -270,9 +267,6 @@ def _evaluate(self, v):
         # perform inverse non-uniformly FFT transformation back to 3D rectangular coordinates
         x = anufft(pf, self._precomp["fourier_pts"], self.sz, real=True)
 
-        # Roll, return the x with the last three dimensions as self.sz
-        # Higher dimensions should be like v.
-        x = x.reshape((*sz_roll, *self.sz))
         return xp.asnumpy(x)
 
     def _evaluate_t(self, x):
@@ -287,9 +281,6 @@ def _evaluate_t(self, x):
             dimensions of `x`.
         """
         x = xp.asarray(x)
-        # roll dimensions
-        sz_roll = x.shape[:-3]
-        x = x.reshape((-1, *self.sz))
 
         n_data = x.shape[0]
         n_r = np.size(self._precomp["radial_wtd"], 0)
@@ -379,7 +370,4 @@ def _evaluate_t(self, x):
             ind = self._indices["ells"] == ell
             v[:, ind] = v_ell
 
-        # Roll dimensions, last dimension should be self.count,
-        # Higher dimensions like x.
-        v = v.reshape((*sz_roll, self.count))
         return xp.asnumpy(v)
diff --git a/src/aspire/utils/matrix.py b/src/aspire/utils/matrix.py
index 38375c745f..bdc46c25eb 100644
--- a/src/aspire/utils/matrix.py
+++ b/src/aspire/utils/matrix.py
@@ -88,12 +88,6 @@ def mdim_mat_fun_conj(X, d1, d2, f):
         and columns of the multidimensional matrix X.
     """
 
-    # Roll up outer dimensions if any.
-    dim = 2 * d1 + 1
-    sz_roll = X.shape[:-dim]
-    shp = X.shape[-dim:]
-    X = X.reshape((-1, *shp))
-
     X = f(X)
 
     # Swap the last d2 axes with the first d1 axes
@@ -111,9 +105,6 @@ def mdim_mat_fun_conj(X, d1, d2, f):
 
     X = np.conj(X)
 
-    # Unroll outer dimensions.
-    X = X.reshape(*sz_roll, *X.shape[1:])
-
     return X
 
 

From f08b8dbc9fdd32787c87a639de0fbd55955a33df Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Thu, 21 Nov 2024 12:42:45 -0500
Subject: [PATCH 054/184] cleanup the fb ffb 2d shapes a little more.

---
 src/aspire/basis/fb_2d.py  | 11 ++++++-----
 src/aspire/basis/fb_3d.py  |  5 ++---
 src/aspire/basis/ffb_2d.py |  4 +---
 3 files changed, 9 insertions(+), 11 deletions(-)

diff --git a/src/aspire/basis/fb_2d.py b/src/aspire/basis/fb_2d.py
index e737f0fa8b..6477d23e5c 100644
--- a/src/aspire/basis/fb_2d.py
+++ b/src/aspire/basis/fb_2d.py
@@ -182,8 +182,9 @@ def _evaluate(self, v):
             This is an array whose last dimensions equal `self.sz` and the remaining
             dimensions correspond to first dimensions of `v`.
         """
-        # Transpose here once, instead of several times below  #RCOPT
-        v = v.T
+        n_data = v.shape[0]
+        # Transpose here once, instead of several times below
+        v = v.reshape(n_data, self.count).T
 
         r_idx = self.basis_coords["r_idx"]
         ang_idx = self.basis_coords["ang_idx"]
@@ -193,7 +194,7 @@ def _evaluate(self, v):
         ind_radial = 0
         ind_ang = 0
 
-        x = np.zeros(shape=tuple([np.prod(self.sz)] + list(v.shape[1:])), dtype=v.dtype)
+        x = np.zeros((np.prod(self.sz), n_data), dtype=v.dtype)
         for ell in range(0, self.ell_max + 1):
             k_max = self.k_max[ell]
             idx_radial = ind_radial + np.arange(0, k_max, dtype=int)
@@ -214,7 +215,7 @@ def _evaluate(self, v):
 
             ind_radial += len(idx_radial)
 
-        x = x.T.reshape(-1, *self.sz)  # RCOPT
+        x = x.T.reshape(n_data, *self.sz)
 
         return x
 
@@ -229,7 +230,7 @@ def _evaluate_t(self, x):
             `self.count` and whose first dimensions correspond to
             first dimensions of `v`.
         """
-        x = x.reshape(x.shape[0], -1)
+        x = x.reshape(x.shape[0], np.prod(self.sz))
 
         r_idx = self.basis_coords["r_idx"]
         ang_idx = self.basis_coords["ang_idx"]
diff --git a/src/aspire/basis/fb_3d.py b/src/aspire/basis/fb_3d.py
index b2910c4262..cd04aa1f3b 100644
--- a/src/aspire/basis/fb_3d.py
+++ b/src/aspire/basis/fb_3d.py
@@ -176,7 +176,7 @@ def _evaluate(self, v):
 
             ind_radial += len(idx_radial)
 
-        return x.reshape(*stack_shape, *self.sz)
+        return x.reshape(v.shape[0], *self.sz)
 
     def _evaluate_t(self, x):
         """
@@ -189,8 +189,7 @@ def _evaluate_t(self, x):
             equals `self.count` and whose remaining dimensions correspond
             to higher dimensions of `v`.
         """
-        stack_shape = x.shape[: -self.ndim]
-        x = x.reshape(-1, np.prod(self.sz))
+        x = x.reshape(x.shape[0], np.prod(self.sz))
         r_idx = self.basis_coords["r_idx"]
         ang_idx = self.basis_coords["ang_idx"]
         mask = self.basis_coords["mask"].flatten()
diff --git a/src/aspire/basis/ffb_2d.py b/src/aspire/basis/ffb_2d.py
index 8806bad78a..f39e89f683 100644
--- a/src/aspire/basis/ffb_2d.py
+++ b/src/aspire/basis/ffb_2d.py
@@ -110,8 +110,6 @@ def _evaluate(self, v):
             and the first dimension correspond to remaining dimension of `v`.
         """
         v = xp.asarray(v)
-        sz_roll = v.shape[:-1]
-        v = v.reshape(-1, self.count)
 
         # number of 2D image samples
         n_data = v.shape[0]
@@ -169,7 +167,7 @@ def _evaluate(self, v):
         x = 2 * anufft(pf, 2 * pi * freqs, self.sz, real=True)
 
         # Return X as Image instance with the last two dimensions as *self.sz
-        x = x.reshape((*sz_roll, *self.sz))
+        x = x.reshape((n_data, *self.sz))
 
         return xp.asnumpy(x)
 

From f173ee1a8d636956397629738e345d06a4a2a1e0 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Thu, 21 Nov 2024 12:44:05 -0500
Subject: [PATCH 055/184] purge roll_dim, unroll_dim

---
 src/aspire/utils/__init__.py |  2 --
 src/aspire/utils/matrix.py   | 24 ------------------------
 tests/test_matrix.py         | 25 -------------------------
 3 files changed, 51 deletions(-)

diff --git a/src/aspire/utils/__init__.py b/src/aspire/utils/__init__.py
index 83a681373d..2f8dfefe9d 100644
--- a/src/aspire/utils/__init__.py
+++ b/src/aspire/utils/__init__.py
@@ -47,10 +47,8 @@
     mat_to_vec,
     mdim_mat_fun_conj,
     nearest_rotations,
-    roll_dim,
     symmat_to_vec,
     symmat_to_vec_iso,
-    unroll_dim,
     vec_to_mat,
     vec_to_symmat,
     vec_to_symmat_iso,
diff --git a/src/aspire/utils/matrix.py b/src/aspire/utils/matrix.py
index bdc46c25eb..e11768218d 100644
--- a/src/aspire/utils/matrix.py
+++ b/src/aspire/utils/matrix.py
@@ -11,30 +11,6 @@
 SQRT2_R = 1 / SQRT2
 
 
-def unroll_dim(X, dim):
-    # TODO: dim is still 1-indexed like in MATLAB to reduce headaches for now
-    # TODO: unroll/roll are great candidates for a context manager since they're always used in conjunction.
-    dim = dim - 1
-    old_shape = X.shape
-    new_shape = old_shape[:dim]
-
-    new_shape += (-1,)
-
-    Y = m_reshape(X, new_shape)
-
-    removed_dims = old_shape[dim:]
-
-    return Y, removed_dims
-
-
-def roll_dim(X, dim):
-    # TODO: dim is still 1-indexed like in MATLAB to reduce headaches for now
-    old_shape = X.shape
-    new_shape = old_shape[:-1] + dim
-    Y = m_reshape(X, new_shape)
-    return Y
-
-
 def vecmat_to_volmat(X):
     """
     Roll up vector matrices into volume matrices
diff --git a/tests/test_matrix.py b/tests/test_matrix.py
index cd101a3654..be1fda9343 100644
--- a/tests/test_matrix.py
+++ b/tests/test_matrix.py
@@ -12,9 +12,7 @@
     mean_aligned_angular_distance,
     nearest_rotations,
     randn,
-    roll_dim,
     symmat_to_vec_iso,
-    unroll_dim,
     utest_tolerance,
     vec_to_symmat,
     vec_to_symmat_iso,
@@ -32,29 +30,6 @@ def setUp(self):
     def tearDown(self):
         pass
 
-    def testUnrollDims(self):
-        m = np.arange(1, 1201).reshape((5, 2, 10, 3, 4), order="F")
-        m2, sz = unroll_dim(
-            m, 2
-        )  # second argument is 1-indexed - all dims including and after this are unrolled
-
-        # m2 will now have shape (5, (2x10x3x4)) = (5, 240)
-        self.assertEqual(m2.shape, (5, 240))
-        # The values should still be filled in with the first axis values changing fastest
-        self.assertTrue(np.allclose(m2[:, 0], np.array([1, 2, 3, 4, 5])))
-
-        # sz are the dimensions that were unrolled
-        self.assertEqual(sz, (2, 10, 3, 4))
-
-    def testRollDims(self):
-        m = np.arange(1, 1201).reshape((5, 2, 120), order="F")
-        m2 = roll_dim(m, (10, 3, 4))
-
-        # m2 will now have shape (5, 2, 10, 3, 4)
-        self.assertEqual(m2.shape, (5, 2, 10, 3, 4))
-        # The values should still be filled in with the first axis values changing fastest
-        self.assertTrue(np.allclose(m2[:, 0, 0, 0, 0], np.array([1, 2, 3, 4, 5])))
-
     def testVecmatToVolmat(self):
         m = np.empty((8, 27, 10))
         m2 = vecmat_to_volmat(m)

From 4bc0d5b83fa57820df896b5bfcc87f682fbc34e0 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Thu, 21 Nov 2024 12:49:17 -0500
Subject: [PATCH 056/184] purge vecmat volmat

---
 src/aspire/utils/__init__.py |  2 --
 src/aspire/utils/matrix.py   | 39 ------------------------------------
 tests/test_matrix.py         | 14 -------------
 3 files changed, 55 deletions(-)

diff --git a/src/aspire/utils/__init__.py b/src/aspire/utils/__init__.py
index 2f8dfefe9d..69e0c7df9f 100644
--- a/src/aspire/utils/__init__.py
+++ b/src/aspire/utils/__init__.py
@@ -52,8 +52,6 @@
     vec_to_mat,
     vec_to_symmat,
     vec_to_symmat_iso,
-    vecmat_to_volmat,
-    volmat_to_vecmat,
 )
 from .multiprocessing import (
     mem_based_cpu_suggestion,
diff --git a/src/aspire/utils/matrix.py b/src/aspire/utils/matrix.py
index e11768218d..2b7373fadb 100644
--- a/src/aspire/utils/matrix.py
+++ b/src/aspire/utils/matrix.py
@@ -11,45 +11,6 @@
 SQRT2_R = 1 / SQRT2
 
 
-def vecmat_to_volmat(X):
-    """
-    Roll up vector matrices into volume matrices
-
-    :param X: A vector matrix of size L1^3-by-L2^3-by-...
-    :return: A volume "matrix" of size L1-by-L1-by-L1-by-L2-by-L2-by-L2-by-...
-    """
-    # TODO: Use context manager?
-    shape = X.shape
-    assert X.ndim >= 2, "Array should have at least 2 dimensions"
-
-    L1 = round(shape[0] ** (1 / 3))
-    L2 = round(shape[1] ** (1 / 3))
-
-    assert L1**3 == shape[0], "First dimension of X must be cubic"
-    assert L2**3 == shape[1], "Second dimension of X must be cubic"
-
-    return m_reshape(X, (L1, L1, L1, L2, L2, L2) + (shape[2:]))
-
-
-def volmat_to_vecmat(X):
-    """
-    Unroll volume matrices to vector matrices
-
-    :param X: A volume "matrix" of size L1-by-L1-by-L1-by-L2-by-L2-by-L2-by-...
-    :return: A vector matrix of size L1^3-by-L2^3-by-...
-    """
-    # TODO: Use context manager?
-    shape = X.shape
-    assert X.ndim >= 6, "Array should have at least 6 dimensions"
-    assert shape[0] == shape[1] == shape[2], "Dimensions 1-3 should be identical"
-    assert shape[3] == shape[4] == shape[5], "Dimensions 4-6 should be identical"
-
-    l1 = shape[0]
-    l2 = shape[3]
-
-    return m_reshape(X, (l1**3, l2**3) + (shape[6:]))
-
-
 def mdim_mat_fun_conj(X, d1, d2, f):
     """
     Conjugate a multidimensional matrix using a linear mapping
diff --git a/tests/test_matrix.py b/tests/test_matrix.py
index be1fda9343..540c86a377 100644
--- a/tests/test_matrix.py
+++ b/tests/test_matrix.py
@@ -16,8 +16,6 @@
     utest_tolerance,
     vec_to_symmat,
     vec_to_symmat_iso,
-    vecmat_to_volmat,
-    volmat_to_vecmat,
 )
 
 DATA_DIR = os.path.join(os.path.dirname(__file__), "saved_test_data")
@@ -30,18 +28,6 @@ def setUp(self):
     def tearDown(self):
         pass
 
-    def testVecmatToVolmat(self):
-        m = np.empty((8, 27, 10))
-        m2 = vecmat_to_volmat(m)
-
-        self.assertEqual(m2.shape, (2, 2, 2, 3, 3, 3, 10))
-
-    def testVolmatToVecmat(self):
-        m = np.empty((3, 3, 3, 2, 2, 2, 5))
-        m2 = volmat_to_vecmat(m)
-
-        self.assertEqual(m2.shape, (27, 8, 5))
-
     def testMatToVec1(self):
         m = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
 

From 2714f40c1337d3e84c9730666c59220c063dd172 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Thu, 21 Nov 2024 14:18:59 -0500
Subject: [PATCH 057/184] purge m_reshape

---
 src/aspire/operators/filters.py   |  1 -
 src/aspire/utils/matlab_compat.py | 20 --------------------
 src/aspire/utils/matrix.py        |  7 ++-----
 src/aspire/utils/random.py        | 14 +++++++-------
 4 files changed, 9 insertions(+), 33 deletions(-)

diff --git a/src/aspire/operators/filters.py b/src/aspire/operators/filters.py
index f290d88eb1..22a3fffa55 100644
--- a/src/aspire/operators/filters.py
+++ b/src/aspire/operators/filters.py
@@ -126,7 +126,6 @@ def evaluate_grid(self, L, *args, dtype=np.float32, **kwargs):
         :return: Filter values at omega's points.
         """
 
-        # Note we can probably unwind the "F"/m_reshape here
         grid2d = grid_2d(L, indexing="yx", dtype=dtype)
         omega = np.pi * np.vstack((grid2d["x"].flatten(), grid2d["y"].flatten()))
         h = self.evaluate(omega, *args, **kwargs)
diff --git a/src/aspire/utils/matlab_compat.py b/src/aspire/utils/matlab_compat.py
index 4b46e2227f..56bc1eecaf 100644
--- a/src/aspire/utils/matlab_compat.py
+++ b/src/aspire/utils/matlab_compat.py
@@ -9,26 +9,6 @@
 import scipy.sparse as sparse
 
 
-def m_reshape(x, new_shape):
-    # This is a somewhat round-about way of saying:
-    #   return x.reshape(new_shape, order='F')
-    # We follow this approach since numba/cupy don't support the 'order'
-    # argument, and we may want to use those decorators in the future
-    # Note that flattening is required before reshaping, because
-    if isinstance(new_shape, tuple):
-        return m_flatten(x).reshape(new_shape[::-1]).T
-    else:
-        return x
-
-
-def m_flatten(x):
-    # This is a somewhat round-about way of saying:
-    #   return x.flatten(order='F')
-    # We follow this approach since numba/cupy don't support the 'order'
-    # argument, and we may want to use those decorators in the future
-    return x.T.flatten()
-
-
 def stable_eigsh(*args, **kwargs):
     """
     A Wrapper function to fix sign problem of eigen-vectors
diff --git a/src/aspire/utils/matrix.py b/src/aspire/utils/matrix.py
index 2b7373fadb..19bd1287f4 100644
--- a/src/aspire/utils/matrix.py
+++ b/src/aspire/utils/matrix.py
@@ -5,8 +5,6 @@
 import numpy as np
 from scipy.linalg import eigh
 
-from aspire.utils.matlab_compat import m_reshape
-
 SQRT2 = np.sqrt(2)
 SQRT2_R = 1 / SQRT2
 
@@ -260,7 +258,7 @@ def eigs(A, k):
     """
     sig_sz = A.shape[: int(A.ndim / 2)]
     sig_len = np.prod(sig_sz)
-    A = m_reshape(A, (sig_len, sig_len))
+    A = A.reshape((sig_len, sig_len))
 
     dtype = A.dtype
     w, v = eigh(
@@ -271,8 +269,7 @@ def eigs(A, k):
     w = w[::-1].astype(dtype)
     v = np.fliplr(v)
 
-    v = m_reshape(v, sig_sz + (k,)).astype(dtype)
-
+    v = v.reshape(sig_sz + (k,)).astype(dtype)
     return v, np.diag(w)
 
 
diff --git a/src/aspire/utils/random.py b/src/aspire/utils/random.py
index f19d3abd93..416ce25650 100644
--- a/src/aspire/utils/random.py
+++ b/src/aspire/utils/random.py
@@ -5,8 +5,6 @@
 import numpy as np
 from scipy.special import erfinv
 
-from aspire.utils.matlab_compat import m_reshape
-
 # A list of random states, used as a stack
 random_states = []
 
@@ -49,19 +47,21 @@ def randn(*args, **kwargs):
     with Random(seed):
         uniform = np.random.rand(*args, **kwargs)
         result = np.sqrt(2) * erfinv(2 * uniform - 1)
-        # TODO: Rearranging elements to get consistent behavior with MATLAB 'randn2'
-        result = m_reshape(result.flatten(), args)
+        # Note, rearranging elements to get consistent behavior with MATLAB 'randn2'
+        result = result.T.reshape(args, order="F")
+
         return result
 
 
 def rand(size, seed=None):
     """
-    Note this is for MATLAB repro (see m_reshape).
+    Wraps numpy.random.random with ASPIRE Random context manager.
+    Note this is re-packed from F order, only for MATLAB repro.
 
-    Other uses prefer use of `random`.
+    For all other uses prefer `random`.
     """
     with Random(seed):
-        return m_reshape(np.random.random(np.prod(size)), size)
+        return np.random.random(size).reshape(size, order="F")
 
 
 def random(*args, **kwargs):

From a8c61c921165766584af6fbc40fe269685dc0c22 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Thu, 21 Nov 2024 15:18:16 -0500
Subject: [PATCH 058/184] remove flat applications of rand

---
 src/aspire/abinitio/commonline_sync3n.py | 4 ++--
 src/aspire/classification/rir_class2d.py | 5 ++---
 src/aspire/source/simulation.py          | 4 ++--
 src/aspire/utils/random.py               | 9 +++++++++
 4 files changed, 15 insertions(+), 7 deletions(-)

diff --git a/src/aspire/abinitio/commonline_sync3n.py b/src/aspire/abinitio/commonline_sync3n.py
index a0db2263ec..6b89e4bce2 100644
--- a/src/aspire/abinitio/commonline_sync3n.py
+++ b/src/aspire/abinitio/commonline_sync3n.py
@@ -12,11 +12,11 @@
     Rotation,
     all_pairs,
     nearest_rotations,
+    random,
     tqdm,
     trange,
 )
 from aspire.utils.matlab_compat import stable_eigsh
-from aspire.utils.random import rand
 
 logger = logging.getLogger(__name__)
 
@@ -1029,7 +1029,7 @@ def _J_sync_power_method(self, Rijs):
 
         # Initialize candidate eigenvectors
         n_Rijs = Rijs.shape[0]
-        vec = rand(n_Rijs, seed=self.seed)
+        vec = random(n_Rijs, seed=self.seed)
         vec = vec / norm(vec)
         residual = 1
         itr = 0
diff --git a/src/aspire/classification/rir_class2d.py b/src/aspire/classification/rir_class2d.py
index 19a3479ffe..a29db41eee 100644
--- a/src/aspire/classification/rir_class2d.py
+++ b/src/aspire/classification/rir_class2d.py
@@ -8,8 +8,7 @@
 from aspire.classification import Class2D
 from aspire.classification.legacy_implementations import bispec_2drot_large, pca_y
 from aspire.numeric import ComplexPCA
-from aspire.utils import trange
-from aspire.utils.random import rand
+from aspire.utils import random, trange
 
 logger = logging.getLogger(__name__)
 
@@ -428,7 +427,7 @@ def _devel_bispectrum(self, coef):
             self.pca_basis.complex_angular_indices != 0
         ]  # filter non_zero_freqs eq 18,19
         pm = m / np.sum(m)
-        x = rand(len(m))
+        x = random(len(m))
         m_mask = x < self.sample_n * pm
 
         M = None
diff --git a/src/aspire/source/simulation.py b/src/aspire/source/simulation.py
index 15dd1827e4..1b2f18a7be 100644
--- a/src/aspire/source/simulation.py
+++ b/src/aspire/source/simulation.py
@@ -17,7 +17,7 @@
     make_symmat,
     uniform_random_angles,
 )
-from aspire.utils.random import rand, randi, randn
+from aspire.utils.random import randi, randn, random
 from aspire.volume import AsymmetricVolume, Volume
 
 logger = logging.getLogger(__name__)
@@ -143,7 +143,7 @@ def __init__(
 
         if amplitudes is None:
             min_, max_ = 2.0 / 3, 3.0 / 2
-            amplitudes = min_ + rand(n, seed=seed).astype(dtype) * (max_ - min_)
+            amplitudes = min_ + random(n, seed=seed).astype(dtype) * (max_ - min_)
 
         self.C = self.vols.n_vols
 
diff --git a/src/aspire/utils/random.py b/src/aspire/utils/random.py
index 416ce25650..4ccc5f74f9 100644
--- a/src/aspire/utils/random.py
+++ b/src/aspire/utils/random.py
@@ -2,6 +2,8 @@
 Utilities for controlling and generating random numbers.
 """
 
+import warnings
+
 import numpy as np
 from scipy.special import erfinv
 
@@ -60,6 +62,13 @@ def rand(size, seed=None):
 
     For all other uses prefer `random`.
     """
+    warnings.warn(
+        "`rand` is deprecated and included only for MATLAB"
+        " reference applications.  It may be removed in future"
+        " releases.  Use `random` for new development.",
+        DeprecationWarning,
+        stacklevel=2
+    )
     with Random(seed):
         return np.random.random(size).reshape(size, order="F")
 

From 0a22e342536505e81d385fd793b7a0850db298dd Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Thu, 21 Nov 2024 16:05:56 -0500
Subject: [PATCH 059/184] resolve rebase conflicts

---
 src/aspire/basis/fb_3d.py  | 4 +---
 src/aspire/utils/random.py | 2 +-
 2 files changed, 2 insertions(+), 4 deletions(-)

diff --git a/src/aspire/basis/fb_3d.py b/src/aspire/basis/fb_3d.py
index cd04aa1f3b..c8aafd2566 100644
--- a/src/aspire/basis/fb_3d.py
+++ b/src/aspire/basis/fb_3d.py
@@ -148,8 +148,6 @@ def _evaluate(self, v):
             This is an array whose first dimensions equal `self.z` and the
             remaining dimensions correspond to dimensions two and higher of `v`.
         """
-        stack_shape = v.shape[:-1]
-        v = v.reshape(-1, v.shape[-1])
         r_idx = self.basis_coords["r_idx"]
         ang_idx = self.basis_coords["ang_idx"]
         mask = self.basis_coords["mask"].flatten()
@@ -216,4 +214,4 @@ def _evaluate_t(self, x):
 
             ind_radial += len(idx_radial)
 
-        return v.reshape(*stack_shape, self.count)
+        return v
diff --git a/src/aspire/utils/random.py b/src/aspire/utils/random.py
index 4ccc5f74f9..035f2f5d61 100644
--- a/src/aspire/utils/random.py
+++ b/src/aspire/utils/random.py
@@ -67,7 +67,7 @@ def rand(size, seed=None):
         " reference applications.  It may be removed in future"
         " releases.  Use `random` for new development.",
         DeprecationWarning,
-        stacklevel=2
+        stacklevel=2,
     )
     with Random(seed):
         return np.random.random(size).reshape(size, order="F")

From ef98e4c84495f9481a680bfa5871fb74c43c0db0 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Tue, 26 Nov 2024 15:06:33 -0500
Subject: [PATCH 060/184] resolve future tuple syntax with Python 3.9

try changing syntax
---
 src/aspire/utils/matrix.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/src/aspire/utils/matrix.py b/src/aspire/utils/matrix.py
index 19bd1287f4..e0322e966e 100644
--- a/src/aspire/utils/matrix.py
+++ b/src/aspire/utils/matrix.py
@@ -96,7 +96,10 @@ def symmat_to_vec(mat):
 
     sz = mat.shape[:-2]
     tri_indices = np.triu_indices(N)
-    vec = mat[..., *tri_indices].reshape(*sz, N * (N + 1) // 2)
+    # Python 3.11 this can change to mat[..., *tri_indices]
+    #   See PEP 646, variadics
+    #   https://peps.python.org/pep-0646/#multiple-unpackings-in-a-tuple-not-allowed
+    vec = mat[..., tri_indices[0], tri_indices[1]].reshape(*sz, N * (N + 1) // 2)
 
     return vec
 

From afb965ecd4b0e4692f7341bc04a90ac0f3f322c1 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Mon, 9 Dec 2024 15:58:40 -0500
Subject: [PATCH 061/184] cleanup cov3d sim eigs trans

---
 src/aspire/source/simulation.py | 13 +++++--------
 1 file changed, 5 insertions(+), 8 deletions(-)

diff --git a/src/aspire/source/simulation.py b/src/aspire/source/simulation.py
index 1b2f18a7be..de5f0957c7 100644
--- a/src/aspire/source/simulation.py
+++ b/src/aspire/source/simulation.py
@@ -395,11 +395,10 @@ def mean_true(self):
 
     def covar_true(self):
         eigs_true, lamdbas_true = self.eigs()
-        eigs_true = eigs_true.T.to_vec()
+        eigs_true = eigs_true.to_vec()
 
         covar_true = eigs_true.T @ lamdbas_true @ eigs_true
-        # Hrmm
-        covar_true = covar_true.reshape((self.L,) * 6).transpose(2, 1, 0, 5, 4, 3)
+        covar_true = covar_true.reshape((self.L,) * 6)
 
         return covar_true
 
@@ -415,7 +414,6 @@ def eigs(self):
         vols_c = self.vols - self.mean_true()
 
         p = np.ones(C) / C
-        # RCOPT, we may be able to do better here if we dig in.
         Q, R = qr(vols_c.to_vec().T, mode="economic")
 
         # Rank is at most C-1, so remove last vector
@@ -423,11 +421,10 @@ def eigs(self):
         R = R[:-1, :]
 
         w, v = eigh(make_symmat(R @ np.diag(p) @ R.T))
-        eigs_true = Volume.from_vec((Q @ v).T)
-
         # Arrange in descending order (flip column order in eigenvector matrix)
-        w = w[::-1]
-        eigs_true = Volume(eigs_true.asnumpy()[::-1])
+        w, v = w[::-1], v[::-1]
+
+        eigs_true = Volume.from_vec((Q @ v).T)
 
         return eigs_true, np.diag(w)
 

From 080ecdb8ecd7ecff98bb74eab99fd48cd58103db Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Tue, 10 Dec 2024 10:19:17 -0500
Subject: [PATCH 062/184] rename legacy `rand` to `matlab_rand`

---
 src/aspire/utils/__init__.py | 2 +-
 src/aspire/utils/random.py   | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/aspire/utils/__init__.py b/src/aspire/utils/__init__.py
index 69e0c7df9f..df4134a6d2 100644
--- a/src/aspire/utils/__init__.py
+++ b/src/aspire/utils/__init__.py
@@ -59,7 +59,7 @@
     physical_core_cpu_suggestion,
     virtual_core_cpu_suggestion,
 )
-from .random import Random, choice, rand, randi, randn, random
+from .random import Random, choice, matlab_rand, randi, randn, random
 from .relion_interop import RelionStarFile, relion_metadata_fields
 from .resolution_estimation import FourierRingCorrelation, FourierShellCorrelation
 from .rotation import Rotation
diff --git a/src/aspire/utils/random.py b/src/aspire/utils/random.py
index 035f2f5d61..22cd0cf9f6 100644
--- a/src/aspire/utils/random.py
+++ b/src/aspire/utils/random.py
@@ -55,7 +55,7 @@ def randn(*args, **kwargs):
         return result
 
 
-def rand(size, seed=None):
+def matlab_rand(size, seed=None):
     """
     Wraps numpy.random.random with ASPIRE Random context manager.
     Note this is re-packed from F order, only for MATLAB repro.
@@ -63,7 +63,7 @@ def rand(size, seed=None):
     For all other uses prefer `random`.
     """
     warnings.warn(
-        "`rand` is deprecated and included only for MATLAB"
+        "`matlab_rand` is deprecated and included only for MATLAB"
         " reference applications.  It may be removed in future"
         " releases.  Use `random` for new development.",
         DeprecationWarning,

From 626237d5180e572dfe14e0442c9b145027a18c7c Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Thu, 12 Dec 2024 08:55:51 -0500
Subject: [PATCH 063/184] Add mrcs extension to ext list for Image.load

---
 src/aspire/image/image.py |  1 +
 tests/test_image.py       | 26 +++++++++++++++++++++++++-
 2 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/src/aspire/image/image.py b/src/aspire/image/image.py
index 1e980e88dd..9b76402ea4 100644
--- a/src/aspire/image/image.py
+++ b/src/aspire/image/image.py
@@ -147,6 +147,7 @@ class Image:
     # Map file extensions to their respective readers
     extensions = {
         ".mrc": load_mrc,
+        ".mrcs": load_mrc,
         ".tif": load_tiff,
         ".tiff": load_tiff,
     }
diff --git a/tests/test_image.py b/tests/test_image.py
index 593f781e7d..4b76d96883 100644
--- a/tests/test_image.py
+++ b/tests/test_image.py
@@ -458,7 +458,7 @@ def test_load_bad_ext():
 
 def test_load_mrc(dtype):
     """
-    Test `Image.load` round-trip.
+    Test `Image.load` round-trip for `mrc` extension.
     """
 
     # `sample.mrc` is single precision
@@ -480,6 +480,30 @@ def test_load_mrc(dtype):
     assert im2.dtype == dtype
 
 
+def test_load_mrcs(dtype):
+    """
+    Test `Image.load` round-trip for `mrcs` extension.
+    """
+
+    # `sample.mrcs` is single precision
+    filepath = os.path.join(DATA_DIR, "sample.mrcs")
+
+    # Load data from file
+    im = Image.load(filepath, dtype=dtype)
+
+    with tempfile.TemporaryDirectory() as tmpdir_name:
+        # tmp filename
+        test_filepath = os.path.join(tmpdir_name, "test.mrcs")
+
+        im.save(test_filepath)
+
+        im2 = Image.load(test_filepath, dtype)
+
+    # Check the single precision round-trip
+    assert np.array_equal(im, im2)
+    assert im2.dtype == dtype
+
+
 def test_load_tiff():
     """
     Test `Image.load` with a TIFF file

From b8f84e2f07d5efa047ecf837f0584f375bab9d7c Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Tue, 7 Jan 2025 13:37:36 -0500
Subject: [PATCH 064/184] Update sph_harm tests for scipy 1.15.0

---
 tests/test_basis_utils.py | 23 +++++++++++++----------
 1 file changed, 13 insertions(+), 10 deletions(-)

diff --git a/tests/test_basis_utils.py b/tests/test_basis_utils.py
index 2e416a5dee..f6f69b5f6b 100644
--- a/tests/test_basis_utils.py
+++ b/tests/test_basis_utils.py
@@ -25,27 +25,33 @@ def test_sph_harm_low_order():
     y = np.linspace(0, 2 * np.pi, 42)
 
     ref = sp_sph_harm(m, j, y, x)  # Note calling convention is different
-    np.testing.assert_allclose(sph_harm(j, m, x, y), ref)
+    # Prescribe an atol because some of the ref values can be very
+    # small, which can impact relative tolerance.
+    np.testing.assert_allclose(sph_harm(j, m, x, y), ref, atol=1e-8)
 
     # negative m
     m *= -1
     ref = sp_sph_harm(m, j, y, x)  # Note calling convention is different
-    np.testing.assert_allclose(sph_harm(j, m, x, y), ref)
+    # Prescribe an atol because some of the ref values can be very
+    # small, which can impact relative tolerance.
+    np.testing.assert_allclose(sph_harm(j, m, x, y), ref, atol=1e-8)
 
 
 def test_sph_harm_high_order():
     """
-    Test we remain finite at higher orders where `scipy.special.sph_harm` overflows.
+    Test we remain finite at higher orders where legacy `scipy.special.sph_harm` overflowed.
     """
+    # Older (<1.15.0) versions of Scipy overflowed with these values.
+    # Scipy>=1.15.0 has better overflow behavior,
+    #   but the method `sph_harm` will be deprecated in 1.17.0.
+
     m = 87
     j = 87
     x = 0.12345
     y = 0.56789
 
-    # If scipy fixed their implementation for higher orders in the future,
-    # this check should fail and we can reconsider that package.
-    ref = sp_sph_harm(m, j, y, x)  # Note calling convention is different
-    assert not np.isfinite(ref)
+    # Check we are finite.
+    assert np.isfinite(sph_harm(j, m, x, y))
 
     # Can manually check against pyshtools,
     # but we are avoiding that package dependency.
@@ -61,9 +67,6 @@ def test_sph_harm_high_order():
     #     normalization="ortho",
     # )
 
-    # Check we are finite.
-    assert np.isfinite(sph_harm(j, m, x, y))
-
 
 class BesselTestCase(TestCase):
     def setUp(self):

From e39f813b62ad0287f096ece8b4008153edd43f82 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Tue, 7 Jan 2025 13:48:54 -0500
Subject: [PATCH 065/184] Fix future deprecation of sph_harm

replace with reference method with sph_harm_y
---
 tests/test_basis_utils.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tests/test_basis_utils.py b/tests/test_basis_utils.py
index f6f69b5f6b..ea669dd054 100644
--- a/tests/test_basis_utils.py
+++ b/tests/test_basis_utils.py
@@ -1,7 +1,7 @@
 from unittest import TestCase
 
 import numpy as np
-from scipy.special import sph_harm as sp_sph_harm
+from scipy.special import sph_harm_y
 
 from aspire.basis.basis_utils import (
     all_besselj_zeros,
@@ -24,14 +24,14 @@ def test_sph_harm_low_order():
     x = np.linspace(0, np.pi, 42)
     y = np.linspace(0, 2 * np.pi, 42)
 
-    ref = sp_sph_harm(m, j, y, x)  # Note calling convention is different
+    ref = sph_harm_y(j, m, x, y)  # Note Scipy calling convention is different
     # Prescribe an atol because some of the ref values can be very
     # small, which can impact relative tolerance.
     np.testing.assert_allclose(sph_harm(j, m, x, y), ref, atol=1e-8)
 
     # negative m
     m *= -1
-    ref = sp_sph_harm(m, j, y, x)  # Note calling convention is different
+    ref = sph_harm_y(j, m, x, y)  # Note Scipy calling convention is different
     # Prescribe an atol because some of the ref values can be very
     # small, which can impact relative tolerance.
     np.testing.assert_allclose(sph_harm(j, m, x, y), ref, atol=1e-8)

From 43f8a3afea1af6350ac8966b0b3c0d9ad46cb672 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Tue, 7 Jan 2025 14:03:44 -0500
Subject: [PATCH 066/184] make unit test compatible with old and new scipy

---
 tests/test_basis_utils.py | 17 ++++++++++++++++-
 1 file changed, 16 insertions(+), 1 deletion(-)

diff --git a/tests/test_basis_utils.py b/tests/test_basis_utils.py
index ea669dd054..e45e61d91f 100644
--- a/tests/test_basis_utils.py
+++ b/tests/test_basis_utils.py
@@ -1,7 +1,22 @@
 from unittest import TestCase
 
 import numpy as np
-from scipy.special import sph_harm_y
+
+# This can be removed when project requires scipy>=1.15.0
+# scipy<1.15.0 provide `sph_harm`
+import scipy
+from packaging.version import Version
+
+if Version(scipy.__version__) < Version("1.15.0"):
+    from scipy.special import sph_harm as sp_sph_harm
+
+    # This has a different convention from upstream sph_harm_y
+    def sph_harm_y(j, m, x, y):
+        return sp_sph_harm(m, j, y, x)
+
+else:
+    # scipy>=1.15.0 provide `sph_harm_y`
+    from scipy.special import sph_harm_y
 
 from aspire.basis.basis_utils import (
     all_besselj_zeros,

From 266e71d09f8deb72c1e6acca93ad082303124735 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Tue, 26 Nov 2024 08:59:36 -0500
Subject: [PATCH 067/184] Add new default and legacy class averager wrappers

---
 src/aspire/denoising/__init__.py  |   2 +-
 src/aspire/denoising/class_avg.py | 181 ++++++++++++++++++++++++++++--
 tests/test_class_src.py           |   9 +-
 3 files changed, 179 insertions(+), 13 deletions(-)

diff --git a/src/aspire/denoising/__init__.py b/src/aspire/denoising/__init__.py
index 920638eb44..ddeb353291 100644
--- a/src/aspire/denoising/__init__.py
+++ b/src/aspire/denoising/__init__.py
@@ -1,5 +1,5 @@
 from .adaptive_support import adaptive_support
-from .class_avg import ClassAvgSource, DebugClassAvgSource, DefaultClassAvgSource
+from .class_avg import ClassAvgSource, DebugClassAvgSource, DefaultClassAvgSource, ClassAvgSourceLegacy
 
 # isort: off
 from .denoiser import Denoiser
diff --git a/src/aspire/denoising/class_avg.py b/src/aspire/denoising/class_avg.py
index a003d0bd1d..d1977b07e1 100644
--- a/src/aspire/denoising/class_avg.py
+++ b/src/aspire/denoising/class_avg.py
@@ -5,13 +5,16 @@
 from aspire.basis import FFBBasis2D
 from aspire.classification import (
     Averager2D,
+    BandedSNRImageQualityFunction,
     BFRAverager2D,
     BFSRAverager2D,
     Class2D,
     ClassSelector,
+    GlobalClassSelector,
     NeighborVarianceWithRepulsionClassSelector,
     RIRClass2D,
     TopClassSelector,
+    VarianceImageQualityFunction,
 )
 from aspire.image import Image
 from aspire.source import ImageSource
@@ -442,6 +445,86 @@ def __init__(
         )
 
 
+class ClassAvgSourceLegacy(ClassAvgSource):
+    """
+    Source for denoised 2D images using class average methods.
+
+    Defaults to using global variance based class selection,
+    and a brute force image alignment (rotational only).
+
+    This is most similar to what was reported for papers using the
+    MATLAB code.
+    """
+
+    def __init__(
+        self,
+        src,
+        n_nbor=50,
+        classifier=None,
+        class_selector=None,
+        averager=None,
+        averager_src=None,
+    ):
+        """
+        Instantiates ClassAvgSourcev132 with the following parameters.
+
+        :param src: Source used for image classification.
+        :param n_nbor: Number of nearest neighbors. Default 50.
+        :param classifier: `Class2D` classifier instance.
+            Default `None` creates `RIRClass2D`.
+            See code for parameter details.
+        :param class_selector: `ClassSelector` instance.
+            Default `None` creates `NeighborVarianceWithRepulsionClassSelector`.
+        :param averager: `Averager2D` instance.
+            Default `None` ceates `BFRAverager2D` instance.
+            See code for parameter details.
+        :param averager_src: Optionally explicitly assign source
+            to BFSRAverager2D during initialization.
+            Raises error when combined with an explicit `averager`
+            argument.
+
+        :return: ClassAvgSource instance.
+        """
+        dtype = src.dtype
+
+        if classifier is None:
+            classifier = RIRClass2D(
+                src,
+                fspca_components=400,
+                bispectrum_components=300,  # Compressed Features after last PCA stage.
+                n_nbor=n_nbor,
+                large_pca_implementation="legacy",
+                nn_implementation="sklearn",  # Note this is different than debug
+                bispectrum_implementation="legacy",
+            )
+
+        if averager is None:
+            if averager_src is None:
+                averager_src = src
+
+            basis_2d = self._get_classifier_basis(classifier)
+
+            averager = BFRAverager2D(
+                composite_basis=basis_2d,
+                src=averager_src,
+                dtype=dtype,
+            )
+        elif averager_src is not None:
+            raise RuntimeError(
+                "When providing an instantiated `averager`, cannot assign `averager_src`."
+            )
+
+        if class_selector is None:
+            class_selector = NeighborVarianceWithRepulsionClassSelector()
+
+        super().__init__(
+            src=src,
+            classifier=classifier,
+            class_selector=class_selector,
+            averager=averager,
+        )
+
+
 def DefaultClassAvgSource(
     src,
     n_nbor=50,
@@ -455,7 +538,8 @@ def DefaultClassAvgSource(
     Source for denoised 2D images using class average methods.
 
     Accepts `version`, to dispatch ClassAvgSource with parameters
-    below.  Default `version` is latest available.
+    below.  Default `version` is latest available.  Different versions
+    may have different defaults.
 
     :param src: Source used for image classification.
     :param n_nbor: Number of nearest neighbors. Default 50.
@@ -463,15 +547,12 @@ def DefaultClassAvgSource(
         Default `None` creates `RIRClass2D`.
         See code for parameter details.
     :param class_selector: `ClassSelector` instance.
-        Default `None` creates `NeighborVarianceWithRepulsionClassSelector`.
     :param averager: `Averager2D` instance.
-        Default `None` ceates `BFSRAverager2D` instance.
-        See code for parameter details.
     :param averager_src: Optionally explicitly assign source
-        to BFSRAverager2D during initialization.
+        to `averager` during initialization.
         Raises error when combined with an explicit `averager`
         argument.
-    :param version: Optionally selects a versioned DefaultClassAvgSource.
+    :param version: Optionally selects a versioned `DefaultClassAvgSource`.
         Defaults to latest available.
     :return: ClassAvgSource instance.
     """
@@ -479,6 +560,7 @@ def DefaultClassAvgSource(
     _versions = {
         None: ClassAvgSourcev110,
         "latest": ClassAvgSourcev110,
+        "0.13.2": ClassAvgSourcev132,
         "0.11.0": ClassAvgSourcev110,
     }
 
@@ -496,6 +578,89 @@ def DefaultClassAvgSource(
     )
 
 
+class ClassAvgSourcev132(ClassAvgSource):
+    """
+    Source for denoised 2D images using class average methods.
+
+    Defaults to using SNR based class selection,
+    avoiding neighbors of previous classes,
+    and a brute force image alignment (rotational only).
+    """
+
+    def __init__(
+        self,
+        src,
+        n_nbor=50,
+        classifier=None,
+        class_selector=None,
+        averager=None,
+        averager_src=None,
+    ):
+        """
+        Instantiates ClassAvgSourcev132 with the following parameters.
+
+        :param src: Source used for image classification.
+        :param n_nbor: Number of nearest neighbors. Default 50.
+        :param classifier: `Class2D` classifier instance.
+            Default `None` creates `RIRClass2D`.
+            See code for parameter details.
+        :param class_selector: `ClassSelector` instance.
+            Default `None` creates `GlobalWithRepulsionClassSelector` with
+            `BandedSNRImageQualityFunction`. This will select the
+            images with the highest banded SNR.
+        :param averager: `Averager2D` instance.
+            Default `None` ceates `BFRAverager2D` instance.
+            See code for parameter details.
+        :param averager_src: Optionally explicitly assign source
+            to BFSRAverager2D during initialization.
+            Raises error when combined with an explicit `averager`
+            argument.
+
+        :return: ClassAvgSource instance.
+        """
+        dtype = src.dtype
+
+        if classifier is None:
+            classifier = RIRClass2D(
+                src,
+                fspca_components=400,
+                bispectrum_components=300,  # Compressed Features after last PCA stage.
+                n_nbor=n_nbor,
+                large_pca_implementation="legacy",
+                nn_implementation="sklearn",  # Note this is different than debug
+                bispectrum_implementation="legacy",
+            )
+
+        if averager is None:
+            if averager_src is None:
+                averager_src = src
+
+            basis_2d = self._get_classifier_basis(classifier)
+
+            averager = BFRAverager2D(
+                composite_basis=basis_2d,
+                src=averager_src,
+                dtype=dtype,
+            )
+        elif averager_src is not None:
+            raise RuntimeError(
+                "When providing an instantiated `averager`, cannot assign `averager_src`."
+            )
+
+        if class_selector is None:
+            quality_function = BandedSNRImageQualityFunction()
+            class_selector = GlobalWithRepulsionClassSelector(
+                averager, quality_function
+            )
+
+        super().__init__(
+            src=src,
+            classifier=classifier,
+            class_selector=class_selector,
+            averager=averager,
+        )
+
+
 class ClassAvgSourcev110(ClassAvgSource):
     """
     Source for denoised 2D images using class average methods.
@@ -503,8 +668,6 @@ class ClassAvgSourcev110(ClassAvgSource):
     Defaults to using Contrast based class selection (on the fly, compressed),
     avoiding neighbors of previous classes,
     and a brute force image alignment.
-
-    Currently this is the most reasonable default for experimental data.
     """
 
     def __init__(
@@ -517,7 +680,7 @@ def __init__(
         averager_src=None,
     ):
         """
-        Instantiates ClassAvgSourcev11 with the following parameters.
+        Instantiates ClassAvgSourcev110 with the following parameters.
 
         :param src: Source used for image classification.
         :param n_nbor: Number of nearest neighbors. Default 50.
diff --git a/tests/test_class_src.py b/tests/test_class_src.py
index 866c08c3e4..4a2b96a620 100644
--- a/tests/test_class_src.py
+++ b/tests/test_class_src.py
@@ -24,7 +24,11 @@
     VarianceImageQualityFunction,
 )
 from aspire.classification.class_selection import _HeapItem
-from aspire.denoising import DebugClassAvgSource, DefaultClassAvgSource
+from aspire.denoising import (
+    ClassAvgSourceLegacy,
+    DebugClassAvgSource,
+    DefaultClassAvgSource,
+)
 from aspire.image import Image
 from aspire.source import RelionSource, Simulation
 from aspire.utils import Rotation
@@ -46,8 +50,7 @@
     np.float64,
     pytest.param(np.float32, marks=pytest.mark.expensive),
 ]
-CLS_SRCS = [DebugClassAvgSource, DefaultClassAvgSource]
-# For very small problems, it usually isn't worth running in parallel.
+CLS_SRCS = [DebugClassAvgSource, DefaultClassAvgSource, ClassAvgSourceLegacy]
 
 
 BASIS = [

From 498bc1a0b04bb0a4de8c7d196b29c3226a24a91d Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Wed, 27 Nov 2024 09:19:03 -0500
Subject: [PATCH 068/184] change 10028 and 10081 to use legacy class avg align

---
 gallery/experiments/experimental_abinitio_pipeline_10028.py | 4 ++--
 gallery/experiments/experimental_abinitio_pipeline_10081.py | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/gallery/experiments/experimental_abinitio_pipeline_10028.py b/gallery/experiments/experimental_abinitio_pipeline_10028.py
index 7892302758..f5c68ea344 100644
--- a/gallery/experiments/experimental_abinitio_pipeline_10028.py
+++ b/gallery/experiments/experimental_abinitio_pipeline_10028.py
@@ -29,7 +29,7 @@
 import numpy as np
 
 from aspire.abinitio import CLSync3N
-from aspire.denoising import DefaultClassAvgSource, DenoisedSource, DenoiserCov2D
+from aspire.denoising import ClassAvgSourceLegacy, DenoisedSource, DenoiserCov2D
 from aspire.noise import AnisotropicNoiseEstimator
 from aspire.reconstruction import MeanEstimator
 from aspire.source import OrientedSource, RelionSource
@@ -137,7 +137,7 @@
 # Now perform classification and averaging for each class.
 # This also demonstrates the potential to use a different source for classification and averaging.
 
-avgs = DefaultClassAvgSource(
+avgs = ClassAvgSourceLegacy(
     classification_src,
     n_nbor=n_nbor,
     averager_src=src,
diff --git a/gallery/experiments/experimental_abinitio_pipeline_10081.py b/gallery/experiments/experimental_abinitio_pipeline_10081.py
index 7d83842c24..37b09515fd 100644
--- a/gallery/experiments/experimental_abinitio_pipeline_10081.py
+++ b/gallery/experiments/experimental_abinitio_pipeline_10081.py
@@ -26,7 +26,7 @@
 from pathlib import Path
 
 from aspire.abinitio import CLSymmetryC3C4
-from aspire.denoising import DefaultClassAvgSource
+from aspire.denoising import ClassAvgSourceLegacy
 from aspire.noise import AnisotropicNoiseEstimator
 from aspire.reconstruction import MeanEstimator
 from aspire.source import OrientedSource, RelionSource
@@ -91,7 +91,7 @@
 
 # Now perform classification and averaging for each class.
 # Automatically configure parallel processing
-avgs = DefaultClassAvgSource(src, n_nbor=n_nbor)
+avgs = ClassAvgSourceLegacy(src, n_nbor=n_nbor)
 
 # We'll continue our pipeline with the first ``n_classes`` from ``avgs``.
 avgs = avgs[:n_classes]

From 59ddef13e38e8021970d42a14a5561bef4f76d9a Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Wed, 27 Nov 2024 09:22:34 -0500
Subject: [PATCH 069/184] use default overwrite in exps

---
 gallery/experiments/experimental_abinitio_pipeline_10028.py | 2 +-
 gallery/experiments/experimental_abinitio_pipeline_10073.py | 2 +-
 gallery/experiments/experimental_abinitio_pipeline_10081.py | 2 +-
 gallery/experiments/simulated_abinitio_pipeline.py          | 2 +-
 4 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/gallery/experiments/experimental_abinitio_pipeline_10028.py b/gallery/experiments/experimental_abinitio_pipeline_10028.py
index f5c68ea344..c9fd0460c1 100644
--- a/gallery/experiments/experimental_abinitio_pipeline_10028.py
+++ b/gallery/experiments/experimental_abinitio_pipeline_10028.py
@@ -146,7 +146,7 @@
 avgs = avgs[:n_classes].cache()
 
 # Save off the set of class average images.
-avgs.save("experimental_10028_class_averages.star", overwrite=True)
+avgs.save("experimental_10028_class_averages.star")
 
 if interactive:
     avgs.images[:10].show()
diff --git a/gallery/experiments/experimental_abinitio_pipeline_10073.py b/gallery/experiments/experimental_abinitio_pipeline_10073.py
index d936a98d65..7e576abf83 100644
--- a/gallery/experiments/experimental_abinitio_pipeline_10073.py
+++ b/gallery/experiments/experimental_abinitio_pipeline_10073.py
@@ -122,7 +122,7 @@
 avgs = avgs[:n_classes].cache()
 
 # Save off the set of class average images.
-avgs.save("experimental_10073_class_averages_global.star", overwrite=True)
+avgs.save("experimental_10073_class_averages_global.star")
 
 
 # %%
diff --git a/gallery/experiments/experimental_abinitio_pipeline_10081.py b/gallery/experiments/experimental_abinitio_pipeline_10081.py
index 37b09515fd..f5dfbcbb42 100644
--- a/gallery/experiments/experimental_abinitio_pipeline_10081.py
+++ b/gallery/experiments/experimental_abinitio_pipeline_10081.py
@@ -97,7 +97,7 @@
 avgs = avgs[:n_classes]
 
 # Save off the set of class average images.
-avgs.save("experimental_10081_class_averages.star", overwrite=True)
+avgs.save("experimental_10081_class_averages.star")
 
 # %%
 # Common Line Estimation
diff --git a/gallery/experiments/simulated_abinitio_pipeline.py b/gallery/experiments/simulated_abinitio_pipeline.py
index e9b81065de..b622970cb9 100644
--- a/gallery/experiments/simulated_abinitio_pipeline.py
+++ b/gallery/experiments/simulated_abinitio_pipeline.py
@@ -166,7 +166,7 @@ def noise_function(x, y):
     avgs.images[:10].show()
 
 # Save off the set of class average images.
-avgs.save("simulated_abinitio_pipeline_class_averages.star", overwrite=True)
+avgs.save("simulated_abinitio_pipeline_class_averages.star")
 
 # %%
 # Common Line Estimation

From edf3ada3b6dede63fffc43b51cf34c4f742b154a Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Mon, 2 Dec 2024 11:34:05 -0500
Subject: [PATCH 070/184] add GlobalVarianceClassSelector and update docs for
 Legacy

---
 docs/source/class_source.rst                  | 18 ++++++++++---
 .../experimental_abinitio_pipeline_10028.py   |  4 +--
 .../experimental_abinitio_pipeline_10081.py   |  4 +--
 src/aspire/classification/__init__.py         |  1 +
 src/aspire/classification/class_selection.py  | 26 +++++++++++++++++++
 src/aspire/denoising/__init__.py              |  7 ++++-
 src/aspire/denoising/class_avg.py             |  8 +++---
 tests/test_class_src.py                       |  4 +--
 8 files changed, 57 insertions(+), 15 deletions(-)

diff --git a/docs/source/class_source.rst b/docs/source/class_source.rst
index c4dd021e14..ed7c5a391b 100644
--- a/docs/source/class_source.rst
+++ b/docs/source/class_source.rst
@@ -49,8 +49,8 @@ required component and assign them here for complete control.
 
 """"""""""
 
-While that allows for full customization, two helper classes are
-provided that supply defaults as a jumping off point.  Both of these
+While that allows for full customization, helper classes are
+provided that supply defaults as a jumping off point.  These
 helper sources only require an input ``Source`` to be instantiated.
 They can still be fully customized, but they are intended to start
 with sensible defaults, so users only need to instantiate the specific
@@ -61,6 +61,7 @@ components they wish to configure.
    classDiagram
       ClassAvgSource <|-- DebugClassAvgSource
       ClassAvgSource <|-- DefaultClassAvgSource
+      ClassAvgSource <|-- LegacyClassAvgSource
       class DebugClassAvgSource{
 	 src: ImageSource
 	 classifier: RIRClass2D
@@ -68,11 +69,19 @@ components they wish to configure.
 	 averager: BFRAverager2D
 	 +images()
       }
+      class LegacyClassAvgSource{
+	 src: ImageSource
+	 classifier: RIRClass2D
+	 class_selector: GlobalVarianceClassSelector
+	 averager: BFRAverager2D
+	 +images()
+      }
       class DefaultClassAvgSource{
-	 version="0.11.0"
+	 version="0.13.2"
 	 src: ImageSource
 	 classifier: RIRClass2D
 	 class_selector: NeighborVarianceWithRepulsionClassSelector
+	 quality_function: BandedSNRImageQualityFunction
 	 averager: BFSRAverager2D
 	 +images()
       }
@@ -86,7 +95,7 @@ mappings etc.
 
 ``DefaultClassAvgSource`` applies the most sensible defaults available
 in the current ASPIRE release.  ``DefaultClassAvgSource`` takes a
-version string, such as ``0.11.0`` which will return a specific
+version string, such as ``0.13.2`` which will return a specific
 configuration.  This version should allow users to perform a similar
 experiment across releases as ASPIRE implements improved methods.
 When a version is not provided, ``DefaultClassAvgSource`` defaults to
@@ -160,6 +169,7 @@ can reduce pipeline run times by an order of magnitude.
        ClassSelector <|-- TopClassSelector
        ClassSelector <|-- RandomClassSelector
        ClassSelector <|-- NeighborVarianceClassSelector
+       ClassSelector <|-- GlobalVarianceClassSelector
        ClassSelector <|-- DistanceClassSelector
        ClassSelector o-- GreedyClassRepulsionMixin
 
diff --git a/gallery/experiments/experimental_abinitio_pipeline_10028.py b/gallery/experiments/experimental_abinitio_pipeline_10028.py
index c9fd0460c1..034d6dfd68 100644
--- a/gallery/experiments/experimental_abinitio_pipeline_10028.py
+++ b/gallery/experiments/experimental_abinitio_pipeline_10028.py
@@ -29,7 +29,7 @@
 import numpy as np
 
 from aspire.abinitio import CLSync3N
-from aspire.denoising import ClassAvgSourceLegacy, DenoisedSource, DenoiserCov2D
+from aspire.denoising import DenoisedSource, DenoiserCov2D, LegacyClassAvgSource
 from aspire.noise import AnisotropicNoiseEstimator
 from aspire.reconstruction import MeanEstimator
 from aspire.source import OrientedSource, RelionSource
@@ -137,7 +137,7 @@
 # Now perform classification and averaging for each class.
 # This also demonstrates the potential to use a different source for classification and averaging.
 
-avgs = ClassAvgSourceLegacy(
+avgs = LegacyClassAvgSource(
     classification_src,
     n_nbor=n_nbor,
     averager_src=src,
diff --git a/gallery/experiments/experimental_abinitio_pipeline_10081.py b/gallery/experiments/experimental_abinitio_pipeline_10081.py
index f5dfbcbb42..0b9eeef9de 100644
--- a/gallery/experiments/experimental_abinitio_pipeline_10081.py
+++ b/gallery/experiments/experimental_abinitio_pipeline_10081.py
@@ -26,7 +26,7 @@
 from pathlib import Path
 
 from aspire.abinitio import CLSymmetryC3C4
-from aspire.denoising import ClassAvgSourceLegacy
+from aspire.denoising import LegacyClassAvgSource
 from aspire.noise import AnisotropicNoiseEstimator
 from aspire.reconstruction import MeanEstimator
 from aspire.source import OrientedSource, RelionSource
@@ -91,7 +91,7 @@
 
 # Now perform classification and averaging for each class.
 # Automatically configure parallel processing
-avgs = ClassAvgSourceLegacy(src, n_nbor=n_nbor)
+avgs = LegacyClassAvgSource(src, n_nbor=n_nbor)
 
 # We'll continue our pipeline with the first ``n_classes`` from ``avgs``.
 avgs = avgs[:n_classes]
diff --git a/src/aspire/classification/__init__.py b/src/aspire/classification/__init__.py
index 16ed250f23..df92d36932 100644
--- a/src/aspire/classification/__init__.py
+++ b/src/aspire/classification/__init__.py
@@ -16,6 +16,7 @@
     ClassSelector,
     DistanceClassSelector,
     GlobalClassSelector,
+    GlobalVarianceClassSelector,
     GlobalWithRepulsionClassSelector,
     NeighborVarianceClassSelector,
     NeighborVarianceWithRepulsionClassSelector,
diff --git a/src/aspire/classification/class_selection.py b/src/aspire/classification/class_selection.py
index d035dc609e..dcc77d1ccd 100644
--- a/src/aspire/classification/class_selection.py
+++ b/src/aspire/classification/class_selection.py
@@ -647,3 +647,29 @@ class RampWeightedVarianceImageQualityFunction(
     """
     Computes the variance of pixels after weighting with Ramp function.
     """
+
+
+class GlobalVarianceClassSelector(GlobalClassSelector):
+    """
+    GlobalClassSelector with VarianceImageQualityFunction.
+
+    Computes per image variance for all images provided by
+    `averager`, and selects for highest variance.
+
+    Requires aligning the entire set of class averages.
+    """
+
+    def __init__(self, averager, heap_size_limit_bytes=2e9):
+        """
+        See `GlobalClassSelector` and `VarianceImageQualityFunction`
+        for additional documentation.
+
+        :param averager: An Averager2D subclass.
+        :param heap_size_limit_bytes: Max heap size in Bytes.
+            Defaults 2GB, 0 will disable.
+        """
+        super().__init__(
+            averager=averager,
+            quality_function=VarianceImageQualityFunction(),
+            heap_size_limit_bytes=heap_size_limit_bytes,
+        )
diff --git a/src/aspire/denoising/__init__.py b/src/aspire/denoising/__init__.py
index ddeb353291..256c135872 100644
--- a/src/aspire/denoising/__init__.py
+++ b/src/aspire/denoising/__init__.py
@@ -1,5 +1,10 @@
 from .adaptive_support import adaptive_support
-from .class_avg import ClassAvgSource, DebugClassAvgSource, DefaultClassAvgSource, ClassAvgSourceLegacy
+from .class_avg import (
+    ClassAvgSource,
+    DebugClassAvgSource,
+    DefaultClassAvgSource,
+    LegacyClassAvgSource,
+)
 
 # isort: off
 from .denoiser import Denoiser
diff --git a/src/aspire/denoising/class_avg.py b/src/aspire/denoising/class_avg.py
index d1977b07e1..e53514c8b4 100644
--- a/src/aspire/denoising/class_avg.py
+++ b/src/aspire/denoising/class_avg.py
@@ -10,11 +10,11 @@
     BFSRAverager2D,
     Class2D,
     ClassSelector,
-    GlobalClassSelector,
+    GlobalVarianceClassSelector,
+    GlobalWithRepulsionClassSelector,
     NeighborVarianceWithRepulsionClassSelector,
     RIRClass2D,
     TopClassSelector,
-    VarianceImageQualityFunction,
 )
 from aspire.image import Image
 from aspire.source import ImageSource
@@ -445,7 +445,7 @@ def __init__(
         )
 
 
-class ClassAvgSourceLegacy(ClassAvgSource):
+class LegacyClassAvgSource(ClassAvgSource):
     """
     Source for denoised 2D images using class average methods.
 
@@ -515,7 +515,7 @@ def __init__(
             )
 
         if class_selector is None:
-            class_selector = NeighborVarianceWithRepulsionClassSelector()
+            class_selector = GlobalVarianceClassSelector(averager=averager)
 
         super().__init__(
             src=src,
diff --git a/tests/test_class_src.py b/tests/test_class_src.py
index 4a2b96a620..f714187739 100644
--- a/tests/test_class_src.py
+++ b/tests/test_class_src.py
@@ -25,9 +25,9 @@
 )
 from aspire.classification.class_selection import _HeapItem
 from aspire.denoising import (
-    ClassAvgSourceLegacy,
     DebugClassAvgSource,
     DefaultClassAvgSource,
+    LegacyClassAvgSource,
 )
 from aspire.image import Image
 from aspire.source import RelionSource, Simulation
@@ -50,7 +50,7 @@
     np.float64,
     pytest.param(np.float32, marks=pytest.mark.expensive),
 ]
-CLS_SRCS = [DebugClassAvgSource, DefaultClassAvgSource, ClassAvgSourceLegacy]
+CLS_SRCS = [DebugClassAvgSource, DefaultClassAvgSource, LegacyClassAvgSource]
 
 
 BASIS = [

From 7c47d68d2dfd456d655663df198d0a2fe71ca3da Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Mon, 2 Dec 2024 14:59:49 -0500
Subject: [PATCH 071/184] address numpy >0d scalar assignment deprecation

---
 src/aspire/classification/class_selection.py | 12 ++++++++++--
 1 file changed, 10 insertions(+), 2 deletions(-)

diff --git a/src/aspire/classification/class_selection.py b/src/aspire/classification/class_selection.py
index dcc77d1ccd..130e559e9d 100644
--- a/src/aspire/classification/class_selection.py
+++ b/src/aspire/classification/class_selection.py
@@ -217,6 +217,10 @@ def __init__(self, value, index, image):
         :param index: Image index
         :param image: Image object
         """
+        # Numpy scalar deprecation
+        if isinstance(value, np.ndarray) and value.ndim > 0:
+            value = value.item()
+
         self.value = float(value)
         self.index = int(index)
         self.image = image
@@ -324,6 +328,10 @@ def _select(self, classes, reflections, distances):
         for i, im in enumerate(self.averager.average(classes, reflections)):
             quality_score = self._quality_function(im)
 
+            # Numpy scalar deprecation
+            if isinstance(quality_score, np.ndarray) and quality_score.ndim > 0:
+                quality_score = quality_score.item()
+
             # Assign in global quality score array
             self._quality_scores[i] = quality_score
 
@@ -512,7 +520,7 @@ def _function(self, img):
 
         :return: Pixel variance.
         """
-        return np.var(img)
+        return np.var(img).item()
 
 
 class BandedSNRImageQualityFunction(ImageQualityFunction):
@@ -547,7 +555,7 @@ def _function(self, img, center_radius=0.5, outer_band_start=0.8, outer_band_end
                 f"Band of ({outer_band_start}, {outer_band_end}) empty for image size {L}, adjust band boundaries."
             )
 
-        return np.var(img[center_mask]) / np.var(img[outer_mask])
+        return (np.var(img[center_mask]) / np.var(img[outer_mask])).item()
 
 
 class BandpassImageQualityFunction(ImageQualityFunction):

From a1d73189c49823a68442fa4a36d6ccf04a93ec7b Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Thu, 2 Jan 2025 10:09:10 -0500
Subject: [PATCH 072/184] self review string and default updates

---
 src/aspire/denoising/class_avg.py | 35 +++++++++++++++++--------------
 1 file changed, 19 insertions(+), 16 deletions(-)

diff --git a/src/aspire/denoising/class_avg.py b/src/aspire/denoising/class_avg.py
index e53514c8b4..9b61864fa9 100644
--- a/src/aspire/denoising/class_avg.py
+++ b/src/aspire/denoising/class_avg.py
@@ -466,7 +466,7 @@ def __init__(
         averager_src=None,
     ):
         """
-        Instantiates ClassAvgSourcev132 with the following parameters.
+        Instantiates `ClassAvgSource` with the following parameters.
 
         :param src: Source used for image classification.
         :param n_nbor: Number of nearest neighbors. Default 50.
@@ -474,14 +474,15 @@ def __init__(
             Default `None` creates `RIRClass2D`.
             See code for parameter details.
         :param class_selector: `ClassSelector` instance.
-            Default `None` creates `NeighborVarianceWithRepulsionClassSelector`.
+            Default `None` creates `GlobalVarianceClassSelector`.
         :param averager: `Averager2D` instance.
             Default `None` ceates `BFRAverager2D` instance.
             See code for parameter details.
-        :param averager_src: Optionally explicitly assign source
-            to BFSRAverager2D during initialization.
-            Raises error when combined with an explicit `averager`
-            argument.
+        :param averager_src: Optionally explicitly assign source to
+            `BFRAverager2D` during initialization.  Allows users to
+             provide distinct sources for classification and
+             averaging.  Raises error when combined with an explicit
+             `averager` argument.
 
         :return: ClassAvgSource instance.
         """
@@ -548,18 +549,19 @@ def DefaultClassAvgSource(
         See code for parameter details.
     :param class_selector: `ClassSelector` instance.
     :param averager: `Averager2D` instance.
-    :param averager_src: Optionally explicitly assign source
-        to `averager` during initialization.
-        Raises error when combined with an explicit `averager`
-        argument.
+    :param averager_src: Optionally explicitly assign source to
+        `averager` during initialization.  Allows users to
+         provide distinct sources for classification and
+         averaging.  Raises error when combined with an explicit
+         `averager` argument.
     :param version: Optionally selects a versioned `DefaultClassAvgSource`.
         Defaults to latest available.
     :return: ClassAvgSource instance.
     """
 
     _versions = {
-        None: ClassAvgSourcev110,
-        "latest": ClassAvgSourcev110,
+        None: ClassAvgSourcev132,
+        "latest": ClassAvgSourcev132,
         "0.13.2": ClassAvgSourcev132,
         "0.11.0": ClassAvgSourcev110,
     }
@@ -611,10 +613,11 @@ def __init__(
         :param averager: `Averager2D` instance.
             Default `None` ceates `BFRAverager2D` instance.
             See code for parameter details.
-        :param averager_src: Optionally explicitly assign source
-            to BFSRAverager2D during initialization.
-            Raises error when combined with an explicit `averager`
-            argument.
+        :param averager_src: Optionally explicitly assign source to
+            `averager` during initialization.  Allows users to
+             provide distinct sources for classification and
+             averaging.  Raises error when combined with an explicit
+             `averager` argument.
 
         :return: ClassAvgSource instance.
         """

From a29670b3a37c4496de9b789561f8eb3561419bb8 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Thu, 2 Jan 2025 10:31:19 -0500
Subject: [PATCH 073/184] cover both v110 and v132 class average defaults

---
 tests/test_class_src.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/tests/test_class_src.py b/tests/test_class_src.py
index f714187739..c105bf0a6d 100644
--- a/tests/test_class_src.py
+++ b/tests/test_class_src.py
@@ -29,6 +29,7 @@
     DefaultClassAvgSource,
     LegacyClassAvgSource,
 )
+from aspire.denoising.class_avg import ClassAvgSourcev110
 from aspire.image import Image
 from aspire.source import RelionSource, Simulation
 from aspire.utils import Rotation
@@ -50,7 +51,12 @@
     np.float64,
     pytest.param(np.float32, marks=pytest.mark.expensive),
 ]
-CLS_SRCS = [DebugClassAvgSource, DefaultClassAvgSource, LegacyClassAvgSource]
+CLS_SRCS = [
+    DebugClassAvgSource,
+    DefaultClassAvgSource,
+    LegacyClassAvgSource,
+    ClassAvgSourcev110,
+]
 
 
 BASIS = [

From e969ca33727421e6ef91d5ef87564afeea2e1913 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Tue, 7 Jan 2025 10:37:34 -0500
Subject: [PATCH 074/184] Correct CTFFilter formula/units

---
 src/aspire/operators/filters.py | 24 +++++++++++++++++-------
 1 file changed, 17 insertions(+), 7 deletions(-)

diff --git a/src/aspire/operators/filters.py b/src/aspire/operators/filters.py
index 22a3fffa55..b20247861b 100644
--- a/src/aspire/operators/filters.py
+++ b/src/aspire/operators/filters.py
@@ -420,12 +420,15 @@ def __init__(
         """
         A CTF (Contrast Transfer Function) Filter
 
+        Note if comparing to legacy MATLAB cryo_CTF_Relion,
+        take care regarding defocus unit conversion to nm.
+
         :param pixel_size:  Pixel size in angstrom, default 1.
         :param voltage:     Electron voltage in kV
         :param defocus_u:   Defocus depth along the u-axis in angstrom
         :param defocus_v:   Defocus depth along the v-axis in angstrom
         :param defocus_ang: Angle between the x-axis and the u-axis in radians
-        :param Cs:          Spherical aberration constant
+        :param Cs:          Spherical aberration constant in mm
         :param alpha:       Amplitude contrast phase in radians
         :param B:           Envelope decay in inverse square angstrom (default 0)
         """
@@ -440,11 +443,13 @@ def __init__(
         self.alpha = alpha
         self.B = B
 
-        self.defocus_mean = 0.5 * (self.defocus_u + self.defocus_v)
-        self.defocus_diff = 0.5 * (self.defocus_u - self.defocus_v)
+        # Convert angstrom to nm and divide by 2
+        self._defocus_mean_nm = 0.05 * (self.defocus_u + self.defocus_v)
+        self._defocus_diff_nm = 0.05 * (self.defocus_u - self.defocus_v)
 
     def _evaluate(self, omega):
-        om_y, om_x = np.vsplit(omega / (2 * np.pi * self.pixel_size), 2)
+        # Note the grid is wrt nm.
+        om_y, om_x = np.vsplit(omega / (2 * np.pi * self.pixel_size / 10), 2)
 
         eps = np.finfo(np.pi).eps
         ind_nz = (np.abs(om_x) > eps) | (np.abs(om_y) > eps)
@@ -452,10 +457,15 @@ def _evaluate(self, omega):
         angles_nz -= self.defocus_ang
 
         defocus = np.zeros_like(om_x)
-        defocus[ind_nz] = self.defocus_mean + self.defocus_diff * np.cos(2 * angles_nz)
+        # Note the division by 2  for _defocus_diff_nm is in `__init__`.
+        defocus[ind_nz] = self._defocus_mean_nm + self._defocus_diff_nm * np.cos(
+            2 * angles_nz
+        )
 
-        c2 = -np.pi * self.wavelength * defocus
-        c4 = 0.5 * np.pi * (self.Cs * 1e7) * self.wavelength**3
+        # Note lambda must be in nm, and `Cs` must be converted from mm to nm.
+        lambda_nm = self.wavelength / 10
+        c2 = -np.pi * lambda_nm * defocus
+        c4 = 0.5 * np.pi * (self.Cs * 1e6) * lambda_nm**3
 
         r2 = om_x**2 + om_y**2
         r4 = r2**2

From a94bf0d16a8196389ed29884c21eff1bb997e0af Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Tue, 7 Jan 2025 11:02:51 -0500
Subject: [PATCH 075/184] Add a hard coded CTFFilter reference from MATLAB

---
 tests/test_filters.py | 33 +++++++++++++++++++++++++++++++++
 1 file changed, 33 insertions(+)

diff --git a/tests/test_filters.py b/tests/test_filters.py
index da30d6c400..bf0bbbc2cb 100644
--- a/tests/test_filters.py
+++ b/tests/test_filters.py
@@ -399,3 +399,36 @@ def test_array_filter_dtype_passthrough(dtype):
     filt_vals = filt.evaluate_grid(L, dtype=dtype)
 
     assert filt_vals.dtype == dtype
+
+
+def test_ctf_reference():
+    """
+    Test CTFFilter against a MATLAB reference.
+    """
+    fltr = CTFFilter(
+        pixel_size=4.56,
+        voltage=200,
+        defocus_u=10000,
+        defocus_v=15000,
+        defocus_ang=1.23,
+        Cs=2.0,
+        alpha=0.1,
+    )
+    h = fltr.evaluate_grid(5)
+
+    # Compare with MATLAB.  Note DF converted to nm
+    # >> n=5; V=200; DF1=1000; DF2=1500; theta=1.23; Cs=2.0; A=0.1; pxA=4.56;
+    # >> ref_h=cryo_CTF_Relion(n,V,DF1,DF2,theta,Cs,pxA,A)
+    ref_h = np.array(
+        [
+            [-0.6152, 0.0299, -0.5638, 0.9327, 0.9736],
+            [-0.9865, 0.2598, -0.7543, 0.9383, 0.1733],
+            [-0.1876, -0.9918, -0.1000, -0.9918, -0.1876],
+            [0.1733, 0.9383, -0.7543, 0.2598, -0.9865],
+            [0.9736, 0.9327, -0.5638, 0.0299, -0.6152],
+        ]
+    )
+
+    # Test we're within 1%.
+    #   There are minor differences in the formulas for wavelength and grids.
+    np.testing.assert_allclose(h, ref_h, rtol=0.01)

From 16d644bcaeb2c0ca9e7f96bcdc97a0bfb47b94a3 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Tue, 3 Dec 2024 16:15:36 -0500
Subject: [PATCH 076/184] stash add batch ca

---
 src/aspire/classification/averager2d.py | 85 +++++++++++++++++++++++++
 1 file changed, 85 insertions(+)

diff --git a/src/aspire/classification/averager2d.py b/src/aspire/classification/averager2d.py
index b447a6d91c..256f7174f5 100644
--- a/src/aspire/classification/averager2d.py
+++ b/src/aspire/classification/averager2d.py
@@ -6,6 +6,7 @@
 from aspire.basis import Coef
 from aspire.classification.reddy_chatterji import reddy_chatterji_register
 from aspire.image import Image, ImageStacker, MeanImageStacker
+from aspire.numeric import xp
 from aspire.utils import trange
 from aspire.utils.coor_trans import grid_2d
 
@@ -721,6 +722,90 @@ def average(
         return AligningAverager2D.average(self, classes, reflections, coefs)
 
 
+class BBFSR(BFSRAverager2D):
+    """
+    Batch Brute Force Shift and Rotational alignment.
+    """
+
+    def align(self, classes, reflections, basis_coefficients):
+        """
+        During this process `rotations`, `reflections`, `shifts` and
+        `correlations` properties will be computed for aligners.
+
+        `rotations` is an (src.n, n_nbor) array of angles,
+        which should represent the rotations needed to align images within
+        that class. `rotations` is measured in CCW radians.
+
+        `shifts` is None or an (src.n, n_nbor) array of 2D shifts
+        which should represent the translation needed to best align the images
+        within that class.
+
+        `correlations` is an (src.n, n_nbor) array representing
+        a correlation like measure between classified images and their base
+        image (image index 0).
+
+        Subclasses of should implement and extend this method.
+
+        :param classes: (src.n, n_nbor) integer array of img indices.
+        :param reflections: (src.n, n_nbor) bool array of corresponding reflections,
+        :param basis_coefficients: (n_img, self.alignment_basis.count) basis coefficients,
+
+        :returns: (rotations, shifts, correlations)
+        """
+        # Construct array of angles to brute force.
+        _angles = xp.linspace(0, -2 * np.pi, self.n_angles, endpoint=False).reshape(
+            self.n_angles, 1
+        )
+        _rot_ops = xp.exp(
+            1j * self.alignment_basis.complex_angular_indices.reshape(-1, 1) * _angles
+        )
+        # xxx ensure shape
+        assert _rot_ops.shape == (self.alignment_basis.complex_count, self.n_angles)
+
+        # Result arrays
+        n_classes, n_nbor = classes.shape
+        rots = np.empty((n_classes, n_nbor), dtype=self.dtype)
+        correlations = np.empty((n_classes, n_nbor), dtype=self.dtype)
+
+        for k in trange(n_classes):
+            # Get the coefs for these neighbors
+            if basis_coefficients is None:
+                # Retrieve relevant images
+                neighbors_imgs = Image(self._cls_images(classes[k]))
+
+                # Evaluate_t into basis
+                nbr_coef = self.composite_basis.evaluate_t(neighbors_imgs)
+            else:
+                nbr_coef = basis_coefficients[classes[k]]
+
+            # Convert to array of complex coef, implicit copy.
+            nbr_coef = xp.array(nbr_coef.to_complex().asnumpy())
+
+            # Handle reflections
+            refl = reflections[k]
+            nbr_coef[refl] = xp.conj(nbr_coef[refl])
+
+            # Generate table of translations for image 0.
+            # Note we invert the translation, applying -rot to image 0
+            #   to avoid translating each member of the class
+            #   for the alignment test (each a dot product,
+            #   performed as a large matmul).
+            base_img = nbr_coef[0].reshape(self.complex_count, 1)
+
+            # (cnt, n_transl) * (cnt, 1) -> (cnt, n_transl)
+            rot_base_imgs = _rot_ops * base_img
+
+            # (n_nbor, cnt) @ (cnt, n_transl) = (n_nbor, n_transl)
+            dots = nbr_coef @ rot_base_imgs
+            idx = np.argmax(dots, axis=1)
+
+            # Assign results for this class
+            correlations[k] = dots[:, idx].asnumpy()
+            rots[k] = -1 * _angles[idx].asnumpy()  # return the reverse rot
+
+        return rots, correlations
+
+
 class EMAverager2D(Averager2D):
     """
     Citation needed.

From d83d061691a4778d258f31d67342377b95d529fe Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Thu, 5 Dec 2024 11:42:36 -0500
Subject: [PATCH 077/184] stash add, repro reals ref

---
 src/aspire/classification/averager2d.py | 30 ++++++++++++++-----------
 1 file changed, 17 insertions(+), 13 deletions(-)

diff --git a/src/aspire/classification/averager2d.py b/src/aspire/classification/averager2d.py
index 256f7174f5..fa9e4fd45b 100644
--- a/src/aspire/classification/averager2d.py
+++ b/src/aspire/classification/averager2d.py
@@ -753,19 +753,17 @@ def align(self, classes, reflections, basis_coefficients):
         :returns: (rotations, shifts, correlations)
         """
         # Construct array of angles to brute force.
-        _angles = xp.linspace(0, -2 * np.pi, self.n_angles, endpoint=False).reshape(
-            self.n_angles, 1
-        )
+        _angles = xp.linspace(0, 2 * np.pi, self.n_angles, endpoint=False)
+
         _rot_ops = xp.exp(
             1j * self.alignment_basis.complex_angular_indices.reshape(-1, 1) * _angles
         )
-        # xxx ensure shape
-        assert _rot_ops.shape == (self.alignment_basis.complex_count, self.n_angles)
 
         # Result arrays
         n_classes, n_nbor = classes.shape
+        print('dbg n_classes, n_nbor', n_classes, n_nbor)
         rots = np.empty((n_classes, n_nbor), dtype=self.dtype)
-        correlations = np.empty((n_classes, n_nbor), dtype=self.dtype)
+        correlations = np.zeros((n_classes, n_nbor), dtype=self.dtype)
 
         for k in trange(n_classes):
             # Get the coefs for these neighbors
@@ -790,20 +788,26 @@ def align(self, classes, reflections, basis_coefficients):
             #   to avoid translating each member of the class
             #   for the alignment test (each a dot product,
             #   performed as a large matmul).
-            base_img = nbr_coef[0].reshape(self.complex_count, 1)
+            base_img = nbr_coef[0].reshape(self.alignment_basis.complex_count, 1)
 
             # (cnt, n_transl) * (cnt, 1) -> (cnt, n_transl)
             rot_base_imgs = _rot_ops * base_img
+            ## try factoring rot_base_imgs.conj() to here
 
             # (n_nbor, cnt) @ (cnt, n_transl) = (n_nbor, n_transl)
-            dots = nbr_coef @ rot_base_imgs
-            idx = np.argmax(dots, axis=1)
+            dots = xp.real(nbr_coef @ rot_base_imgs.conj())
+            idx = xp.argmax(dots, axis=1)
+            idx[0] = 0  # Force base image, just in case.
 
             # Assign results for this class
-            correlations[k] = dots[:, idx].asnumpy()
-            rots[k] = -1 * _angles[idx].asnumpy()  # return the reverse rot
-
-        return rots, correlations
+            correlations[k,:] = xp.take_along_axis(dots, idx.reshape(n_nbor,1), axis=1).flatten()
+            correlations[k,0] = 1  # Force base correlation, just in case
+            #  Todo, make an actual correlation, (normalize)
+            
+            # Assign the reverse rotation
+            rots[k] = -1 * xp.asnumpy( _angles[idx]) 
+
+        return rots, None, correlations
 
 
 class EMAverager2D(Averager2D):

From 3cd8698bcef4387e5654d8bedb822212cb6df356 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Thu, 5 Dec 2024 13:59:06 -0500
Subject: [PATCH 078/184] move some of the conjugation outside of loop

---
 src/aspire/classification/averager2d.py | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/src/aspire/classification/averager2d.py b/src/aspire/classification/averager2d.py
index fa9e4fd45b..2c4823ad09 100644
--- a/src/aspire/classification/averager2d.py
+++ b/src/aspire/classification/averager2d.py
@@ -755,9 +755,9 @@ def align(self, classes, reflections, basis_coefficients):
         # Construct array of angles to brute force.
         _angles = xp.linspace(0, 2 * np.pi, self.n_angles, endpoint=False)
 
-        _rot_ops = xp.exp(
+        _rot_ops_conj = xp.exp(
             1j * self.alignment_basis.complex_angular_indices.reshape(-1, 1) * _angles
-        )
+        ).conj()
 
         # Result arrays
         n_classes, n_nbor = classes.shape
@@ -791,11 +791,10 @@ def align(self, classes, reflections, basis_coefficients):
             base_img = nbr_coef[0].reshape(self.alignment_basis.complex_count, 1)
 
             # (cnt, n_transl) * (cnt, 1) -> (cnt, n_transl)
-            rot_base_imgs = _rot_ops * base_img
-            ## try factoring rot_base_imgs.conj() to here
+            rot_base_imgs_conj = _rot_ops_conj * base_img.conj()
 
             # (n_nbor, cnt) @ (cnt, n_transl) = (n_nbor, n_transl)
-            dots = xp.real(nbr_coef @ rot_base_imgs.conj())
+            dots = xp.real(nbr_coef @ rot_base_imgs_conj)
             idx = xp.argmax(dots, axis=1)
             idx[0] = 0  # Force base image, just in case.
 

From f95ee652ac59cdbe89bccb4e9b269c48379af2b7 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Mon, 9 Dec 2024 11:12:48 -0500
Subject: [PATCH 079/184] Fix complex conversion dot products

---
 src/aspire/basis/steerable.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/aspire/basis/steerable.py b/src/aspire/basis/steerable.py
index a2b9872886..4065869116 100644
--- a/src/aspire/basis/steerable.py
+++ b/src/aspire/basis/steerable.py
@@ -418,8 +418,8 @@ def to_real(self, complex_coef):
             idx_neg = idx_pos + self.k_max[ell]
 
             c = complex_coef[..., idx]
-            coef[..., idx_pos] = 2.0 * np.real(c)
-            coef[..., idx_neg] = -2.0 * np.imag(c)
+            coef[..., idx_pos] = np.real(c)
+            coef[..., idx_neg] = -np.imag(c)
 
             ind += np.size(idx)
             ind_pos += 2 * self.k_max[ell]
@@ -469,7 +469,7 @@ def to_complex(self, coef):
             idx = ind + np.arange(self.k_max[ell], dtype=int)
             complex_coef[..., idx] = (
                 coef[..., self._pos[idx]] - imaginary * coef[..., self._neg[idx]]
-            ) / 2.0
+            )
 
             ind += np.size(idx)
 

From eed952a3596870b94922cb7e1a6dc012b362c0e3 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Mon, 9 Dec 2024 11:43:22 -0500
Subject: [PATCH 080/184] implement BFR with more optimal linalg

---
 src/aspire/classification/averager2d.py | 162 ++++++------------------
 1 file changed, 36 insertions(+), 126 deletions(-)

diff --git a/src/aspire/classification/averager2d.py b/src/aspire/classification/averager2d.py
index 2c4823ad09..cf84232e2c 100644
--- a/src/aspire/classification/averager2d.py
+++ b/src/aspire/classification/averager2d.py
@@ -272,23 +272,19 @@ def align(self, classes, reflections, basis_coefficients):
         Performs the actual rotational alignment estimation,
         returning parameters needed for averaging.
         """
+        # Construct array of angles to brute force.
+        _angles = xp.linspace(0, 2 * np.pi, self.n_angles, endpoint=False)
 
-        # Admit simple case of single case alignment
-        classes = np.atleast_2d(classes)
-        reflections = np.atleast_2d(reflections)
+        _rot_ops_conj = xp.exp(
+            1j * self.alignment_basis.complex_angular_indices.reshape(-1, 1) * _angles
+        ).conj()
 
+        # Result arrays
         n_classes, n_nbor = classes.shape
+        rots = np.zeros((n_classes, n_nbor), dtype=self.dtype)
+        correlations = np.zeros((n_classes, n_nbor), dtype=self.dtype)
 
-        # Construct array of angles to brute force.
-        test_angles = np.linspace(0, -2 * np.pi, self.n_angles, endpoint=False)
-
-        # Instantiate matrices for results
-        rotations = np.empty(classes.shape, dtype=self.dtype)
-        correlations = np.empty(classes.shape, dtype=self.dtype)
-
-        def _innerloop(k):
-            _correlations = np.full((n_nbor, self.n_angles), fill_value=-np.inf)
-            _correlations[0, 0] = 1  # Set this now so we can skip it in the loop.
+        for k in trange(n_classes):
             # Get the coefs for these neighbors
             if basis_coefficients is None:
                 # Retrieve relevant images
@@ -308,38 +304,39 @@ def _innerloop(k):
             else:
                 nbr_coef = basis_coefficients[classes[k]]
 
-            norm_0 = np.linalg.norm(nbr_coef[0])
-            for i, angle in enumerate(test_angles):
-                # Rotate the set of neighbors by angle,
-                rotated_nbrs = self.alignment_basis.rotate(
-                    nbr_coef, angle, reflections[k]
-                )
+            # Convert to array of complex coef, implicit copy.
+            nbr_coef = xp.array(nbr_coef.to_complex().asnumpy())
 
-                # then store dot between class base image (0) and each nbor
-                for j, nbor in enumerate(rotated_nbrs.asnumpy()):
-                    # Skip the base image.
-                    if j == 0:
-                        continue
-                    norm_nbor = np.linalg.norm(nbor)
-                    _correlations[j, i] = np.dot(nbr_coef.asnumpy()[0], nbor) / (
-                        norm_nbor * norm_0
-                    )
+            # Handle reflections
+            refl = reflections[k]
+            nbr_coef[refl] = xp.conj(nbr_coef[refl])
 
-            # Now find the index of the angle reporting highest correlation
-            angle_idx = np.argmax(_correlations, axis=1)
+            # Generate table of rotations for image 0.
+            # Note we invert the rotation, applying -rot to image 0
+            #   to avoid rotating each member of the class
+            #   for the argmax alignment test (each a dot product,
+            #   performed as a large matmul).
+            base_img = nbr_coef[0].reshape(self.alignment_basis.complex_count, 1)
 
-            # Take the correlation corresponding to angle_idx
-            _correlations = np.take_along_axis(
-                _correlations, np.expand_dims(angle_idx, axis=1), axis=1
-            ).reshape(n_nbor)
+            # (cnt, n_transl) * (cnt, 1) -> (cnt, n_transl)
+            rot_base_imgs_conj = _rot_ops_conj * base_img.conj()
 
-            return test_angles[angle_idx], _correlations
+            # (n_nbor, cnt) @ (cnt, n_transl) = (n_nbor, n_transl)
+            dots = xp.real(nbr_coef @ rot_base_imgs_conj)
+            idx = xp.argmax(dots, axis=1)
+            idx[0] = 0  # Force base image, just in case.
 
-        for k in trange(n_classes):
-            # Store angles and correlations for this class
-            rotations[k], correlations[k] = _innerloop(k)
+            # Assign results for this class
+            correlations[k, :] = xp.take_along_axis(
+                dots, idx.reshape(n_nbor, 1), axis=1
+            ).flatten()
+            # correlations[k,0] = 1  # Force base correlation, just in case
+            #  Todo, do we care to spend the compute to make an actual correlation? (normalize)
+
+            # Assign the reverse rotation
+            rots[k] = -1 * xp.asnumpy(_angles[idx])
 
-        return rotations, None, correlations
+        return rots, None, correlations
 
 
 class BFSRAverager2D(BFRAverager2D):
@@ -722,93 +719,6 @@ def average(
         return AligningAverager2D.average(self, classes, reflections, coefs)
 
 
-class BBFSR(BFSRAverager2D):
-    """
-    Batch Brute Force Shift and Rotational alignment.
-    """
-
-    def align(self, classes, reflections, basis_coefficients):
-        """
-        During this process `rotations`, `reflections`, `shifts` and
-        `correlations` properties will be computed for aligners.
-
-        `rotations` is an (src.n, n_nbor) array of angles,
-        which should represent the rotations needed to align images within
-        that class. `rotations` is measured in CCW radians.
-
-        `shifts` is None or an (src.n, n_nbor) array of 2D shifts
-        which should represent the translation needed to best align the images
-        within that class.
-
-        `correlations` is an (src.n, n_nbor) array representing
-        a correlation like measure between classified images and their base
-        image (image index 0).
-
-        Subclasses of should implement and extend this method.
-
-        :param classes: (src.n, n_nbor) integer array of img indices.
-        :param reflections: (src.n, n_nbor) bool array of corresponding reflections,
-        :param basis_coefficients: (n_img, self.alignment_basis.count) basis coefficients,
-
-        :returns: (rotations, shifts, correlations)
-        """
-        # Construct array of angles to brute force.
-        _angles = xp.linspace(0, 2 * np.pi, self.n_angles, endpoint=False)
-
-        _rot_ops_conj = xp.exp(
-            1j * self.alignment_basis.complex_angular_indices.reshape(-1, 1) * _angles
-        ).conj()
-
-        # Result arrays
-        n_classes, n_nbor = classes.shape
-        print('dbg n_classes, n_nbor', n_classes, n_nbor)
-        rots = np.empty((n_classes, n_nbor), dtype=self.dtype)
-        correlations = np.zeros((n_classes, n_nbor), dtype=self.dtype)
-
-        for k in trange(n_classes):
-            # Get the coefs for these neighbors
-            if basis_coefficients is None:
-                # Retrieve relevant images
-                neighbors_imgs = Image(self._cls_images(classes[k]))
-
-                # Evaluate_t into basis
-                nbr_coef = self.composite_basis.evaluate_t(neighbors_imgs)
-            else:
-                nbr_coef = basis_coefficients[classes[k]]
-
-            # Convert to array of complex coef, implicit copy.
-            nbr_coef = xp.array(nbr_coef.to_complex().asnumpy())
-
-            # Handle reflections
-            refl = reflections[k]
-            nbr_coef[refl] = xp.conj(nbr_coef[refl])
-
-            # Generate table of translations for image 0.
-            # Note we invert the translation, applying -rot to image 0
-            #   to avoid translating each member of the class
-            #   for the alignment test (each a dot product,
-            #   performed as a large matmul).
-            base_img = nbr_coef[0].reshape(self.alignment_basis.complex_count, 1)
-
-            # (cnt, n_transl) * (cnt, 1) -> (cnt, n_transl)
-            rot_base_imgs_conj = _rot_ops_conj * base_img.conj()
-
-            # (n_nbor, cnt) @ (cnt, n_transl) = (n_nbor, n_transl)
-            dots = xp.real(nbr_coef @ rot_base_imgs_conj)
-            idx = xp.argmax(dots, axis=1)
-            idx[0] = 0  # Force base image, just in case.
-
-            # Assign results for this class
-            correlations[k,:] = xp.take_along_axis(dots, idx.reshape(n_nbor,1), axis=1).flatten()
-            correlations[k,0] = 1  # Force base correlation, just in case
-            #  Todo, make an actual correlation, (normalize)
-            
-            # Assign the reverse rotation
-            rots[k] = -1 * xp.asnumpy( _angles[idx]) 
-
-        return rots, None, correlations
-
-
 class EMAverager2D(Averager2D):
     """
     Citation needed.

From b33e4751d40556b13b3d1ba00adfd01e6c794267 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Tue, 10 Dec 2024 10:12:36 -0500
Subject: [PATCH 081/184] fixup CuPy wrapping of BFR

---
 src/aspire/classification/averager2d.py | 11 +++++------
 1 file changed, 5 insertions(+), 6 deletions(-)

diff --git a/src/aspire/classification/averager2d.py b/src/aspire/classification/averager2d.py
index cf84232e2c..f2c9a8aded 100644
--- a/src/aspire/classification/averager2d.py
+++ b/src/aspire/classification/averager2d.py
@@ -275,9 +275,8 @@ def align(self, classes, reflections, basis_coefficients):
         # Construct array of angles to brute force.
         _angles = xp.linspace(0, 2 * np.pi, self.n_angles, endpoint=False)
 
-        _rot_ops_conj = xp.exp(
-            1j * self.alignment_basis.complex_angular_indices.reshape(-1, 1) * _angles
-        ).conj()
+        ks = xp.asarray(self.alignment_basis.complex_angular_indices).reshape(-1, 1)
+        _rot_ops_conj = xp.exp(1j * ks * _angles).conj()
 
         # Result arrays
         n_classes, n_nbor = classes.shape
@@ -327,9 +326,9 @@ def align(self, classes, reflections, basis_coefficients):
             idx[0] = 0  # Force base image, just in case.
 
             # Assign results for this class
-            correlations[k, :] = xp.take_along_axis(
-                dots, idx.reshape(n_nbor, 1), axis=1
-            ).flatten()
+            correlations[k, :] = xp.asnumpy(
+                xp.take_along_axis(dots, idx.reshape(n_nbor, 1), axis=1).flatten()
+            )
             # correlations[k,0] = 1  # Force base correlation, just in case
             #  Todo, do we care to spend the compute to make an actual correlation? (normalize)
 

From 84a0cad7b509f1ae300036be5b8b5e5df332d2e5 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Wed, 11 Dec 2024 08:48:59 -0500
Subject: [PATCH 082/184] stuff batched BFR into BFSR

---
 src/aspire/classification/averager2d.py | 167 ++++++++++++++----------
 1 file changed, 95 insertions(+), 72 deletions(-)

diff --git a/src/aspire/classification/averager2d.py b/src/aspire/classification/averager2d.py
index f2c9a8aded..f69bd1facf 100644
--- a/src/aspire/classification/averager2d.py
+++ b/src/aspire/classification/averager2d.py
@@ -260,7 +260,6 @@ def __init__(
         super().__init__(composite_basis, src, alignment_basis, dtype=dtype)
 
         self.n_angles = n_angles
-        self._base_image_shift = None
 
         if not hasattr(self.alignment_basis, "rotate"):
             raise RuntimeError(
@@ -287,19 +286,10 @@ def align(self, classes, reflections, basis_coefficients):
             # Get the coefs for these neighbors
             if basis_coefficients is None:
                 # Retrieve relevant images
-                neighbors_imgs = self._cls_images(classes[k]).copy()
-
-                # We optionally can shift the base image by `_base_image_shift`
-                # Shift in real space to avoid extra conversions
-                if self._base_image_shift is not None:
-                    neighbors_imgs[0] = (
-                        Image(neighbors_imgs[0])
-                        .shift(self._base_image_shift)
-                        .asnumpy()[0]
-                    )
+                neighbors_imgs = Image(self._cls_images(classes[k]))
 
                 # Evaluate_t into basis
-                nbr_coef = self.composite_basis.evaluate_t(Image(neighbors_imgs))
+                nbr_coef = self.alignment_basis.evaluate_t(neighbors_imgs)
             else:
                 nbr_coef = basis_coefficients[classes[k]]
 
@@ -392,7 +382,22 @@ def align(self, classes, reflections, basis_coefficients):
         classes = np.atleast_2d(classes)
         reflections = np.atleast_2d(reflections)
 
-        n_classes = classes.shape[0]
+        # Result arrays
+        # These arrays will incrementally store our best alignment.
+        n_classes, n_nbor = classes.shape
+        rotations = np.zeros((n_classes, n_nbor), dtype=self.dtype)
+        correlations = np.ones((n_classes, n_nbor), dtype=self.dtype) * -np.inf
+        shifts = np.empty((*classes.shape, 2), dtype=int)
+
+        # Work arrays
+        _rotations = np.zeros((1, n_nbor), dtype=self.dtype)
+        _correlations = np.ones((1, n_nbor), dtype=self.dtype) * -np.inf
+
+        # Construct array of angles to brute force.
+        _angles = xp.linspace(0, 2 * np.pi, self.n_angles, endpoint=False)
+
+        ks = xp.asarray(self.alignment_basis.complex_angular_indices).reshape(-1, 1)
+        _rot_ops_conj = xp.exp(1j * ks * _angles).conj()
 
         # Create a search grid and force initial pair to (0,0)
         # This is done primarily in case of a tie later, we would take unshifted.
@@ -400,68 +405,86 @@ def align(self, classes, reflections, basis_coefficients):
             self.src.L, self.radius, roll_zero=True
         )
 
-        # These arrays will incrementally store our best alignment.
-        rotations = np.empty(classes.shape, dtype=self.dtype)
-        correlations = np.ones(classes.shape, dtype=self.dtype) * -np.inf
-        shifts = np.empty((*classes.shape, 2), dtype=int)
-
-        # We want to maintain the original coefs for the base images,
-        #  because we will mutate them with shifts in the loop.
-        if basis_coefficients is None:
-            original_coef = self.composite_basis.evaluate_t(
-                Image(self._cls_images(classes[:, 0], src=self.src))
-            )
-        else:
-            original_coef = basis_coefficients[classes[:, 0], :].copy()
-
-        # Sanity check the original_coef shape
-        assert original_coef.shape == (n_classes, self.alignment_basis.count)
-
-        # Loop over shift search space, updating best result
-        for x, y in zip(x_shifts, y_shifts):
-            shift = np.array([x, y], dtype=int)
-            logger.debug(f"Computing rotational alignment after shift ({x},{y}).")
-
-            # Shift the coef representing the first (base) entry in each class
-            #   by the negation of the shift
-            # Shifting one image is more efficient than shifting every neighbor
-            if basis_coefficients is not None:
-                basis_coefficients[classes[:, 0], :] = self.alignment_basis.shift(
-                    original_coef, -shift
-                )
-            else:
-                # Store the latest shift so that super class can access it.
-                # This allows us to retrieve and shift coefficients on the fly,
-                #   instead of storing them all.
-                self._base_image_shift = -shift
-
-            _rotations, _, _correlations = self._bfr_align(
-                classes, reflections, basis_coefficients
-            )
-
-            # Each class-neighbor pair may have a best shift-rot from a different shift.
-            # Test and update
-            improved_indices = _correlations > correlations
-            rotations[improved_indices] = _rotations[improved_indices]
-            correlations[improved_indices] = _correlations[improved_indices]
-            shifts[improved_indices] = shift
-
-            # Cleanup/Restore unshifted base coefs
+        for k in trange(n_classes):
+            # We want to locally cache the original images,
+            #  because we will mutate them with shifts in the next loop.
+            #  This avoids recomputing them before each shift
             if basis_coefficients is None:
-                # Reset this flag
-                self._base_image_shift = None
+                original_images = Image(self._cls_images(classes[k], src=self.src))
             else:
-                basis_coefficients[classes[:, 0], :] = original_coef
-
-            if (x, y) == (0, 0):
-                logger.debug("Initial rotational alignment complete (shift (0,0))")
-                assert np.sum(improved_indices) == np.size(
-                    classes
-                ), f"{np.sum(improved_indices)} =?= {np.size(classes)}"
-            else:
-                logger.debug(
-                    f"Shift ({x},{y}) complete. Improved {np.sum(improved_indices)} alignments."
+                original_coef = basis_coefficients[classes[k], :]
+                # batch here? or maybe just force always passing images....?
+                original_images = self.alignment_basis.evaluate(original_coef)
+
+            # Working copy
+            _images = original_images.asnumpy().copy()
+
+            # Loop over shift search space, updating best result
+            for x, y in zip(x_shifts, y_shifts):
+                shift = np.array([x, y], dtype=int)
+                logger.debug(f"Computing rotational alignment after shift ({x},{y}).")
+
+                # For each shift, the set of neighbor images is shifted.
+                #   This order is chosen because:
+                #   i) allows concatenation of shifts and rotation
+                #   operations after orientation estimation
+                #   ii) because generally the number of neighbors << the
+                #   number of test rotations.
+
+                # Note the base image[0] is never shifted.
+                _images[1:] = original_images[1:].shift(shift).asnumpy()
+
+                # Convert to array of complex coef, implicit copy.
+                _coef = self.alignment_basis.evaluate_t(_images)
+                _coef = xp.array(_coef.to_complex().asnumpy())
+
+                # Handle reflections
+                # XXXX should we do this earlier via Images? (flipud?)
+                # How would we communicate this reflection to other areas of the code...
+                refl = reflections[k]
+                _coef[refl] = xp.conj(_coef[refl])
+
+                # Generate table of rotations for image 0.
+                # Note we invert the rotation, applying -rot to image 0
+                #   to avoid rotating each member of the class
+                #   for the argmax alignment test (each a dot product,
+                #   performed as a large matmul).
+                base_img = _coef[0].reshape(self.alignment_basis.complex_count, 1)
+
+                # (cnt, n_transl) * (cnt, 1) -> (cnt, n_transl)
+                rot_base_imgs_conj = _rot_ops_conj * base_img.conj()
+
+                # (n_nbor, cnt) @ (cnt, n_transl) = (n_nbor, n_transl)
+                dots = xp.real(_coef @ rot_base_imgs_conj)
+                idx = xp.argmax(dots, axis=1)
+                idx[0] = 0  # Force base image, just in case.
+
+                # Assign results for this class
+                _correlations[:] = xp.asnumpy(
+                    xp.take_along_axis(dots, idx.reshape(n_nbor, 1), axis=1).flatten()
                 )
+                # _correlations[0] = 1  # Force base correlation, just in case
+                #  Todo, do we care to spend the compute to make an actual correlation? (normalize)
+
+                # Assign the reverse rotation
+                _rotations[:] = -1 * xp.asnumpy(_angles[idx])
+
+                # Each class-neighbor pair may have a best shift-rot from a different shift.
+                # Test and update
+                improved_indices = _correlations[0] > correlations[k]
+                rotations[k, improved_indices] = _rotations[0, improved_indices]
+                correlations[k, improved_indices] = _correlations[0, improved_indices]
+                shifts[k, improved_indices] = shift
+
+                if (x, y) == (0, 0):
+                    logger.debug("Initial rotational alignment complete (shift (0,0))")
+                    assert np.sum(improved_indices) == np.size(
+                        classes[0]
+                    ), f"{np.sum(improved_indices)} =?= {np.size(classes)}"
+                else:
+                    logger.debug(
+                        f"Shift ({x},{y}) complete. Improved {np.sum(improved_indices)} alignments."
+                    )
 
         return rotations, shifts, correlations
 

From ef56c69318183a838483adf2947c2bf7fde38e63 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Thu, 12 Dec 2024 14:31:23 -0500
Subject: [PATCH 083/184] invert BFR BFSR relationship

---
 src/aspire/classification/averager2d.py | 38 +++++++++++++++++++------
 1 file changed, 29 insertions(+), 9 deletions(-)

diff --git a/src/aspire/classification/averager2d.py b/src/aspire/classification/averager2d.py
index f69bd1facf..0f9d4cdc57 100644
--- a/src/aspire/classification/averager2d.py
+++ b/src/aspire/classification/averager2d.py
@@ -235,7 +235,7 @@ def _shift_search_grid(self, L, radius, roll_zero=False):
         return X, Y
 
 
-class BFRAverager2D(AligningAverager2D):
+class _BFRAverager2D(AligningAverager2D):
     """
     This perfoms a Brute Force Rotational alignment.
 
@@ -328,7 +328,7 @@ def align(self, classes, reflections, basis_coefficients):
         return rots, None, correlations
 
 
-class BFSRAverager2D(BFRAverager2D):
+class BFSRAverager2D(AligningAverager2D):
     """
     This perfoms a Brute Force Shift and Rotational alignment.
     It is potentially expensive to brute force this search space.
@@ -359,20 +359,25 @@ def __init__(
             composite_basis,
             src,
             alignment_basis,
-            n_angles,
             dtype=dtype,
         )
 
-        self.radius = radius if radius is not None else src.L // 16
-
-        # Each shift will require calling the parent BFRAverager2D.align
-        self._bfr_align = super().align
+        self.n_angles = n_angles
 
-        if not hasattr(self.alignment_basis, "shift"):
+        if not hasattr(self.alignment_basis, "rotate"):
             raise RuntimeError(
-                f"{self.__class__.__name__}'s alignment_basis {self.alignment_basis} must provide a `shift` method."
+                f"{self.__class__.__name__}'s alignment_basis {self.alignment_basis} must provide a `rotate` method."
             )
 
+        self.radius = radius if radius is not None else src.L // 16
+
+        if self.radius != 0:
+
+            if not hasattr(self.alignment_basis, "shift"):
+                raise RuntimeError(
+                    f"{self.__class__.__name__}'s alignment_basis {self.alignment_basis} must provide a `shift` method."
+                )
+
     def align(self, classes, reflections, basis_coefficients):
         """
         See `AligningAverager2D.align`
@@ -489,6 +494,21 @@ def align(self, classes, reflections, basis_coefficients):
         return rotations, shifts, correlations
 
 
+class BFRAverager2D(BFSRAverager2D):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, radius=0, **kwargs)
+
+    def align(self, *args, **kwargs):
+        """
+        See `AligningAverager2D.align`
+        """
+        # BFR shifts should all be zeros.
+        # Replace with `None` to induce short ciruit shifting during stacking.
+        rotations, _, correlations = super().align(*args, **kwargs)
+
+        return rotations, None, correlations
+
+
 class ReddyChatterjiAverager2D(AligningAverager2D):
     """
     Attempts rotational estimation using Reddy Chatterji log polar Fourier cross correlation.

From 1520d42b48d5aba813c184e0c6a939bfdf7d9e1e Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Fri, 13 Dec 2024 11:46:11 -0500
Subject: [PATCH 084/184] fixing pbars

---
 src/aspire/classification/averager2d.py | 57 +++++++++++++++++--------
 src/aspire/utils/logging.py             | 12 ++++++
 2 files changed, 51 insertions(+), 18 deletions(-)

diff --git a/src/aspire/classification/averager2d.py b/src/aspire/classification/averager2d.py
index 0f9d4cdc57..316ecda66d 100644
--- a/src/aspire/classification/averager2d.py
+++ b/src/aspire/classification/averager2d.py
@@ -7,7 +7,7 @@
 from aspire.classification.reddy_chatterji import reddy_chatterji_register
 from aspire.image import Image, ImageStacker, MeanImageStacker
 from aspire.numeric import xp
-from aspire.utils import trange
+from aspire.utils import tqdm, trange
 from aspire.utils.coor_trans import grid_2d
 
 logger = logging.getLogger(__name__)
@@ -131,7 +131,7 @@ def __init__(
             )
 
     @abstractmethod
-    def align(self, classes, reflections, basis_coefficients):
+    def align(self, classes, reflections, basis_coefficients=None):
         """
         During this process `rotations`, `reflections`, `shifts` and
         `correlations` properties will be computed for aligners.
@@ -206,7 +206,7 @@ def _innerloop(i):
             # Averaging in composite_basis
             return self.image_stacker(neighbors_coefs.asnumpy())
 
-        for i in trange(n_classes):
+        for i in trange(n_classes, desc="Stacking class averages"):
             b_avgs[i] = _innerloop(i)
 
         # Now we convert the averaged images from Basis to Cartesian.
@@ -266,7 +266,7 @@ def __init__(
                 f"{self.__class__.__name__}'s alignment_basis {self.alignment_basis} must provide a `rotate` method."
             )
 
-    def align(self, classes, reflections, basis_coefficients):
+    def align(self, classes, reflections, basis_coefficients=None):
         """
         Performs the actual rotational alignment estimation,
         returning parameters needed for averaging.
@@ -282,7 +282,7 @@ def align(self, classes, reflections, basis_coefficients):
         rots = np.zeros((n_classes, n_nbor), dtype=self.dtype)
         correlations = np.zeros((n_classes, n_nbor), dtype=self.dtype)
 
-        for k in trange(n_classes):
+        for k in trange(n_classes, desc="Rotationally aligning classes"):
             # Get the coefs for these neighbors
             if basis_coefficients is None:
                 # Retrieve relevant images
@@ -349,7 +349,7 @@ def __init__(
         dtype=None,
     ):
         """
-        See AligningAverager2D and BFRAverager2D, adds: radius
+        See AligningAverager2D adds `n_angles` and `radius`.
 
         :params n_angles: Number of brute force rotations to attempt, defaults 360.
         :param radius: Brute force translation search radius.
@@ -378,7 +378,7 @@ def __init__(
                     f"{self.__class__.__name__}'s alignment_basis {self.alignment_basis} must provide a `shift` method."
                 )
 
-    def align(self, classes, reflections, basis_coefficients):
+    def align(self, classes, reflections, basis_coefficients=None):
         """
         See `AligningAverager2D.align`
         """
@@ -410,7 +410,7 @@ def align(self, classes, reflections, basis_coefficients):
             self.src.L, self.radius, roll_zero=True
         )
 
-        for k in trange(n_classes):
+        for k in trange(n_classes, desc="Rotationally aligning classes"):
             # We want to locally cache the original images,
             #  because we will mutate them with shifts in the next loop.
             #  This avoids recomputing them before each shift
@@ -425,7 +425,13 @@ def align(self, classes, reflections, basis_coefficients):
             _images = original_images.asnumpy().copy()
 
             # Loop over shift search space, updating best result
-            for x, y in zip(x_shifts, y_shifts):
+            for x, y in tqdm(
+                zip(x_shifts, y_shifts),
+                total=len(x_shifts),
+                desc="\tmaximizing over shifts",
+                disable=len(x_shifts) == 1,
+                leave=False,
+            ):
                 shift = np.array([x, y], dtype=int)
                 logger.debug(f"Computing rotational alignment after shift ({x},{y}).")
 
@@ -440,7 +446,7 @@ def align(self, classes, reflections, basis_coefficients):
                 _images[1:] = original_images[1:].shift(shift).asnumpy()
 
                 # Convert to array of complex coef, implicit copy.
-                _coef = self.alignment_basis.evaluate_t(_images)
+                _coef = self.alignment_basis.evaluate_t(Image(_images))
                 _coef = xp.array(_coef.to_complex().asnumpy())
 
                 # Handle reflections
@@ -500,11 +506,19 @@ def __init__(self, *args, **kwargs):
 
     def align(self, *args, **kwargs):
         """
-        See `AligningAverager2D.align`
+        See `BFSRAverager2D.align`
         """
         # BFR shifts should all be zeros.
         # Replace with `None` to induce short ciruit shifting during stacking.
-        rotations, _, correlations = super().align(*args, **kwargs)
+        rotations, shifts, correlations = super().align(*args, **kwargs)
+
+        # RM?
+        if not np.all(shifts.flatten() == 0):
+            logger.error(
+                "BFR should return zero shifts."
+                "  BFSR returned non zero shifts."
+                "  Forcing to `None`."
+            )
 
         return rotations, None, correlations
 
@@ -555,7 +569,7 @@ def __init__(
 
         super().__init__(composite_basis, src, composite_basis, dtype=dtype)
 
-    def align(self, classes, reflections, basis_coefficients):
+    def align(self, classes, reflections, basis_coefficients=None):
         """
         Performs the actual rotational alignment estimation,
         returning parameters needed for averaging.
@@ -583,7 +597,7 @@ def _innerloop(k):
                 dtype=self.dtype,
             )
 
-        for k in trange(n_classes):
+        for k in trange(n_classes, desc="Rotationally aligning classes"):
             rotations[k], shifts[k], correlations[k] = _innerloop(k)
 
         return rotations, shifts, correlations
@@ -632,7 +646,7 @@ def _innerloop(i):
             # Averaging in composite_basis
             return self.image_stacker(neighbors_coefs.asnumpy())
 
-        for i in trange(n_classes):
+        for i in trange(n_classes, desc="Stacking class averages"):
             b_avgs[i] = _innerloop(i)
 
         # Now we convert the averaged images from Basis to Cartesian.
@@ -686,7 +700,7 @@ def __init__(
         # Assign search radius
         self.radius = radius if radius is not None else src.L // 8
 
-    def align(self, classes, reflections, basis_coefficients):
+    def align(self, classes, reflections, basis_coefficients=None):
         """
         Performs the actual rotational alignment estimation,
         returning parameters needed for averaging.
@@ -712,7 +726,14 @@ def _innerloop(k):
             _correlations = np.ones(classes.shape[1:], dtype=self.dtype) * -np.inf
             _shifts = np.zeros((*classes.shape[1:], 2), dtype=int)
 
-            for xs, ys in zip(X, Y):
+            for xs, ys in tqdm(
+                zip(X, Y),
+                total=len(X),
+                desc="\tmaximizing over shifts",
+                disable=len(X) == 1,
+                leave=False,
+            ):
+
                 s = np.array([xs, ys])
                 # Get the array of images for this class
 
@@ -741,7 +762,7 @@ def _innerloop(k):
 
             return _rotations, _shifts, _correlations
 
-        for k in trange(n_classes):
+        for k in trange(n_classes, desc="Rotationally aligning classes"):
             rotations[k], shifts[k], correlations[k] = _innerloop(k)
 
         return rotations, shifts, correlations
diff --git a/src/aspire/utils/logging.py b/src/aspire/utils/logging.py
index 7a3151a3ea..0b719d07c9 100644
--- a/src/aspire/utils/logging.py
+++ b/src/aspire/utils/logging.py
@@ -72,10 +72,16 @@ def tqdm(*args, **kwargs):
     Currently setting `aspire.config['logging']['tqdm_disable']`
     true/false will disable/enable tqdm progress bars.
     """
+    # Calling code may use disable kwarg
+    caller_disable = kwargs.pop("disable", False)
 
     disable = config["logging"]["tqdm_disable"] or (
         getConsoleLoggingLevel() not in ["DEBUG", "INFO"]
     )
+
+    # Disable if calling code flag is true or when configured to disable.
+    disable = disable or caller_disable
+
     return _tqdm.tqdm(*args, **kwargs, disable=disable)
 
 
@@ -86,10 +92,16 @@ def trange(*args, **kwargs):
     Currently setting `aspire.config['logging']['tqdm_disable']`
     true/false will disable/enable tqdm progress bars.
     """
+    # Calling code may use disable kwarg
+    caller_disable = kwargs.pop("disable", False)
 
     disable = config["logging"]["tqdm_disable"] or (
         getConsoleLoggingLevel() not in ["DEBUG", "INFO"]
     )
+
+    # Disable if calling code flag is true or when configured to disable.
+    disable = disable or caller_disable
+
     return _tqdm.trange(*args, **kwargs, disable=disable)
 
 

From 92bcd94895fbd39faaf0296c104ea1da8e54cb39 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Thu, 2 Jan 2025 14:48:17 -0500
Subject: [PATCH 085/184] remove temp BFR only work

---
 src/aspire/classification/averager2d.py | 128 ++++--------------------
 1 file changed, 19 insertions(+), 109 deletions(-)

diff --git a/src/aspire/classification/averager2d.py b/src/aspire/classification/averager2d.py
index 316ecda66d..a9c6383ee5 100644
--- a/src/aspire/classification/averager2d.py
+++ b/src/aspire/classification/averager2d.py
@@ -235,99 +235,6 @@ def _shift_search_grid(self, L, radius, roll_zero=False):
         return X, Y
 
 
-class _BFRAverager2D(AligningAverager2D):
-    """
-    This perfoms a Brute Force Rotational alignment.
-
-    For each class,
-        constructs n_angles rotations of all class members,
-        and then identifies angle yielding largest correlation(dot).
-    """
-
-    def __init__(
-        self,
-        composite_basis,
-        src,
-        alignment_basis=None,
-        n_angles=360,
-        dtype=None,
-    ):
-        """
-        See AligningAverager2D, adds:
-
-        :param n_angles: Number of brute force rotations to attempt, defaults 360.
-        """
-        super().__init__(composite_basis, src, alignment_basis, dtype=dtype)
-
-        self.n_angles = n_angles
-
-        if not hasattr(self.alignment_basis, "rotate"):
-            raise RuntimeError(
-                f"{self.__class__.__name__}'s alignment_basis {self.alignment_basis} must provide a `rotate` method."
-            )
-
-    def align(self, classes, reflections, basis_coefficients=None):
-        """
-        Performs the actual rotational alignment estimation,
-        returning parameters needed for averaging.
-        """
-        # Construct array of angles to brute force.
-        _angles = xp.linspace(0, 2 * np.pi, self.n_angles, endpoint=False)
-
-        ks = xp.asarray(self.alignment_basis.complex_angular_indices).reshape(-1, 1)
-        _rot_ops_conj = xp.exp(1j * ks * _angles).conj()
-
-        # Result arrays
-        n_classes, n_nbor = classes.shape
-        rots = np.zeros((n_classes, n_nbor), dtype=self.dtype)
-        correlations = np.zeros((n_classes, n_nbor), dtype=self.dtype)
-
-        for k in trange(n_classes, desc="Rotationally aligning classes"):
-            # Get the coefs for these neighbors
-            if basis_coefficients is None:
-                # Retrieve relevant images
-                neighbors_imgs = Image(self._cls_images(classes[k]))
-
-                # Evaluate_t into basis
-                nbr_coef = self.alignment_basis.evaluate_t(neighbors_imgs)
-            else:
-                nbr_coef = basis_coefficients[classes[k]]
-
-            # Convert to array of complex coef, implicit copy.
-            nbr_coef = xp.array(nbr_coef.to_complex().asnumpy())
-
-            # Handle reflections
-            refl = reflections[k]
-            nbr_coef[refl] = xp.conj(nbr_coef[refl])
-
-            # Generate table of rotations for image 0.
-            # Note we invert the rotation, applying -rot to image 0
-            #   to avoid rotating each member of the class
-            #   for the argmax alignment test (each a dot product,
-            #   performed as a large matmul).
-            base_img = nbr_coef[0].reshape(self.alignment_basis.complex_count, 1)
-
-            # (cnt, n_transl) * (cnt, 1) -> (cnt, n_transl)
-            rot_base_imgs_conj = _rot_ops_conj * base_img.conj()
-
-            # (n_nbor, cnt) @ (cnt, n_transl) = (n_nbor, n_transl)
-            dots = xp.real(nbr_coef @ rot_base_imgs_conj)
-            idx = xp.argmax(dots, axis=1)
-            idx[0] = 0  # Force base image, just in case.
-
-            # Assign results for this class
-            correlations[k, :] = xp.asnumpy(
-                xp.take_along_axis(dots, idx.reshape(n_nbor, 1), axis=1).flatten()
-            )
-            # correlations[k,0] = 1  # Force base correlation, just in case
-            #  Todo, do we care to spend the compute to make an actual correlation? (normalize)
-
-            # Assign the reverse rotation
-            rots[k] = -1 * xp.asnumpy(_angles[idx])
-
-        return rots, None, correlations
-
-
 class BFSRAverager2D(AligningAverager2D):
     """
     This perfoms a Brute Force Shift and Rotational alignment.
@@ -395,8 +302,8 @@ def align(self, classes, reflections, basis_coefficients=None):
         shifts = np.empty((*classes.shape, 2), dtype=int)
 
         # Work arrays
-        _rotations = np.zeros((1, n_nbor), dtype=self.dtype)
-        _correlations = np.ones((1, n_nbor), dtype=self.dtype) * -np.inf
+        _rotations = np.zeros((n_nbor), dtype=self.dtype)
+        _correlations = np.ones((n_nbor), dtype=self.dtype) * -np.inf
 
         # Construct array of angles to brute force.
         _angles = xp.linspace(0, 2 * np.pi, self.n_angles, endpoint=False)
@@ -418,7 +325,6 @@ def align(self, classes, reflections, basis_coefficients=None):
                 original_images = Image(self._cls_images(classes[k], src=self.src))
             else:
                 original_coef = basis_coefficients[classes[k], :]
-                # batch here? or maybe just force always passing images....?
                 original_images = self.alignment_basis.evaluate(original_coef)
 
             # Working copy
@@ -442,29 +348,32 @@ def align(self, classes, reflections, basis_coefficients=None):
                 #   ii) because generally the number of neighbors << the
                 #   number of test rotations.
 
-                # Note the base image[0] is never shifted.
-                _images[1:] = original_images[1:].shift(shift).asnumpy()
+                # Skip zero shifting.
+                if np.any(shift != 0):
+                    # Note the base image[0] is never shifted.
+                    _images[1:] = original_images[1:].shift(shift).asnumpy()
 
                 # Convert to array of complex coef, implicit copy.
                 _coef = self.alignment_basis.evaluate_t(Image(_images))
                 _coef = xp.array(_coef.to_complex().asnumpy())
 
                 # Handle reflections
-                # XXXX should we do this earlier via Images? (flipud?)
-                # How would we communicate this reflection to other areas of the code...
                 refl = reflections[k]
                 _coef[refl] = xp.conj(_coef[refl])
 
                 # Generate table of rotations for image 0.
-                # Note we invert the rotation, applying -rot to image 0
-                #   to avoid rotating each member of the class
-                #   for the argmax alignment test (each a dot product,
-                #   performed as a large matmul).
+                # Note we invert the rotations later.
+                #   Applying rot to image 0
+                #   avoids rotating each member of the class
+                #   for the argmax alignment test.
                 base_img = _coef[0].reshape(self.alignment_basis.complex_count, 1)
 
                 # (cnt, n_transl) * (cnt, 1) -> (cnt, n_transl)
                 rot_base_imgs_conj = _rot_ops_conj * base_img.conj()
 
+                # Compute dot product of each base-neighbor pair.
+                #   The collection of dots is performed in bulk
+                #   as a large matmul.
                 # (n_nbor, cnt) @ (cnt, n_transl) = (n_nbor, n_transl)
                 dots = xp.real(_coef @ rot_base_imgs_conj)
                 idx = xp.argmax(dots, axis=1)
@@ -475,16 +384,17 @@ def align(self, classes, reflections, basis_coefficients=None):
                     xp.take_along_axis(dots, idx.reshape(n_nbor, 1), axis=1).flatten()
                 )
                 # _correlations[0] = 1  # Force base correlation, just in case
-                #  Todo, do we care to spend the compute to make an actual correlation? (normalize)
+                #  Todo, do we care to spend the compute to make an actual correlation? (normalizing).
+                #  The matlab code did do something like this, but only for diagnostic purposes.
 
                 # Assign the reverse rotation
                 _rotations[:] = -1 * xp.asnumpy(_angles[idx])
 
-                # Each class-neighbor pair may have a best shift-rot from a different shift.
                 # Test and update
-                improved_indices = _correlations[0] > correlations[k]
-                rotations[k, improved_indices] = _rotations[0, improved_indices]
-                correlations[k, improved_indices] = _correlations[0, improved_indices]
+                # Each base-neighbor pair may have a best shift+rot from a different shift iteration.
+                improved_indices = _correlations > correlations[k]
+                rotations[k, improved_indices] = _rotations[improved_indices]
+                correlations[k, improved_indices] = _correlations[improved_indices]
                 shifts[k, improved_indices] = shift
 
                 if (x, y) == (0, 0):

From a449e827c31b04b8f983a57ff0e95295a168a050 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Thu, 2 Jan 2025 16:03:42 -0500
Subject: [PATCH 086/184] minor cleanup

---
 src/aspire/classification/averager2d.py | 17 ++++++++++-------
 1 file changed, 10 insertions(+), 7 deletions(-)

diff --git a/src/aspire/classification/averager2d.py b/src/aspire/classification/averager2d.py
index a9c6383ee5..eea06df36d 100644
--- a/src/aspire/classification/averager2d.py
+++ b/src/aspire/classification/averager2d.py
@@ -384,8 +384,13 @@ def align(self, classes, reflections, basis_coefficients=None):
                     xp.take_along_axis(dots, idx.reshape(n_nbor, 1), axis=1).flatten()
                 )
                 # _correlations[0] = 1  # Force base correlation, just in case
-                #  Todo, do we care to spend the compute to make an actual correlation? (normalizing).
-                #  The matlab code did do something like this, but only for diagnostic purposes.
+                #  Todo, discuss:
+                #      If we care to compute the an actual correlation?
+                #      This involves a lot of somewhat careful normalizing.
+                #      The MATLAB code did do something like this,
+                #      but only for diagnostic purposes.
+                #      In our code anything beyond a relative calculation
+                #      would currently be superflous.
 
                 # Assign the reverse rotation
                 _rotations[:] = -1 * xp.asnumpy(_angles[idx])
@@ -422,12 +427,10 @@ def align(self, *args, **kwargs):
         # Replace with `None` to induce short ciruit shifting during stacking.
         rotations, shifts, correlations = super().align(*args, **kwargs)
 
-        # RM?
+        # Sanity check the results did not indicate shifts.
         if not np.all(shifts.flatten() == 0):
-            logger.error(
-                "BFR should return zero shifts."
-                "  BFSR returned non zero shifts."
-                "  Forcing to `None`."
+            raise RuntimeError(
+                "BFR should return zero shifts." "  BFSR returned non zero shifts."
             )
 
         return rotations, None, correlations

From e56f4129a4f694384125c49fc77e7689e601d392 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Thu, 16 Jan 2025 07:45:06 -0500
Subject: [PATCH 087/184] change n_transl comments to n_rot

---
 src/aspire/classification/averager2d.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/aspire/classification/averager2d.py b/src/aspire/classification/averager2d.py
index eea06df36d..ef56dc604a 100644
--- a/src/aspire/classification/averager2d.py
+++ b/src/aspire/classification/averager2d.py
@@ -368,13 +368,13 @@ def align(self, classes, reflections, basis_coefficients=None):
                 #   for the argmax alignment test.
                 base_img = _coef[0].reshape(self.alignment_basis.complex_count, 1)
 
-                # (cnt, n_transl) * (cnt, 1) -> (cnt, n_transl)
+                # (cnt, n_rot) * (cnt, 1) -> (cnt, n_rot)
                 rot_base_imgs_conj = _rot_ops_conj * base_img.conj()
 
                 # Compute dot product of each base-neighbor pair.
                 #   The collection of dots is performed in bulk
                 #   as a large matmul.
-                # (n_nbor, cnt) @ (cnt, n_transl) = (n_nbor, n_transl)
+                # (n_nbor, cnt) @ (cnt, n_rot) = (n_nbor, n_rot)
                 dots = xp.real(_coef @ rot_base_imgs_conj)
                 idx = xp.argmax(dots, axis=1)
                 idx[0] = 0  # Force base image, just in case.

From e198ef5a3170b575b8c6b35b90e6b14feb4a6999 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Thu, 16 Jan 2025 07:56:56 -0500
Subject: [PATCH 088/184] change correlations ~~> dot_products in code and
 related strings/docs/tests

---
 .../tutorials/tutorials/class_averaging.py    |  4 +-
 src/aspire/classification/averager2d.py       | 61 +++++++++----------
 tests/test_averager2d.py                      |  4 +-
 3 files changed, 32 insertions(+), 37 deletions(-)

diff --git a/gallery/tutorials/tutorials/class_averaging.py b/gallery/tutorials/tutorials/class_averaging.py
index b609c59985..bf4ba66f9a 100644
--- a/gallery/tutorials/tutorials/class_averaging.py
+++ b/gallery/tutorials/tutorials/class_averaging.py
@@ -224,11 +224,11 @@
 
 est_rotations = avgs.averager.rotations
 est_shifts = avgs.averager.shifts
-est_correlations = avgs.averager.correlations
+est_dot_products = avgs.averager.dot_products
 
 print(f"Estimated Rotations: {est_rotations}")
 print(f"Estimated Shifts: {est_shifts}")
-print(f"Estimated Correlations: {est_correlations}")
+print(f"Estimated Dot Products: {est_dot_products}")
 
 # Compare the original unaligned images with the estimated alignment.
 # Get the indices from the classification results.
diff --git a/src/aspire/classification/averager2d.py b/src/aspire/classification/averager2d.py
index ef56dc604a..4d12641ad4 100644
--- a/src/aspire/classification/averager2d.py
+++ b/src/aspire/classification/averager2d.py
@@ -134,7 +134,7 @@ def __init__(
     def align(self, classes, reflections, basis_coefficients=None):
         """
         During this process `rotations`, `reflections`, `shifts` and
-        `correlations` properties will be computed for aligners.
+        `dot_products` properties will be computed for aligners.
 
         `rotations` is an (src.n, n_nbor) array of angles,
         which should represent the rotations needed to align images within
@@ -144,8 +144,8 @@ def align(self, classes, reflections, basis_coefficients=None):
         which should represent the translation needed to best align the images
         within that class.
 
-        `correlations` is an (src.n, n_nbor) array representing
-        a correlation like measure between classified images and their base
+        `dot_products` is an (src.n, n_nbor) array representing
+        the dot between classified images and their base
         image (image index 0).
 
         Subclasses of should implement and extend this method.
@@ -154,7 +154,7 @@ def align(self, classes, reflections, basis_coefficients=None):
         :param reflections: (src.n, n_nbor) bool array of corresponding reflections,
         :param basis_coefficients: (n_img, self.alignment_basis.count) basis coefficients,
 
-        :returns: (rotations, shifts, correlations)
+        :returns: (rotations, shifts, dot_products)
         """
 
     def average(
@@ -170,7 +170,7 @@ def average(
         classes = np.atleast_2d(classes)
         reflections = np.atleast_2d(reflections)
 
-        self.rotations, self.shifts, self.correlations = self.align(
+        self.rotations, self.shifts, self.dot_products = self.align(
             classes, reflections, coefs
         )
 
@@ -298,12 +298,12 @@ def align(self, classes, reflections, basis_coefficients=None):
         # These arrays will incrementally store our best alignment.
         n_classes, n_nbor = classes.shape
         rotations = np.zeros((n_classes, n_nbor), dtype=self.dtype)
-        correlations = np.ones((n_classes, n_nbor), dtype=self.dtype) * -np.inf
+        dot_products = np.ones((n_classes, n_nbor), dtype=self.dtype) * -np.inf
         shifts = np.empty((*classes.shape, 2), dtype=int)
 
         # Work arrays
         _rotations = np.zeros((n_nbor), dtype=self.dtype)
-        _correlations = np.ones((n_nbor), dtype=self.dtype) * -np.inf
+        _dot_products = np.ones((n_nbor), dtype=self.dtype) * -np.inf
 
         # Construct array of angles to brute force.
         _angles = xp.linspace(0, 2 * np.pi, self.n_angles, endpoint=False)
@@ -380,26 +380,21 @@ def align(self, classes, reflections, basis_coefficients=None):
                 idx[0] = 0  # Force base image, just in case.
 
                 # Assign results for this class
-                _correlations[:] = xp.asnumpy(
+                _dot_products[:] = xp.asnumpy(
                     xp.take_along_axis(dots, idx.reshape(n_nbor, 1), axis=1).flatten()
                 )
-                # _correlations[0] = 1  # Force base correlation, just in case
-                #  Todo, discuss:
-                #      If we care to compute the an actual correlation?
-                #      This involves a lot of somewhat careful normalizing.
-                #      The MATLAB code did do something like this,
-                #      but only for diagnostic purposes.
-                #      In our code anything beyond a relative calculation
-                #      would currently be superflous.
+                # Note, legacy codes would normalize to form correlations.
+                #   These were only used for diagnostic purposes.
+                #   Normalizing is skipped here to save computation.
 
                 # Assign the reverse rotation
                 _rotations[:] = -1 * xp.asnumpy(_angles[idx])
 
                 # Test and update
                 # Each base-neighbor pair may have a best shift+rot from a different shift iteration.
-                improved_indices = _correlations > correlations[k]
+                improved_indices = _dot_products > dot_products[k]
                 rotations[k, improved_indices] = _rotations[improved_indices]
-                correlations[k, improved_indices] = _correlations[improved_indices]
+                dot_products[k, improved_indices] = _dot_products[improved_indices]
                 shifts[k, improved_indices] = shift
 
                 if (x, y) == (0, 0):
@@ -412,7 +407,7 @@ def align(self, classes, reflections, basis_coefficients=None):
                         f"Shift ({x},{y}) complete. Improved {np.sum(improved_indices)} alignments."
                     )
 
-        return rotations, shifts, correlations
+        return rotations, shifts, dot_products
 
 
 class BFRAverager2D(BFSRAverager2D):
@@ -425,7 +420,7 @@ def align(self, *args, **kwargs):
         """
         # BFR shifts should all be zeros.
         # Replace with `None` to induce short ciruit shifting during stacking.
-        rotations, shifts, correlations = super().align(*args, **kwargs)
+        rotations, shifts, dot_products = super().align(*args, **kwargs)
 
         # Sanity check the results did not indicate shifts.
         if not np.all(shifts.flatten() == 0):
@@ -433,7 +428,7 @@ def align(self, *args, **kwargs):
                 "BFR should return zero shifts." "  BFSR returned non zero shifts."
             )
 
-        return rotations, None, correlations
+        return rotations, None, dot_products
 
 
 class ReddyChatterjiAverager2D(AligningAverager2D):
@@ -496,7 +491,7 @@ def align(self, classes, reflections, basis_coefficients=None):
 
         # Instantiate matrices for results
         rotations = np.zeros(classes.shape, dtype=self.dtype)
-        correlations = np.zeros(classes.shape, dtype=self.dtype)
+        dot_products = np.zeros(classes.shape, dtype=self.dtype)
         shifts = np.zeros((*classes.shape, 2), dtype=int)
 
         def _innerloop(k):
@@ -511,9 +506,9 @@ def _innerloop(k):
             )
 
         for k in trange(n_classes, desc="Rotationally aligning classes"):
-            rotations[k], shifts[k], correlations[k] = _innerloop(k)
+            rotations[k], shifts[k], dot_products[k] = _innerloop(k)
 
-        return rotations, shifts, correlations
+        return rotations, shifts, dot_products
 
     def average(
         self,
@@ -526,7 +521,7 @@ def average(
         Otherwise is similar to `AligningAverager2D.average`.
         """
 
-        self.rotations, self.shifts, self.correlations = self.align(
+        self.rotations, self.shifts, self.dot_products = self.align(
             classes, reflections, coefs
         )
 
@@ -627,7 +622,7 @@ def align(self, classes, reflections, basis_coefficients=None):
 
         # Instantiate matrices for inner loop, and best results.
         rotations = np.zeros(classes.shape, dtype=self.dtype)
-        correlations = np.ones(classes.shape, dtype=self.dtype) * -np.inf
+        dot_products = np.ones(classes.shape, dtype=self.dtype) * -np.inf
         shifts = np.zeros((*classes.shape, 2), dtype=int)
 
         X, Y = self._shift_search_grid(self.alignment_src.L, self.radius)
@@ -636,7 +631,7 @@ def _innerloop(k):
             unshifted_images = self._cls_images(classes[k])
             # Instantiate matrices for inner loop, and best results.
             _rotations = np.zeros(classes.shape[1:], dtype=self.dtype)
-            _correlations = np.ones(classes.shape[1:], dtype=self.dtype) * -np.inf
+            _dot_products = np.ones(classes.shape[1:], dtype=self.dtype) * -np.inf
             _shifts = np.zeros((*classes.shape[1:], 2), dtype=int)
 
             for xs, ys in tqdm(
@@ -657,7 +652,7 @@ def _innerloop(k):
                 images[1:] = Image(unshifted_images[1:]).shift(s).asnumpy()
 
                 # returned shifts ignored since we are forcing shift of `s` above
-                __rotations, _, __correlations = reddy_chatterji_register(
+                __rotations, _, __dot_products = reddy_chatterji_register(
                     images,
                     reflections[k],
                     mask=self.mask,
@@ -667,18 +662,18 @@ def _innerloop(k):
 
                 # Where corr has improved
                 #  update our rolling best results with this loop.
-                improved = __correlations > _correlations
-                _correlations = np.where(improved, __correlations, _correlations)
+                improved = __dot_products > _dot_products
+                _dot_products = np.where(improved, __dot_products, _dot_products)
                 _rotations = np.where(improved, __rotations, _rotations)
                 _shifts = np.where(improved[..., np.newaxis], s, _shifts)
                 logger.debug(f"Shift {s} has improved {np.sum(improved)} results")
 
-            return _rotations, _shifts, _correlations
+            return _rotations, _shifts, _dot_products
 
         for k in trange(n_classes, desc="Rotationally aligning classes"):
-            rotations[k], shifts[k], correlations[k] = _innerloop(k)
+            rotations[k], shifts[k], dot_products[k] = _innerloop(k)
 
-        return rotations, shifts, correlations
+        return rotations, shifts, dot_products
 
     def average(
         self,
diff --git a/tests/test_averager2d.py b/tests/test_averager2d.py
index 1bbb7ca18b..becf0ad2d1 100644
--- a/tests/test_averager2d.py
+++ b/tests/test_averager2d.py
@@ -123,7 +123,7 @@ class AligningAverager2DBase(Averager2DBase):
 
     `.rotations`
     `.shifts`
-    `.correlations`
+    `.dot_products`
     """
 
     averager = AligningAverager2D
@@ -155,7 +155,7 @@ def test_attributes(self):
 
         self.assertTrue(hasattr(avgr, "shifts"))
 
-        self.assertTrue(hasattr(avgr, "correlations"))
+        self.assertTrue(hasattr(avgr, "dot_products"))
 
     def _getSrc(self):
         if not hasattr(self, "shifts"):

From 1860ef21c6252f1ffc311a157b861d3dc6b07855 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Thu, 16 Jan 2025 10:12:59 -0500
Subject: [PATCH 089/184] Migrate base image rotation table out of shift loop

---
 src/aspire/classification/averager2d.py | 24 ++++++++++++++----------
 1 file changed, 14 insertions(+), 10 deletions(-)

diff --git a/src/aspire/classification/averager2d.py b/src/aspire/classification/averager2d.py
index 4d12641ad4..8649a7c021 100644
--- a/src/aspire/classification/averager2d.py
+++ b/src/aspire/classification/averager2d.py
@@ -321,15 +321,29 @@ def align(self, classes, reflections, basis_coefficients=None):
             # We want to locally cache the original images,
             #  because we will mutate them with shifts in the next loop.
             #  This avoids recomputing them before each shift
+            # The coefficient for the base images are also computed here.
             if basis_coefficients is None:
                 original_images = Image(self._cls_images(classes[k], src=self.src))
+                _coef0 = self.alignment_basis.evaluate_t(original_images[0])
             else:
                 original_coef = basis_coefficients[classes[k], :]
                 original_images = self.alignment_basis.evaluate(original_coef)
+                _coef0 = original_coef[0]
 
             # Working copy
             _images = original_images.asnumpy().copy()
 
+            # Generate table of rotations for image 0.
+            # Note we invert the rotations later.
+            #   Applying rot to image 0
+            #   avoids rotating each member of the class
+            #   for the argmax alignment test.
+            # Convert to array of complex coef, implicit copy.
+            _coef0 = xp.array(_coef0.to_complex().asnumpy())
+            base_img = _coef0.reshape(self.alignment_basis.complex_count, 1)
+            # (cnt, n_rot) * (cnt, 1) -> (cnt, n_rot)
+            rot_base_imgs_conj = _rot_ops_conj * base_img.conj()
+
             # Loop over shift search space, updating best result
             for x, y in tqdm(
                 zip(x_shifts, y_shifts),
@@ -361,16 +375,6 @@ def align(self, classes, reflections, basis_coefficients=None):
                 refl = reflections[k]
                 _coef[refl] = xp.conj(_coef[refl])
 
-                # Generate table of rotations for image 0.
-                # Note we invert the rotations later.
-                #   Applying rot to image 0
-                #   avoids rotating each member of the class
-                #   for the argmax alignment test.
-                base_img = _coef[0].reshape(self.alignment_basis.complex_count, 1)
-
-                # (cnt, n_rot) * (cnt, 1) -> (cnt, n_rot)
-                rot_base_imgs_conj = _rot_ops_conj * base_img.conj()
-
                 # Compute dot product of each base-neighbor pair.
                 #   The collection of dots is performed in bulk
                 #   as a large matmul.

From ec6a5d7f2fa909ee8f6c70fc353df0adaf960f81 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Thu, 30 Jan 2025 15:37:34 -0500
Subject: [PATCH 090/184] remove cufinufft 3d1 method workaround

---
 src/aspire/nufft/cufinufft.py | 10 ----------
 1 file changed, 10 deletions(-)

diff --git a/src/aspire/nufft/cufinufft.py b/src/aspire/nufft/cufinufft.py
index fd869aacfd..40b5914182 100644
--- a/src/aspire/nufft/cufinufft.py
+++ b/src/aspire/nufft/cufinufft.py
@@ -62,15 +62,6 @@ def __init__(self, sz, fourier_pts, epsilon=1e-8, ntransforms=1, **kwargs):
             2, self.sz, self.ntransforms, self.epsilon, -1, dtype=self.complex_dtype
         )
 
-        self.adjoint_opts = dict()
-        if self.dtype is np.float64 and self.dim == 3 and self.epsilon < 1e3:
-            # Note this is an algorithmic implementation dictated by shmem.
-            logger.info(
-                "Converting cufinufft gpu_method=1 from default of 2 for 3D1 transform,"
-                f"to support computation in double precision with tol={self.epsilon}."
-            )
-            self.adjoint_opts["gpu_method"] = 1
-
         self._adjoint_plan = cufPlan(
             1,
             self.sz,
@@ -78,7 +69,6 @@ def __init__(self, sz, fourier_pts, epsilon=1e-8, ntransforms=1, **kwargs):
             self.epsilon,
             1,
             dtype=self.complex_dtype,
-            **self.adjoint_opts,
         )
 
         self._transform_plan.setpts(*self.fourier_pts)

From bb674800a1f4de8ed95bb786941fde9632c26686 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Thu, 30 Jan 2025 10:21:58 -0500
Subject: [PATCH 091/184] plumb batch_size through class averager code

---
 src/aspire/classification/averager2d.py | 45 +++++++++++++++++++++----
 src/aspire/denoising/class_avg.py       | 22 ++++++++++++
 2 files changed, 60 insertions(+), 7 deletions(-)

diff --git a/src/aspire/classification/averager2d.py b/src/aspire/classification/averager2d.py
index 8649a7c021..e25187d5a2 100644
--- a/src/aspire/classification/averager2d.py
+++ b/src/aspire/classification/averager2d.py
@@ -18,15 +18,18 @@ class Averager2D(ABC):
     Base class for 2D Image Averaging methods.
     """
 
-    def __init__(self, composite_basis, src, dtype=None):
+    def __init__(self, composite_basis, src, batch_size=512, dtype=None):
         """
         :param composite_basis:  Basis to be used during class average composition (eg FFB2D)
         :param src: Source of original images.
+        :param batch_size: Integer size of batches used for basis conversion.
         :param dtype: Numpy dtype to be used during alignment.
         """
 
         self.composite_basis = composite_basis
         self.src = src
+        self.batch_size = int(batch_size)
+
         if dtype is None:
             if self.composite_basis:
                 self.dtype = self.composite_basis.dtype
@@ -97,6 +100,7 @@ def __init__(
         src,
         alignment_basis=None,
         image_stacker=None,
+        batch_size=512,
         dtype=None,
     ):
         """
@@ -105,12 +109,14 @@ def __init__(
         :param alignment_basis: Optional, basis to be used only during alignment (eg FSPCA).
         :param image_stacker: Optional, provide a user defined `ImageStacker` instance,
             used during image stacking (averaging).  Defaults to MeanImageStacker.
+        :param batch_size: Integer size of batches used for basis conversion.
         :param dtype: Numpy dtype to be used during alignment.
         """
 
         super().__init__(
             composite_basis=composite_basis,
             src=src,
+            batch_size=batch_size,
             dtype=dtype,
         )
         # If alignment_basis is None, use composite_basis
@@ -176,7 +182,12 @@ def average(
 
         n_classes, n_nbor = classes.shape
 
-        b_avgs = np.empty((n_classes, self.composite_basis.count), dtype=self.src.dtype)
+        # Result (image) array
+        avgs = np.empty((n_classes, *self.composite_basis.sz), dtype=self.src.dtype)
+        # Tmp (basis) batch result array)
+        b_avgs = np.empty(
+            (self.batch_size, self.composite_basis.count), dtype=self.src.dtype
+        )
 
         def _innerloop(i):
             # Get coefs in Composite_Basis if not provided as an argument.
@@ -206,11 +217,22 @@ def _innerloop(i):
             # Averaging in composite_basis
             return self.image_stacker(neighbors_coefs.asnumpy())
 
-        for i in trange(n_classes, desc="Stacking class averages"):
-            b_avgs[i] = _innerloop(i)
+        desc = f"Stacking and evaluating class averages from {self.composite_basis.__class__.__name__} to Cartesian"
+        for start in trange(0, n_classes, self.batch_size, desc=desc):
+            end = min(start + self.batch_size, n_classes)
+            for i, cls in enumerate(
+                trange(start, end, desc="Stacking batch", leave=False)
+            ):
+                b_avgs[i] = _innerloop(cls)  # average stacked in basis
 
-        # Now we convert the averaged images from Basis to Cartesian.
-        return Coef(self.composite_basis, b_avgs).evaluate()
+            # Now we convert the averaged images from Basis to Cartesian,
+            #   assigning to result array.
+            # Note i should usually be batch_size, but may be less on final batch.
+            avgs[start:end] = (
+                Coef(self.composite_basis, b_avgs[: i + 1]).evaluate().asnumpy()
+            )
+
+        return Image(avgs)
 
     def _shift_search_grid(self, L, radius, roll_zero=False):
         """
@@ -253,6 +275,7 @@ def __init__(
         alignment_basis=None,
         n_angles=360,
         radius=None,
+        batch_size=512,
         dtype=None,
     ):
         """
@@ -266,6 +289,7 @@ def __init__(
             composite_basis,
             src,
             alignment_basis,
+            batch_size=batch_size,
             dtype=dtype,
         )
 
@@ -458,6 +482,7 @@ def __init__(
         composite_basis,
         src,
         alignment_src=None,
+        batch_size=512,
         dtype=None,
     ):
         """
@@ -465,6 +490,7 @@ def __init__(
         :param src: Source of original images.
         :param alignment_src: Optional, source to be used during class average alignment.
             Must be the same resolution as `src`.
+        :param batch_size: Integer size of batches used for basis conversion.
         :param dtype: Numpy dtype to be used during alignment.
         """
 
@@ -479,7 +505,9 @@ def __init__(
 
         self.mask = grid_2d(src.L, normalized=False)["r"] < src.L // 2
 
-        super().__init__(composite_basis, src, composite_basis, dtype=dtype)
+        super().__init__(
+            composite_basis, src, composite_basis, batch_size=batch_size, dtype=dtype
+        )
 
     def align(self, classes, reflections, basis_coefficients=None):
         """
@@ -587,6 +615,7 @@ def __init__(
         src,
         alignment_src=None,
         radius=None,
+        batch_size=512,
         dtype=None,
     ):
         """
@@ -599,6 +628,7 @@ def __init__(
             Must be the same resolution as `src`.
         :param radius: Brute force translation search radius.
             Defaults to src.L//8.
+        :param batch_size: Integer size of batches used for basis conversion.
         :param dtype: Numpy dtype to be used during alignment.
         """
 
@@ -606,6 +636,7 @@ def __init__(
             composite_basis,
             src,
             alignment_src,
+            batch_size=batch_size,
             dtype=dtype,
         )
 
diff --git a/src/aspire/denoising/class_avg.py b/src/aspire/denoising/class_avg.py
index 9b61864fa9..8aff4c23e4 100644
--- a/src/aspire/denoising/class_avg.py
+++ b/src/aspire/denoising/class_avg.py
@@ -33,6 +33,7 @@ def __init__(
         classifier,
         class_selector,
         averager,
+        batch_size=512,
     ):
         """
         Constructor of an object for denoising 2D images using class averaging methods.
@@ -42,8 +43,10 @@ def __init__(
             Example, RIRClass2D.
         :param class_selector: A ClassSelector subclass.
         :param averager: An Averager2D subclass.
+        :param batch_size: Integer size for batched operations.
         """
         self.src = src
+        self.batch_size = int(batch_size)
         if not isinstance(self.src, ImageSource):
             raise ValueError(
                 f"`src` should be subclass of `ImageSource`, found {self.src}."
@@ -397,6 +400,7 @@ def __init__(
         classifier=None,
         class_selector=None,
         averager=None,
+        batch_size=512,
     ):
         """
         Instantiates with default debug paramaters.
@@ -411,6 +415,7 @@ def __init__(
         :param averager: `Averager2D` instance.
             Default `None` ceates `BFRAverager2D` instance.
             See code for parameter details.
+        :param batch_size: Integer size for batched operations.
 
         :return: ClassAvgSource instance.
         """
@@ -432,6 +437,7 @@ def __init__(
                 self._get_classifier_basis(classifier),
                 src,
                 dtype=dtype,
+                batch_size=batch_size,
             )
 
         if class_selector is None:
@@ -442,6 +448,7 @@ def __init__(
             classifier=classifier,
             class_selector=class_selector,
             averager=averager,
+            batch_size=batch_size,
         )
 
 
@@ -464,6 +471,7 @@ def __init__(
         class_selector=None,
         averager=None,
         averager_src=None,
+        batch_size=512,
     ):
         """
         Instantiates `ClassAvgSource` with the following parameters.
@@ -483,6 +491,7 @@ def __init__(
              provide distinct sources for classification and
              averaging.  Raises error when combined with an explicit
              `averager` argument.
+        :param batch_size: Integer size for batched operations.
 
         :return: ClassAvgSource instance.
         """
@@ -508,6 +517,7 @@ def __init__(
             averager = BFRAverager2D(
                 composite_basis=basis_2d,
                 src=averager_src,
+                batch_size=batch_size,
                 dtype=dtype,
             )
         elif averager_src is not None:
@@ -523,6 +533,7 @@ def __init__(
             classifier=classifier,
             class_selector=class_selector,
             averager=averager,
+            batch_size=batch_size,
         )
 
 
@@ -533,6 +544,7 @@ def DefaultClassAvgSource(
     class_selector=None,
     averager=None,
     averager_src=None,
+    batch_size=512,
     version=None,
 ):
     """
@@ -554,6 +566,7 @@ def DefaultClassAvgSource(
          provide distinct sources for classification and
          averaging.  Raises error when combined with an explicit
          `averager` argument.
+    :param batch_size: Integer size for batched operations.
     :param version: Optionally selects a versioned `DefaultClassAvgSource`.
         Defaults to latest available.
     :return: ClassAvgSource instance.
@@ -577,6 +590,7 @@ def DefaultClassAvgSource(
         class_selector=class_selector,
         averager=averager,
         averager_src=averager_src,
+        batch_size=batch_size,
     )
 
 
@@ -597,6 +611,7 @@ def __init__(
         class_selector=None,
         averager=None,
         averager_src=None,
+        batch_size=512,
     ):
         """
         Instantiates ClassAvgSourcev132 with the following parameters.
@@ -618,6 +633,7 @@ def __init__(
              provide distinct sources for classification and
              averaging.  Raises error when combined with an explicit
              `averager` argument.
+        :param batch_size: Integer size for batched operations.
 
         :return: ClassAvgSource instance.
         """
@@ -643,6 +659,7 @@ def __init__(
             averager = BFRAverager2D(
                 composite_basis=basis_2d,
                 src=averager_src,
+                batch_size=batch_size,
                 dtype=dtype,
             )
         elif averager_src is not None:
@@ -661,6 +678,7 @@ def __init__(
             classifier=classifier,
             class_selector=class_selector,
             averager=averager,
+            batch_size=batch_size,
         )
 
 
@@ -681,6 +699,7 @@ def __init__(
         class_selector=None,
         averager=None,
         averager_src=None,
+        batch_size=512,
     ):
         """
         Instantiates ClassAvgSourcev110 with the following parameters.
@@ -699,6 +718,7 @@ def __init__(
             to BFSRAverager2D during initialization.
             Raises error when combined with an explicit `averager`
             argument.
+        :param batch_size: Integer size for batched operations.
 
         :return: ClassAvgSource instance.
         """
@@ -724,6 +744,7 @@ def __init__(
             averager = BFSRAverager2D(
                 composite_basis=basis_2d,
                 src=averager_src,
+                batch_size=batch_size,
                 dtype=dtype,
             )
         elif averager_src is not None:
@@ -739,4 +760,5 @@ def __init__(
             classifier=classifier,
             class_selector=class_selector,
             averager=averager,
+            batch_size=batch_size,
         )

From 9f919401ba618e75002f7686e54952cc5872762d Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Fri, 10 Jan 2025 15:43:40 -0500
Subject: [PATCH 092/184] stashing epsdR port of isotropic noise est

---
 src/aspire/noise/__init__.py |   1 +
 src/aspire/noise/noise.py    | 133 +++++++++++++++++++++++++++++++++++
 2 files changed, 134 insertions(+)

diff --git a/src/aspire/noise/__init__.py b/src/aspire/noise/__init__.py
index 1c5c45e08d..67978bcd2f 100644
--- a/src/aspire/noise/__init__.py
+++ b/src/aspire/noise/__init__.py
@@ -2,6 +2,7 @@
     AnisotropicNoiseEstimator,
     BlueNoiseAdder,
     CustomNoiseAdder,
+    IsotropicNoiseEstimator,
     NoiseAdder,
     NoiseEstimator,
     PinkNoiseAdder,
diff --git a/src/aspire/noise/noise.py b/src/aspire/noise/noise.py
index 642ef0f5c5..4b07aa1019 100644
--- a/src/aspire/noise/noise.py
+++ b/src/aspire/noise/noise.py
@@ -345,3 +345,136 @@ def estimate_noise_psd(self):
         noise_psd_est[mid, mid] -= mean_est**2
 
         return noise_psd_est
+
+
+class IsotropicNoiseEstimator(NoiseEstimator):
+    """
+    Isotropic Noise Estimator.
+    """
+
+    def estimate(self):
+        """
+        :return: The estimated noise variance of the images.
+        """
+
+        # AnisotropicNoiseEstimator.filter is an ArrayFilter.
+        #   We average the variance over all frequencies,
+
+        return np.mean(self.filter.evaluate_grid(self.src.L))
+
+    def _create_filter(self, noise_psd=None):
+        """
+        :param noise_psd: Noise PSD of images
+        :return: The estimated noise power spectral distribution (PSD) of the images in the form of a filter object.
+        """
+        if noise_psd is None:
+            noise_psd = self.estimate_noise_psd()
+        return ArrayFilter(noise_psd)
+
+    def estimate_noise_psd(self):
+        """
+        :return: The estimated noise variance of the images in the Source used to create this estimator.
+        TODO: How's this initial estimate of variance different from the 'estimate' method?
+        """
+
+        return noise_psd_est
+
+    @staticmethod
+    def epsdR(images, samples_idx, max_d=None):
+        """
+        Estimate the 1D isotropic autocorrelation function of `images`.
+        The samples to use in each image are given by `samples_idx` mask.
+        The correlation is computed up to a maximal distance of `max_d`.
+
+        :param images: Images as a Numpy array shaped (n_img,L,L).
+        :param samples_idx: Boolean mask shaped (L,L).
+        :param max_d: Max computed correlation distance in pixels.
+        :return: Tuple radial PSD, distances map, count of nonzero correlations.
+        """
+
+        n, L, L2 = images.shape
+        if L != L2:
+            raise RuntimeError(f"Images must be square, received {images.shape}")
+
+        # Correlations more than `max_d` pixels apart are not computed.
+        if max_d is None:
+            max_d = L - 1
+        if max_d > L - 1:
+            logger.info(
+                f"`max_d` value {max_d}greater than number of image pixels {L}, clipping to {L-1}."
+            )
+        max_d = min(max_d, L - 1)
+
+        # Compute distances
+        # Note grid_2d['r'] is not used because we always want zero centered integer grid,
+        #   yielding integer dists (radius**2) values.
+        J, I = np.mgrid[0:max_d, 0:max_d]
+        dists = I * I + J * J
+        dsquare = np.sort(np.unique(dists[dists <= max_d**2]))
+        x = np.sqrt(dsquare)  # actual distance
+
+        # corrs[i] is the sum of all x[j]x[j+d] where d = x[i]
+        corrs = np.zeros_like(dsquare)
+        # corrcount[i] is the number of pairs summed in corr[i]
+        corrcount = np.zeros_like(dsquare, dtype=int)
+
+        # distmap maps [i,j] to k where dsquare[k] = i**2 + j**2.
+        #   -1 indicates distance is larger than max_d
+        distmap = np.full(shape=dists.shape, fill_value=-1)
+
+        # This differs from the MATLAB code because Numpy does not directly provide `bsearch`.
+        for i, d in enumerate(dsquare):
+            inds = dists == d  # locations having distance `d`
+            distmap[inds] = i  # assign index into dsquare `i`
+        # # Mapped distance indices where i**2+j**2 <= max_d**2
+        # validdists = np.where(distmap != -1)  # Note this is a 2-tuple
+
+        # Compute Ncorr using a constant unit image.
+        mask = np.zeros((L, L))
+        mask[samples_idx] = 1
+        tmp = np.zeros((2 * L + 1, 2 * L + 1))  # pad
+        tmp[:L, :L] = mask
+        ftmp = fft.fft2(tmp)
+        Ncorr = fft.ifft2(ftmp * ftmp.conj())
+        Ncorr = Ncorr[:max_d, :max_d]  # crop
+        Ncorr = np.round(Ncorr)
+
+        # Values of isotropic autocorrelation function
+        # R[i] is value of ACF at distance x[i]
+        R = np.zeros(len(corrs))
+
+        samples = np.zeros((L, L))
+        tmp[:, :] = 0  # reset tmp
+        for k in trange(n, desc="Processing image autocorrelations"):
+            # Mask unused pixels (note, think can merge these lines later)
+            samples[samples_idx] = images[k][samples_idx]
+
+            # Compute non-preiodic autocorrelation
+            tmp[:L, :L] = samples  # pad
+            ftmp = fft.fft2(tmp)
+            s = fft.ifft2(ftmp * ftmp.conj()).real  # matlab code does not cast here
+            s = s[:max_d, :max_d]  # crop
+
+            # I didn't port this correctly.  (MATLAB auto (un/)raveling)
+            # # Accumulate all autocorrelation values R[k1,k2] such that
+            # # k1^2+k2^2=const (all autocorrelations of a certain distance).
+            # for j in range(np.size(validdists)):
+            #     currdist = validdists[j]
+            #     dmidx = distmap[currdist]
+            #     corrs[dmidx] = corrs[dmidx] + s[currdist]
+            #     corrcount[dmidx] = corrcount[dmidx] + Ncorr[currdist]
+
+            for i in range(max_d):
+                for j in range(max_d):
+                    idx = distmap[i, j]
+                    if idx != -1:
+                        corrs[idx] = corrs[idx] + s[i, j]
+                        corrcount[idx] = corrcount[idx] + Ncorr[i, j]
+
+        # Remove distances which had no samples
+        idx = np.where(corrcount != 0)  # [0]
+        R = corrs[idx] / corrcount[idx]
+        x = x[idx]
+        cnt = corrcount[idx]
+
+        return R, x, cnt

From 460f6a0b70617c8c82827329446591495c089838 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Tue, 14 Jan 2025 09:28:49 -0500
Subject: [PATCH 093/184] epsdR port of isotropic noise est matching in dbgr

---
 src/aspire/noise/noise.py | 41 +++++++++++++++++++++++----------------
 1 file changed, 24 insertions(+), 17 deletions(-)

diff --git a/src/aspire/noise/noise.py b/src/aspire/noise/noise.py
index 4b07aa1019..b2a27a58d4 100644
--- a/src/aspire/noise/noise.py
+++ b/src/aspire/noise/noise.py
@@ -398,23 +398,23 @@ def epsdR(images, samples_idx, max_d=None):
 
         # Correlations more than `max_d` pixels apart are not computed.
         if max_d is None:
-            max_d = L - 1
+            max_d = np.floor(L/3)
         if max_d > L - 1:
             logger.info(
                 f"`max_d` value {max_d}greater than number of image pixels {L}, clipping to {L-1}."
             )
-        max_d = min(max_d, L - 1)
+        max_d = int(min(max_d, L - 1))
 
         # Compute distances
         # Note grid_2d['r'] is not used because we always want zero centered integer grid,
         #   yielding integer dists (radius**2) values.
-        J, I = np.mgrid[0:max_d, 0:max_d]
+        J, I = np.mgrid[0:max_d+1, 0:max_d+1]
         dists = I * I + J * J
         dsquare = np.sort(np.unique(dists[dists <= max_d**2]))
         x = np.sqrt(dsquare)  # actual distance
 
         # corrs[i] is the sum of all x[j]x[j+d] where d = x[i]
-        corrs = np.zeros_like(dsquare)
+        corrs = np.zeros_like(dsquare,dtype=np.float64)
         # corrcount[i] is the number of pairs summed in corr[i]
         corrcount = np.zeros_like(dsquare, dtype=int)
 
@@ -435,8 +435,8 @@ def epsdR(images, samples_idx, max_d=None):
         tmp = np.zeros((2 * L + 1, 2 * L + 1))  # pad
         tmp[:L, :L] = mask
         ftmp = fft.fft2(tmp)
-        Ncorr = fft.ifft2(ftmp * ftmp.conj())
-        Ncorr = Ncorr[:max_d, :max_d]  # crop
+        Ncorr = fft.ifft2(ftmp * ftmp.conj()).real  # matlab code does not cast here, but internally detects conj sym...
+        Ncorr = Ncorr[:max_d+1, :max_d+1]  # crop
         Ncorr = np.round(Ncorr)
 
         # Values of isotropic autocorrelation function
@@ -452,25 +452,29 @@ def epsdR(images, samples_idx, max_d=None):
             # Compute non-preiodic autocorrelation
             tmp[:L, :L] = samples  # pad
             ftmp = fft.fft2(tmp)
-            s = fft.ifft2(ftmp * ftmp.conj()).real  # matlab code does not cast here
-            s = s[:max_d, :max_d]  # crop
-
-            # I didn't port this correctly.  (MATLAB auto (un/)raveling)
+            s = fft.ifft2(ftmp * ftmp.conj()).real  # matlab code does not cast here, but internally detects conj sym...
+            s = s[0:max_d+1, 0:max_d+1]  # crop
+            
             # # Accumulate all autocorrelation values R[k1,k2] such that
             # # k1^2+k2^2=const (all autocorrelations of a certain distance).
+            for i in range(max_d+1):
+                for j in range(max_d+1):
+                    idx = distmap[i, j]
+                    if idx != -1:
+                        corrs[idx] = corrs[idx] + s[i, j]
+                        corrcount[idx] = corrcount[idx] + Ncorr[i, j]
+
+            # TODO, fix this MATLAB optmized implementation and compare with the clearer code above.
+            # I didn't port this validdist slice optimized version correctly(yet).
+            # it uses implicit (MATLAB auto flat (un/)raveling)
+            # im not sure the speedup would be similar in python anyway.
             # for j in range(np.size(validdists)):
             #     currdist = validdists[j]
             #     dmidx = distmap[currdist]
             #     corrs[dmidx] = corrs[dmidx] + s[currdist]
             #     corrcount[dmidx] = corrcount[dmidx] + Ncorr[currdist]
 
-            for i in range(max_d):
-                for j in range(max_d):
-                    idx = distmap[i, j]
-                    if idx != -1:
-                        corrs[idx] = corrs[idx] + s[i, j]
-                        corrcount[idx] = corrcount[idx] + Ncorr[i, j]
-
+                        
         # Remove distances which had no samples
         idx = np.where(corrcount != 0)  # [0]
         R = corrs[idx] / corrcount[idx]
@@ -478,3 +482,6 @@ def epsdR(images, samples_idx, max_d=None):
         cnt = corrcount[idx]
 
         return R, x, cnt
+
+    @staticmethod
+    def epsdS(images, samples_idx, max_d=None):

From dc88863efba7918e892e0ef87822ff3ed86e0fc3 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Tue, 14 Jan 2025 10:45:58 -0500
Subject: [PATCH 094/184] initial port of epsdS

need gwindow
---
 src/aspire/noise/noise.py | 85 ++++++++++++++++++++++++++++++++++++++-
 1 file changed, 83 insertions(+), 2 deletions(-)

diff --git a/src/aspire/noise/noise.py b/src/aspire/noise/noise.py
index b2a27a58d4..b628bdf850 100644
--- a/src/aspire/noise/noise.py
+++ b/src/aspire/noise/noise.py
@@ -392,7 +392,7 @@ def epsdR(images, samples_idx, max_d=None):
         :return: Tuple radial PSD, distances map, count of nonzero correlations.
         """
 
-        n, L, L2 = images.shape
+        n_img, L, L2 = images.shape
         if L != L2:
             raise RuntimeError(f"Images must be square, received {images.shape}")
 
@@ -445,9 +445,10 @@ def epsdR(images, samples_idx, max_d=None):
 
         samples = np.zeros((L, L))
         tmp[:, :] = 0  # reset tmp
-        for k in trange(n, desc="Processing image autocorrelations"):
+        for k in trange(n_img, desc="Processing image autocorrelations"):
             # Mask unused pixels (note, think can merge these lines later)
             samples[samples_idx] = images[k][samples_idx]
+            # Note, we can also compute the noise energy estimate used later at this time to avoid looping over images twice.
 
             # Compute non-preiodic autocorrelation
             tmp[:L, :L] = samples  # pad
@@ -485,3 +486,83 @@ def epsdR(images, samples_idx, max_d=None):
 
     @staticmethod
     def epsdS(images, samples_idx, max_d=None):
+        """
+        Estimate the 2D isotropic power spectrum of `images`.
+        The samples to use in each image are given by `samples_idx` mask.
+        The correlation is computed up to a maximal distance of `max_d`.
+        """
+        R, x, _ = IsotropicNoiseEstimator.epsdR(images=images, samples_idx=samples_idx, max_d=max_d)
+
+        n_img, L, L2 = images.shape
+        if L != L2:
+            raise RuntimeError(f"Images must be square, received {images.shape}")
+
+        # Correlations more than `max_d` pixels apart are not computed.
+        if max_d is None:
+            max_d = np.floor(L/3)
+        if max_d > L - 1:
+            logger.info(
+                f"`max_d` value {max_d}greater than number of image pixels {L}, clipping to {L-1}."
+            )
+        max_d = int(min(max_d, L - 1))
+
+        # Use the 1D autocorrelation estimted above to populate an
+        # array of the 2D isotropic autocorrelction. This
+        # autocorrelation is later Fourier transformed to get the
+        # power spectrum.
+        R2 = np.zeros((2*L-1, 2*L-1), dtype=np.float64)
+
+        J, I = np.mgrid[-L+1:L, -L+1:L]
+        dists2 = I * I + J * J
+        dsquare2 = np.sort(np.unique(dists2[dists2 <= max_d**2]))
+        for i,d in enumerate(dsquare2):
+            idx = dists2==d
+            R2[idx] = R[i]
+
+        # R2 seems okay here.
+        #breakpoint()
+
+        # Window te 2D autocorrelation and Fourier transform it to get the power
+        # spectrum. Always use the Gaussian window, as it has positive Fourier
+        # transform.
+        w = gwrindor(L, max_d)
+        P2 = fft.centered_fft2(R2*w)
+        if (err := np.linalg.norm(P2.imag) / np.linalg.norm(P2)) > 1e-12:
+            logger.warning(f'Large imaginary components in P2 {err}.')
+        P2 = P2.real
+        
+        # Normalize the power spectrum P2. The power spectrum is normalized such
+        # that its energy is equal to the average energy of the noise samples used
+        # to estimate it.
+
+        E=0  # Total energy of the noise samples used to estimate the power spectrum.
+        samples = np.zeros((L, L))        
+        for k in trange(n_img, desc="Estimating image noise energy"):
+            samples[samples_idx] = images[k][samples_idx]            
+            E += np.sum( (samples - np.mean(samples))**2)
+        # Mean energy of the noise samples
+        meanE = E / (samples.size * n_img)
+
+        # Normalize P2 such that its mean energy is preserved and is equal to
+        # meanE, that is, mean(P2)==meanE. That way the mean energy does not
+        # go down if the number of pixels is artifically changed (say be
+        # upsampling, downsampling, or cropping). Note that P2 is already in
+        # units of energy, and so the total energy is given by sum(P2) and
+        # not by norm(P2).
+        P2 = P2 / np.sum(P2) * meanE * P2.size
+
+        # Check that P2 has no negative values.
+        # Due to the truncation of the Gaussian window, we get small negative
+        # values. So unless they are very big, we just ignore them.
+        negidx = P2 < 0
+        if np.count_nonzero(negidx):
+            maxnegerr = np.max(np.abs(P2[negidx]))
+            logger.debug(f'Maximal negative P2 value = {maxnegerr}')
+            if maxnegerr > 1e-2:
+                negnorm = np.linalg.norm(P2[negidx])
+                logger.warning(f'Power spectrum P2 has negative values with energy {negnorm}.')
+            P2[negidx] = 0  # zero out negative estimates
+
+        return P2, R, R2, x
+            
+        

From dd0b9346a89a4d589f323178f02a408649ecaf3d Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Tue, 14 Jan 2025 11:34:27 -0500
Subject: [PATCH 095/184] Add gwindow.  P2 comparable to MATLAB.

[skip ci]
---
 src/aspire/noise/noise.py | 85 ++++++++++++++++++++++++---------------
 1 file changed, 52 insertions(+), 33 deletions(-)

diff --git a/src/aspire/noise/noise.py b/src/aspire/noise/noise.py
index b628bdf850..6fc6b17caa 100644
--- a/src/aspire/noise/noise.py
+++ b/src/aspire/noise/noise.py
@@ -398,7 +398,7 @@ def epsdR(images, samples_idx, max_d=None):
 
         # Correlations more than `max_d` pixels apart are not computed.
         if max_d is None:
-            max_d = np.floor(L/3)
+            max_d = np.floor(L / 3)
         if max_d > L - 1:
             logger.info(
                 f"`max_d` value {max_d}greater than number of image pixels {L}, clipping to {L-1}."
@@ -408,13 +408,13 @@ def epsdR(images, samples_idx, max_d=None):
         # Compute distances
         # Note grid_2d['r'] is not used because we always want zero centered integer grid,
         #   yielding integer dists (radius**2) values.
-        J, I = np.mgrid[0:max_d+1, 0:max_d+1]
+        J, I = np.mgrid[0 : max_d + 1, 0 : max_d + 1]
         dists = I * I + J * J
         dsquare = np.sort(np.unique(dists[dists <= max_d**2]))
         x = np.sqrt(dsquare)  # actual distance
 
         # corrs[i] is the sum of all x[j]x[j+d] where d = x[i]
-        corrs = np.zeros_like(dsquare,dtype=np.float64)
+        corrs = np.zeros_like(dsquare, dtype=np.float64)
         # corrcount[i] is the number of pairs summed in corr[i]
         corrcount = np.zeros_like(dsquare, dtype=int)
 
@@ -435,8 +435,10 @@ def epsdR(images, samples_idx, max_d=None):
         tmp = np.zeros((2 * L + 1, 2 * L + 1))  # pad
         tmp[:L, :L] = mask
         ftmp = fft.fft2(tmp)
-        Ncorr = fft.ifft2(ftmp * ftmp.conj()).real  # matlab code does not cast here, but internally detects conj sym...
-        Ncorr = Ncorr[:max_d+1, :max_d+1]  # crop
+        Ncorr = fft.ifft2(
+            ftmp * ftmp.conj()
+        ).real  # matlab code does not cast here, but internally detects conj sym...
+        Ncorr = Ncorr[: max_d + 1, : max_d + 1]  # crop
         Ncorr = np.round(Ncorr)
 
         # Values of isotropic autocorrelation function
@@ -453,13 +455,15 @@ def epsdR(images, samples_idx, max_d=None):
             # Compute non-preiodic autocorrelation
             tmp[:L, :L] = samples  # pad
             ftmp = fft.fft2(tmp)
-            s = fft.ifft2(ftmp * ftmp.conj()).real  # matlab code does not cast here, but internally detects conj sym...
-            s = s[0:max_d+1, 0:max_d+1]  # crop
-            
+            s = fft.ifft2(
+                ftmp * ftmp.conj()
+            ).real  # matlab code does not cast here, but internally detects conj sym...
+            s = s[0 : max_d + 1, 0 : max_d + 1]  # crop
+
             # # Accumulate all autocorrelation values R[k1,k2] such that
             # # k1^2+k2^2=const (all autocorrelations of a certain distance).
-            for i in range(max_d+1):
-                for j in range(max_d+1):
+            for i in range(max_d + 1):
+                for j in range(max_d + 1):
                     idx = distmap[i, j]
                     if idx != -1:
                         corrs[idx] = corrs[idx] + s[i, j]
@@ -475,7 +479,6 @@ def epsdR(images, samples_idx, max_d=None):
             #     corrs[dmidx] = corrs[dmidx] + s[currdist]
             #     corrcount[dmidx] = corrcount[dmidx] + Ncorr[currdist]
 
-                        
         # Remove distances which had no samples
         idx = np.where(corrcount != 0)  # [0]
         R = corrs[idx] / corrcount[idx]
@@ -484,6 +487,22 @@ def epsdR(images, samples_idx, max_d=None):
 
         return R, x, cnt
 
+    @staticmethod
+    def gwindow(L, max_d, alpha=3.0):
+        """
+        Create a 2D gaussian window used for 2D power spectrum estimation.
+
+        Given `L` retuns a `(2L-1, 2L-1)` Gaussian window where
+        `max_d` is the width of the Gaussian and `alpha` is the
+        reciprical of the standard deviation of the Gaussian window.
+        See Harris 78.
+        """
+
+        X, Y = np.mgrid[-(L - 1) : L, -(L - 1) : L]  # -(L-1) to (L-1) inclusive
+        W = np.exp(-alpha * (X**2 + Y**2) / (2 * max_d**2))
+
+        return W
+
     @staticmethod
     def epsdS(images, samples_idx, max_d=None):
         """
@@ -491,7 +510,9 @@ def epsdS(images, samples_idx, max_d=None):
         The samples to use in each image are given by `samples_idx` mask.
         The correlation is computed up to a maximal distance of `max_d`.
         """
-        R, x, _ = IsotropicNoiseEstimator.epsdR(images=images, samples_idx=samples_idx, max_d=max_d)
+        R, x, _ = IsotropicNoiseEstimator.epsdR(
+            images=images, samples_idx=samples_idx, max_d=max_d
+        )
 
         n_img, L, L2 = images.shape
         if L != L2:
@@ -499,7 +520,7 @@ def epsdS(images, samples_idx, max_d=None):
 
         # Correlations more than `max_d` pixels apart are not computed.
         if max_d is None:
-            max_d = np.floor(L/3)
+            max_d = np.floor(L / 3)
         if max_d > L - 1:
             logger.info(
                 f"`max_d` value {max_d}greater than number of image pixels {L}, clipping to {L-1}."
@@ -510,38 +531,36 @@ def epsdS(images, samples_idx, max_d=None):
         # array of the 2D isotropic autocorrelction. This
         # autocorrelation is later Fourier transformed to get the
         # power spectrum.
-        R2 = np.zeros((2*L-1, 2*L-1), dtype=np.float64)
+        R2 = np.zeros((2 * L - 1, 2 * L - 1), dtype=np.float64)
 
-        J, I = np.mgrid[-L+1:L, -L+1:L]
+        J, I = np.mgrid[-L + 1 : L, -L + 1 : L]
         dists2 = I * I + J * J
         dsquare2 = np.sort(np.unique(dists2[dists2 <= max_d**2]))
-        for i,d in enumerate(dsquare2):
-            idx = dists2==d
+        for i, d in enumerate(dsquare2):
+            idx = dists2 == d
             R2[idx] = R[i]
 
-        # R2 seems okay here.
-        #breakpoint()
-
         # Window te 2D autocorrelation and Fourier transform it to get the power
         # spectrum. Always use the Gaussian window, as it has positive Fourier
         # transform.
-        w = gwrindor(L, max_d)
-        P2 = fft.centered_fft2(R2*w)
+        w = IsotropicNoiseEstimator.gwindow(L, max_d)
+        P2 = fft.centered_fft2(R2 * w)
         if (err := np.linalg.norm(P2.imag) / np.linalg.norm(P2)) > 1e-12:
-            logger.warning(f'Large imaginary components in P2 {err}.')
+            logger.warning(f"Large imaginary components in P2 {err}.")
         P2 = P2.real
-        
+
         # Normalize the power spectrum P2. The power spectrum is normalized such
         # that its energy is equal to the average energy of the noise samples used
         # to estimate it.
 
-        E=0  # Total energy of the noise samples used to estimate the power spectrum.
-        samples = np.zeros((L, L))        
+        E = 0.0  # Total energy of the noise samples used to estimate the power spectrum.
+        samples = np.zeros((L, L))
         for k in trange(n_img, desc="Estimating image noise energy"):
-            samples[samples_idx] = images[k][samples_idx]            
-            E += np.sum( (samples - np.mean(samples))**2)
+            samples[samples_idx] = images[k][samples_idx]
+            E += np.sum((samples - np.mean(samples)) ** 2)
         # Mean energy of the noise samples
-        meanE = E / (samples.size * n_img)
+        n_samples_per_img = np.count_nonzero(samples_idx)
+        meanE = E / (n_samples_per_img * n_img)
 
         # Normalize P2 such that its mean energy is preserved and is equal to
         # meanE, that is, mean(P2)==meanE. That way the mean energy does not
@@ -557,12 +576,12 @@ def epsdS(images, samples_idx, max_d=None):
         negidx = P2 < 0
         if np.count_nonzero(negidx):
             maxnegerr = np.max(np.abs(P2[negidx]))
-            logger.debug(f'Maximal negative P2 value = {maxnegerr}')
+            logger.debug(f"Maximal negative P2 value = {maxnegerr}")
             if maxnegerr > 1e-2:
                 negnorm = np.linalg.norm(P2[negidx])
-                logger.warning(f'Power spectrum P2 has negative values with energy {negnorm}.')
+                logger.warning(
+                    f"Power spectrum P2 has negative values with energy {negnorm}."
+                )
             P2[negidx] = 0  # zero out negative estimates
 
         return P2, R, R2, x
-            
-        

From 7dcaecf6c9412b89e98c9983b437b81480b3ec29 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Tue, 14 Jan 2025 13:49:12 -0500
Subject: [PATCH 096/184] migrate gaussian_window to utils

---
 src/aspire/noise/noise.py    | 20 ++------------------
 src/aspire/utils/__init__.py |  1 +
 src/aspire/utils/misc.py     | 24 ++++++++++++++++++++++++
 3 files changed, 27 insertions(+), 18 deletions(-)

diff --git a/src/aspire/noise/noise.py b/src/aspire/noise/noise.py
index 6fc6b17caa..0398dee077 100644
--- a/src/aspire/noise/noise.py
+++ b/src/aspire/noise/noise.py
@@ -13,7 +13,7 @@
     PowerFilter,
     ScalarFilter,
 )
-from aspire.utils import grid_2d, randn, trange
+from aspire.utils import gaussian_window, grid_2d, randn, trange
 
 logger = logging.getLogger(__name__)
 
@@ -487,22 +487,6 @@ def epsdR(images, samples_idx, max_d=None):
 
         return R, x, cnt
 
-    @staticmethod
-    def gwindow(L, max_d, alpha=3.0):
-        """
-        Create a 2D gaussian window used for 2D power spectrum estimation.
-
-        Given `L` retuns a `(2L-1, 2L-1)` Gaussian window where
-        `max_d` is the width of the Gaussian and `alpha` is the
-        reciprical of the standard deviation of the Gaussian window.
-        See Harris 78.
-        """
-
-        X, Y = np.mgrid[-(L - 1) : L, -(L - 1) : L]  # -(L-1) to (L-1) inclusive
-        W = np.exp(-alpha * (X**2 + Y**2) / (2 * max_d**2))
-
-        return W
-
     @staticmethod
     def epsdS(images, samples_idx, max_d=None):
         """
@@ -543,7 +527,7 @@ def epsdS(images, samples_idx, max_d=None):
         # Window te 2D autocorrelation and Fourier transform it to get the power
         # spectrum. Always use the Gaussian window, as it has positive Fourier
         # transform.
-        w = IsotropicNoiseEstimator.gwindow(L, max_d)
+        w = gaussian_window(L, max_d)
         P2 = fft.centered_fft2(R2 * w)
         if (err := np.linalg.norm(P2.imag) / np.linalg.norm(P2)) > 1e-12:
             logger.warning(f"Large imaginary components in P2 {err}.")
diff --git a/src/aspire/utils/__init__.py b/src/aspire/utils/__init__.py
index df4134a6d2..74e13c8b82 100644
--- a/src/aspire/utils/__init__.py
+++ b/src/aspire/utils/__init__.py
@@ -24,6 +24,7 @@
     gaussian_1d,
     gaussian_2d,
     gaussian_3d,
+    gaussian_window,
     importlib_path,
     inverse_r,
     J_conjugate,
diff --git a/src/aspire/utils/misc.py b/src/aspire/utils/misc.py
index 7ce6602f86..b4248bc67e 100644
--- a/src/aspire/utils/misc.py
+++ b/src/aspire/utils/misc.py
@@ -359,6 +359,30 @@ def fuzzy_mask(L, dtype, r0=None, risetime=None):
     return m
 
 
+def gaussian_window(L, max_d, alpha=3.0):
+    """
+    Create a 2D gaussian window used for 2D power spectrum estimation.
+
+    Given `L` retuns a `(2L-1, 2L-1)` array where
+    `max_d` is the width of the Gaussian and `alpha` is the
+    reciprical of the standard deviation of the Gaussian window.
+    See Harris 78.
+
+    When `alpha=1`, this function should be equivalent to
+    `gaussian_2d(L=2*L-1, sigma=max_d)`.
+
+    :param L: Number of radial pixels
+    :param max_d: Width of Gaussian (stddev)
+    :param alpha: Reciprical of stddev of window
+    :return: Numpy array with shape `(2L-1, 2L-1)`x
+    """
+
+    X, Y = np.mgrid[-(L - 1) : L, -(L - 1) : L]  # -(L-1) to (L-1) inclusive
+    W = np.exp(-alpha * (X**2 + Y**2) / (2 * max_d**2))
+
+    return W
+
+
 def all_pairs(n, return_map=False):
     """
     All pairs indexing (i,j) for i<j and a pairs-to-linear index mapping.

From 19af6fc4b4bc25cde4b679627d8122637580596d Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Tue, 14 Jan 2025 14:54:02 -0500
Subject: [PATCH 097/184] Begin cleanup and add test for gaussian_window

---
 src/aspire/noise/noise.py | 11 +++++------
 src/aspire/utils/misc.py  |  6 ++++--
 tests/test_noise.py       | 38 ++++++++++++++++++++++++++++++++++----
 3 files changed, 43 insertions(+), 12 deletions(-)

diff --git a/src/aspire/noise/noise.py b/src/aspire/noise/noise.py
index 0398dee077..7eaaee9457 100644
--- a/src/aspire/noise/noise.py
+++ b/src/aspire/noise/noise.py
@@ -376,8 +376,7 @@ def estimate_noise_psd(self):
         :return: The estimated noise variance of the images in the Source used to create this estimator.
         TODO: How's this initial estimate of variance different from the 'estimate' method?
         """
-
-        return noise_psd_est
+        pass
 
     @staticmethod
     def epsdR(images, samples_idx, max_d=None):
@@ -408,8 +407,8 @@ def epsdR(images, samples_idx, max_d=None):
         # Compute distances
         # Note grid_2d['r'] is not used because we always want zero centered integer grid,
         #   yielding integer dists (radius**2) values.
-        J, I = np.mgrid[0 : max_d + 1, 0 : max_d + 1]
-        dists = I * I + J * J
+        X, Y = np.mgrid[0 : max_d + 1, 0 : max_d + 1]
+        dists = X * X + Y * Y
         dsquare = np.sort(np.unique(dists[dists <= max_d**2]))
         x = np.sqrt(dsquare)  # actual distance
 
@@ -517,8 +516,8 @@ def epsdS(images, samples_idx, max_d=None):
         # power spectrum.
         R2 = np.zeros((2 * L - 1, 2 * L - 1), dtype=np.float64)
 
-        J, I = np.mgrid[-L + 1 : L, -L + 1 : L]
-        dists2 = I * I + J * J
+        X, Y = np.mgrid[-L + 1 : L, -L + 1 : L]
+        dists2 = X * X + Y * Y
         dsquare2 = np.sort(np.unique(dists2[dists2 <= max_d**2]))
         for i, d in enumerate(dsquare2):
             idx = dists2 == d
diff --git a/src/aspire/utils/misc.py b/src/aspire/utils/misc.py
index b4248bc67e..3cae04ee1a 100644
--- a/src/aspire/utils/misc.py
+++ b/src/aspire/utils/misc.py
@@ -359,7 +359,7 @@ def fuzzy_mask(L, dtype, r0=None, risetime=None):
     return m
 
 
-def gaussian_window(L, max_d, alpha=3.0):
+def gaussian_window(L, max_d, alpha=3.0, dtype=np.float64):
     """
     Create a 2D gaussian window used for 2D power spectrum estimation.
 
@@ -369,7 +369,7 @@ def gaussian_window(L, max_d, alpha=3.0):
     See Harris 78.
 
     When `alpha=1`, this function should be equivalent to
-    `gaussian_2d(L=2*L-1, sigma=max_d)`.
+    `gaussian_2d(size=2*L-1, sigma=max_d)`.
 
     :param L: Number of radial pixels
     :param max_d: Width of Gaussian (stddev)
@@ -378,6 +378,8 @@ def gaussian_window(L, max_d, alpha=3.0):
     """
 
     X, Y = np.mgrid[-(L - 1) : L, -(L - 1) : L]  # -(L-1) to (L-1) inclusive
+    X = X.astype(dtype, copy=False)
+    Y = Y.astype(dtype, copy=False)
     W = np.exp(-alpha * (X**2 + Y**2) / (2 * max_d**2))
 
     return W
diff --git a/tests/test_noise.py b/tests/test_noise.py
index 803deb1b57..eecded6dfa 100644
--- a/tests/test_noise.py
+++ b/tests/test_noise.py
@@ -1,4 +1,3 @@
-import itertools
 import logging
 import os.path
 
@@ -16,6 +15,7 @@
 )
 from aspire.operators import FunctionFilter, ScalarFilter
 from aspire.source.simulation import Simulation
+from aspire.utils import gaussian_2d, gaussian_window, utest_tolerance
 from aspire.volume import AsymmetricVolume
 
 DATA_DIR = os.path.join(os.path.dirname(__file__), "saved_test_data")
@@ -42,9 +42,19 @@ def sim_fixture_id(params):
     return f"res={res}, dtype={dtype.__name__}"
 
 
-@pytest.fixture(params=itertools.product(RESOLUTIONS, DTYPES), ids=sim_fixture_id)
-def sim_fixture(request):
-    resolution, dtype = request.param
+@pytest.fixture(params=DTYPES, ids=lambda x: f"dtype={x}")
+def dtype(request):
+    return request.param
+
+
+@pytest.fixture(params=RESOLUTIONS, ids=lambda x: f"resolution={x}")
+def resolution(request):
+    return request.param
+
+
+@pytest.fixture
+def sim_fixture(resolution, dtype):
+    # resolution, dtype = request.param
     # Setup a sim with no noise, no ctf, no shifts,
     #   using a compactly supported volume.
     # ie, clean centered projections.
@@ -285,3 +295,23 @@ def aniso_spectrum(x, y):
 
     # Check we're within 5%
     assert np.isclose(est_noise_variance, target_noise_variance, rtol=0.05)
+
+
+def test_gaussian_window(resolution, dtype):
+    """
+    Tests `gaussian_window` by comparing with `gaussian_2d`.
+    """
+
+    # Used by both tests below
+    max_d = resolution // 3
+    g2d = gaussian_2d(size=2 * resolution - 1, sigma=max_d, dtype=dtype)
+
+    # Test unit alpha
+    w = gaussian_window(L=resolution, max_d=max_d, dtype=dtype, alpha=1)
+    np.testing.assert_allclose(w, g2d, atol=utest_tolerance(dtype))
+
+    # Test default alpha=3, e**(alpha * ...) == (e**(...))**alpha
+    #   where (e**(...)) is provided by g2d.
+    a = 3.0
+    w = gaussian_window(L=resolution, max_d=max_d, alpha=a, dtype=dtype)
+    np.testing.assert_allclose(w, g2d**a, atol=utest_tolerance(dtype))

From 5135721b37dc76b540c0003dbe45e88165ebb254 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Wed, 15 Jan 2025 10:34:11 -0500
Subject: [PATCH 098/184] make ImageAccessor more compatible with Image

---
 src/aspire/source/image.py | 27 ++++++++++++++-------------
 1 file changed, 14 insertions(+), 13 deletions(-)

diff --git a/src/aspire/source/image.py b/src/aspire/source/image.py
index 0b0b909950..d9d3184f7f 100644
--- a/src/aspire/source/image.py
+++ b/src/aspire/source/image.py
@@ -39,20 +39,21 @@ class _ImageAccessor:
     Helper class for accessing images from an ImageSource as slices via the `src.images[start:stop:step]` API.
     """
 
-    def __init__(self, fun, num_imgs):
+    def __init__(self, fun, n_images):
         """
         :param fun: The private image-accessing method specific to the ImageSource associated with this ImageAccessor.
                     Generally _images() but can be substituted with a custom method.
-        :param num_imgs: The max number of images that this ImageAccessor can load (generally ImageSource.n).
+        :param n_images: The max number of images that this ImageAccessor can load (generally ImageSource.n).
         """
         self.fun = fun
-        self.num_imgs = num_imgs
+        self.n_images = n_images
+
 
     def __getitem__(self, indices):
         """
         ImageAccessor can be indexed via Python slice object, 1-D NumPy array, list, or a single integer,
         corresponding to the indices of the requested images. By default, slices default to a start of 0,
-        an end of self.num_imgs, and a step of 1.
+        an end of self.n_images, and a step of 1.
 
         :return: An Image object containing the requested images.
         """
@@ -71,11 +72,11 @@ def __getitem__(self, indices):
             if start < 0 and stop is None:
                 # slice(-10, None, None) -> slice(-10, *0* ,1)
                 stop = 0
-            # All other cases, limit to num_imgs
-            #   slice(s, None, None) -> slice(s, num_imgs, 1)
-            #   slice(s, 10**10, None) -> slice(0, num_imgs, 1)
-            elif not stop or stop > self.num_imgs:
-                stop = self.num_imgs
+            # All other cases, limit to n_images
+            #   slice(s, None, None) -> slice(s, n_images, 1)
+            #   slice(s, 10**10, None) -> slice(0, n_images, 1)
+            elif not stop or stop > self.n_images:
+                stop = self.n_images
 
             if not step:
                 step = 1
@@ -92,12 +93,12 @@ def __getitem__(self, indices):
             raise KeyError("Only one-dimensional indexing is allowed for images.")
 
         # final check for out-of-range indices
-        out_of_range = indices >= self.num_imgs
+        out_of_range = indices >= self.n_images
         if out_of_range.any():
             raise KeyError(f"Out-of-range indices: {list(indices[out_of_range])}")
 
         # check for negative indices and flip to positive
-        indices = indices % self.num_imgs
+        indices = indices % self.n_images
 
         return self.fun(indices)
 
@@ -262,7 +263,7 @@ def __getitem__(self, indices):
 
         :param indices: Requested indices as a Python slice object,
             1-D NumPy array, list, or a single integer. Slices default
-            to a start of 0, an end of self.num_imgs, and a step of 1.
+            to a start of 0, an end of self.n_images, and a step of 1.
             See _ImageAccessor.
         :return: Source composed of the images and metadata at `indices`.
         """
@@ -388,7 +389,7 @@ def _set_n(self, n):
             raise TypeError("`n` must be an integer")
         n = int(n)
 
-        self._img_accessor.num_imgs = n
+        self._img_accessor.n_images = n
         self._n = n
 
     @property

From df8d532981a18b2717f6b9f906669170ebbde4c5 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Wed, 15 Jan 2025 10:34:24 -0500
Subject: [PATCH 099/184] make ImageAccessor more compatible with Image

---
 src/aspire/source/image.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/src/aspire/source/image.py b/src/aspire/source/image.py
index d9d3184f7f..e315baa84e 100644
--- a/src/aspire/source/image.py
+++ b/src/aspire/source/image.py
@@ -48,7 +48,6 @@ def __init__(self, fun, n_images):
         self.fun = fun
         self.n_images = n_images
 
-
     def __getitem__(self, indices):
         """
         ImageAccessor can be indexed via Python slice object, 1-D NumPy array, list, or a single integer,

From d7471046741044fb5f6d11332e99a73439fa7098 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Wed, 15 Jan 2025 13:19:29 -0500
Subject: [PATCH 100/184] Migrate epsdS towards use of `Image` and into class
 method

---
 src/aspire/noise/noise.py | 71 ++++++++++++++++++++++++++++-----------
 1 file changed, 52 insertions(+), 19 deletions(-)

diff --git a/src/aspire/noise/noise.py b/src/aspire/noise/noise.py
index 7eaaee9457..75684fcb0b 100644
--- a/src/aspire/noise/noise.py
+++ b/src/aspire/noise/noise.py
@@ -352,23 +352,39 @@ class IsotropicNoiseEstimator(NoiseEstimator):
     Isotropic Noise Estimator.
     """
 
+    def __init__(self, src, bgRadius=None, max_d=None, batchSize=512):
+        """
+        Any additional args/kwargs are passed on to the Source's 'images' method
+
+        :param src: A Source object which can give us images on demand
+        :param bgRadius: The radius of the disk whose complement is used to estimate the noise.
+            Radius is relative proportion, where `1` represents
+            the radius of disc inscribing a `(src.L, src.L)` image.
+            Default of `None` uses `(np.floor(src.L / 2) - 1) / L`
+        :param max_d: Max computed correlation distance as a proportion of `src.L`.
+            Default of `None` uses `np.floor(src.L/3) / L`.
+        :param batchSize:  The size of the batches in which to compute the variance estimate.
+        """
+        self.src = src
+        self.dtype = self.src.dtype
+        self.bgRadius = (np.floor(self.src.L / 2) - 1) / self.src.L
+        self.batchSize = batchSize
+        self.max_d = np.floor(self.src.L / 3) / self.src.L
+
+        self.filter = self._create_filter()
+
     def estimate(self):
         """
         :return: The estimated noise variance of the images.
         """
 
-        # AnisotropicNoiseEstimator.filter is an ArrayFilter.
-        #   We average the variance over all frequencies,
-
         return np.mean(self.filter.evaluate_grid(self.src.L))
 
-    def _create_filter(self, noise_psd=None):
+    def _create_filter(self):
         """
-        :param noise_psd: Noise PSD of images
         :return: The estimated noise power spectral distribution (PSD) of the images in the form of a filter object.
         """
-        if noise_psd is None:
-            noise_psd = self.estimate_noise_psd()
+        noise_psd = self.estimate_noise_psd()
         return ArrayFilter(noise_psd)
 
     def estimate_noise_psd(self):
@@ -376,7 +392,13 @@ def estimate_noise_psd(self):
         :return: The estimated noise variance of the images in the Source used to create this estimator.
         TODO: How's this initial estimate of variance different from the 'estimate' method?
         """
-        pass
+        # Setup
+        samples_idx = grid_2d(self.src.L, normalized=True)["r"] >= self.bgRadius
+        max_d_pixels = round(self.max_d * self.src.L)
+
+        psd = IsotropicNoiseEstimator.epsdS(self.src.images, samples_idx, max_d_pixels)[0]
+
+        return psd
 
     @staticmethod
     def epsdR(images, samples_idx, max_d=None):
@@ -385,15 +407,17 @@ def epsdR(images, samples_idx, max_d=None):
         The samples to use in each image are given by `samples_idx` mask.
         The correlation is computed up to a maximal distance of `max_d`.
 
-        :param images: Images as a Numpy array shaped (n_img,L,L).
+        :param images: Images instance
         :param samples_idx: Boolean mask shaped (L,L).
         :param max_d: Max computed correlation distance in pixels.
-        :return: Tuple radial PSD, distances map, count of nonzero correlations.
+        :return:
+            - Radial PSD array of shape
+            - Distances map array
+            - Count of nonzero correlations array
         """
 
-        n_img, L, L2 = images.shape
-        if L != L2:
-            raise RuntimeError(f"Images must be square, received {images.shape}")
+        n_img = images.n_images
+        L = samples_idx.shape[-1]
 
         # Correlations more than `max_d` pixels apart are not computed.
         if max_d is None:
@@ -448,7 +472,7 @@ def epsdR(images, samples_idx, max_d=None):
         tmp[:, :] = 0  # reset tmp
         for k in trange(n_img, desc="Processing image autocorrelations"):
             # Mask unused pixels (note, think can merge these lines later)
-            samples[samples_idx] = images[k][samples_idx]
+            samples[samples_idx] = images[k].asnumpy()[0][samples_idx]
             # Note, we can also compute the noise energy estimate used later at this time to avoid looping over images twice.
 
             # Compute non-preiodic autocorrelation
@@ -492,15 +516,24 @@ def epsdS(images, samples_idx, max_d=None):
         Estimate the 2D isotropic power spectrum of `images`.
         The samples to use in each image are given by `samples_idx` mask.
         The correlation is computed up to a maximal distance of `max_d`.
+
+        :param images: Images instance
+        :param samples_idx: Boolean mask shaped (L,L).
+        :param max_d: Max computed correlation distance in pixels.
+        :return:
+            - 2D PSD array
+            - Radial PSD array
+            - Distances map array
+            - Count of nonzero correlations array
         """
+
+        n_img = images.n_images
+        L = samples_idx.shape[-1]
+
         R, x, _ = IsotropicNoiseEstimator.epsdR(
             images=images, samples_idx=samples_idx, max_d=max_d
         )
 
-        n_img, L, L2 = images.shape
-        if L != L2:
-            raise RuntimeError(f"Images must be square, received {images.shape}")
-
         # Correlations more than `max_d` pixels apart are not computed.
         if max_d is None:
             max_d = np.floor(L / 3)
@@ -539,7 +572,7 @@ def epsdS(images, samples_idx, max_d=None):
         E = 0.0  # Total energy of the noise samples used to estimate the power spectrum.
         samples = np.zeros((L, L))
         for k in trange(n_img, desc="Estimating image noise energy"):
-            samples[samples_idx] = images[k][samples_idx]
+            samples[samples_idx] = images[k].asnumpy()[0][samples_idx]
             E += np.sum((samples - np.mean(samples)) ** 2)
         # Mean energy of the noise samples
         n_samples_per_img = np.count_nonzero(samples_idx)

From ae0243f017d3aee89ec068d957676237eefd0b26 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Thu, 16 Jan 2025 15:04:13 -0500
Subject: [PATCH 101/184] Simplify NoiseEstimator classes

---
 src/aspire/noise/noise.py       | 67 +++++++++++++++++++--------------
 tests/test_anisotropic_noise.py |  2 +-
 2 files changed, 40 insertions(+), 29 deletions(-)

diff --git a/src/aspire/noise/noise.py b/src/aspire/noise/noise.py
index 75684fcb0b..2fac22e39b 100644
--- a/src/aspire/noise/noise.py
+++ b/src/aspire/noise/noise.py
@@ -1,5 +1,6 @@
 import abc
 import logging
+from functools import cached_property
 
 import numpy as np
 
@@ -241,13 +242,28 @@ def __init__(self, src, bgRadius=1, batchSize=512):
         self.bgRadius = bgRadius
         self.batchSize = batchSize
 
-        self.filter = self._create_filter()
+    @cached_property
+    def filter(self):
+        """
+        Property returning `Filter` object for this noise estimator.
+        This property will be computed and cached on first call.
+
+        :return: `Filter` object.
+        """
 
+        return self._create_filter()
+
+    @abc.abstractmethod
     def estimate(self):
         """
         :return: The estimated noise variance of the images.
         """
-        raise NotImplementedError("Subclasses implement the `estimate` method.")
+
+    @abc.abstractmethod
+    def _create_filter(self):
+        """
+        Private method for computing and returning `Filter` object.
+        """
 
 
 class WhiteNoiseEstimator(NoiseEstimator):
@@ -263,21 +279,17 @@ def estimate(self):
         #   so we only evaluate for the zero frequencies.
         return self.filter.evaluate(np.zeros((2, 1), dtype=self.dtype)).item()
 
-    def _create_filter(self, noise_variance=None):
+    def _create_filter(self):
         """
-        :param noise_variance: Noise variance of images
         :return: The estimated noise power spectral distribution (PSD) of the images in the form of a filter object.
         """
-        if noise_variance is None:
-            logger.info(f"Determining Noise variance in batches of {self.batchSize}")
-            noise_variance = self._estimate_noise_variance()
-            logger.info(f"Noise variance = {noise_variance}")
+        logger.info(f"Determining Noise variance in batches of {self.batchSize}")
+        noise_variance = self._estimate_noise_variance()
+        logger.info(f"Noise variance = {noise_variance}")
         return ScalarFilter(dim=2, value=noise_variance)
 
     def _estimate_noise_variance(self):
         """
-        Any additional arguments/keyword-arguments are passed on to the Source's 'images' method
-
         :return: The estimated noise variance of the images in the Source used to create this estimator.
         TODO: How's this initial estimate of variance different from the 'estimate' method?
         """
@@ -312,16 +324,13 @@ def estimate(self):
 
         return np.mean(self.filter.evaluate_grid(self.src.L))
 
-    def _create_filter(self, noise_psd=None):
+    def _create_filter(self):
         """
-        :param noise_psd: Noise PSD of images
         :return: The estimated noise power spectral distribution (PSD) of the images in the form of a filter object.
         """
-        if noise_psd is None:
-            noise_psd = self.estimate_noise_psd()
-        return ArrayFilter(noise_psd)
+        return ArrayFilter(self._estimate_noise_psd())
 
-    def estimate_noise_psd(self):
+    def _estimate_noise_psd(self):
         """
         :return: The estimated noise variance of the images in the Source used to create this estimator.
         TODO: How's this initial estimate of variance different from the 'estimate' method?
@@ -350,9 +359,11 @@ def estimate_noise_psd(self):
 class IsotropicNoiseEstimator(NoiseEstimator):
     """
     Isotropic Noise Estimator.
+
+    Ported from MATLAB `cryo_noise_estimation`.
     """
 
-    def __init__(self, src, bgRadius=None, max_d=None, batchSize=512):
+    def __init__(self, src, bgRadius=None, max_d=None):
         """
         Any additional args/kwargs are passed on to the Source's 'images' method
 
@@ -363,15 +374,16 @@ def __init__(self, src, bgRadius=None, max_d=None, batchSize=512):
             Default of `None` uses `(np.floor(src.L / 2) - 1) / L`
         :param max_d: Max computed correlation distance as a proportion of `src.L`.
             Default of `None` uses `np.floor(src.L/3) / L`.
-        :param batchSize:  The size of the batches in which to compute the variance estimate.
         """
-        self.src = src
-        self.dtype = self.src.dtype
-        self.bgRadius = (np.floor(self.src.L / 2) - 1) / self.src.L
-        self.batchSize = batchSize
-        self.max_d = np.floor(self.src.L / 3) / self.src.L
 
-        self.filter = self._create_filter()
+        if bgRadius is None:
+            bgRadius = (np.floor(src.L / 2) - 1) / src.L
+
+        super().__init__(src=src, bgRadius=bgRadius, batchSize=1)
+
+        self.max_d = max_d
+        if self.max_d is None:
+            self.max_d = np.floor(src.L / 3) / src.L
 
     def estimate(self):
         """
@@ -384,10 +396,9 @@ def _create_filter(self):
         """
         :return: The estimated noise power spectral distribution (PSD) of the images in the form of a filter object.
         """
-        noise_psd = self.estimate_noise_psd()
-        return ArrayFilter(noise_psd)
+        return ArrayFilter(self._estimate_noise_psd())
 
-    def estimate_noise_psd(self):
+    def _estimate_noise_psd(self):
         """
         :return: The estimated noise variance of the images in the Source used to create this estimator.
         TODO: How's this initial estimate of variance different from the 'estimate' method?
@@ -396,7 +407,7 @@ def estimate_noise_psd(self):
         samples_idx = grid_2d(self.src.L, normalized=True)["r"] >= self.bgRadius
         max_d_pixels = round(self.max_d * self.src.L)
 
-        psd = IsotropicNoiseEstimator.epsdS(self.src.images, samples_idx, max_d_pixels)[0]
+        psd = self.epsdS(self.src.images, samples_idx, max_d_pixels)[0]
 
         return psd
 
diff --git a/tests/test_anisotropic_noise.py b/tests/test_anisotropic_noise.py
index 4b06fcef9a..797c27709c 100644
--- a/tests/test_anisotropic_noise.py
+++ b/tests/test_anisotropic_noise.py
@@ -37,7 +37,7 @@ def tearDown(self):
 
     def testAnisotropicNoisePSD(self):
         noise_estimator = AnisotropicNoiseEstimator(self.sim, batchSize=512)
-        noise_psd = noise_estimator.estimate_noise_psd()
+        noise_psd = noise_estimator._estimate_noise_psd()
         self.assertTrue(
             np.allclose(
                 noise_psd,

From e7c06d4399d36ed392294e9b2be1673757397086 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Fri, 17 Jan 2025 10:38:58 -0500
Subject: [PATCH 102/184] Begin adding into unit test framework

some issues

[skip ci]
---
 tests/test_noise.py | 53 ++++++++++++++++++++++++++-------------------
 1 file changed, 31 insertions(+), 22 deletions(-)

diff --git a/tests/test_noise.py b/tests/test_noise.py
index eecded6dfa..3ad2a4546b 100644
--- a/tests/test_noise.py
+++ b/tests/test_noise.py
@@ -9,6 +9,7 @@
     AnisotropicNoiseEstimator,
     BlueNoiseAdder,
     CustomNoiseAdder,
+    IsotropicNoiseEstimator,
     PinkNoiseAdder,
     WhiteNoiseAdder,
     WhiteNoiseEstimator,
@@ -28,6 +29,7 @@
 VARS = [0.1] + [
     pytest.param(10 ** (-x), marks=pytest.mark.expensive) for x in range(2, 5)
 ]
+NOISE_ESTIMATORS = [AnisotropicNoiseEstimator, IsotropicNoiseEstimator]
 
 
 def _noise_function(x, y):
@@ -42,12 +44,12 @@ def sim_fixture_id(params):
     return f"res={res}, dtype={dtype.__name__}"
 
 
-@pytest.fixture(params=DTYPES, ids=lambda x: f"dtype={x}")
+@pytest.fixture(params=DTYPES, ids=lambda x: f"dtype={x}", scope="module")
 def dtype(request):
     return request.param
 
 
-@pytest.fixture(params=RESOLUTIONS, ids=lambda x: f"resolution={x}")
+@pytest.fixture(params=RESOLUTIONS, ids=lambda x: f"resolution={x}", scope="module")
 def resolution(request):
     return request.param
 
@@ -67,6 +69,11 @@ def sim_fixture(resolution, dtype):
     )
 
 
+@pytest.fixture(params=NOISE_ESTIMATORS, ids=lambda x: f"noise_estimator={x.__name__}")
+def noise_estimator_fixture(request):
+    return request.param
+
+
 @pytest.fixture(
     params=[
         WhiteNoiseAdder(var=1),
@@ -82,7 +89,7 @@ def test_white_noise_estimator_clean_corners(sim_fixture):
     """
     Tests that a clean image yields a noise estimate that is virtually zero.
     """
-    noise_estimator = WhiteNoiseEstimator(sim_fixture, batchSize=512)
+    noise_estimator = WhiteNoiseEstimator(sim_fixture)
     noise_variance = noise_estimator.estimate()
     # Using a compactly supported volume should yield
     #   virtually no noise in the image corners.
@@ -120,10 +127,12 @@ def test_white_noise_adder(sim_fixture, target_noise_variance):
     assert sim_fixture.noise_adder.noise_var == target_noise_variance
 
     # Create an estimator from the source
-    noise_estimator = WhiteNoiseEstimator(sim_fixture, batchSize=512)
+    noise_estimator = WhiteNoiseEstimator(sim_fixture)
 
     # Match estimate within 1%
-    assert np.isclose(target_noise_variance, noise_estimator.estimate(), rtol=0.01)
+    np.testing.assert_allclose(
+        target_noise_variance, noise_estimator.estimate(), rtol=0.01
+    )
 
 
 @pytest.mark.parametrize(
@@ -159,7 +168,7 @@ def test_custom_noise_adder(sim_fixture, target_noise_variance):
     sampled_noise_var = np.var(im_noise_sample.asnumpy())
 
     logger.debug(f"Sampled Noise Variance {sampled_noise_var}")
-    assert np.isclose(sampled_noise_var, target_noise_variance, rtol=0.1)
+    np.testing.assert_allclose(sampled_noise_var, target_noise_variance, rtol=0.1)
 
 
 @pytest.mark.parametrize(
@@ -202,7 +211,7 @@ def test_from_snr_white(sim_fixture, target_noise_variance):
     )
 
     # Compare with WhiteNoiseEstimator consuming sim_from_snr
-    noise_estimator = WhiteNoiseEstimator(sim_from_snr, batchSize=512)
+    noise_estimator = WhiteNoiseEstimator(sim_from_snr)
     est_noise_variance = noise_estimator.estimate()
     logger.info(
         "est_noise_variance, target_noise_variance ="
@@ -210,13 +219,15 @@ def test_from_snr_white(sim_fixture, target_noise_variance):
     )
 
     # Check we're within 5%
-    assert np.isclose(est_noise_variance, target_noise_variance, rtol=0.05)
+    np.testing.assert_allclose(est_noise_variance, target_noise_variance, rtol=0.05)
 
 
 @pytest.mark.parametrize(
     "target_noise_variance", VARS, ids=lambda param: f"var={param}"
 )
-def test_blue_iso_noise_estimation(sim_fixture, target_noise_variance):
+def test_blue_iso_noise_estimation(
+    sim_fixture, target_noise_variance, noise_estimator_fixture
+):
     """
     Test that prescribing isotropic blue-ish noise
     is close to target for a variety of paramaters.
@@ -225,23 +236,23 @@ def test_blue_iso_noise_estimation(sim_fixture, target_noise_variance):
     # Create the CustomNoiseAdder
     sim_fixture.noise_adder = BlueNoiseAdder(var=target_noise_variance)
 
-    # TODO, potentially remove or change to Isotropic after #842
-    # Compare with AnisotropicNoiseEstimator consuming sim_from_snr
-    noise_estimator = AnisotropicNoiseEstimator(sim_fixture, batchSize=512)
+    noise_estimator = noise_estimator_fixture(sim_fixture)
     est_noise_variance = noise_estimator.estimate()
     logger.info(
         "est_noise_variance, target_noise_variance ="
         f" {est_noise_variance}, {target_noise_variance}"
     )
 
-    # Check we're within 5%
-    assert np.isclose(est_noise_variance, target_noise_variance, rtol=0.05)
+    # Check
+    np.testing.assert_allclose(est_noise_variance, target_noise_variance, rtol=0.20)
 
 
 @pytest.mark.parametrize(
     "target_noise_variance", VARS, ids=lambda param: f"var={param}"
 )
-def test_pink_iso_noise_estimation(sim_fixture, target_noise_variance):
+def test_pink_iso_noise_estimation(
+    sim_fixture, target_noise_variance, noise_estimator_fixture
+):
     """
     Test that prescribing isotropic pink-ish noise
     is close to target for a variety of paramaters.
@@ -250,17 +261,15 @@ def test_pink_iso_noise_estimation(sim_fixture, target_noise_variance):
     # Create the CustomNoiseAdder
     sim_fixture.noise_adder = PinkNoiseAdder(var=target_noise_variance)
 
-    # TODO, potentially remove or change to Isotropic after #842
-    # Compare with AnisotropicNoiseEstimator consuming sim_from_snr
-    noise_estimator = AnisotropicNoiseEstimator(sim_fixture, batchSize=512)
+    noise_estimator = noise_estimator_fixture(sim_fixture)
     est_noise_variance = noise_estimator.estimate()
     logger.info(
         "est_noise_variance, target_noise_variance ="
         f" {est_noise_variance}, {target_noise_variance}"
     )
 
-    # Check we're within 5%
-    assert np.isclose(est_noise_variance, target_noise_variance, rtol=0.05)
+    # Check
+    np.testing.assert_allclose(est_noise_variance, target_noise_variance, rtol=0.20)
 
 
 @pytest.mark.parametrize(
@@ -286,7 +295,7 @@ def aniso_spectrum(x, y):
 
     # TODO, potentially remove after #842
     # Compare with AnisotropicNoiseEstimator consuming sim_from_snr
-    noise_estimator = AnisotropicNoiseEstimator(sim_fixture, batchSize=512)
+    noise_estimator = AnisotropicNoiseEstimator(sim_fixture)
     est_noise_variance = noise_estimator.estimate()
     logger.info(
         "est_noise_variance, target_noise_variance ="
@@ -294,7 +303,7 @@ def aniso_spectrum(x, y):
     )
 
     # Check we're within 5%
-    assert np.isclose(est_noise_variance, target_noise_variance, rtol=0.05)
+    np.testing.assert_allclose(est_noise_variance, target_noise_variance, rtol=0.05)
 
 
 def test_gaussian_window(resolution, dtype):

From 4ec718f8207d0de41d0e627749a64eaba1e4a9b3 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Fri, 24 Jan 2025 14:36:12 -0500
Subject: [PATCH 103/184] Add initial implementation of legacy_whiten in Image
 and Xform

---
 src/aspire/image/image.py    | 70 ++++++++++++++++++++++++++++++++++++
 src/aspire/image/xform.py    | 29 +++++++++++++++
 src/aspire/noise/__init__.py |  2 +-
 src/aspire/noise/noise.py    |  8 ++---
 src/aspire/source/image.py   | 36 ++++++++++++++++++-
 tests/test_noise.py          |  4 +--
 6 files changed, 140 insertions(+), 9 deletions(-)

diff --git a/src/aspire/image/image.py b/src/aspire/image/image.py
index 9b76402ea4..8ad716805f 100644
--- a/src/aspire/image/image.py
+++ b/src/aspire/image/image.py
@@ -416,6 +416,76 @@ def shift(self, shifts):
 
         return self._im_translate(shifts)
 
+    def legacy_whiten(self, psd, delta):
+        """
+        Apply the legacy MATLAB whitening transformation to `im`.
+
+        :param psd: PSD (as computed by LegacyNoiseEstimator)
+        :param delta: Threshold used to determine which frequencies to whiten
+            and which to set to zero. By default all sqrt(PSD) values in the `noise_estimate`
+            less than eps(self.dtype) are zeroed out in the whitening filter.
+        """
+        n = self.n_images
+        L = self.resolution
+        j = L // 2
+        K = psd.shape[-1]
+        k = int(np.ceil(K / 2))
+
+        # Create result array
+        res = np.empty((n, L, L), dtype=self.dtype)
+
+        # The whitening filter is the sqrt of the power spectrum of the noise, normalized to unit energy.
+        psd = xp.asarray(psd, dtype=np.float64)
+        fltr = xp.sqrt(psd)
+        fltr = fltr / xp.linalg.norm(fltr)
+
+        assert xp.linalg.norm(fltr.imag) < 10 * delta
+        assert xp.linalg.norm(fltr - xp.flipud(fltr)) < 10 * delta
+        assert xp.linalg.norm(fltr - xp.fliplr(fltr)) < 10 * delta
+
+        # Enforce symmetry
+        fltr = (fltr + xp.flipud(fltr)) / 2
+        fltr = (fltr + xp.fliplr(fltr)) / 2
+
+        # The filter may have very small values or even zeros.
+        # We don't want to process these, so make a list of all large entries.
+        nzidx = fltr > 100 * delta
+        fnz = fltr[nzidx]
+
+        pp = xp.zeros((K, K), dtype=np.float64)
+        p = xp.zeros((K, K), dtype=np.complex128)
+        for i, proj in enumerate(self.asnumpy()):
+
+            # Zero pad the image to twice the size
+            if L % 2 == 1:
+                pp[k - j - 1 : k + j, k - j - 1 : k + j] = xp.asarray(proj)
+            else:
+                pp[k - j - 2 : k + j, k - j - 1 : k + j - 1] = xp.asarray(proj)
+
+            # Take the Fourier Transform of the padded image.
+            fp = fft.centered_fft2(pp)
+
+            # Divide the image by the whitening filter but only in
+            # places where the filter is large.  In frequencies where
+            # the filter is tiny we cannot whiten so we just put
+            # zeros.
+            p[nzidx] = fp[nzidx] / fnz
+            p2 = fft.centered_ifft2(p)
+
+            # The resulting image should be real.
+            if xp.linalg.norm(p2.imag) / xp.linalg.norm(p2) > 1e-13:
+                raise RuntimeError("Whitened image has strong imaginary component.")
+
+            if L % 2 == 1:
+                p2 = p2[k - j - 1 : k + j, k - j - 1 : k + j].real
+            else:
+                p2 = p2[k - j - 1 : k + j - 1, k - j - 1 : k + j - 1].real
+
+            # Assign the resulting image.
+            res[i] = xp.asnumpy(p2)
+
+        return Image(res)
+
     def downsample(self, ds_res, zero_nyquist=True):
         """
         Downsample Image to a specific resolution. This method returns a new Image.
diff --git a/src/aspire/image/xform.py b/src/aspire/image/xform.py
index 8003b4a1c5..fbe898618b 100644
--- a/src/aspire/image/xform.py
+++ b/src/aspire/image/xform.py
@@ -214,6 +214,35 @@ def __str__(self):
         return f"Downsample (Resolution {self.resolution})"
 
 
+class LegacyWhiten(Xform):
+    """
+    A Xform that implements MATLAB legacy whitening.
+    """
+
+    def __init__(self, psd, delta):
+        """
+        Initialize LegacyWhiten Xform.
+
+        :param psd: PSD (as computed by LegacyNoiseEstimator)
+        :param delta: Threshold used to determine which frequencies to whiten
+            and which to set to zero. By default all sqrt(PSD) values in the `noise_estimate`
+            less than eps(self.dtype) are zeroed out in the whitening filter.
+        """
+
+        self.psd = psd
+        self.delta = delta
+        super().__init__()
+
+    def _forward(self, im, indices):
+        """
+        Apply the legacy MATLAB whitening transformation to `im`.
+        """
+        return im.legacy_whiten(self.psd, self.delta)
+
+    def __str__(self):
+        return "Legacy Whitening Xform."
+
+
 class FilterXform(SymmetricXform):
     """
     A `Xform` that applies a single `Filter` object to a stack of 2D images (as an Image object).
diff --git a/src/aspire/noise/__init__.py b/src/aspire/noise/__init__.py
index 67978bcd2f..48b5684727 100644
--- a/src/aspire/noise/__init__.py
+++ b/src/aspire/noise/__init__.py
@@ -2,7 +2,7 @@
     AnisotropicNoiseEstimator,
     BlueNoiseAdder,
     CustomNoiseAdder,
-    IsotropicNoiseEstimator,
+    LegacyNoiseEstimator,
     NoiseAdder,
     NoiseEstimator,
     PinkNoiseAdder,
diff --git a/src/aspire/noise/noise.py b/src/aspire/noise/noise.py
index 2fac22e39b..77feda71b5 100644
--- a/src/aspire/noise/noise.py
+++ b/src/aspire/noise/noise.py
@@ -356,11 +356,9 @@ def _estimate_noise_psd(self):
         return noise_psd_est
 
 
-class IsotropicNoiseEstimator(NoiseEstimator):
+class LegacyNoiseEstimator(NoiseEstimator):
     """
-    Isotropic Noise Estimator.
-
-    Ported from MATLAB `cryo_noise_estimation`.
+    Isotropic noise estimator ported from MATLAB `cryo_noise_estimation`.
     """
 
     def __init__(self, src, bgRadius=None, max_d=None):
@@ -541,7 +539,7 @@ def epsdS(images, samples_idx, max_d=None):
         n_img = images.n_images
         L = samples_idx.shape[-1]
 
-        R, x, _ = IsotropicNoiseEstimator.epsdR(
+        R, x, _ = LegacyNoiseEstimator.epsdR(
             images=images, samples_idx=samples_idx, max_d=max_d
         )
 
diff --git a/src/aspire/source/image.py b/src/aspire/source/image.py
index e315baa84e..064834f28a 100644
--- a/src/aspire/source/image.py
+++ b/src/aspire/source/image.py
@@ -16,10 +16,11 @@
     FilterXform,
     IndexedXform,
     LambdaXform,
+    LegacyWhiten,
     Multiply,
     Pipeline,
 )
-from aspire.noise import NoiseEstimator, WhiteNoiseEstimator
+from aspire.noise import LegacyNoiseEstimator, NoiseEstimator, WhiteNoiseEstimator
 from aspire.operators import (
     CTFFilter,
     Filter,
@@ -824,6 +825,39 @@ def whiten(self, noise_estimate=None, epsilon=None):
         logger.info("Adding Whitening Filter Xform to end of generation pipeline")
         self.generation_pipeline.add_xform(FilterXform(whiten_filter))
 
+    @_as_copy
+    def legacy_whiten(self, noise_response=None, delta=None):
+        """
+        Reproduce the legacy MATLAB whitening process.
+
+        :param noise_response: Noise response is provided either
+            directly as an array, or a `LegacyNoiseEstimator` instance.
+        :param delta: Threshold used to determine which frequencies to whiten
+            and which to set to zero. By default all sqrt(PSD) values in the `noise_estimate`
+            less than eps(self.dtype) are zeroed out in the whitening filter.
+        """
+
+        if noise_response is None:
+            logger.info("Computing noise response.")
+            psd = LegacyNoiseEstimator(self).filter.xfer_fn_array
+        elif isinstance(noise_response, LegacyNoiseEstimator):
+            psd = noise_response.filter.xfer_fn_array
+        elif isinstance(noise_response, np.ndarray):
+            if not noise_response.shape == (self.L * 2 - 1,) * 2:
+                raise ValueError(
+                    f"Unexepected `noise_response` array shape {noise_response.shape}."
+                )
+            # Take the array directly
+            psd = noise_response
+        else:
+            raise ValueError("Unexepected `noise_response` type.")
+
+        if delta is None:
+            delta = np.finfo(np.float32).eps
+
+        logger.info("Adding LegacyWhiten Filter Xform to end of generation pipeline")
+        self.generation_pipeline.add_xform(LegacyWhiten(psd, delta))
+
     @_as_copy
     def phase_flip(self):
         """
diff --git a/tests/test_noise.py b/tests/test_noise.py
index 3ad2a4546b..01fae0ae10 100644
--- a/tests/test_noise.py
+++ b/tests/test_noise.py
@@ -9,7 +9,7 @@
     AnisotropicNoiseEstimator,
     BlueNoiseAdder,
     CustomNoiseAdder,
-    IsotropicNoiseEstimator,
+    LegacyNoiseEstimator,
     PinkNoiseAdder,
     WhiteNoiseAdder,
     WhiteNoiseEstimator,
@@ -29,7 +29,7 @@
 VARS = [0.1] + [
     pytest.param(10 ** (-x), marks=pytest.mark.expensive) for x in range(2, 5)
 ]
-NOISE_ESTIMATORS = [AnisotropicNoiseEstimator, IsotropicNoiseEstimator]
+NOISE_ESTIMATORS = [AnisotropicNoiseEstimator, LegacyNoiseEstimator]
 
 
 def _noise_function(x, y):

From 47db4b6230e799b0dadf6e3fa3f0fba1e553374e Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Mon, 27 Jan 2025 10:35:08 -0500
Subject: [PATCH 104/184] Cleanup legacy whitening methods/comments/asserts

---
 src/aspire/image/image.py | 29 ++++++++++----
 src/aspire/image/xform.py |  8 ++--
 src/aspire/noise/noise.py | 80 +++++++++++++++++++--------------------
 3 files changed, 64 insertions(+), 53 deletions(-)

diff --git a/src/aspire/image/image.py b/src/aspire/image/image.py
index 8ad716805f..5e75fdd3de 100644
--- a/src/aspire/image/image.py
+++ b/src/aspire/image/image.py
@@ -418,12 +418,17 @@ def shift(self, shifts):
 
     def legacy_whiten(self, psd, delta):
         """
-        Apply the legacy MATLAB whitening transformation to `im`.
+        Apply the legacy MATLAB whitening transformation.
 
-        :param psd: PSD (as computed by LegacyNoiseEstimator)
+        Note that this legacy method will compute the convolution in
+        (complex)double precision regardless of this instances
+        `dtype`.  However, the resulting image stack will be cast to
+        the instance `dtype`.
+
+        :param psd: PSD as computed by `LegacyNoiseEstimator`.
         :param delta: Threshold used to determine which frequencies to whiten
-            and which to set to zero. By default all sqrt(PSD) values in the `noise_estimate`
-            less than eps(self.dtype) are zeroed out in the whitening filter.
+            and which to set to zero. By default all `sqrt(psd)` values
+            less than `delta` are zeroed out in the whitening filter.
         """
         n = self.n_images
         L = self.resolution
@@ -439,9 +444,17 @@ def legacy_whiten(self, psd, delta):
         fltr = xp.sqrt(psd)
         fltr = fltr / xp.linalg.norm(fltr)
 
-        assert xp.linalg.norm(fltr.imag) < 10 * delta
-        assert xp.linalg.norm(fltr - xp.flipud(fltr)) < 10 * delta
-        assert xp.linalg.norm(fltr - xp.fliplr(fltr)) < 10 * delta
+        # Error checking
+        if (err := xp.linalg.norm(fltr.imag)) < 10 * delta:
+            raise RuntimeError(
+                f"Whitening filter has non trivial imaginary component {err}."
+            )
+        err_ud = xp.linalg.norm(fltr - xp.flipud(fltr))
+        err_lr = xp.linalg.norm(fltr - xp.fliplr(fltr))
+        if (err_ud < 10 * delta) or (err_lr < 10 * delta):
+            raise RuntimeError(
+                f"Whitening filter has non trivial symmetry lr {err_lr}, ud {err_ud}."
+            )
 
         # Enforce symmetry
         fltr = (fltr + xp.flipud(fltr)) / 2
@@ -467,7 +480,7 @@ def legacy_whiten(self, psd, delta):
 
             # Divide the image by the whitening filter but only in
             # places where the filter is large.  In frequencies where
-            # the filter is tiny we cannot whiten so we just put
+            # the filter is tiny we cannot whiten so we just use
             # zeros.
             p[nzidx] = fp[nzidx] / fnz
             p2 = fft.centered_ifft2(p)
diff --git a/src/aspire/image/xform.py b/src/aspire/image/xform.py
index fbe898618b..4710f02d56 100644
--- a/src/aspire/image/xform.py
+++ b/src/aspire/image/xform.py
@@ -223,10 +223,10 @@ def __init__(self, psd, delta):
         """
         Initialize LegacyWhiten Xform.
 
-        :param psd: PSD (as computed by LegacyNoiseEstimator)
+        :param psd: PSD (as computed by LegacyNoiseEstimator).
         :param delta: Threshold used to determine which frequencies to whiten
-            and which to set to zero. By default all sqrt(PSD) values in the `noise_estimate`
-            less than eps(self.dtype) are zeroed out in the whitening filter.
+            and which to set to zero. By default all `sqrt(psd)` values
+            less than `delta` are zeroed out in the whitening filter.
         """
 
         self.psd = psd
@@ -236,6 +236,8 @@ def __init__(self, psd, delta):
     def _forward(self, im, indices):
         """
         Apply the legacy MATLAB whitening transformation to `im`.
+
+        The tranform is applied to all of `im`, `indices` are unused.
         """
         return im.legacy_whiten(self.psd, self.delta)
 
diff --git a/src/aspire/noise/noise.py b/src/aspire/noise/noise.py
index 77feda71b5..026ed2f70e 100644
--- a/src/aspire/noise/noise.py
+++ b/src/aspire/noise/noise.py
@@ -363,9 +363,9 @@ class LegacyNoiseEstimator(NoiseEstimator):
 
     def __init__(self, src, bgRadius=None, max_d=None):
         """
-        Any additional args/kwargs are passed on to the Source's 'images' method
+        Given an `ImageSource`, constructs
 
-        :param src: A Source object which can give us images on demand
+        :param src: A `ImageSource` object.
         :param bgRadius: The radius of the disk whose complement is used to estimate the noise.
             Radius is relative proportion, where `1` represents
             the radius of disc inscribing a `(src.L, src.L)` image.
@@ -399,9 +399,8 @@ def _create_filter(self):
     def _estimate_noise_psd(self):
         """
         :return: The estimated noise variance of the images in the Source used to create this estimator.
-        TODO: How's this initial estimate of variance different from the 'estimate' method?
         """
-        # Setup
+        # Setup parameters
         samples_idx = grid_2d(self.src.L, normalized=True)["r"] >= self.bgRadius
         max_d_pixels = round(self.max_d * self.src.L)
 
@@ -416,9 +415,10 @@ def epsdR(images, samples_idx, max_d=None):
         The samples to use in each image are given by `samples_idx` mask.
         The correlation is computed up to a maximal distance of `max_d`.
 
-        :param images: Images instance
-        :param samples_idx: Boolean mask shaped (L,L).
+        :param images: `Image` instance
+        :param samples_idx: Boolean mask shaped `(L,L)`.
         :param max_d: Max computed correlation distance in pixels.
+           Default of `None` yields `np.floor(L / 3)`.
         :return:
             - Radial PSD array of shape
             - Distances map array
@@ -438,28 +438,26 @@ def epsdR(images, samples_idx, max_d=None):
         max_d = int(min(max_d, L - 1))
 
         # Compute distances
-        # Note grid_2d['r'] is not used because we always want zero centered integer grid,
-        #   yielding integer dists (radius**2) values.
+        # Note grid_2d['r'] is not used because we want an integer grid directly;
+        #   yields integer dists (radius**2) values.
         X, Y = np.mgrid[0 : max_d + 1, 0 : max_d + 1]
         dists = X * X + Y * Y
         dsquare = np.sort(np.unique(dists[dists <= max_d**2]))
-        x = np.sqrt(dsquare)  # actual distance
+        x = np.sqrt(dsquare)  # actual distances
 
         # corrs[i] is the sum of all x[j]x[j+d] where d = x[i]
         corrs = np.zeros_like(dsquare, dtype=np.float64)
         # corrcount[i] is the number of pairs summed in corr[i]
-        corrcount = np.zeros_like(dsquare, dtype=int)
+        corrcount = np.zeros_like(dsquare, dtype=np.int64)
 
         # distmap maps [i,j] to k where dsquare[k] = i**2 + j**2.
         #   -1 indicates distance is larger than max_d
         distmap = np.full(shape=dists.shape, fill_value=-1)
 
-        # This differs from the MATLAB code because Numpy does not directly provide `bsearch`.
+        # This differs from the MATLAB code, avoids `bisect`.
         for i, d in enumerate(dsquare):
             inds = dists == d  # locations having distance `d`
             distmap[inds] = i  # assign index into dsquare `i`
-        # # Mapped distance indices where i**2+j**2 <= max_d**2
-        # validdists = np.where(distmap != -1)  # Note this is a 2-tuple
 
         # Compute Ncorr using a constant unit image.
         mask = np.zeros((L, L))
@@ -467,9 +465,10 @@ def epsdR(images, samples_idx, max_d=None):
         tmp = np.zeros((2 * L + 1, 2 * L + 1))  # pad
         tmp[:L, :L] = mask
         ftmp = fft.fft2(tmp)
-        Ncorr = fft.ifft2(
-            ftmp * ftmp.conj()
-        ).real  # matlab code does not cast here, but internally detects conj sym...
+        # MATLAB code internally detects sym and implicitly casts,
+        #   we explicitly cast.
+        # Optimization note: This and the later fft call could be optimized with real/herm sym FFT calls.
+        Ncorr = fft.ifft2(ftmp * ftmp.conj()).real
         Ncorr = Ncorr[: max_d + 1, : max_d + 1]  # crop
         Ncorr = np.round(Ncorr)
 
@@ -480,20 +479,26 @@ def epsdR(images, samples_idx, max_d=None):
         samples = np.zeros((L, L))
         tmp[:, :] = 0  # reset tmp
         for k in trange(n_img, desc="Processing image autocorrelations"):
-            # Mask unused pixels (note, think can merge these lines later)
+            # Mask off unused pixels
             samples[samples_idx] = images[k].asnumpy()[0][samples_idx]
-            # Note, we can also compute the noise energy estimate used later at this time to avoid looping over images twice.
+            # Optimization note: We could also compute the noise
+            # energy estimate used later at this time to avoid looping
+            # over images twice.
 
             # Compute non-preiodic autocorrelation
             tmp[:L, :L] = samples  # pad
             ftmp = fft.fft2(tmp)
-            s = fft.ifft2(
-                ftmp * ftmp.conj()
-            ).real  # matlab code does not cast here, but internally detects conj sym...
+            # MATLAB code internally detects conj sym and implicitly casts,
+            #   we explicitly cast.
+            s = fft.ifft2(ftmp * ftmp.conj()).real
             s = s[0 : max_d + 1, 0 : max_d + 1]  # crop
 
-            # # Accumulate all autocorrelation values R[k1,k2] such that
-            # # k1^2+k2^2=const (all autocorrelations of a certain distance).
+            # Accumulate all autocorrelation values R[k1,k2] such that
+            # k1**2 + k2**2 = dist (all autocorrelations of a certain distance).
+            # Optimization note: The MATLAB code used another map
+            #   layer `validdists` to remove one loop layer here, but
+            #   it relied primarily on MATLABs implicit ravel/unravel
+            #   and would be less clear in Python. Simpler code was ported.
             for i in range(max_d + 1):
                 for j in range(max_d + 1):
                     idx = distmap[i, j]
@@ -501,18 +506,8 @@ def epsdR(images, samples_idx, max_d=None):
                         corrs[idx] = corrs[idx] + s[i, j]
                         corrcount[idx] = corrcount[idx] + Ncorr[i, j]
 
-            # TODO, fix this MATLAB optmized implementation and compare with the clearer code above.
-            # I didn't port this validdist slice optimized version correctly(yet).
-            # it uses implicit (MATLAB auto flat (un/)raveling)
-            # im not sure the speedup would be similar in python anyway.
-            # for j in range(np.size(validdists)):
-            #     currdist = validdists[j]
-            #     dmidx = distmap[currdist]
-            #     corrs[dmidx] = corrs[dmidx] + s[currdist]
-            #     corrcount[dmidx] = corrcount[dmidx] + Ncorr[currdist]
-
         # Remove distances which had no samples
-        idx = np.where(corrcount != 0)  # [0]
+        idx = np.where(corrcount != 0)
         R = corrs[idx] / corrcount[idx]
         x = x[idx]
         cnt = corrcount[idx]
@@ -529,6 +524,7 @@ def epsdS(images, samples_idx, max_d=None):
         :param images: Images instance
         :param samples_idx: Boolean mask shaped (L,L).
         :param max_d: Max computed correlation distance in pixels.
+           Default of `None` yields `np.floor(L / 3)`.
         :return:
             - 2D PSD array
             - Radial PSD array
@@ -539,10 +535,6 @@ def epsdS(images, samples_idx, max_d=None):
         n_img = images.n_images
         L = samples_idx.shape[-1]
 
-        R, x, _ = LegacyNoiseEstimator.epsdR(
-            images=images, samples_idx=samples_idx, max_d=max_d
-        )
-
         # Correlations more than `max_d` pixels apart are not computed.
         if max_d is None:
             max_d = np.floor(L / 3)
@@ -552,7 +544,11 @@ def epsdS(images, samples_idx, max_d=None):
             )
         max_d = int(min(max_d, L - 1))
 
-        # Use the 1D autocorrelation estimted above to populate an
+        R, x, _ = LegacyNoiseEstimator.epsdR(
+            images=images, samples_idx=samples_idx, max_d=max_d
+        )
+
+        # Use the 1D autocorrelation estimated above to populate an
         # array of the 2D isotropic autocorrelction. This
         # autocorrelation is later Fourier transformed to get the
         # power spectrum.
@@ -565,7 +561,7 @@ def epsdS(images, samples_idx, max_d=None):
             idx = dists2 == d
             R2[idx] = R[i]
 
-        # Window te 2D autocorrelation and Fourier transform it to get the power
+        # Window the 2D autocorrelation and Fourier transform it to get the power
         # spectrum. Always use the Gaussian window, as it has positive Fourier
         # transform.
         w = gaussian_window(L, max_d)
@@ -579,7 +575,7 @@ def epsdS(images, samples_idx, max_d=None):
         # to estimate it.
 
         E = 0.0  # Total energy of the noise samples used to estimate the power spectrum.
-        samples = np.zeros((L, L))
+        samples = np.zeros((L, L), dtype=np.float64)
         for k in trange(n_img, desc="Estimating image noise energy"):
             samples[samples_idx] = images[k].asnumpy()[0][samples_idx]
             E += np.sum((samples - np.mean(samples)) ** 2)
@@ -605,7 +601,7 @@ def epsdS(images, samples_idx, max_d=None):
             if maxnegerr > 1e-2:
                 negnorm = np.linalg.norm(P2[negidx])
                 logger.warning(
-                    f"Power spectrum P2 has negative values with energy {negnorm}."
+                    f"Power spectrum P2 has non trivial negative values with energy {negnorm}."
                 )
             P2[negidx] = 0  # zero out negative estimates
 

From f10f58701287459cfef75831fc8e51df119fd942 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Mon, 27 Jan 2025 10:36:06 -0500
Subject: [PATCH 105/184] Change batchSize to batch_size

---
 .../tutorials/tutorials/cov3d_simulation.py    |  2 +-
 src/aspire/commands/cov3d.py                   |  2 +-
 src/aspire/noise/noise.py                      | 18 +++++++++---------
 src/aspire/source/image.py                     |  2 +-
 tests/test_anisotropic_noise.py                |  8 ++++----
 5 files changed, 16 insertions(+), 16 deletions(-)

diff --git a/gallery/tutorials/tutorials/cov3d_simulation.py b/gallery/tutorials/tutorials/cov3d_simulation.py
index 8bb72906b7..fdec8f601e 100644
--- a/gallery/tutorials/tutorials/cov3d_simulation.py
+++ b/gallery/tutorials/tutorials/cov3d_simulation.py
@@ -56,7 +56,7 @@
 basis = FBBasis3D(img_size)
 
 # Estimate the noise variance. This is needed for the covariance estimation step below.
-noise_estimator = WhiteNoiseEstimator(sim, batchSize=500)
+noise_estimator = WhiteNoiseEstimator(sim, batch_size=500)
 noise_variance = noise_estimator.estimate()
 print(f"Noise Variance = {noise_variance}")
 
diff --git a/src/aspire/commands/cov3d.py b/src/aspire/commands/cov3d.py
index 9f2ff0c45f..3a0752740e 100644
--- a/src/aspire/commands/cov3d.py
+++ b/src/aspire/commands/cov3d.py
@@ -55,7 +55,7 @@ def cov3d(
     mean_estimator = MeanEstimator(source, basis, batch_size=8192)
     mean_est = mean_estimator.estimate()
 
-    noise_estimator = WhiteNoiseEstimator(source, batchSize=500)
+    noise_estimator = WhiteNoiseEstimator(source, batch_size=500)
     # Estimate the noise variance. This is needed for the covariance estimation step below.
     noise_variance = noise_estimator.estimate()
     logger.info(f"Noise Variance = {noise_variance}")
diff --git a/src/aspire/noise/noise.py b/src/aspire/noise/noise.py
index 026ed2f70e..6436efa9e1 100644
--- a/src/aspire/noise/noise.py
+++ b/src/aspire/noise/noise.py
@@ -227,7 +227,7 @@ class NoiseEstimator:
     Noise Estimator base class.
     """
 
-    def __init__(self, src, bgRadius=1, batchSize=512):
+    def __init__(self, src, bgRadius=1, batch_size=512):
         """
         Any additional args/kwargs are passed on to the Source's 'images' method
 
@@ -235,12 +235,12 @@ def __init__(self, src, bgRadius=1, batchSize=512):
         :param bgRadius: The radius of the disk whose complement is used to estimate the noise.
             Radius is relative proportion, where `1` represents
             the radius of disc inscribing a `(src.L, src.L)` image.
-        :param batchSize:  The size of the batches in which to compute the variance estimate.
+        :param batch_size:  The size of the batches in which to compute the variance estimate.
         """
         self.src = src
         self.dtype = self.src.dtype
         self.bgRadius = bgRadius
-        self.batchSize = batchSize
+        self.batch_size = batch_size
 
     @cached_property
     def filter(self):
@@ -283,7 +283,7 @@ def _create_filter(self):
         """
         :return: The estimated noise power spectral distribution (PSD) of the images in the form of a filter object.
         """
-        logger.info(f"Determining Noise variance in batches of {self.batchSize}")
+        logger.info(f"Determining Noise variance in batches of {self.batch_size}")
         noise_variance = self._estimate_noise_variance()
         logger.info(f"Noise variance = {noise_variance}")
         return ScalarFilter(dim=2, value=noise_variance)
@@ -299,8 +299,8 @@ def _estimate_noise_variance(self):
 
         first_moment = 0
         second_moment = 0
-        for i in trange(0, self.src.n, self.batchSize):
-            images = self.src.images[i : i + self.batchSize].asnumpy()
+        for i in trange(0, self.src.n, self.batch_size):
+            images = self.src.images[i : i + self.batch_size].asnumpy()
             images_masked = images * mask
 
             _denominator = self.src.n * np.sum(mask)
@@ -341,8 +341,8 @@ def _estimate_noise_psd(self):
 
         mean_est = 0
         noise_psd_est = np.zeros((self.src.L, self.src.L)).astype(self.src.dtype)
-        for i in trange(0, self.src.n, self.batchSize):
-            images = self.src.images[i : i + self.batchSize].asnumpy()
+        for i in trange(0, self.src.n, self.batch_size):
+            images = self.src.images[i : i + self.batch_size].asnumpy()
             images_masked = images * mask
 
             _denominator = self.src.n * np.sum(mask)
@@ -377,7 +377,7 @@ def __init__(self, src, bgRadius=None, max_d=None):
         if bgRadius is None:
             bgRadius = (np.floor(src.L / 2) - 1) / src.L
 
-        super().__init__(src=src, bgRadius=bgRadius, batchSize=1)
+        super().__init__(src=src, bgRadius=bgRadius, batch_size=1)
 
         self.max_d = max_d
         if self.max_d is None:
diff --git a/src/aspire/source/image.py b/src/aspire/source/image.py
index 064834f28a..1d16acd95c 100644
--- a/src/aspire/source/image.py
+++ b/src/aspire/source/image.py
@@ -1412,7 +1412,7 @@ def estimate_noise_power(
             support_radius_proportion = support_radius / (self.L // 2)
 
         est = WhiteNoiseEstimator(
-            src=self, bgRadius=support_radius_proportion, batchSize=batch_size
+            src=self, bgRadius=support_radius_proportion, batch_size=batch_size
         )
 
         return est.estimate()
diff --git a/tests/test_anisotropic_noise.py b/tests/test_anisotropic_noise.py
index 797c27709c..516e1652d8 100644
--- a/tests/test_anisotropic_noise.py
+++ b/tests/test_anisotropic_noise.py
@@ -36,7 +36,7 @@ def tearDown(self):
         pass
 
     def testAnisotropicNoisePSD(self):
-        noise_estimator = AnisotropicNoiseEstimator(self.sim, batchSize=512)
+        noise_estimator = AnisotropicNoiseEstimator(self.sim, batch_size=512)
         noise_psd = noise_estimator._estimate_noise_psd()
         self.assertTrue(
             np.allclose(
@@ -128,7 +128,7 @@ def testAnisotropicNoisePSD(self):
         )
 
     def testAnisotropicNoiseVariance(self):
-        noise_estimator = AnisotropicNoiseEstimator(self.sim, batchSize=512)
+        noise_estimator = AnisotropicNoiseEstimator(self.sim, batch_size=512)
         noise_variance = noise_estimator.estimate()
         self.assertTrue(
             np.allclose(
@@ -154,9 +154,9 @@ def testParseval(self):
         wht_noise = np.random.randn(1024, 128, 128).astype(self.dtype)
         src = ArrayImageSource(wht_noise)
 
-        wht_noise_estimator = WhiteNoiseEstimator(src, batchSize=512)
+        wht_noise_estimator = WhiteNoiseEstimator(src, batch_size=512)
         wht_noise_variance = wht_noise_estimator.estimate()
-        noise_estimator = AnisotropicNoiseEstimator(src, batchSize=512)
+        noise_estimator = AnisotropicNoiseEstimator(src, batch_size=512)
         noise_variance = noise_estimator.estimate()
 
         self.assertTrue(np.allclose(noise_variance, wht_noise_variance))

From 2a8cdaea7b655f827978b1160c111d0c8a0ac5e9 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Mon, 27 Jan 2025 10:36:40 -0500
Subject: [PATCH 106/184] Change bgRadius to bg_radius

---
 src/aspire/noise/noise.py  | 22 +++++++++++-----------
 src/aspire/source/image.py |  2 +-
 2 files changed, 12 insertions(+), 12 deletions(-)

diff --git a/src/aspire/noise/noise.py b/src/aspire/noise/noise.py
index 6436efa9e1..1605612d02 100644
--- a/src/aspire/noise/noise.py
+++ b/src/aspire/noise/noise.py
@@ -227,19 +227,19 @@ class NoiseEstimator:
     Noise Estimator base class.
     """
 
-    def __init__(self, src, bgRadius=1, batch_size=512):
+    def __init__(self, src, bg_radius=1, batch_size=512):
         """
         Any additional args/kwargs are passed on to the Source's 'images' method
 
         :param src: A Source object which can give us images on demand
-        :param bgRadius: The radius of the disk whose complement is used to estimate the noise.
+        :param bg_radius: The radius of the disk whose complement is used to estimate the noise.
             Radius is relative proportion, where `1` represents
             the radius of disc inscribing a `(src.L, src.L)` image.
         :param batch_size:  The size of the batches in which to compute the variance estimate.
         """
         self.src = src
         self.dtype = self.src.dtype
-        self.bgRadius = bgRadius
+        self.bg_radius = bg_radius
         self.batch_size = batch_size
 
     @cached_property
@@ -295,7 +295,7 @@ def _estimate_noise_variance(self):
         """
         # Run estimate using saved parameters
         g2d = grid_2d(self.src.L, indexing="yx", dtype=self.dtype)
-        mask = g2d["r"] >= self.bgRadius
+        mask = g2d["r"] >= self.bg_radius
 
         first_moment = 0
         second_moment = 0
@@ -337,7 +337,7 @@ def _estimate_noise_psd(self):
         """
         # Run estimate using saved parameters
         g2d = grid_2d(self.src.L, indexing="yx", dtype=self.dtype)
-        mask = g2d["r"] >= self.bgRadius
+        mask = g2d["r"] >= self.bg_radius
 
         mean_est = 0
         noise_psd_est = np.zeros((self.src.L, self.src.L)).astype(self.src.dtype)
@@ -361,12 +361,12 @@ class LegacyNoiseEstimator(NoiseEstimator):
     Isotropic noise estimator ported from MATLAB `cryo_noise_estimation`.
     """
 
-    def __init__(self, src, bgRadius=None, max_d=None):
+    def __init__(self, src, bg_radius=None, max_d=None):
         """
         Given an `ImageSource`, constructs
 
         :param src: A `ImageSource` object.
-        :param bgRadius: The radius of the disk whose complement is used to estimate the noise.
+        :param bg_radius: The radius of the disk whose complement is used to estimate the noise.
             Radius is relative proportion, where `1` represents
             the radius of disc inscribing a `(src.L, src.L)` image.
             Default of `None` uses `(np.floor(src.L / 2) - 1) / L`
@@ -374,10 +374,10 @@ def __init__(self, src, bgRadius=None, max_d=None):
             Default of `None` uses `np.floor(src.L/3) / L`.
         """
 
-        if bgRadius is None:
-            bgRadius = (np.floor(src.L / 2) - 1) / src.L
+        if bg_radius is None:
+            bg_radius = (np.floor(src.L / 2) - 1) / src.L
 
-        super().__init__(src=src, bgRadius=bgRadius, batch_size=1)
+        super().__init__(src=src, bg_radius=bg_radius, batch_size=1)
 
         self.max_d = max_d
         if self.max_d is None:
@@ -401,7 +401,7 @@ def _estimate_noise_psd(self):
         :return: The estimated noise variance of the images in the Source used to create this estimator.
         """
         # Setup parameters
-        samples_idx = grid_2d(self.src.L, normalized=True)["r"] >= self.bgRadius
+        samples_idx = grid_2d(self.src.L, normalized=True)["r"] >= self.bg_radius
         max_d_pixels = round(self.max_d * self.src.L)
 
         psd = self.epsdS(self.src.images, samples_idx, max_d_pixels)[0]
diff --git a/src/aspire/source/image.py b/src/aspire/source/image.py
index 1d16acd95c..8711dc97d1 100644
--- a/src/aspire/source/image.py
+++ b/src/aspire/source/image.py
@@ -1412,7 +1412,7 @@ def estimate_noise_power(
             support_radius_proportion = support_radius / (self.L // 2)
 
         est = WhiteNoiseEstimator(
-            src=self, bgRadius=support_radius_proportion, batch_size=batch_size
+            src=self, bg_radius=support_radius_proportion, batch_size=batch_size
         )
 
         return est.estimate()

From dcdf14e1aff546ec86154cd62c6f1abe2ce126c0 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Mon, 27 Jan 2025 10:48:30 -0500
Subject: [PATCH 107/184] More minor LegacyWhiten cleanup

---
 src/aspire/source/image.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/src/aspire/source/image.py b/src/aspire/source/image.py
index 8711dc97d1..5e2481d68d 100644
--- a/src/aspire/source/image.py
+++ b/src/aspire/source/image.py
@@ -833,8 +833,9 @@ def legacy_whiten(self, noise_response=None, delta=None):
         :param noise_response: Noise response is provided either
             directly as an array, or a `LegacyNoiseEstimator` instance.
         :param delta: Threshold used to determine which frequencies to whiten
-            and which to set to zero. By default all sqrt(PSD) values in the `noise_estimate`
-            less than eps(self.dtype) are zeroed out in the whitening filter.
+            and which to set to zero. By default all `sqrt(psd)` values
+            less than `delta` are zeroed out in the whitening filter.
+            Default of `None` yields `np.finfo(np.float32).eps`.
         """
 
         if noise_response is None:
@@ -855,7 +856,7 @@ def legacy_whiten(self, noise_response=None, delta=None):
         if delta is None:
             delta = np.finfo(np.float32).eps
 
-        logger.info("Adding LegacyWhiten Filter Xform to end of generation pipeline")
+        logger.info("Adding LegacyWhiten Filter Xform to end of generation pipeline.")
         self.generation_pipeline.add_xform(LegacyWhiten(psd, delta))
 
     @_as_copy

From 61918f8bcc0acbc747c3beae2b114e27c3d130a1 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Mon, 27 Jan 2025 12:49:15 -0500
Subject: [PATCH 108/184] More xp interop for legacy whitening

---
 src/aspire/noise/noise.py | 65 ++++++++++++++++++++++-----------------
 1 file changed, 36 insertions(+), 29 deletions(-)

diff --git a/src/aspire/noise/noise.py b/src/aspire/noise/noise.py
index 1605612d02..5b650bca47 100644
--- a/src/aspire/noise/noise.py
+++ b/src/aspire/noise/noise.py
@@ -440,19 +440,19 @@ def epsdR(images, samples_idx, max_d=None):
         # Compute distances
         # Note grid_2d['r'] is not used because we want an integer grid directly;
         #   yields integer dists (radius**2) values.
-        X, Y = np.mgrid[0 : max_d + 1, 0 : max_d + 1]
+        X, Y = xp.mgrid[0 : max_d + 1, 0 : max_d + 1]
         dists = X * X + Y * Y
-        dsquare = np.sort(np.unique(dists[dists <= max_d**2]))
-        x = np.sqrt(dsquare)  # actual distances
+        dsquare = xp.sort(xp.unique(dists[dists <= max_d**2]))
+        x = xp.sqrt(dsquare)  # actual distances
 
         # corrs[i] is the sum of all x[j]x[j+d] where d = x[i]
-        corrs = np.zeros_like(dsquare, dtype=np.float64)
+        corrs = xp.zeros_like(dsquare, dtype=np.float64)
         # corrcount[i] is the number of pairs summed in corr[i]
-        corrcount = np.zeros_like(dsquare, dtype=np.int64)
+        corrcount = xp.zeros_like(dsquare, dtype=np.int64)
 
         # distmap maps [i,j] to k where dsquare[k] = i**2 + j**2.
         #   -1 indicates distance is larger than max_d
-        distmap = np.full(shape=dists.shape, fill_value=-1)
+        distmap = xp.full(shape=dists.shape, fill_value=-1)
 
         # This differs from the MATLAB code, avoids `bisect`.
         for i, d in enumerate(dsquare):
@@ -460,9 +460,9 @@ def epsdR(images, samples_idx, max_d=None):
             distmap[inds] = i  # assign index into dsquare `i`
 
         # Compute Ncorr using a constant unit image.
-        mask = np.zeros((L, L))
+        mask = xp.zeros((L, L))
         mask[samples_idx] = 1
-        tmp = np.zeros((2 * L + 1, 2 * L + 1))  # pad
+        tmp = xp.zeros((2 * L + 1, 2 * L + 1))  # pad
         tmp[:L, :L] = mask
         ftmp = fft.fft2(tmp)
         # MATLAB code internally detects sym and implicitly casts,
@@ -470,13 +470,13 @@ def epsdR(images, samples_idx, max_d=None):
         # Optimization note: This and the later fft call could be optimized with real/herm sym FFT calls.
         Ncorr = fft.ifft2(ftmp * ftmp.conj()).real
         Ncorr = Ncorr[: max_d + 1, : max_d + 1]  # crop
-        Ncorr = np.round(Ncorr)
+        Ncorr = xp.round(Ncorr)
 
         # Values of isotropic autocorrelation function
         # R[i] is value of ACF at distance x[i]
-        R = np.zeros(len(corrs))
+        R = xp.zeros(len(corrs))
 
-        samples = np.zeros((L, L))
+        samples = xp.zeros((L, L))
         tmp[:, :] = 0  # reset tmp
         for k in trange(n_img, desc="Processing image autocorrelations"):
             # Mask off unused pixels
@@ -507,10 +507,10 @@ def epsdR(images, samples_idx, max_d=None):
                         corrcount[idx] = corrcount[idx] + Ncorr[i, j]
 
         # Remove distances which had no samples
-        idx = np.where(corrcount != 0)
-        R = corrs[idx] / corrcount[idx]
-        x = x[idx]
-        cnt = corrcount[idx]
+        idx = xp.where(corrcount != 0)
+        R = xp.asnumpy(corrs[idx] / corrcount[idx])
+        x = xp.asnumpy(x[idx])
+        cnt = xp.asnumpy(corrcount[idx])
 
         return R, x, cnt
 
@@ -535,6 +535,9 @@ def epsdS(images, samples_idx, max_d=None):
         n_img = images.n_images
         L = samples_idx.shape[-1]
 
+        # Migrate mask to GPU as needed
+        _samples_idx = xp.asarray(samples_idx)
+
         # Correlations more than `max_d` pixels apart are not computed.
         if max_d is None:
             max_d = np.floor(L / 3)
@@ -547,26 +550,27 @@ def epsdS(images, samples_idx, max_d=None):
         R, x, _ = LegacyNoiseEstimator.epsdR(
             images=images, samples_idx=samples_idx, max_d=max_d
         )
+        _R = xp.asarray(R)  # Migrate to GPU for assignments below
 
         # Use the 1D autocorrelation estimated above to populate an
         # array of the 2D isotropic autocorrelction. This
         # autocorrelation is later Fourier transformed to get the
         # power spectrum.
-        R2 = np.zeros((2 * L - 1, 2 * L - 1), dtype=np.float64)
+        R2 = xp.zeros((2 * L - 1, 2 * L - 1), dtype=np.float64)
 
-        X, Y = np.mgrid[-L + 1 : L, -L + 1 : L]
+        X, Y = xp.mgrid[-L + 1 : L, -L + 1 : L]
         dists2 = X * X + Y * Y
-        dsquare2 = np.sort(np.unique(dists2[dists2 <= max_d**2]))
+        dsquare2 = xp.sort(xp.unique(dists2[dists2 <= max_d**2]))
         for i, d in enumerate(dsquare2):
             idx = dists2 == d
-            R2[idx] = R[i]
+            R2[idx] = _R[i]
 
         # Window the 2D autocorrelation and Fourier transform it to get the power
         # spectrum. Always use the Gaussian window, as it has positive Fourier
         # transform.
-        w = gaussian_window(L, max_d)
+        w = xp.asarray(gaussian_window(L, max_d))
         P2 = fft.centered_fft2(R2 * w)
-        if (err := np.linalg.norm(P2.imag) / np.linalg.norm(P2)) > 1e-12:
+        if (err := xp.linalg.norm(P2.imag) / xp.linalg.norm(P2)) > 1e-12:
             logger.warning(f"Large imaginary components in P2 {err}.")
         P2 = P2.real
 
@@ -575,12 +579,12 @@ def epsdS(images, samples_idx, max_d=None):
         # to estimate it.
 
         E = 0.0  # Total energy of the noise samples used to estimate the power spectrum.
-        samples = np.zeros((L, L), dtype=np.float64)
+        samples = xp.zeros((L, L), dtype=np.float64)
         for k in trange(n_img, desc="Estimating image noise energy"):
-            samples[samples_idx] = images[k].asnumpy()[0][samples_idx]
-            E += np.sum((samples - np.mean(samples)) ** 2)
+            samples[_samples_idx] = images[k].asnumpy()[0][samples_idx]
+            E += xp.sum((samples - xp.mean(samples)) ** 2)
         # Mean energy of the noise samples
-        n_samples_per_img = np.count_nonzero(samples_idx)
+        n_samples_per_img = xp.count_nonzero(_samples_idx)
         meanE = E / (n_samples_per_img * n_img)
 
         # Normalize P2 such that its mean energy is preserved and is equal to
@@ -589,20 +593,23 @@ def epsdS(images, samples_idx, max_d=None):
         # upsampling, downsampling, or cropping). Note that P2 is already in
         # units of energy, and so the total energy is given by sum(P2) and
         # not by norm(P2).
-        P2 = P2 / np.sum(P2) * meanE * P2.size
+        P2 = P2 / xp.sum(P2) * meanE * P2.size
 
         # Check that P2 has no negative values.
         # Due to the truncation of the Gaussian window, we get small negative
         # values. So unless they are very big, we just ignore them.
         negidx = P2 < 0
-        if np.count_nonzero(negidx):
-            maxnegerr = np.max(np.abs(P2[negidx]))
+        if xp.count_nonzero(negidx):
+            maxnegerr = xp.max(xp.abs(P2[negidx]))
             logger.debug(f"Maximal negative P2 value = {maxnegerr}")
             if maxnegerr > 1e-2:
-                negnorm = np.linalg.norm(P2[negidx])
+                negnorm = xp.linalg.norm(P2[negidx])
                 logger.warning(
                     f"Power spectrum P2 has non trivial negative values with energy {negnorm}."
                 )
             P2[negidx] = 0  # zero out negative estimates
 
+        P2 = xp.asnumpy(P2)
+        R2 = xp.asnumpy(R2)
+        # R, x already on host
         return P2, R, R2, x

From cff6de85e74bc42995a18e67a01c7b5bc47c3da5 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Mon, 27 Jan 2025 13:30:17 -0500
Subject: [PATCH 109/184] typos found in testing

---
 src/aspire/image/image.py | 4 ++--
 src/aspire/noise/noise.py | 2 +-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/aspire/image/image.py b/src/aspire/image/image.py
index 5e75fdd3de..658591880c 100644
--- a/src/aspire/image/image.py
+++ b/src/aspire/image/image.py
@@ -445,13 +445,13 @@ def legacy_whiten(self, psd, delta):
         fltr = fltr / xp.linalg.norm(fltr)
 
         # Error checking
-        if (err := xp.linalg.norm(fltr.imag)) < 10 * delta:
+        if (err := xp.linalg.norm(fltr.imag)) > 10 * delta:
             raise RuntimeError(
                 f"Whitening filter has non trivial imaginary component {err}."
             )
         err_ud = xp.linalg.norm(fltr - xp.flipud(fltr))
         err_lr = xp.linalg.norm(fltr - xp.fliplr(fltr))
-        if (err_ud < 10 * delta) or (err_lr < 10 * delta):
+        if (err_ud > 10 * delta) or (err_lr > 10 * delta):
             raise RuntimeError(
                 f"Whitening filter has non trivial symmetry lr {err_lr}, ud {err_ud}."
             )
diff --git a/src/aspire/noise/noise.py b/src/aspire/noise/noise.py
index 5b650bca47..9d4433c816 100644
--- a/src/aspire/noise/noise.py
+++ b/src/aspire/noise/noise.py
@@ -485,7 +485,7 @@ def epsdR(images, samples_idx, max_d=None):
             # energy estimate used later at this time to avoid looping
             # over images twice.
 
-            # Compute non-preiodic autocorrelation
+            # Compute non-periodic autocorrelation
             tmp[:L, :L] = samples  # pad
             ftmp = fft.fft2(tmp)
             # MATLAB code internally detects conj sym and implicitly casts,

From 26a8226a1103aa1054afa34c15cbf085b8b8ab3b Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Mon, 27 Jan 2025 14:59:18 -0500
Subject: [PATCH 110/184] Sensitive to the sample grid

---
 src/aspire/noise/noise.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/src/aspire/noise/noise.py b/src/aspire/noise/noise.py
index 9d4433c816..d0cac53284 100644
--- a/src/aspire/noise/noise.py
+++ b/src/aspire/noise/noise.py
@@ -401,7 +401,9 @@ def _estimate_noise_psd(self):
         :return: The estimated noise variance of the images in the Source used to create this estimator.
         """
         # Setup parameters
-        samples_idx = grid_2d(self.src.L, normalized=True)["r"] >= self.bg_radius
+        samples_idx = grid_2d(self.src.L, normalized=False)["r"] >= (
+            self.bg_radius * self.src.L
+        )
         max_d_pixels = round(self.max_d * self.src.L)
 
         psd = self.epsdS(self.src.images, samples_idx, max_d_pixels)[0]

From 584e5fd05b2a5161d170f6f067e503a420a87b57 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Tue, 28 Jan 2025 08:13:44 -0500
Subject: [PATCH 111/184] repro even samples_idx grid for whitening

---
 src/aspire/image/image.py | 2 +-
 src/aspire/noise/noise.py | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/aspire/image/image.py b/src/aspire/image/image.py
index 658591880c..0562c5e36e 100644
--- a/src/aspire/image/image.py
+++ b/src/aspire/image/image.py
@@ -473,7 +473,7 @@ def legacy_whiten(self, psd, delta):
             if L % 2 == 1:
                 pp[k - j - 1 : k + j, k - j - 1 : k + j] = xp.asarray(proj)
             else:
-                pp[k - j - 2 : k + j, k - j - 1 : k + j - 1] = xp.asarray(proj)
+                pp[k - j - 1 : k + j - 1, k - j - 1 : k + j - 1] = xp.asarray(proj)
 
             # Take the Fourier Transform of the padded image.
             fp = fft.centered_fft2(pp)
diff --git a/src/aspire/noise/noise.py b/src/aspire/noise/noise.py
index d0cac53284..3a5ae9e3c4 100644
--- a/src/aspire/noise/noise.py
+++ b/src/aspire/noise/noise.py
@@ -401,7 +401,7 @@ def _estimate_noise_psd(self):
         :return: The estimated noise variance of the images in the Source used to create this estimator.
         """
         # Setup parameters
-        samples_idx = grid_2d(self.src.L, normalized=False)["r"] >= (
+        samples_idx = grid_2d(self.src.L, normalized=False, shifted=True)["r"] >= (
             self.bg_radius * self.src.L
         )
         max_d_pixels = round(self.max_d * self.src.L)

From 0a7b9bdadcda880bf08a68e403e52959adae5067 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Tue, 28 Jan 2025 10:20:51 -0500
Subject: [PATCH 112/184] Add pair of `legacy_whiten` unit tests to preprocess
 pipeline suite

---
 tests/test_preprocess_pipeline.py | 49 ++++++++++++++++++++++++++++++-
 1 file changed, 48 insertions(+), 1 deletion(-)

diff --git a/tests/test_preprocess_pipeline.py b/tests/test_preprocess_pipeline.py
index c0ea0b7b6a..e53322a266 100644
--- a/tests/test_preprocess_pipeline.py
+++ b/tests/test_preprocess_pipeline.py
@@ -4,7 +4,11 @@
 import numpy as np
 import pytest
 
-from aspire.noise import AnisotropicNoiseEstimator, CustomNoiseAdder
+from aspire.noise import (
+    AnisotropicNoiseEstimator,
+    CustomNoiseAdder,
+    LegacyNoiseEstimator,
+)
 from aspire.operators.filters import FunctionFilter, RadialCTFFilter
 from aspire.source import ArrayImageSource
 from aspire.source.simulation import Simulation
@@ -182,6 +186,49 @@ def test_whiten_safeguard_default(dtype):
     assert np.count_nonzero(whiten_filt == 0) == 0
 
 
+@pytest.mark.parametrize(
+    "dtype", [np.float32, pytest.param(np.float64, marks=pytest.mark.expensive)]
+)
+def test_legacy_whiten(dtype):
+    """
+    Test `legacy_whiten` method.
+    """
+    L = 64
+    sim = get_sim_object(L, dtype)
+    sim = sim.legacy_whiten()
+    imgs_wt = sim.images[:num_images].asnumpy()
+
+    # calculate correlation between two neighboring pixels from background
+    corr_coef = np.corrcoef(imgs_wt[:, L - 1, L - 1], imgs_wt[:, L - 2, L - 1])
+
+    # correlation matrix should be close to identity
+    np.testing.assert_allclose(np.eye(2), corr_coef, atol=1e-1)
+    # dtype of returned images should be the same
+    assert dtype == imgs_wt.dtype
+
+
+@pytest.mark.parametrize(
+    "dtype", [np.float32, pytest.param(np.float64, marks=pytest.mark.expensive)]
+)
+def test_legacy_whiten_2(dtype):
+    """
+    Test `legacy_whiten` method with alternate invocation.
+    """
+    L = 63
+    sim = get_sim_object(L, dtype)
+    noise_estimator = LegacyNoiseEstimator(sim)
+    sim = sim.legacy_whiten(noise_estimator)
+    imgs_wt = sim.images[:num_images].asnumpy()
+
+    # calculate correlation between two neighboring pixels from background
+    corr_coef = np.corrcoef(imgs_wt[:, L - 1, L - 1], imgs_wt[:, L - 2, L - 1])
+
+    # correlation matrix should be close to identity
+    np.testing.assert_allclose(np.eye(2), corr_coef, atol=1e-1)
+    # dtype of returned images should be the same
+    assert dtype == imgs_wt.dtype
+
+
 @pytest.mark.parametrize("L, dtype", params)
 def testInvertContrast(L, dtype):
     sim1 = get_sim_object(L, dtype)

From de8a6b292b0d54143f7b8b2b6d2242aa3fb8ec37 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Wed, 29 Jan 2025 11:10:48 -0500
Subject: [PATCH 113/184] update exp 10028 pipeline to use legacy_whitening

---
 .../experimental_abinitio_pipeline_10028.py   | 31 +++++++------------
 1 file changed, 11 insertions(+), 20 deletions(-)

diff --git a/gallery/experiments/experimental_abinitio_pipeline_10028.py b/gallery/experiments/experimental_abinitio_pipeline_10028.py
index 034d6dfd68..9e1570bb13 100644
--- a/gallery/experiments/experimental_abinitio_pipeline_10028.py
+++ b/gallery/experiments/experimental_abinitio_pipeline_10028.py
@@ -30,7 +30,6 @@
 
 from aspire.abinitio import CLSync3N
 from aspire.denoising import DenoisedSource, DenoiserCov2D, LegacyClassAvgSource
-from aspire.noise import AnisotropicNoiseEstimator
 from aspire.reconstruction import MeanEstimator
 from aspire.source import OrientedSource, RelionSource
 
@@ -43,11 +42,11 @@
 # Example simulation configuration.
 
 interactive = False  # Draw blocking interactive plots?
-do_cov2d = True  # Use CWF coefficients
+do_cov2d = False  # Use CWF coefficients
 n_imgs = None  # Set to None for all images in starfile, can set smaller for tests.
 img_size = 32  # Downsample the images/reconstruction to a desired resolution
-n_classes = 2000  # How many class averages to compute.
-n_nbor = 100  # How many neighbors to stack
+n_classes = 3000  # How many class averages to compute.
+n_nbor = 50  # How many neighbors to stack
 starfile_in = "10028/data/shiny_2sets_fixed9.star"
 data_folder = "."  # This depends on the specific starfile entries.
 volume_output_filename = f"10028_abinitio_c{n_classes}_m{n_nbor}_{img_size}.mrc"
@@ -68,8 +67,9 @@
 )
 
 # Downsample the images
+# Caching is used for speeding up large datasets on high memory machines.
 logger.info(f"Set the resolution to {img_size} X {img_size}")
-src = src.downsample(img_size)
+src = src.downsample(img_size).cache()
 
 # Peek
 if interactive:
@@ -77,28 +77,19 @@
 
 # Use phase_flip to attempt correcting for CTF.
 logger.info("Perform phase flip to input images.")
-src = src.phase_flip()
+src = src.phase_flip().cache()
 
 # Estimate the noise and `Whiten` based on the estimated noise
-aiso_noise_estimator = AnisotropicNoiseEstimator(src)
-src = src.whiten(aiso_noise_estimator)
-
-# Plot the noise profile for inspection
-if interactive:
-    plt.imshow(aiso_noise_estimator.filter.evaluate_grid(img_size))
-    plt.show()
+src = src.legacy_whiten().cache()
 
 # Peek, what do the whitened images look like...
 if interactive:
     src.images[:10].show()
 
-# # Optionally invert image contrast, depends on data convention.
-# # This is not needed for 10028, but included anyway.
-# logger.info("Invert the global density contrast")
-# src = src.invert_contrast()
+# Optionally invert image contrast, depends on data conventions.
+logger.info("Invert the global density contrast")
+src = src.invert_contrast().cache()
 
-# Caching is used for speeding up large datasets on high memory machines.
-src = src.cache()
 
 # %%
 # Optional: CWF Denoising
@@ -163,7 +154,7 @@
 
 # Create a custom orientation estimation object for ``avgs``.
 # This is done to customize the ``n_theta`` value.
-orient_est = CLSync3N(avgs, n_theta=72)
+orient_est = CLSync3N(avgs, n_theta=360)
 
 # Create an ``OrientedSource`` class instance that performs orientation
 # estimation in a lazy fashion upon request of images or rotations.

From 56f325bdf575e94ed0154904dfeb448b75115340 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Fri, 31 Jan 2025 12:29:43 -0500
Subject: [PATCH 114/184] Name change

---
 src/aspire/noise/noise.py | 14 ++++++++++----
 1 file changed, 10 insertions(+), 4 deletions(-)

diff --git a/src/aspire/noise/noise.py b/src/aspire/noise/noise.py
index 3a5ae9e3c4..ee7f0ccbda 100644
--- a/src/aspire/noise/noise.py
+++ b/src/aspire/noise/noise.py
@@ -406,17 +406,21 @@ def _estimate_noise_psd(self):
         )
         max_d_pixels = round(self.max_d * self.src.L)
 
-        psd = self.epsdS(self.src.images, samples_idx, max_d_pixels)[0]
+        psd = self.estimate_power_spectrum_distribution_2d(
+            self.src.images, samples_idx, max_d_pixels
+        )[0]
 
         return psd
 
     @staticmethod
-    def epsdR(images, samples_idx, max_d=None):
+    def estimate_power_spectrum_distribution_1d(images, samples_idx, max_d=None):
         """
         Estimate the 1D isotropic autocorrelation function of `images`.
         The samples to use in each image are given by `samples_idx` mask.
         The correlation is computed up to a maximal distance of `max_d`.
 
+        Port of MATLAB `cryo_epsdR`.
+
         :param images: `Image` instance
         :param samples_idx: Boolean mask shaped `(L,L)`.
         :param max_d: Max computed correlation distance in pixels.
@@ -517,12 +521,14 @@ def epsdR(images, samples_idx, max_d=None):
         return R, x, cnt
 
     @staticmethod
-    def epsdS(images, samples_idx, max_d=None):
+    def estimate_power_spectrum_distribution_2d(images, samples_idx, max_d=None):
         """
         Estimate the 2D isotropic power spectrum of `images`.
         The samples to use in each image are given by `samples_idx` mask.
         The correlation is computed up to a maximal distance of `max_d`.
 
+        Port of MATLAB `cryo_epsdS`.
+
         :param images: Images instance
         :param samples_idx: Boolean mask shaped (L,L).
         :param max_d: Max computed correlation distance in pixels.
@@ -549,7 +555,7 @@ def epsdS(images, samples_idx, max_d=None):
             )
         max_d = int(min(max_d, L - 1))
 
-        R, x, _ = LegacyNoiseEstimator.epsdR(
+        R, x, _ = LegacyNoiseEstimator.estimate_power_spectrum_distribution_1d(
             images=images, samples_idx=samples_idx, max_d=max_d
         )
         _R = xp.asarray(R)  # Migrate to GPU for assignments below

From 31964267761488a144ef60441e0f8e44d4a1b9c9 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Fri, 31 Jan 2025 12:33:50 -0500
Subject: [PATCH 115/184] Complete missing docstring

---
 src/aspire/noise/noise.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/src/aspire/noise/noise.py b/src/aspire/noise/noise.py
index ee7f0ccbda..efd577366d 100644
--- a/src/aspire/noise/noise.py
+++ b/src/aspire/noise/noise.py
@@ -363,7 +363,9 @@ class LegacyNoiseEstimator(NoiseEstimator):
 
     def __init__(self, src, bg_radius=None, max_d=None):
         """
-        Given an `ImageSource`, constructs
+        Given an `ImageSource`, prepares for estimation of noise spectrum.
+
+        Estimate is delayed and computed on first access of `filter` attribute.
 
         :param src: A `ImageSource` object.
         :param bg_radius: The radius of the disk whose complement is used to estimate the noise.

From d48d0c24a29bc18bf301649b8ea5d3e270c1ef27 Mon Sep 17 00:00:00 2001
From: Garrett Wright <47759732+garrettwrong@users.noreply.github.com>
Date: Wed, 5 Feb 2025 08:34:20 -0500
Subject: [PATCH 116/184] WIP legacy_whiten speedup (#1227)

* batch legacy_whiten 2d
* fft opts legacy_whitening
* profile and implement Yoels validdist opt
---
 src/aspire/noise/noise.py       | 103 ++++++++++++++++++++------------
 src/aspire/numeric/base_fft.py  |  12 ++++
 src/aspire/numeric/cupy_fft.py  |   8 +++
 src/aspire/numeric/scipy_fft.py |   6 ++
 src/aspire/source/image.py      |   4 +-
 5 files changed, 92 insertions(+), 41 deletions(-)

diff --git a/src/aspire/noise/noise.py b/src/aspire/noise/noise.py
index efd577366d..d25dcf2068 100644
--- a/src/aspire/noise/noise.py
+++ b/src/aspire/noise/noise.py
@@ -361,7 +361,7 @@ class LegacyNoiseEstimator(NoiseEstimator):
     Isotropic noise estimator ported from MATLAB `cryo_noise_estimation`.
     """
 
-    def __init__(self, src, bg_radius=None, max_d=None):
+    def __init__(self, src, bg_radius=None, max_d=None, batch_size=512):
         """
         Given an `ImageSource`, prepares for estimation of noise spectrum.
 
@@ -374,12 +374,13 @@ def __init__(self, src, bg_radius=None, max_d=None):
             Default of `None` uses `(np.floor(src.L / 2) - 1) / L`
         :param max_d: Max computed correlation distance as a proportion of `src.L`.
             Default of `None` uses `np.floor(src.L/3) / L`.
+        :param batch_size:  The size of the batches in which to compute the variance estimate.
         """
 
         if bg_radius is None:
             bg_radius = (np.floor(src.L / 2) - 1) / src.L
 
-        super().__init__(src=src, bg_radius=bg_radius, batch_size=1)
+        super().__init__(src=src, bg_radius=bg_radius, batch_size=batch_size)
 
         self.max_d = max_d
         if self.max_d is None:
@@ -409,13 +410,18 @@ def _estimate_noise_psd(self):
         max_d_pixels = round(self.max_d * self.src.L)
 
         psd = self.estimate_power_spectrum_distribution_2d(
-            self.src.images, samples_idx, max_d_pixels
+            self.src.images,
+            samples_idx,
+            max_d_pixels,
+            batch_size=self.batch_size,
         )[0]
 
         return psd
 
     @staticmethod
-    def estimate_power_spectrum_distribution_1d(images, samples_idx, max_d=None):
+    def estimate_power_spectrum_distribution_1d(
+        images, samples_idx, max_d=None, batch_size=512
+    ):
         """
         Estimate the 1D isotropic autocorrelation function of `images`.
         The samples to use in each image are given by `samples_idx` mask.
@@ -427,6 +433,7 @@ def estimate_power_spectrum_distribution_1d(images, samples_idx, max_d=None):
         :param samples_idx: Boolean mask shaped `(L,L)`.
         :param max_d: Max computed correlation distance in pixels.
            Default of `None` yields `np.floor(L / 3)`.
+        :param batch_size:  The size of the batches in which to compute the variance estimate.
         :return:
             - Radial PSD array of shape
             - Distances map array
@@ -435,6 +442,7 @@ def estimate_power_spectrum_distribution_1d(images, samples_idx, max_d=None):
 
         n_img = images.n_images
         L = samples_idx.shape[-1]
+        batch_size = min(batch_size, n_img)
 
         # Correlations more than `max_d` pixels apart are not computed.
         if max_d is None:
@@ -466,17 +474,19 @@ def estimate_power_spectrum_distribution_1d(images, samples_idx, max_d=None):
         for i, d in enumerate(dsquare):
             inds = dists == d  # locations having distance `d`
             distmap[inds] = i  # assign index into dsquare `i`
+        # From here on, distmap will be accessed with flat indices
+        distmap = distmap.flatten()
+        validdists = xp.argwhere(distmap != -1)
 
         # Compute Ncorr using a constant unit image.
         mask = xp.zeros((L, L))
         mask[samples_idx] = 1
-        tmp = xp.zeros((2 * L + 1, 2 * L + 1))  # pad
-        tmp[:L, :L] = mask
-        ftmp = fft.fft2(tmp)
-        # MATLAB code internally detects sym and implicitly casts,
-        #   we explicitly cast.
-        # Optimization note: This and the later fft call could be optimized with real/herm sym FFT calls.
-        Ncorr = fft.ifft2(ftmp * ftmp.conj()).real
+        tmp = xp.zeros((batch_size, 2 * L + 1, 2 * L + 1))  # pad
+        tmp[0, :L, :L] = mask
+        # MATLAB code internally detects/implicitly casts,
+        #   we explicitly call rfft2/irfft2.
+        ftmp = fft.rfft2(tmp[0])
+        Ncorr = fft.irfft2(ftmp * ftmp.conj(), s=tmp.shape[1:])
         Ncorr = Ncorr[: max_d + 1, : max_d + 1]  # crop
         Ncorr = xp.round(Ncorr)
 
@@ -484,46 +494,53 @@ def estimate_power_spectrum_distribution_1d(images, samples_idx, max_d=None):
         # R[i] is value of ACF at distance x[i]
         R = xp.zeros(len(corrs))
 
-        samples = xp.zeros((L, L))
-        tmp[:, :] = 0  # reset tmp
-        for k in trange(n_img, desc="Processing image autocorrelations"):
+        samples = xp.zeros((batch_size, L, L))
+        tmp[0, :, :] = 0  # reset tmp
+        corrs = corrs.flatten()
+        corrcount = corrcount.flatten()
+        for start in trange(
+            0, n_img, batch_size, desc="Processing image autocorrelations"
+        ):
+            end = min(n_img, start + batch_size)
+            count = end - start
             # Mask off unused pixels
-            samples[samples_idx] = images[k].asnumpy()[0][samples_idx]
+            samples[:count, samples_idx] = images[start:end].asnumpy()[:, samples_idx]
             # Optimization note: We could also compute the noise
             # energy estimate used later at this time to avoid looping
             # over images twice.
 
             # Compute non-periodic autocorrelation
-            tmp[:L, :L] = samples  # pad
-            ftmp = fft.fft2(tmp)
-            # MATLAB code internally detects conj sym and implicitly casts,
-            #   we explicitly cast.
-            s = fft.ifft2(ftmp * ftmp.conj()).real
-            s = s[0 : max_d + 1, 0 : max_d + 1]  # crop
+            tmp[:count, :L, :L] = samples[:count]  # pad
+            # MATLAB code internally detects/implicitly casts,
+            #   we explicitly call rfft2/irfft2.
+            ftmp = fft.rfft2(tmp[:count])
+            s = fft.irfft2(ftmp * ftmp.conj(), s=tmp.shape[1:])
+            s = s[:, 0 : max_d + 1, 0 : max_d + 1]  # crop
 
             # Accumulate all autocorrelation values R[k1,k2] such that
             # k1**2 + k2**2 = dist (all autocorrelations of a certain distance).
-            # Optimization note: The MATLAB code used another map
-            #   layer `validdists` to remove one loop layer here, but
-            #   it relied primarily on MATLABs implicit ravel/unravel
-            #   and would be less clear in Python. Simpler code was ported.
-            for i in range(max_d + 1):
-                for j in range(max_d + 1):
-                    idx = distmap[i, j]
-                    if idx != -1:
-                        corrs[idx] = corrs[idx] + s[i, j]
-                        corrcount[idx] = corrcount[idx] + Ncorr[i, j]
+            s = xp.sum(s, axis=0).flatten()
+            _Ncorr = Ncorr.flatten() * count
+            for d in validdists:
+                idx = distmap[d]
+                corrs[idx] = corrs[idx] + s[d]
+                corrcount[idx] = corrcount[idx] + _Ncorr[d]
 
         # Remove distances which had no samples
         idx = xp.where(corrcount != 0)
-        R = xp.asnumpy(corrs[idx] / corrcount[idx])
+        R = corrs[idx] / corrcount[idx]
         x = xp.asnumpy(x[idx])
-        cnt = xp.asnumpy(corrcount[idx])
+        cnt = corrcount[idx]
+
+        R = xp.asnumpy(R)
+        cnt = xp.asnumpy(cnt)
 
         return R, x, cnt
 
     @staticmethod
-    def estimate_power_spectrum_distribution_2d(images, samples_idx, max_d=None):
+    def estimate_power_spectrum_distribution_2d(
+        images, samples_idx, max_d=None, batch_size=512
+    ):
         """
         Estimate the 2D isotropic power spectrum of `images`.
         The samples to use in each image are given by `samples_idx` mask.
@@ -535,6 +552,7 @@ def estimate_power_spectrum_distribution_2d(images, samples_idx, max_d=None):
         :param samples_idx: Boolean mask shaped (L,L).
         :param max_d: Max computed correlation distance in pixels.
            Default of `None` yields `np.floor(L / 3)`.
+        :param batch_size:  The size of the batches in which to compute the variance estimate.
         :return:
             - 2D PSD array
             - Radial PSD array
@@ -544,6 +562,7 @@ def estimate_power_spectrum_distribution_2d(images, samples_idx, max_d=None):
 
         n_img = images.n_images
         L = samples_idx.shape[-1]
+        batch_size = min(batch_size, n_img)
 
         # Migrate mask to GPU as needed
         _samples_idx = xp.asarray(samples_idx)
@@ -558,7 +577,7 @@ def estimate_power_spectrum_distribution_2d(images, samples_idx, max_d=None):
         max_d = int(min(max_d, L - 1))
 
         R, x, _ = LegacyNoiseEstimator.estimate_power_spectrum_distribution_1d(
-            images=images, samples_idx=samples_idx, max_d=max_d
+            images=images, samples_idx=samples_idx, max_d=max_d, batch_size=batch_size
         )
         _R = xp.asarray(R)  # Migrate to GPU for assignments below
 
@@ -589,10 +608,16 @@ def estimate_power_spectrum_distribution_2d(images, samples_idx, max_d=None):
         # to estimate it.
 
         E = 0.0  # Total energy of the noise samples used to estimate the power spectrum.
-        samples = xp.zeros((L, L), dtype=np.float64)
-        for k in trange(n_img, desc="Estimating image noise energy"):
-            samples[_samples_idx] = images[k].asnumpy()[0][samples_idx]
-            E += xp.sum((samples - xp.mean(samples)) ** 2)
+        samples = xp.zeros((batch_size, L, L), dtype=np.float64)
+        for start in trange(0, n_img, batch_size, desc="Estimating image noise energy"):
+            end = min(n_img, start + batch_size)
+            cnt = end - start
+
+            samples[:cnt, _samples_idx] = images[start:end].asnumpy()[0][samples_idx]
+            E += xp.sum(
+                (samples[:cnt] - xp.mean(samples[:cnt], axis=(1, 2)).reshape(cnt, 1, 1))
+                ** 2
+            )
         # Mean energy of the noise samples
         n_samples_per_img = xp.count_nonzero(_samples_idx)
         meanE = E / (n_samples_per_img * n_img)
diff --git a/src/aspire/numeric/base_fft.py b/src/aspire/numeric/base_fft.py
index d80ca780d1..6ce029f518 100644
--- a/src/aspire/numeric/base_fft.py
+++ b/src/aspire/numeric/base_fft.py
@@ -18,6 +18,18 @@ def fft2(self, x, axes=(-2, -1), workers=-1):
     def ifft2(self, x, axes=(-2, -1), workers=-1):
         raise NotImplementedError("subclasses must implement this")
 
+    def rfft(self, x, axis=-1, workers=-1):
+        raise NotImplementedError("subclasses must implement this")
+
+    def irfft(self, x, axis=-1, workers=-1):
+        raise NotImplementedError("subclasses must implement this")
+
+    def rfft2(self, x, axes=(-2, -1), workers=-1):
+        raise NotImplementedError("subclasses must implement this")
+
+    def irfft2(self, x, axes=(-2, -1), workers=-1):
+        raise NotImplementedError("subclasses must implement this")
+
     def fftn(self, x, axes=None, workers=-1):
         raise NotImplementedError("subclasses must implement this")
 
diff --git a/src/aspire/numeric/cupy_fft.py b/src/aspire/numeric/cupy_fft.py
index b491a0dcd1..2c5179ed72 100644
--- a/src/aspire/numeric/cupy_fft.py
+++ b/src/aspire/numeric/cupy_fft.py
@@ -111,3 +111,11 @@ def irfft(self, x, **kwargs):
     @_preserve_host
     def rfft(self, x, **kwargs):
         return cufft.rfft(x, **kwargs)
+
+    @_preserve_host
+    def irfft2(self, x, **kwargs):
+        return cufft.irfft2(x, **kwargs)
+
+    @_preserve_host
+    def rfft2(self, x, **kwargs):
+        return cufft.rfft2(x, **kwargs)
diff --git a/src/aspire/numeric/scipy_fft.py b/src/aspire/numeric/scipy_fft.py
index 0ef5c95f16..0d9af75309 100644
--- a/src/aspire/numeric/scipy_fft.py
+++ b/src/aspire/numeric/scipy_fft.py
@@ -48,3 +48,9 @@ def irfft(self, x, **kwargs):
 
     def rfft(self, x, **kwargs):
         return sp.fft.rfft(x, **kwargs)
+
+    def irfft2(self, x, **kwargs):
+        return sp.fft.irfft2(x, **kwargs)
+
+    def rfft2(self, x, **kwargs):
+        return sp.fft.rfft2(x, **kwargs)
diff --git a/src/aspire/source/image.py b/src/aspire/source/image.py
index 5e2481d68d..6eff83d23d 100644
--- a/src/aspire/source/image.py
+++ b/src/aspire/source/image.py
@@ -826,7 +826,7 @@ def whiten(self, noise_estimate=None, epsilon=None):
         self.generation_pipeline.add_xform(FilterXform(whiten_filter))
 
     @_as_copy
-    def legacy_whiten(self, noise_response=None, delta=None):
+    def legacy_whiten(self, noise_response=None, delta=None, batch_size=512):
         """
         Reproduce the legacy MATLAB whitening process.
 
@@ -840,7 +840,7 @@ def legacy_whiten(self, noise_response=None, delta=None):
 
         if noise_response is None:
             logger.info("Computing noise response.")
-            psd = LegacyNoiseEstimator(self).filter.xfer_fn_array
+            psd = LegacyNoiseEstimator(self, batch_size=batch_size).filter.xfer_fn_array
         elif isinstance(noise_response, LegacyNoiseEstimator):
             psd = noise_response.filter.xfer_fn_array
         elif isinstance(noise_response, np.ndarray):

From 69bc86f483676e3cd6300af8d969b70ac2d33209 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Wed, 5 Feb 2025 10:21:15 -0500
Subject: [PATCH 117/184] extend other FFT packages to cover additional real
 functions

---
 src/aspire/numeric/base_fft.py   |  9 +++++++++
 src/aspire/numeric/mkl_fft.py    | 24 ++++++++++++++++++++++++
 src/aspire/numeric/pyfftw_fft.py | 15 +++++++++++++++
 3 files changed, 48 insertions(+)

diff --git a/src/aspire/numeric/base_fft.py b/src/aspire/numeric/base_fft.py
index 6ce029f518..b317f20396 100644
--- a/src/aspire/numeric/base_fft.py
+++ b/src/aspire/numeric/base_fft.py
@@ -42,6 +42,15 @@ def fftshift(self, x, axes=None):
     def ifftshift(self, x, axes=None):
         raise NotImplementedError("subclasses must implement this")
 
+    def dct(self, x, **kwargs):
+        raise NotImplementedError("subclasses must implement this")
+
+    def idct(self, x, **kwargs):
+        raise NotImplementedError("subclasses must implement this")
+
+    def rfftfreq(self, x, **kwargs):
+        raise NotImplementedError("subclasses must implement this")
+
     def centered_ifft(self, x, axis=-1, workers=-1):
         x = self.ifftshift(x, axes=axis)
         x = self.ifft(x, axis=axis, workers=workers)
diff --git a/src/aspire/numeric/mkl_fft.py b/src/aspire/numeric/mkl_fft.py
index 57dfff1e0d..29e9ebfc73 100644
--- a/src/aspire/numeric/mkl_fft.py
+++ b/src/aspire/numeric/mkl_fft.py
@@ -1,5 +1,6 @@
 import mkl_fft
 import numpy as np
+import scipy as sp
 
 from aspire.numeric.base_fft import FFT
 
@@ -55,3 +56,26 @@ def fftshift(self, x, axes=None):
     def ifftshift(self, x, axes=None):
         # N/A in mkl_fft, use np
         return np.fft.ifftshift(x, axes=axes)
+
+    def rfft(self, x, **kwargs):
+        return mkl_fft._numpy_fft.rfft(x, **kwargs)
+
+    def irfft(self, x, **kwargs):
+        return mkl_fft._numpy_fft.irfft(x, **kwargs)
+
+    def rfft2(self, x, **kwargs):
+        return mkl_fft._numpy_fft.rfft2(x, **kwargs)
+
+    def irfft2(self, x, **kwargs):
+        return mkl_fft._numpy_fft.irfft2(x, **kwargs)
+
+    # These are not currently exposed in mkl_fft,
+    #   fall back to scipy.
+    def dct(self, x, **kwargs):
+        return sp.fft.dct(x, **kwargs)
+
+    def idct(self, x, **kwargs):
+        return sp.fft.idct(x, **kwargs)
+
+    def rfftfreq(self, x, **kwargs):
+        return sp.fft.rfftfreq(x, **kwargs)
diff --git a/src/aspire/numeric/pyfftw_fft.py b/src/aspire/numeric/pyfftw_fft.py
index 95a8ea80f7..0a100eca50 100644
--- a/src/aspire/numeric/pyfftw_fft.py
+++ b/src/aspire/numeric/pyfftw_fft.py
@@ -154,6 +154,18 @@ def ifftn(self, a, axes=None, workers=-1):
 
         return b
 
+    def rfft(self, x, **kwargs):
+        return pyfftw.interfaces.numpy_fft.rfft(x, **kwargs)
+
+    def irfft(self, x, **kwargs):
+        return pyfftw.interfaces.numpy_fft.irfft(x, **kwargs)
+
+    def rfft2(self, x, **kwargs):
+        return pyfftw.interfaces.numpy_fft.rfft2(x, **kwargs)
+
+    def irfft2(self, x, **kwargs):
+        return pyfftw.interfaces.numpy_fft.irfft2(x, **kwargs)
+
     def fftshift(self, a, axes=None):
         return scipy_fft.fftshift(a, axes=axes)
 
@@ -165,3 +177,6 @@ def dct(self, x, **kwargs):
 
     def idct(self, x, **kwargs):
         return scipy_fft.idct(x, **kwargs)
+
+    def rfftfreq(self, x, **kwargs):
+        return scipy_fft.rfftfreq(x, **kwargs)

From 9f4043083eba39e7beea0888a0b9abd79c73b988 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Thu, 6 Feb 2025 07:28:42 -0500
Subject: [PATCH 118/184] private method names

---
 src/aspire/noise/noise.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/aspire/noise/noise.py b/src/aspire/noise/noise.py
index d25dcf2068..194677352b 100644
--- a/src/aspire/noise/noise.py
+++ b/src/aspire/noise/noise.py
@@ -409,7 +409,7 @@ def _estimate_noise_psd(self):
         )
         max_d_pixels = round(self.max_d * self.src.L)
 
-        psd = self.estimate_power_spectrum_distribution_2d(
+        psd = self._estimate_power_spectrum_distribution_2d(
             self.src.images,
             samples_idx,
             max_d_pixels,
@@ -419,7 +419,7 @@ def _estimate_noise_psd(self):
         return psd
 
     @staticmethod
-    def estimate_power_spectrum_distribution_1d(
+    def _estimate_power_spectrum_distribution_1d(
         images, samples_idx, max_d=None, batch_size=512
     ):
         """
@@ -538,7 +538,7 @@ def estimate_power_spectrum_distribution_1d(
         return R, x, cnt
 
     @staticmethod
-    def estimate_power_spectrum_distribution_2d(
+    def _estimate_power_spectrum_distribution_2d(
         images, samples_idx, max_d=None, batch_size=512
     ):
         """
@@ -576,7 +576,7 @@ def estimate_power_spectrum_distribution_2d(
             )
         max_d = int(min(max_d, L - 1))
 
-        R, x, _ = LegacyNoiseEstimator.estimate_power_spectrum_distribution_1d(
+        R, x, _ = LegacyNoiseEstimator._estimate_power_spectrum_distribution_1d(
             images=images, samples_idx=samples_idx, max_d=max_d, batch_size=batch_size
         )
         _R = xp.asarray(R)  # Migrate to GPU for assignments below

From 296747845c18394866e491b03c54f591010386e7 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Thu, 6 Feb 2025 07:50:17 -0500
Subject: [PATCH 119/184] Ncorr and tmp name changes

---
 src/aspire/noise/noise.py | 29 +++++++++++++++--------------
 1 file changed, 15 insertions(+), 14 deletions(-)

diff --git a/src/aspire/noise/noise.py b/src/aspire/noise/noise.py
index 194677352b..4f361b2cc5 100644
--- a/src/aspire/noise/noise.py
+++ b/src/aspire/noise/noise.py
@@ -476,26 +476,27 @@ def _estimate_power_spectrum_distribution_1d(
             distmap[inds] = i  # assign index into dsquare `i`
         # From here on, distmap will be accessed with flat indices
         distmap = distmap.flatten()
-        validdists = xp.argwhere(distmap != -1)
+        valid_dists = xp.argwhere(distmap != -1)
 
         # Compute Ncorr using a constant unit image.
         mask = xp.zeros((L, L))
         mask[samples_idx] = 1
-        tmp = xp.zeros((batch_size, 2 * L + 1, 2 * L + 1))  # pad
-        tmp[0, :L, :L] = mask
+        mask_padded = xp.zeros((batch_size, 2 * L + 1, 2 * L + 1))  # pad
+        mask_padded[0, :L, :L] = mask
         # MATLAB code internally detects/implicitly casts,
         #   we explicitly call rfft2/irfft2.
-        ftmp = fft.rfft2(tmp[0])
-        Ncorr = fft.irfft2(ftmp * ftmp.conj(), s=tmp.shape[1:])
-        Ncorr = Ncorr[: max_d + 1, : max_d + 1]  # crop
-        Ncorr = xp.round(Ncorr)
+        fmask_padded = fft.rfft2(mask_padded[0])
+        n_mask_pairs = fft.irfft2(fmask_padded * fmask_padded.conj(), s=mask_padded.shape[1:])
+        n_mask_pairs = n_mask_pairs[: max_d + 1, : max_d + 1]  # crop
+        breakpoint()
+        n_mask_pairs = xp.round(n_mask_pairs)
 
         # Values of isotropic autocorrelation function
         # R[i] is value of ACF at distance x[i]
         R = xp.zeros(len(corrs))
 
         samples = xp.zeros((batch_size, L, L))
-        tmp[0, :, :] = 0  # reset tmp
+        mask_padded[0, :, :] = 0  # reset mask_padded
         corrs = corrs.flatten()
         corrcount = corrcount.flatten()
         for start in trange(
@@ -510,21 +511,21 @@ def _estimate_power_spectrum_distribution_1d(
             # over images twice.
 
             # Compute non-periodic autocorrelation
-            tmp[:count, :L, :L] = samples[:count]  # pad
+            mask_padded[:count, :L, :L] = samples[:count]  # pad
             # MATLAB code internally detects/implicitly casts,
             #   we explicitly call rfft2/irfft2.
-            ftmp = fft.rfft2(tmp[:count])
-            s = fft.irfft2(ftmp * ftmp.conj(), s=tmp.shape[1:])
+            fmask_padded = fft.rfft2(mask_padded[:count])
+            s = fft.irfft2(fmask_padded * fmask_padded.conj(), s=mask_padded.shape[1:])
             s = s[:, 0 : max_d + 1, 0 : max_d + 1]  # crop
 
             # Accumulate all autocorrelation values R[k1,k2] such that
             # k1**2 + k2**2 = dist (all autocorrelations of a certain distance).
             s = xp.sum(s, axis=0).flatten()
-            _Ncorr = Ncorr.flatten() * count
-            for d in validdists:
+            _n_mask_pairs = n_mask_pairs.flatten() * count
+            for d in valid_dists:
                 idx = distmap[d]
                 corrs[idx] = corrs[idx] + s[d]
-                corrcount[idx] = corrcount[idx] + _Ncorr[d]
+                corrcount[idx] = corrcount[idx] + _n_mask_pairs[d]
 
         # Remove distances which had no samples
         idx = xp.where(corrcount != 0)

From 270ae882717b7b446337dfe27248a9477dd9b794 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Thu, 6 Feb 2025 07:59:19 -0500
Subject: [PATCH 120/184] dists dsquare renaming

---
 src/aspire/noise/noise.py | 28 ++++++++++++++--------------
 1 file changed, 14 insertions(+), 14 deletions(-)

diff --git a/src/aspire/noise/noise.py b/src/aspire/noise/noise.py
index 4f361b2cc5..fc504b058a 100644
--- a/src/aspire/noise/noise.py
+++ b/src/aspire/noise/noise.py
@@ -457,23 +457,23 @@ def _estimate_power_spectrum_distribution_1d(
         # Note grid_2d['r'] is not used because we want an integer grid directly;
         #   yields integer dists (radius**2) values.
         X, Y = xp.mgrid[0 : max_d + 1, 0 : max_d + 1]
-        dists = X * X + Y * Y
-        dsquare = xp.sort(xp.unique(dists[dists <= max_d**2]))
-        x = xp.sqrt(dsquare)  # actual distances
+        dsquare = X * X + Y * Y
+        uniq_dsquare = xp.sort(xp.unique(dsquare[dsquare <= max_d**2]))
+        x = xp.sqrt(uniq_dsquare)  # actual distances
 
         # corrs[i] is the sum of all x[j]x[j+d] where d = x[i]
-        corrs = xp.zeros_like(dsquare, dtype=np.float64)
+        corrs = xp.zeros_like(uniq_dsquare, dtype=np.float64)
         # corrcount[i] is the number of pairs summed in corr[i]
-        corrcount = xp.zeros_like(dsquare, dtype=np.int64)
+        corrcount = xp.zeros_like(uniq_dsquare, dtype=np.int64)
 
-        # distmap maps [i,j] to k where dsquare[k] = i**2 + j**2.
+        # distmap maps [i,j] to k where uniq_dsquare[k] = i**2 + j**2.
         #   -1 indicates distance is larger than max_d
-        distmap = xp.full(shape=dists.shape, fill_value=-1)
+        distmap = xp.full(shape=dsquare.shape, fill_value=-1)
 
         # This differs from the MATLAB code, avoids `bisect`.
-        for i, d in enumerate(dsquare):
-            inds = dists == d  # locations having distance `d`
-            distmap[inds] = i  # assign index into dsquare `i`
+        for i, d in enumerate(uniq_dsquare):
+            inds = dsquare == d  # locations having distance `d`
+            distmap[inds] = i  # assign index into uniq_dsquare `i`
         # From here on, distmap will be accessed with flat indices
         distmap = distmap.flatten()
         valid_dists = xp.argwhere(distmap != -1)
@@ -589,10 +589,10 @@ def _estimate_power_spectrum_distribution_2d(
         R2 = xp.zeros((2 * L - 1, 2 * L - 1), dtype=np.float64)
 
         X, Y = xp.mgrid[-L + 1 : L, -L + 1 : L]
-        dists2 = X * X + Y * Y
-        dsquare2 = xp.sort(xp.unique(dists2[dists2 <= max_d**2]))
-        for i, d in enumerate(dsquare2):
-            idx = dists2 == d
+        dists = X * X + Y * Y
+        uniq_dsquare = xp.sort(xp.unique(dists[dists <= max_d**2]))
+        for i, d in enumerate(uniq_dsquare):
+            idx = dists == d
             R2[idx] = _R[i]
 
         # Window the 2D autocorrelation and Fourier transform it to get the power

From fa6bca62df394b0050e16ec47e539d924cb039e7 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Thu, 6 Feb 2025 08:04:48 -0500
Subject: [PATCH 121/184] remove unnecessary flattens from code

---
 src/aspire/noise/noise.py | 9 ++-------
 1 file changed, 2 insertions(+), 7 deletions(-)

diff --git a/src/aspire/noise/noise.py b/src/aspire/noise/noise.py
index fc504b058a..af11bf8d64 100644
--- a/src/aspire/noise/noise.py
+++ b/src/aspire/noise/noise.py
@@ -488,17 +488,10 @@ def _estimate_power_spectrum_distribution_1d(
         fmask_padded = fft.rfft2(mask_padded[0])
         n_mask_pairs = fft.irfft2(fmask_padded * fmask_padded.conj(), s=mask_padded.shape[1:])
         n_mask_pairs = n_mask_pairs[: max_d + 1, : max_d + 1]  # crop
-        breakpoint()
         n_mask_pairs = xp.round(n_mask_pairs)
 
-        # Values of isotropic autocorrelation function
-        # R[i] is value of ACF at distance x[i]
-        R = xp.zeros(len(corrs))
-
         samples = xp.zeros((batch_size, L, L))
         mask_padded[0, :, :] = 0  # reset mask_padded
-        corrs = corrs.flatten()
-        corrcount = corrcount.flatten()
         for start in trange(
             0, n_img, batch_size, desc="Processing image autocorrelations"
         ):
@@ -527,6 +520,8 @@ def _estimate_power_spectrum_distribution_1d(
                 corrs[idx] = corrs[idx] + s[d]
                 corrcount[idx] = corrcount[idx] + _n_mask_pairs[d]
 
+        # Values of isotropic autocorrelation function
+        # R[i] is value of ACF at distance x[i]
         # Remove distances which had no samples
         idx = xp.where(corrcount != 0)
         R = corrs[idx] / corrcount[idx]

From bd38e8aafefd389e0fe64a373ba0c14120473b24 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Thu, 6 Feb 2025 08:16:47 -0500
Subject: [PATCH 122/184] improve comment

---
 src/aspire/noise/noise.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/aspire/noise/noise.py b/src/aspire/noise/noise.py
index af11bf8d64..b01e1739fa 100644
--- a/src/aspire/noise/noise.py
+++ b/src/aspire/noise/noise.py
@@ -497,7 +497,7 @@ def _estimate_power_spectrum_distribution_1d(
         ):
             end = min(n_img, start + batch_size)
             count = end - start
-            # Mask off unused pixels
+            # Pack masked `sample_idx` pixels from `images` batch into `samples`
             samples[:count, samples_idx] = images[start:end].asnumpy()[:, samples_idx]
             # Optimization note: We could also compute the noise
             # energy estimate used later at this time to avoid looping

From 21f374682e63ac23f85b4bb88e2f36a17b9ffacf Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Thu, 6 Feb 2025 08:43:21 -0500
Subject: [PATCH 123/184] renaming style cleanup

---
 src/aspire/noise/noise.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/src/aspire/noise/noise.py b/src/aspire/noise/noise.py
index b01e1739fa..4f097e2b20 100644
--- a/src/aspire/noise/noise.py
+++ b/src/aspire/noise/noise.py
@@ -486,7 +486,9 @@ def _estimate_power_spectrum_distribution_1d(
         # MATLAB code internally detects/implicitly casts,
         #   we explicitly call rfft2/irfft2.
         fmask_padded = fft.rfft2(mask_padded[0])
-        n_mask_pairs = fft.irfft2(fmask_padded * fmask_padded.conj(), s=mask_padded.shape[1:])
+        n_mask_pairs = fft.irfft2(
+            fmask_padded * fmask_padded.conj(), s=mask_padded.shape[1:]
+        )
         n_mask_pairs = n_mask_pairs[: max_d + 1, : max_d + 1]  # crop
         n_mask_pairs = xp.round(n_mask_pairs)
 

From 723c340f6de9b3238f4c2300e604d8d7bba5e620 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Thu, 6 Feb 2025 08:43:36 -0500
Subject: [PATCH 124/184] renaming and slicing refactor

---
 src/aspire/image/image.py | 39 ++++++++++++++++++++++-----------------
 1 file changed, 22 insertions(+), 17 deletions(-)

diff --git a/src/aspire/image/image.py b/src/aspire/image/image.py
index 0562c5e36e..7d30a86752 100644
--- a/src/aspire/image/image.py
+++ b/src/aspire/image/image.py
@@ -432,7 +432,7 @@ def legacy_whiten(self, psd, delta):
         """
         n = self.n_images
         L = self.resolution
-        j = L // 2
+        L_half = L // 2
         K = psd.shape[-1]
         k = int(np.ceil(K / 2))
 
@@ -463,39 +463,44 @@ def legacy_whiten(self, psd, delta):
         # The filter may have very small values or even zeros.
         # We don't want to process these, so make a list of all large entries.
         nzidx = fltr > 100 * delta
-        fnz = fltr[nzidx]
+        fltr_nz = fltr[nzidx]
+
+        padded_proj = xp.zeros((K, K), dtype=np.float64)
+        filtered_fpadded_proj = xp.zeros((K, K), dtype=np.complex128)
+
+        # Precompute the slices
+        if L % 2 == 1:
+            slc = slice(k - L_half - 1, k + L_half)
+        else:
+            slc = slice(k - L_half - 1, k + L_half - 1)
 
-        pp = xp.zeros((K, K), dtype=np.float64)
-        p = xp.zeros((K, K), dtype=np.complex128)
         for i, proj in enumerate(self.asnumpy()):
 
             # Zero pad the image to twice the size
-            if L % 2 == 1:
-                pp[k - j - 1 : k + j, k - j - 1 : k + j] = xp.asarray(proj)
-            else:
-                pp[k - j - 1 : k + j - 1, k - j - 1 : k + j - 1] = xp.asarray(proj)
+            padded_proj[slc, slc] = xp.asarray(proj)
 
             # Take the Fourier Transform of the padded image.
-            fp = fft.centered_fft2(pp)
+            fpadded_proj = fft.centered_fft2(padded_proj)
 
             # Divide the image by the whitening filter but only in
             # places where the filter is large.  In frequencies where
             # the filter is tiny we cannot whiten so we just use
             # zeros.
-            p[nzidx] = fp[nzidx] / fnz
-            p2 = fft.centered_ifft2(p)
+            filtered_fpadded_proj[nzidx] = fpadded_proj[nzidx] / fltr_nz
+            # `filtered_proj` is still padded and complex. Masked and cast below.
+            filtered_proj = fft.centered_ifft2(filtered_fpadded_proj)
 
             # The resulting image should be real.
-            if xp.linalg.norm(p2.imag) / xp.linalg.norm(p2) > 1e-13:
+            if (
+                xp.linalg.norm(filtered_proj.imag) / xp.linalg.norm(filtered_proj)
+                > 1e-13
+            ):
                 raise RuntimeError("Whitened image has strong imaginary component.")
 
-            if L % 2 == 1:
-                p2 = p2[k - j - 1 : k + j, k - j - 1 : k + j].real
-            else:
-                p2 = p2[k - j - 1 : k + j - 1, k - j - 1 : k + j - 1].real
+            filtered_proj = filtered_proj[slc, slc].real
 
             # Assign the resulting image.
-            res[i] = xp.asnumpy(p2)
+            res[i] = xp.asnumpy(filtered_proj)
 
         return Image(res)
 

From 0a3c0691405e41d689635fbe64e1924b361686ba Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Thu, 6 Feb 2025 10:47:57 -0500
Subject: [PATCH 125/184] Use gauss 2d for gauss window

---
 src/aspire/utils/misc.py | 8 +-------
 1 file changed, 1 insertion(+), 7 deletions(-)

diff --git a/src/aspire/utils/misc.py b/src/aspire/utils/misc.py
index 3cae04ee1a..c7d0637cc8 100644
--- a/src/aspire/utils/misc.py
+++ b/src/aspire/utils/misc.py
@@ -368,19 +368,13 @@ def gaussian_window(L, max_d, alpha=3.0, dtype=np.float64):
     reciprical of the standard deviation of the Gaussian window.
     See Harris 78.
 
-    When `alpha=1`, this function should be equivalent to
-    `gaussian_2d(size=2*L-1, sigma=max_d)`.
-
     :param L: Number of radial pixels
     :param max_d: Width of Gaussian (stddev)
     :param alpha: Reciprical of stddev of window
     :return: Numpy array with shape `(2L-1, 2L-1)`x
     """
 
-    X, Y = np.mgrid[-(L - 1) : L, -(L - 1) : L]  # -(L-1) to (L-1) inclusive
-    X = X.astype(dtype, copy=False)
-    Y = Y.astype(dtype, copy=False)
-    W = np.exp(-alpha * (X**2 + Y**2) / (2 * max_d**2))
+    W = gaussian_2d(size=2 * L - 1, sigma=max_d / np.sqrt(alpha), dtype=dtype)
 
     return W
 

From 04161b00851c37e2467b31e36fdaaa394ccc7e1e Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Thu, 6 Feb 2025 11:01:05 -0500
Subject: [PATCH 126/184] Add PSD size check

---
 src/aspire/image/image.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/src/aspire/image/image.py b/src/aspire/image/image.py
index 7d30a86752..e03de5b237 100644
--- a/src/aspire/image/image.py
+++ b/src/aspire/image/image.py
@@ -426,6 +426,7 @@ def legacy_whiten(self, psd, delta):
         the instance `dtype`.
 
         :param psd: PSD as computed by `LegacyNoiseEstimator`.
+            `psd` in this case is shape (2 * self.src.L - 1, 2 * self.src.L - 1).
         :param delta: Threshold used to determine which frequencies to whiten
             and which to set to zero. By default all `sqrt(psd)` values
             less than `delta` are zeroed out in the whitening filter.
@@ -436,6 +437,11 @@ def legacy_whiten(self, psd, delta):
         K = psd.shape[-1]
         k = int(np.ceil(K / 2))
 
+        # Check PSD
+        shp = (2 * L - 1, 2 * L - 1)
+        if psd.shape != shp:
+            raise RuntimeError(f"Incorrect PSD shape {psd.shape}, expectect {shp}.")
+
         # Create result array
         res = np.empty((n, L, L), dtype=self.dtype)
 

From 6c7f65afb8a6d898120500bd833ebca62e565c9b Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Thu, 6 Feb 2025 11:13:49 -0500
Subject: [PATCH 127/184] Make PSD normalization optional

---
 src/aspire/noise/noise.py | 67 +++++++++++++++++++++++----------------
 1 file changed, 40 insertions(+), 27 deletions(-)

diff --git a/src/aspire/noise/noise.py b/src/aspire/noise/noise.py
index 4f097e2b20..a049544545 100644
--- a/src/aspire/noise/noise.py
+++ b/src/aspire/noise/noise.py
@@ -537,7 +537,7 @@ def _estimate_power_spectrum_distribution_1d(
 
     @staticmethod
     def _estimate_power_spectrum_distribution_2d(
-        images, samples_idx, max_d=None, batch_size=512
+        images, samples_idx, max_d=None, batch_size=512, normalize_psd=False
     ):
         """
         Estimate the 2D isotropic power spectrum of `images`.
@@ -551,6 +551,11 @@ def _estimate_power_spectrum_distribution_2d(
         :param max_d: Max computed correlation distance in pixels.
            Default of `None` yields `np.floor(L / 3)`.
         :param batch_size:  The size of the batches in which to compute the variance estimate.
+        :normalize_psd:  Optionally normalize returned PSD.
+            Disabled by default because it will typiccally be
+            renormalized later in preperation for the convolution
+            application in `Image.legacy_whiten`.
+            Enable to reproduce legacy PSD.
         :return:
             - 2D PSD array
             - Radial PSD array
@@ -601,32 +606,40 @@ def _estimate_power_spectrum_distribution_2d(
             logger.warning(f"Large imaginary components in P2 {err}.")
         P2 = P2.real
 
-        # Normalize the power spectrum P2. The power spectrum is normalized such
-        # that its energy is equal to the average energy of the noise samples used
-        # to estimate it.
-
-        E = 0.0  # Total energy of the noise samples used to estimate the power spectrum.
-        samples = xp.zeros((batch_size, L, L), dtype=np.float64)
-        for start in trange(0, n_img, batch_size, desc="Estimating image noise energy"):
-            end = min(n_img, start + batch_size)
-            cnt = end - start
-
-            samples[:cnt, _samples_idx] = images[start:end].asnumpy()[0][samples_idx]
-            E += xp.sum(
-                (samples[:cnt] - xp.mean(samples[:cnt], axis=(1, 2)).reshape(cnt, 1, 1))
-                ** 2
-            )
-        # Mean energy of the noise samples
-        n_samples_per_img = xp.count_nonzero(_samples_idx)
-        meanE = E / (n_samples_per_img * n_img)
-
-        # Normalize P2 such that its mean energy is preserved and is equal to
-        # meanE, that is, mean(P2)==meanE. That way the mean energy does not
-        # go down if the number of pixels is artifically changed (say be
-        # upsampling, downsampling, or cropping). Note that P2 is already in
-        # units of energy, and so the total energy is given by sum(P2) and
-        # not by norm(P2).
-        P2 = P2 / xp.sum(P2) * meanE * P2.size
+        if normalize_psd:
+            # Normalize the power spectrum P2. The power spectrum is normalized such
+            # that its energy is equal to the average energy of the noise samples used
+            # to estimate it.
+
+            E = 0.0  # Total energy of the noise samples used to estimate the power spectrum.
+            samples = xp.zeros((batch_size, L, L), dtype=np.float64)
+            for start in trange(
+                0, n_img, batch_size, desc="Estimating image noise energy"
+            ):
+                end = min(n_img, start + batch_size)
+                cnt = end - start
+
+                samples[:cnt, _samples_idx] = images[start:end].asnumpy()[0][
+                    samples_idx
+                ]
+                E += xp.sum(
+                    (
+                        samples[:cnt]
+                        - xp.mean(samples[:cnt], axis=(1, 2)).reshape(cnt, 1, 1)
+                    )
+                    ** 2
+                )
+            # Mean energy of the noise samples
+            n_samples_per_img = xp.count_nonzero(_samples_idx)
+            meanE = E / (n_samples_per_img * n_img)
+
+            # Normalize P2 such that its mean energy is preserved and is equal to
+            # meanE, that is, mean(P2)==meanE. That way the mean energy does not
+            # go down if the number of pixels is artifically changed (say be
+            # upsampling, downsampling, or cropping). Note that P2 is already in
+            # units of energy, and so the total energy is given by sum(P2) and
+            # not by norm(P2).
+            P2 = P2 / xp.sum(P2) * meanE * P2.size
 
         # Check that P2 has no negative values.
         # Due to the truncation of the Gaussian window, we get small negative

From 1c633c8cd923ba9683b4dca953dcb5b736cffdb0 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Mon, 10 Feb 2025 15:15:25 -0500
Subject: [PATCH 128/184] changed tmp var from mask_padded to buf_padded

---
 src/aspire/noise/noise.py | 16 ++++++++--------
 1 file changed, 8 insertions(+), 8 deletions(-)

diff --git a/src/aspire/noise/noise.py b/src/aspire/noise/noise.py
index a049544545..eb3aed3856 100644
--- a/src/aspire/noise/noise.py
+++ b/src/aspire/noise/noise.py
@@ -481,19 +481,19 @@ def _estimate_power_spectrum_distribution_1d(
         # Compute Ncorr using a constant unit image.
         mask = xp.zeros((L, L))
         mask[samples_idx] = 1
-        mask_padded = xp.zeros((batch_size, 2 * L + 1, 2 * L + 1))  # pad
-        mask_padded[0, :L, :L] = mask
+        buf_padded = xp.zeros((batch_size, 2 * L + 1, 2 * L + 1))  # pad
+        buf_padded[0, :L, :L] = mask
         # MATLAB code internally detects/implicitly casts,
         #   we explicitly call rfft2/irfft2.
-        fmask_padded = fft.rfft2(mask_padded[0])
+        fbuf_padded = fft.rfft2(buf_padded[0])
         n_mask_pairs = fft.irfft2(
-            fmask_padded * fmask_padded.conj(), s=mask_padded.shape[1:]
+            fbuf_padded * fbuf_padded.conj(), s=buf_padded.shape[1:]
         )
         n_mask_pairs = n_mask_pairs[: max_d + 1, : max_d + 1]  # crop
         n_mask_pairs = xp.round(n_mask_pairs)
 
         samples = xp.zeros((batch_size, L, L))
-        mask_padded[0, :, :] = 0  # reset mask_padded
+        buf_padded[0, :, :] = 0  # reset buf_padded
         for start in trange(
             0, n_img, batch_size, desc="Processing image autocorrelations"
         ):
@@ -506,11 +506,11 @@ def _estimate_power_spectrum_distribution_1d(
             # over images twice.
 
             # Compute non-periodic autocorrelation
-            mask_padded[:count, :L, :L] = samples[:count]  # pad
+            buf_padded[:count, :L, :L] = samples[:count]  # pad
             # MATLAB code internally detects/implicitly casts,
             #   we explicitly call rfft2/irfft2.
-            fmask_padded = fft.rfft2(mask_padded[:count])
-            s = fft.irfft2(fmask_padded * fmask_padded.conj(), s=mask_padded.shape[1:])
+            fbuf_padded = fft.rfft2(buf_padded[:count])
+            s = fft.irfft2(fbuf_padded * fbuf_padded.conj(), s=buf_padded.shape[1:])
             s = s[:, 0 : max_d + 1, 0 : max_d + 1]  # crop
 
             # Accumulate all autocorrelation values R[k1,k2] such that

From 054ddf48691f60a16ef9f042ca13a5ffe5138332 Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Thu, 16 Jan 2025 15:43:40 -0500
Subject: [PATCH 129/184] cache emdb_2660 in CI

---
 .github/workflows/workflow.yml | 48 +++++++++++++++++++++++++++++++++-
 1 file changed, 47 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/workflow.yml b/.github/workflows/workflow.yml
index fc31e18968..ee795c7de4 100644
--- a/.github/workflows/workflow.yml
+++ b/.github/workflows/workflow.yml
@@ -58,8 +58,40 @@ jobs:
       with:
         token: ${{ secrets.CODECOV_TOKEN }}
 
-  conda-build:
+  data-cache:
     needs: check
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - name: Setup Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: '3.9'
+      - name: Install Dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install -e ".[dev]"
+      - name: Create Custom Cache Directory Config
+        run: |
+          CACHE_DIR=/tmp/ASPIRE-data
+          mkdir -p ${CACHE_DIR}
+          echo "common:" > ${CACHE_DIR}/config.yaml
+          echo "    cache_dir: ${CACHE_DIR}" >> ${CACHE_DIR}/config.yaml
+          echo "Logging config.yaml for verification:"
+          cat ${CACHE_DIR}/config.yaml
+          echo "POOCH_CACHE_DIR=${CACHE_DIR}" >> $GITHUB_ENV
+      - name: Download EMDB 2660
+        run: |
+          python -c \
+          "from aspire.downloader import emdb_2660; emdb_2660()"
+      - name: Create Cache
+        uses: actions/cache@v4
+        with:
+          path: /tmp/ASPIRE-data
+          key: ${{ runner.os }}-cached-data
+
+  conda-build:
+    needs: [check, data-cache]
     runs-on: ${{ matrix.os }}
     # Only run on review ready pull_requests
     if: ${{ github.event_name == 'pull_request' && github.event.pull_request.draft == false }}
@@ -90,6 +122,11 @@ jobs:
         activate-environment: aspire
         environment-file: environment-${{ matrix.backend }}.yml
         auto-activate-base: false
+    - name: Restore Cache
+      uses: actions/cache@v4
+      with:
+        path: /tmp/ASPIRE-data
+        key: ${{ runner.os }}-cached-data
     - name: Complete Install and Log Environment ${{ matrix.os }} Python ${{ matrix.python-version }}
       run: |
         conda info
@@ -97,6 +134,15 @@ jobs:
         pip install -e ".[dev]"
         pip freeze
         python -c "import numpy; numpy.show_config()"
+    - name: Restore Cache
+      uses: actions/cache@v4
+      with:
+        path: /tmp/ASPIRE-data
+        key: ${{ runner.os }}-cached-data
+    - name: Set Cache Directory
+      run: |
+        mkdir -p $HOME/.config/ASPIRE
+        echo -e "common:\n  cache_dir: /tmp/ASPIRE-data" > $HOME/.config/ASPIRE/config.yaml
     - name: Execute Pytest Conda ${{ matrix.os }} Python ${{ matrix.python-version }}
       run: |
         export OMP_NUM_THREADS=2

From 0b2009c523edf0ab73548959e3ab0e0c380b5459 Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Thu, 16 Jan 2025 15:49:33 -0500
Subject: [PATCH 130/184] Finalized workflow. Squashed commits.

---
 .github/workflows/workflow.yml | 200 +++++++++++++++++++++------------
 1 file changed, 128 insertions(+), 72 deletions(-)

diff --git a/.github/workflows/workflow.yml b/.github/workflows/workflow.yml
index ee795c7de4..a53d6f8ffc 100644
--- a/.github/workflows/workflow.yml
+++ b/.github/workflows/workflow.yml
@@ -21,9 +21,69 @@ jobs:
     - name: Run Tox Check
       run: tox -e check
 
-  build:
+  data-cache:
     needs: check
     runs-on: ubuntu-latest
+    outputs:
+      cache_hash: ${{ steps.compute-cache-hash.outputs.cache_hash }}
+    steps:
+      - uses: actions/checkout@v4
+      - name: Setup Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: '3.9'
+      - name: Install Dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install -e ".[dev]"
+      - name: Restore Cache
+        uses: actions/cache@v4
+        with:
+          key: cached-data-${{ hashFiles('**/registry.py', '**/*workflow.yml') }}
+          restore-keys: |
+            cached-data-
+          path: .github_cache/ASPIRE-data
+          enableCrossOsArchive: true
+      - name: Create Cache Directory
+        run: |
+          # Create cache and config directories
+          mkdir -p .github_cache/ASPIRE-data
+          chmod -R 777 .github_cache/ASPIRE-data
+          echo "cache:" > config.yaml
+          echo "    cache_dir: .github_cache/ASPIRE-data" >> config.yaml
+          echo "Logging config.yaml for verification:"
+          cat config.yaml
+      - name: Download Cache Files
+        run: |
+          export ASPIREDIR=.
+          python -c "
+          from aspire.downloader import emdb_2660, simulated_channelspin
+          emdb_2660()
+          simulated_channelspin()
+          "
+      - name: Compute Cache Directory Hash
+        id: compute-cache-hash
+        run: |
+          echo "Computing hash for .github_cache/ASPIRE-data..."
+          # Compute a hash on the sorted file listing.
+          cache_hash=$(ls -1 .github_cache/ASPIRE-data/** | md5sum)
+          echo "Computed cache hash: $cache_hash"
+          # Expose the computed hash to subsequent steps/jobs.
+          echo "cache_hash=${cache_hash}" >> $GITHUB_OUTPUT
+      - name: Verify Cache Directory Before Saving
+        run: |
+          ls -lhR .github_cache/ASPIRE-data
+          [ -f config.yaml ]
+      - name: Save Cache
+        uses: actions/cache@v4
+        with:
+          key: cached-data-${{ hashFiles('**/registry.py', '**/*workflow.yml') }}
+          path: .github_cache/ASPIRE-data
+          enableCrossOsArchive: true
+
+  build:
+    needs: [check, data-cache]
+    runs-on: ubuntu-latest
     # Run on every code push, but only on review ready PRs
     if: ${{ github.event_name == 'push' || github.event.pull_request.draft == false }}    
     strategy:
@@ -39,7 +99,6 @@ jobs:
           - python-version: '3.9'
             pyenv: pip,docs
 
-
     steps:
     - uses: actions/checkout@v4
     - name: Set up Python ${{ matrix.python-version }} ${{ matrix.pyenv }}
@@ -51,45 +110,41 @@ jobs:
         pip install tox tox-gh-actions
         # Optional packages
         pip install pyfftw  # `test_fft` runs for pyfftw when installed
+    - name: Restore Cache
+      uses: actions/cache@v4
+      with:
+        key: cached-data-${{ hashFiles('**/registry.py', '**/*workflow.yml') }}
+        restore-keys: |
+          cached-data-
+        path: .github_cache/ASPIRE-data
+        enableCrossOsArchive: true
+    - name: Set Cache Directory
+      run: |
+        echo "cache:" > config.yaml
+        echo "    cache_dir: .github_cache/ASPIRE-data" >> config.yaml
+    - name: Verify Restored Cache Directory
+      run: |
+        ls -lhR .github_cache/ASPIRE-data
+        [ -f config.yaml ]
     - name: Test with tox
-      run: tox --skip-missing-interpreters false -e py${{ matrix.python-version }}-${{ matrix.pyenv }}
+      run: |
+        export ASPIREDIR=.
+        tox --skip-missing-interpreters false -e py${{ matrix.python-version }}-${{ matrix.pyenv }}
+    - name: Validate Cache Directory Hash
+      run: |
+        echo "Computing hash for .github_cache/ASPIRE-data..."
+        new_hash=$(ls -1 .github_cache/ASPIRE-data/** | md5sum)
+        echo "Hash from data-cache job: ${{ needs.data-cache.outputs.cache_hash }}"
+        echo "Computed hash now: $new_hash"
+        if [ "${{ needs.data-cache.outputs.cache_hash }}" != "$new_hash" ]; then
+          echo "Error: Cache directory hash has changed!"
+          exit 1
+        fi
     - name: Upload Coverage to CodeCov
       uses: codecov/codecov-action@v4
       with:
         token: ${{ secrets.CODECOV_TOKEN }}
 
-  data-cache:
-    needs: check
-    runs-on: ubuntu-latest
-    steps:
-      - uses: actions/checkout@v4
-      - name: Setup Python
-        uses: actions/setup-python@v5
-        with:
-          python-version: '3.9'
-      - name: Install Dependencies
-        run: |
-          python -m pip install --upgrade pip
-          pip install -e ".[dev]"
-      - name: Create Custom Cache Directory Config
-        run: |
-          CACHE_DIR=/tmp/ASPIRE-data
-          mkdir -p ${CACHE_DIR}
-          echo "common:" > ${CACHE_DIR}/config.yaml
-          echo "    cache_dir: ${CACHE_DIR}" >> ${CACHE_DIR}/config.yaml
-          echo "Logging config.yaml for verification:"
-          cat ${CACHE_DIR}/config.yaml
-          echo "POOCH_CACHE_DIR=${CACHE_DIR}" >> $GITHUB_ENV
-      - name: Download EMDB 2660
-        run: |
-          python -c \
-          "from aspire.downloader import emdb_2660; emdb_2660()"
-      - name: Create Cache
-        uses: actions/cache@v4
-        with:
-          path: /tmp/ASPIRE-data
-          key: ${{ runner.os }}-cached-data
-
   conda-build:
     needs: [check, data-cache]
     runs-on: ${{ matrix.os }}
@@ -122,11 +177,6 @@ jobs:
         activate-environment: aspire
         environment-file: environment-${{ matrix.backend }}.yml
         auto-activate-base: false
-    - name: Restore Cache
-      uses: actions/cache@v4
-      with:
-        path: /tmp/ASPIRE-data
-        key: ${{ runner.os }}-cached-data
     - name: Complete Install and Log Environment ${{ matrix.os }} Python ${{ matrix.python-version }}
       run: |
         conda info
@@ -137,14 +187,18 @@ jobs:
     - name: Restore Cache
       uses: actions/cache@v4
       with:
-        path: /tmp/ASPIRE-data
-        key: ${{ runner.os }}-cached-data
+        key: cached-data-${{ hashFiles('**/registry.py', '**/*workflow.yml') }}
+        restore-keys: |
+          cached-data-
+        path: .github_cache/ASPIRE-data
+        enableCrossOsArchive: true
     - name: Set Cache Directory
       run: |
-        mkdir -p $HOME/.config/ASPIRE
-        echo -e "common:\n  cache_dir: /tmp/ASPIRE-data" > $HOME/.config/ASPIRE/config.yaml
+        echo "cache:" > config.yaml
+        echo "    cache_dir: .github_cache/ASPIRE-data" >> config.yaml
     - name: Execute Pytest Conda ${{ matrix.os }} Python ${{ matrix.python-version }}
       run: |
+        export ASPIREDIR=.
         export OMP_NUM_THREADS=2
         # -n runs test in parallel using pytest-xdist
         pytest -n2 --durations=50 -s
@@ -215,30 +269,6 @@ jobs:
     - name: Cleanup
       run: rm -rf ${{ env.WORK_DIR }}
 
-  # Create cache and download data for Github Actions CI.
-  data-cache:
-    needs: check
-    runs-on: ubuntu-latest
-    steps:
-      - uses: actions/checkout@v4
-      - name: Setup Python
-        uses: actions/setup-python@v5
-        with:
-          python-version: '3.9'
-      - name: Install Dependencies
-        run: |
-          python -m pip install --upgrade pip
-          pip install -e ".[dev]"
-      - name: Create Cache
-        uses: actions/cache@v4
-        with:
-          path: ~/.cache/ASPIRE-data
-          key: ${{ runner.os }}-cached-data
-      - name: Download EMDB 2660
-        run: |
-          python -c \
-          "from aspire.downloader import emdb_2660; emdb_2660()"
-
   # Build branch's docs and gallery.
   docs:
     needs: [check, data-cache]
@@ -256,10 +286,18 @@ jobs:
       - name: Restore Cache
         uses: actions/cache@v4
         with:
-          path: ~/.cache/ASPIRE-data
-          key: ${{ runner.os }}-cached-data
+          key: cached-data-${{ hashFiles('**/registry.py', '**/*workflow.yml') }}
+          restore-keys: |
+            cached-data-
+          path: .github_cache/ASPIRE-data
+          enableCrossOsArchive: true
+      - name: Set Cache Directory
+        run: |
+          echo "cache:" > config.yaml
+          echo "    cache_dir: .github_cache/ASPIRE-data" >> config.yaml
       - name: Build Sphinx docs
         run: |
+          export ASPIREDIR=.
           make distclean
           sphinx-apidoc -f -o ./source ../src -H Modules
           make html
@@ -272,7 +310,7 @@ jobs:
           retention-days: 7
 
   osx_arm:
-    needs: check
+    needs: [check, data-cache]
     runs-on: macos-14
     # Run on every code push, but only on review ready PRs
     if: ${{ github.event_name == 'push' || github.event.pull_request.draft == false }}
@@ -286,5 +324,23 @@ jobs:
         python --version
         pip install -e ".[dev]"  # install aspire
         pip freeze
+    - name: Restore Cache
+      uses: actions/cache@v4
+      with:
+        key: cached-data-${{ hashFiles('**/registry.py', '**/*workflow.yml') }}
+        restore-keys: |
+          cached-data-
+        path: .github_cache/ASPIRE-data
+        enableCrossOsArchive: true
+    - name: Set Cache Directory
+      run: |
+        echo "cache:" > config.yaml
+        echo "    cache_dir: .github_cache/ASPIRE-data" >> config.yaml
+    - name: Verify Restored Cache Directory
+      run: |
+        ls -lhR .github_cache/ASPIRE-data
+        [ -f config.yaml ]
     - name: Test
-      run: python -m pytest -n3 --durations=50
+      run: |
+        export ASPIREDIR=.
+        python -m pytest -n3 --durations=50

From 39fcaa3c3b071935007825b1dfab05606ae6f926 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Tue, 18 Feb 2025 13:14:23 -0500
Subject: [PATCH 131/184] Attempt bug fix enforcing pure rotation for CL
 methods

---
 src/aspire/abinitio/commonline_sync3n.py | 2 +-
 src/aspire/utils/matrix.py               | 7 ++++---
 2 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/src/aspire/abinitio/commonline_sync3n.py b/src/aspire/abinitio/commonline_sync3n.py
index 6b89e4bce2..841f9dfceb 100644
--- a/src/aspire/abinitio/commonline_sync3n.py
+++ b/src/aspire/abinitio/commonline_sync3n.py
@@ -264,7 +264,7 @@ def _sync3n_S_to_rot(self, S, W=None, n_eigs=4):
         rotations = v.reshape(self.n_img, 3, 3).transpose(0, 2, 1)
 
         # Enforce we are returning actual rotations
-        rotations = nearest_rotations(rotations, allow_reflection=True)
+        rotations = nearest_rotations(rotations, allow_reflection=False)
 
         return rotations.astype(self.dtype, copy=False)
 
diff --git a/src/aspire/utils/matrix.py b/src/aspire/utils/matrix.py
index e0322e966e..5b683c9e68 100644
--- a/src/aspire/utils/matrix.py
+++ b/src/aspire/utils/matrix.py
@@ -321,10 +321,11 @@ def nearest_rotations(A, allow_reflection=False):
     U, _, V = np.linalg.svd(A)
 
     if not allow_reflection:
-        # If det(U)*det(V) = -1, we negate the third singular value to
-        # ensure we have a rotation.
+        # If det(U)*det(V) = -1, apply reflection about the origin
         neg_det_idx = np.linalg.det(U) * np.linalg.det(V) < 0
-        U[neg_det_idx] = U[neg_det_idx] @ np.diag((1, 1, -1)).astype(dtype, copy=False)
+        U[neg_det_idx] = U[neg_det_idx] @ np.diag((-1, -1, -1)).astype(
+            dtype, copy=False
+        )
 
     rots = U @ V
 

From 65a7ec76dc4c0840e1741cbd9c310515b4975258 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Wed, 19 Feb 2025 09:35:47 -0500
Subject: [PATCH 132/184] alter Kabsch for correct 2D projections of random
 ortho matrix

---
 src/aspire/utils/matrix.py | 16 ++++++++++------
 1 file changed, 10 insertions(+), 6 deletions(-)

diff --git a/src/aspire/utils/matrix.py b/src/aspire/utils/matrix.py
index 5b683c9e68..1464ebeced 100644
--- a/src/aspire/utils/matrix.py
+++ b/src/aspire/utils/matrix.py
@@ -316,18 +316,22 @@ def nearest_rotations(A, allow_reflection=False):
             f"Array must be of shape (3, 3) or (n, 3, 3). Found shape {A.shape}."
         )
 
-    # For the singular value decomposition A = U @ S @ V,
-    # we compute the nearest rotation matrices R = U @ V.
-    U, _, V = np.linalg.svd(A)
+    # For the singular value decomposition A = U @ S @ VT,
+    # we compute the nearest rotation matrices R = U @ VT.
+    U, _, VT = np.linalg.svd(A)
 
     if not allow_reflection:
-        # If det(U)*det(V) = -1, apply reflection about the origin
-        neg_det_idx = np.linalg.det(U) * np.linalg.det(V) < 0
+        # If det(U)*det(V) = -1,
+        #   apply reflection about the origin,
+        #   rotate around projection axis.
+        r_proj = np.array([[-1, 0, 0], [0, -1, 0], [0, 0, 1]], dtype=dtype)
+        neg_det_idx = np.linalg.det(U) * np.linalg.det(VT) < 0
         U[neg_det_idx] = U[neg_det_idx] @ np.diag((-1, -1, -1)).astype(
             dtype, copy=False
         )
+        VT[neg_det_idx] = VT[neg_det_idx] @ r_proj
 
-    rots = U @ V
+    rots = U @ VT
 
     return rots.reshape(og_shape)
 

From d644c491c6b1e34412883e6f016fd7bbd9ec3052 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Wed, 19 Feb 2025 10:21:39 -0500
Subject: [PATCH 133/184] cleanup Kabsch correction for 2D projections of
 random ortho matrix

---
 src/aspire/utils/matrix.py | 23 ++++++++++++++++-------
 1 file changed, 16 insertions(+), 7 deletions(-)

diff --git a/src/aspire/utils/matrix.py b/src/aspire/utils/matrix.py
index 1464ebeced..d5d495365c 100644
--- a/src/aspire/utils/matrix.py
+++ b/src/aspire/utils/matrix.py
@@ -322,14 +322,23 @@ def nearest_rotations(A, allow_reflection=False):
 
     if not allow_reflection:
         # If det(U)*det(V) = -1,
-        #   apply reflection about the origin,
-        #   rotate around projection axis.
-        r_proj = np.array([[-1, 0, 0], [0, -1, 0], [0, 0, 1]], dtype=dtype)
+        #   we want to find a pure rotation R that is closest to the
+        #   preserving the 2D projection induced by the  orthogonal transform A.
+        #
+        #   This can be done by reflecting about the origin,
+        #   then rotating around projection axis.
+        #
+        #     R = (U @ diag([-1,-1,-1]) @ VT) @ r_projection
+        #
+        # This is accomplished by the single application of d to elements of VT.
+        #
+        #     R = (U @ diag([-1,-1,-1]) @ VT) @ diag([-1,-1,1])
+        #     R = (U * -1 @ VT) * [-1,-1,1]
+        #     R =  (U @ VT) * (-1 * [-1,-1,1])
+        #     R =  U @ (VT * [1,1,-1])
+        d = np.array([1, 1, -1], dtype=dtype)
         neg_det_idx = np.linalg.det(U) * np.linalg.det(VT) < 0
-        U[neg_det_idx] = U[neg_det_idx] @ np.diag((-1, -1, -1)).astype(
-            dtype, copy=False
-        )
-        VT[neg_det_idx] = VT[neg_det_idx] @ r_proj
+        VT[neg_det_idx] = VT[neg_det_idx] * d
 
     rots = U @ VT
 

From 5613fb1ff6abfeaf072e56f1b452fcf933133fd5 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Wed, 19 Feb 2025 10:21:50 -0500
Subject: [PATCH 134/184] add test for nearest_rotation

---
 tests/test_rotation.py | 51 +++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 50 insertions(+), 1 deletion(-)

diff --git a/tests/test_rotation.py b/tests/test_rotation.py
index e02e650bd5..ee87b11d09 100644
--- a/tests/test_rotation.py
+++ b/tests/test_rotation.py
@@ -1,10 +1,12 @@
 import logging
+from itertools import product
 
 import numpy as np
 import pytest
 from scipy.spatial.transform import Rotation as sp_rot
 
-from aspire.utils import Rotation, utest_tolerance
+from aspire.downloader import emdb_2660
+from aspire.utils import Rotation, nearest_rotations, utest_tolerance
 
 logger = logging.getLogger(__name__)
 
@@ -184,3 +186,50 @@ def test_mean_angular_distance(dtype):
     mean_ang_dist = Rotation.mean_angular_distance(rots_z, rots_id)
 
     assert np.allclose(mean_ang_dist, np.pi / 4)
+
+
+def test_rot_with_refl(dtype):
+    """
+    Given random orthogonal matrices, test nearest_rotation returns a
+    pure rotation which preserves our 2D projection operation.
+    """
+
+    N = 10  # Number of random rots
+    SEED = 123
+    L = 128  # Pixel size for projections
+    atol = 1e-7  # Picked to cover both singles and doubles
+
+    # Generate a sample of random rotations
+    random_rot_mats = Rotation.generate_random_rotations(
+        N, seed=SEED, dtype=dtype
+    ).matrices
+
+    # Sanity check we are starting with pure rotations
+    np.testing.assert_allclose(np.linalg.det(random_rot_mats), 1, atol=atol)
+
+    # Reflection operations
+    refls = np.array([[1, 1, -1], [1, -1, 1], [-1, 1, 1], [-1, -1, -1]], dtype=dtype)
+
+    # Enumerate every combination of rot and refl into M
+    M = np.empty((len(random_rot_mats) * len(refls), 3, 3), dtype=dtype)
+    for i, (r, s) in enumerate(product(random_rot_mats, refls)):
+        M[i] = r * s
+
+    # Sanity check all entries are not pure rotations
+    np.testing.assert_allclose(np.linalg.det(M), -1, atol=atol)
+
+    R = nearest_rotations(M, allow_reflection=False)
+    # Sanity check all entries are pure rotations
+    np.testing.assert_allclose(np.linalg.det(R), 1, atol=atol)
+
+    # Create a volume to use for projections
+    v = emdb_2660().astype(dtype).downsample(L)
+
+    # Project using set of transforms M
+    ref_projections = v.project(M)
+
+    # Project using set of rotations R
+    rot_projections = v.project(R)
+
+    # Validate we are equivalent
+    np.testing.assert_allclose(rot_projections, ref_projections, atol=atol)

From 442de93305341ae5688dc1efa530e9e96a78ee64 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Fri, 14 Feb 2025 14:47:22 -0500
Subject: [PATCH 135/184] add filter/preprocess specific test of indexed source

---
 tests/test_indexed_source.py | 99 ++++++++++++++++++++++++++++++++++++
 1 file changed, 99 insertions(+)

diff --git a/tests/test_indexed_source.py b/tests/test_indexed_source.py
index 3092ed16f4..4350e99d7a 100644
--- a/tests/test_indexed_source.py
+++ b/tests/test_indexed_source.py
@@ -3,7 +3,10 @@
 import numpy as np
 import pytest
 
+from aspire.downloader import emdb_8012
+from aspire.operators import CTFFilter
 from aspire.source import Simulation
+from aspire.utils import Rotation
 
 logger = logging.getLogger(__name__)
 
@@ -49,3 +52,99 @@ def test_repr(sim_fixture):
 
     # Check index counts are mentioned in the repr
     assert f"{sim2.n} of {sim.n}" in repr(sim2)
+
+
+@pytest.mark.expensive
+def test_filter_mapping():
+    """
+    This test is designed to ensure that `unique_filters` and `filter_indices`
+    are being remapped correctly upon slicing.
+
+    Additionally it tests that a realistic preprocessing pipeline is equivalent.
+    """
+
+    # Generate N projection images,
+    #   using N//2 rotations and
+    #   N//2 ctf filters such that images[0::2] == images[1::2].
+    N = 100
+    SEED = 1234
+    DT = np.float64
+    DS = 129
+
+    v = emdb_8012().astype(DT)
+
+    # Generate N//2 rotations
+    rots = Rotation.generate_random_rotations(N // 2, dtype=DT, seed=SEED)
+    angles = Rotation(np.repeat(rots, 2, axis=0)).angles
+
+    # Generate N//2 rotations and repeat indices
+    defoci = np.linspace(1000, 25000, N // 2)
+    ctf_filters = [
+        CTFFilter(
+            v.pixel_size,
+            200,
+            defocus_u=defoci[d],
+            defocus_v=defoci[-d],
+            defocus_ang=np.pi / (N // 2) * d,
+            Cs=2.0,
+            alpha=0.1,
+        )
+        for d in range(N // 2)
+    ]
+    ctf_indices = np.repeat(np.arange(N // 2), 2)
+
+    # Construct the source
+    src = Simulation(
+        vols=v,
+        n=N,
+        dtype=DT,
+        seed=SEED,
+        unique_filters=ctf_filters,
+        filter_indices=ctf_indices,
+        angles=angles,
+        offsets=0,
+        amplitudes=1,
+    ).cache()
+
+    srcA = src[0::2]
+    srcB = src[1::2]
+
+    # Sanity check the images before proceeding
+    np.testing.assert_allclose(srcA.images[:], src.images[::2])
+    np.testing.assert_allclose(srcB.images[:], src.images[1::2])
+    # Confirm the intention of the test
+    np.testing.assert_allclose(srcB.images[:], srcA.images[:])
+
+    # Preprocess the `src` stack
+    pp = (
+        src.phase_flip()
+        .downsample(DS)
+        .normalize_background()
+        .legacy_whiten()
+        .invert_contrast()
+        .cache()
+    )
+
+    # Preprocess the indexed sources
+    ppA = (
+        srcA.phase_flip()
+        .downsample(DS)
+        .normalize_background()
+        .legacy_whiten()
+        .invert_contrast()
+        .cache()
+    )
+    ppB = (
+        srcB.phase_flip()
+        .downsample(DS)
+        .normalize_background()
+        .legacy_whiten()
+        .invert_contrast()
+        .cache()
+    )
+
+    # Confirm we match the original images
+    np.testing.assert_allclose(ppA.images[:], pp.images[::2])
+    np.testing.assert_allclose(ppB.images[:], pp.images[1::2])
+    # Confirm A and B are equivalent
+    np.testing.assert_allclose(ppB.images[:], ppA.images[:])

From 38d1014bbd2b3b47d03dfb5afdd94f127913b64a Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Fri, 14 Feb 2025 13:48:52 -0500
Subject: [PATCH 136/184] fix the indexed source filter mapping

---
 src/aspire/source/image.py | 17 +++++++++++++----
 1 file changed, 13 insertions(+), 4 deletions(-)

diff --git a/src/aspire/source/image.py b/src/aspire/source/image.py
index 6eff83d23d..6be0a3e41e 100644
--- a/src/aspire/source/image.py
+++ b/src/aspire/source/image.py
@@ -1519,10 +1519,19 @@ def __init__(self, src, indices, memory=None):
             pixel_size=src.pixel_size,
         )
 
-        # Create filter indices, these are required to pass unharmed through filter eval code
-        #   that is potentially called by other methods later.
-        self.filter_indices = np.zeros(self.n, dtype=int)
-        self.unique_filters = [IdentityFilter()]
+        if src.unique_filters:
+            # Remap the filter indices to be unique.
+            #   Removes duplicates and filters that are unused in new source.
+            _filter_indices = src.filter_indices[self.index_map]
+            # _unq[_inv] reconstructs _filter_indices
+            _unq, _inv = np.unique(_filter_indices, return_inverse=True)
+            # Repack unique_filters
+            self.filter_indices = _inv
+            self.unique_filters = [copy.copy(src.unique_filters[i]) for i in _unq]
+        else:
+            # Pass through the None case
+            self.unique_filters = src.unique_filters
+            self.filter_indices = np.zeros(self.n)
 
         # Any further operations should not mutate this instance.
         self._mutable = False

From c681a7ae39585b75ed03c70f1b70adebcc2c6b5a Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Fri, 14 Feb 2025 16:08:09 -0500
Subject: [PATCH 137/184] fix inds dtype

---
 src/aspire/source/image.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/aspire/source/image.py b/src/aspire/source/image.py
index 6be0a3e41e..509bf4eb10 100644
--- a/src/aspire/source/image.py
+++ b/src/aspire/source/image.py
@@ -1531,7 +1531,7 @@ def __init__(self, src, indices, memory=None):
         else:
             # Pass through the None case
             self.unique_filters = src.unique_filters
-            self.filter_indices = np.zeros(self.n)
+            self.filter_indices = np.zeros(self.n, dtype=int)
 
         # Any further operations should not mutate this instance.
         self._mutable = False

From 9435cc8e9f4e333c841feec995927ae1b8cb1e0c Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Tue, 18 Feb 2025 08:54:32 -0500
Subject: [PATCH 138/184] tmp unmark test, test CI

---
 tests/test_indexed_source.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/test_indexed_source.py b/tests/test_indexed_source.py
index 4350e99d7a..34a0fbb8b7 100644
--- a/tests/test_indexed_source.py
+++ b/tests/test_indexed_source.py
@@ -54,7 +54,7 @@ def test_repr(sim_fixture):
     assert f"{sim2.n} of {sim.n}" in repr(sim2)
 
 
-@pytest.mark.expensive
+# @pytest.mark.expensive
 def test_filter_mapping():
     """
     This test is designed to ensure that `unique_filters` and `filter_indices`

From 134cb1258d6a0e2f2902a2cb5243951c5659a93e Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Tue, 18 Feb 2025 09:24:31 -0500
Subject: [PATCH 139/184] revert

---
 tests/test_indexed_source.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/test_indexed_source.py b/tests/test_indexed_source.py
index 34a0fbb8b7..4350e99d7a 100644
--- a/tests/test_indexed_source.py
+++ b/tests/test_indexed_source.py
@@ -54,7 +54,7 @@ def test_repr(sim_fixture):
     assert f"{sim2.n} of {sim.n}" in repr(sim2)
 
 
-# @pytest.mark.expensive
+@pytest.mark.expensive
 def test_filter_mapping():
     """
     This test is designed to ensure that `unique_filters` and `filter_indices`

From 9b39c17691bb4a0d39bba9e8f0c847f743cf51f1 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Wed, 19 Feb 2025 14:33:20 -0500
Subject: [PATCH 140/184] rescale internal pixel size after downsample

---
 src/aspire/source/image.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/aspire/source/image.py b/src/aspire/source/image.py
index 509bf4eb10..5ca5ba8374 100644
--- a/src/aspire/source/image.py
+++ b/src/aspire/source/image.py
@@ -775,6 +775,7 @@ def downsample(self, L):
         ds_factor = self.L / L
         self.unique_filters = [f.scale(ds_factor) for f in self.unique_filters]
         self.offsets /= ds_factor
+        self.pixel_size *= ds_factor
 
         self.L = L
 

From c7e6f7cb372413b3db4c4ecc7818bb8c1d1b03ae Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Wed, 19 Feb 2025 14:34:39 -0500
Subject: [PATCH 141/184] add more source saving tests

---
 tests/test_indexed_source.py | 85 ++++++++++++++++++++++++++++++++++--
 1 file changed, 82 insertions(+), 3 deletions(-)

diff --git a/tests/test_indexed_source.py b/tests/test_indexed_source.py
index 4350e99d7a..413bea77be 100644
--- a/tests/test_indexed_source.py
+++ b/tests/test_indexed_source.py
@@ -1,11 +1,13 @@
 import logging
+import os
+import tempfile
 
 import numpy as np
 import pytest
 
 from aspire.downloader import emdb_8012
 from aspire.operators import CTFFilter
-from aspire.source import Simulation
+from aspire.source import RelionSource, Simulation
 from aspire.utils import Rotation
 
 logger = logging.getLogger(__name__)
@@ -110,7 +112,7 @@ def test_filter_mapping():
     srcB = src[1::2]
 
     # Sanity check the images before proceeding
-    np.testing.assert_allclose(srcA.images[:], src.images[::2])
+    np.testing.assert_allclose(srcA.images[:], src.images[0::2])
     np.testing.assert_allclose(srcB.images[:], src.images[1::2])
     # Confirm the intention of the test
     np.testing.assert_allclose(srcB.images[:], srcA.images[:])
@@ -144,7 +146,84 @@ def test_filter_mapping():
     )
 
     # Confirm we match the original images
-    np.testing.assert_allclose(ppA.images[:], pp.images[::2])
+    np.testing.assert_allclose(ppA.images[:], pp.images[0::2])
     np.testing.assert_allclose(ppB.images[:], pp.images[1::2])
     # Confirm A and B are equivalent
     np.testing.assert_allclose(ppB.images[:], ppA.images[:])
+
+    # Create a tmp dir for this test output
+    with tempfile.TemporaryDirectory() as tmpdir_name:
+        # Save the initial images
+        src.save(os.path.join(tmpdir_name, "src.star"))
+        srcA.save(os.path.join(tmpdir_name, "srcA.star"))
+        srcB.save(os.path.join(tmpdir_name, "srcB.star"))
+
+        # Save the preprocessed images.
+        pp.save(os.path.join(tmpdir_name, "pp.star"))
+        ppA.save(os.path.join(tmpdir_name, "ppA.star"))
+        ppB.save(os.path.join(tmpdir_name, "ppB.star"))
+
+        # Reload, assigning `pixel_size`.
+        _src = RelionSource(
+            os.path.join(tmpdir_name, "src.star"), pixel_size=src.pixel_size
+        )
+        _srcA = RelionSource(
+            os.path.join(tmpdir_name, "srcA.star"), pixel_size=srcA.pixel_size
+        )
+        _srcB = RelionSource(
+            os.path.join(tmpdir_name, "srcB.star"), pixel_size=srcB.pixel_size
+        )
+        _pp = RelionSource(
+            os.path.join(tmpdir_name, "pp.star"), pixel_size=pp.pixel_size
+        )
+        _ppA = RelionSource(
+            os.path.join(tmpdir_name, "ppA.star"), pixel_size=ppA.pixel_size
+        )
+        _ppB = RelionSource(
+            os.path.join(tmpdir_name, "ppB.star"), pixel_size=ppB.pixel_size
+        )
+
+        # Confirm reloaded sources match the source it was saved from.
+        #   This implies the equalities in the next section translate
+        #   to the original source as well
+        #   Ideally, _if everything is working_, many of these are redundant.
+        np.testing.assert_allclose(_src.images[:], src.images[:])
+        np.testing.assert_allclose(_srcA.images[:], srcA.images[:])
+        np.testing.assert_allclose(_srcB.images[:], srcB.images[:])
+        np.testing.assert_allclose(_pp.images[:], pp.images[:])
+        np.testing.assert_allclose(_ppA.images[:], ppA.images[:])
+        np.testing.assert_allclose(_ppB.images[:], ppB.images[:])
+
+        # Confirm reloading slices matches the reloading saved stack of images.
+        np.testing.assert_allclose(_srcA.images[:], _src.images[0::2])
+        np.testing.assert_allclose(_srcB.images[:], _src.images[1::2])
+        np.testing.assert_allclose(_ppA.images[:], _pp.images[0::2])
+        np.testing.assert_allclose(_ppB.images[:], _pp.images[1::2])
+        # Confirm A and B are still equivalent
+        np.testing.assert_allclose(_srcB.images[:], _srcA.images[:])
+        np.testing.assert_allclose(_ppB.images[:], _ppA.images[:])
+
+        # # Confirm pre-processing the reloaded sources matches
+        # # reloading the pre-processed sources.
+        # pp_A = (
+        #     _srcA.phase_flip()
+        #     .downsample(DS)
+        #     .normalize_background()
+        #     .legacy_whiten()
+        #     .invert_contrast()
+        #     .cache()
+        # )
+        # pp_B = (
+        #     _srcB.phase_flip()
+        #     .downsample(DS)
+        #     .normalize_background()
+        #     .legacy_whiten()
+        #     .invert_contrast()
+        #     .cache()
+        # )
+        # hrmm, that's not good
+        # breakpoint()
+        # np.testing.assert_allclose(pp_A.images[:], pp.images[0::2])
+        # np.testing.assert_allclose(pp_B.images[:], pp.images[1::2])
+        # np.testing.assert_allclose(pp_A.images[:], _pp.images[0::2])
+        # np.testing.assert_allclose(pp_B.images[:], _pp.images[1::2])

From e77783ddb042f3ec4f0993541c527c286d71a4de Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Wed, 19 Feb 2025 14:50:26 -0500
Subject: [PATCH 142/184] pass through pixel_size=None

---
 src/aspire/source/image.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/aspire/source/image.py b/src/aspire/source/image.py
index 5ca5ba8374..96dbaea30f 100644
--- a/src/aspire/source/image.py
+++ b/src/aspire/source/image.py
@@ -775,7 +775,8 @@ def downsample(self, L):
         ds_factor = self.L / L
         self.unique_filters = [f.scale(ds_factor) for f in self.unique_filters]
         self.offsets /= ds_factor
-        self.pixel_size *= ds_factor
+        if self.pixel_size is not None:
+            self.pixel_size *= ds_factor
 
         self.L = L
 

From bfe7bf8975d68ea651f2438abf720b70c9aebbb5 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Thu, 20 Feb 2025 13:00:14 -0500
Subject: [PATCH 143/184] exhaustive index-save-load testing

---
 tests/test_indexed_source.py | 64 ++++++++++++++++++------------------
 1 file changed, 32 insertions(+), 32 deletions(-)

diff --git a/tests/test_indexed_source.py b/tests/test_indexed_source.py
index 413bea77be..0ad3540706 100644
--- a/tests/test_indexed_source.py
+++ b/tests/test_indexed_source.py
@@ -146,10 +146,10 @@ def test_filter_mapping():
     )
 
     # Confirm we match the original images
-    np.testing.assert_allclose(ppA.images[:], pp.images[0::2])
-    np.testing.assert_allclose(ppB.images[:], pp.images[1::2])
+    np.testing.assert_allclose(ppA.images[:], pp.images[0::2], atol=1e-6)
+    np.testing.assert_allclose(ppB.images[:], pp.images[1::2], atol=1e-6)
     # Confirm A and B are equivalent
-    np.testing.assert_allclose(ppB.images[:], ppA.images[:])
+    np.testing.assert_allclose(ppB.images[:], ppA.images[:], atol=1e-6)
 
     # Create a tmp dir for this test output
     with tempfile.TemporaryDirectory() as tmpdir_name:
@@ -190,40 +190,40 @@ def test_filter_mapping():
         np.testing.assert_allclose(_src.images[:], src.images[:])
         np.testing.assert_allclose(_srcA.images[:], srcA.images[:])
         np.testing.assert_allclose(_srcB.images[:], srcB.images[:])
-        np.testing.assert_allclose(_pp.images[:], pp.images[:])
-        np.testing.assert_allclose(_ppA.images[:], ppA.images[:])
-        np.testing.assert_allclose(_ppB.images[:], ppB.images[:])
+        np.testing.assert_allclose(_pp.images[:], pp.images[:], atol=1e-6)
+        np.testing.assert_allclose(_ppA.images[:], ppA.images[:], atol=1e-6)
+        np.testing.assert_allclose(_ppB.images[:], ppB.images[:], atol=1e-6)
 
         # Confirm reloading slices matches the reloading saved stack of images.
         np.testing.assert_allclose(_srcA.images[:], _src.images[0::2])
         np.testing.assert_allclose(_srcB.images[:], _src.images[1::2])
-        np.testing.assert_allclose(_ppA.images[:], _pp.images[0::2])
-        np.testing.assert_allclose(_ppB.images[:], _pp.images[1::2])
+        np.testing.assert_allclose(_ppA.images[:], _pp.images[0::2], atol=1e-5)
+        np.testing.assert_allclose(_ppB.images[:], _pp.images[1::2], atol=1e-5)
         # Confirm A and B are still equivalent
         np.testing.assert_allclose(_srcB.images[:], _srcA.images[:])
         np.testing.assert_allclose(_ppB.images[:], _ppA.images[:])
 
-        # # Confirm pre-processing the reloaded sources matches
-        # # reloading the pre-processed sources.
-        # pp_A = (
-        #     _srcA.phase_flip()
-        #     .downsample(DS)
-        #     .normalize_background()
-        #     .legacy_whiten()
-        #     .invert_contrast()
-        #     .cache()
-        # )
-        # pp_B = (
-        #     _srcB.phase_flip()
-        #     .downsample(DS)
-        #     .normalize_background()
-        #     .legacy_whiten()
-        #     .invert_contrast()
-        #     .cache()
-        # )
-        # hrmm, that's not good
-        # breakpoint()
-        # np.testing.assert_allclose(pp_A.images[:], pp.images[0::2])
-        # np.testing.assert_allclose(pp_B.images[:], pp.images[1::2])
-        # np.testing.assert_allclose(pp_A.images[:], _pp.images[0::2])
-        # np.testing.assert_allclose(pp_B.images[:], _pp.images[1::2])
+        # Confirm pre-processing the reloaded sources matches
+        # reloading the pre-processed sources.
+        pp_A = (
+            _srcA.phase_flip()
+            .downsample(DS)
+            .normalize_background()
+            .legacy_whiten()
+            .invert_contrast()
+            .cache()
+        )
+        pp_B = (
+            _srcB.phase_flip()
+            .downsample(DS)
+            .normalize_background()
+            .legacy_whiten()
+            .invert_contrast()
+            .cache()
+        )
+
+
+        np.testing.assert_allclose(pp_A.images[:], pp.images[0::2], atol=0.006)
+        np.testing.assert_allclose(pp_B.images[:], pp.images[1::2], atol=0.006)
+        np.testing.assert_allclose(pp_A.images[:], _pp.images[0::2], atol=0.006)
+        np.testing.assert_allclose(pp_B.images[:], _pp.images[1::2], atol=0.006)

From 0cda2ba1f4c36a8abcc50756f67c8f7e9b6e11a4 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Thu, 20 Feb 2025 14:32:10 -0500
Subject: [PATCH 144/184] cleanup indexed source test, prep for review

---
 tests/test_indexed_source.py | 29 +++--------------------------
 1 file changed, 3 insertions(+), 26 deletions(-)

diff --git a/tests/test_indexed_source.py b/tests/test_indexed_source.py
index 0ad3540706..c7b7161a1b 100644
--- a/tests/test_indexed_source.py
+++ b/tests/test_indexed_source.py
@@ -62,7 +62,8 @@ def test_filter_mapping():
     This test is designed to ensure that `unique_filters` and `filter_indices`
     are being remapped correctly upon slicing.
 
-    Additionally it tests that a realistic preprocessing pipeline is equivalent.
+    Additionally it tests that a realistic preprocessing pipeline is equivalent
+    and can be saved then reloaded.
     """
 
     # Generate N projection images,
@@ -187,6 +188,7 @@ def test_filter_mapping():
         #   This implies the equalities in the next section translate
         #   to the original source as well
         #   Ideally, _if everything is working_, many of these are redundant.
+        #   If something fails to work, they may help pinpoint the fault.
         np.testing.assert_allclose(_src.images[:], src.images[:])
         np.testing.assert_allclose(_srcA.images[:], srcA.images[:])
         np.testing.assert_allclose(_srcB.images[:], srcB.images[:])
@@ -202,28 +204,3 @@ def test_filter_mapping():
         # Confirm A and B are still equivalent
         np.testing.assert_allclose(_srcB.images[:], _srcA.images[:])
         np.testing.assert_allclose(_ppB.images[:], _ppA.images[:])
-
-        # Confirm pre-processing the reloaded sources matches
-        # reloading the pre-processed sources.
-        pp_A = (
-            _srcA.phase_flip()
-            .downsample(DS)
-            .normalize_background()
-            .legacy_whiten()
-            .invert_contrast()
-            .cache()
-        )
-        pp_B = (
-            _srcB.phase_flip()
-            .downsample(DS)
-            .normalize_background()
-            .legacy_whiten()
-            .invert_contrast()
-            .cache()
-        )
-
-
-        np.testing.assert_allclose(pp_A.images[:], pp.images[0::2], atol=0.006)
-        np.testing.assert_allclose(pp_B.images[:], pp.images[1::2], atol=0.006)
-        np.testing.assert_allclose(pp_A.images[:], _pp.images[0::2], atol=0.006)
-        np.testing.assert_allclose(pp_B.images[:], _pp.images[1::2], atol=0.006)

From 86eada5fbb9271e301d40b211ab79ea6bc3876cf Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Wed, 26 Feb 2025 15:49:20 -0500
Subject: [PATCH 145/184] Test: verify pooch path in workflow

---
 .github/workflows/workflow.yml | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/.github/workflows/workflow.yml b/.github/workflows/workflow.yml
index a53d6f8ffc..f57ee4c8d7 100644
--- a/.github/workflows/workflow.yml
+++ b/.github/workflows/workflow.yml
@@ -129,6 +129,8 @@ jobs:
     - name: Test with tox
       run: |
         export ASPIREDIR=.
+        python -c \
+        "from aspire.downloader.data_fetcher import _data_fetcher; print(_data_fetcher.path)"
         tox --skip-missing-interpreters false -e py${{ matrix.python-version }}-${{ matrix.pyenv }}
     - name: Validate Cache Directory Hash
       run: |

From f703821a999cd24fe04312dd8fec8780c7f7069e Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Wed, 26 Feb 2025 16:08:27 -0500
Subject: [PATCH 146/184] Set env variable to pas to tox. squashed commits.

---
 .github/workflows/workflow.yml | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/.github/workflows/workflow.yml b/.github/workflows/workflow.yml
index f57ee4c8d7..e515a0de46 100644
--- a/.github/workflows/workflow.yml
+++ b/.github/workflows/workflow.yml
@@ -128,10 +128,9 @@ jobs:
         [ -f config.yaml ]
     - name: Test with tox
       run: |
-        export ASPIREDIR=.
-        python -c \
-        "from aspire.downloader.data_fetcher import _data_fetcher; print(_data_fetcher.path)"
-        tox --skip-missing-interpreters false -e py${{ matrix.python-version }}-${{ matrix.pyenv }}
+        tox --override testenv.set_env=ASPIREDIR=${{ github.workspace }} \
+            --skip-missing-interpreters false \
+            -e py${{ matrix.python-version }}-${{ matrix.pyenv }}
     - name: Validate Cache Directory Hash
       run: |
         echo "Computing hash for .github_cache/ASPIRE-data..."

From b553087d4924f41fc81967c35b5d35b470bc648b Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Thu, 20 Feb 2025 14:32:10 -0500
Subject: [PATCH 147/184] Experimental pipeline updates, add halfset example.

[skip ci]
---
 .../example_halfset_pipeline_10028.py         | 150 +++++++++++++++
 .../experimental_abinitio_pipeline_10028.py   | 182 ------------------
 ...xperimental_abinitio_pipeline_10028_jsb.py | 154 +++++++++++++++
 .../experimental_abinitio_pipeline_10073.py   | 158 ---------------
 ...xperimental_abinitio_pipeline_10073_jsb.py | 168 ++++++++++++++++
 5 files changed, 472 insertions(+), 340 deletions(-)
 create mode 100644 gallery/experiments/example_halfset_pipeline_10028.py
 delete mode 100644 gallery/experiments/experimental_abinitio_pipeline_10028.py
 create mode 100644 gallery/experiments/experimental_abinitio_pipeline_10028_jsb.py
 delete mode 100644 gallery/experiments/experimental_abinitio_pipeline_10073.py
 create mode 100644 gallery/experiments/experimental_abinitio_pipeline_10073_jsb.py

diff --git a/gallery/experiments/example_halfset_pipeline_10028.py b/gallery/experiments/example_halfset_pipeline_10028.py
new file mode 100644
index 0000000000..3e35a736b2
--- /dev/null
+++ b/gallery/experiments/example_halfset_pipeline_10028.py
@@ -0,0 +1,150 @@
+"""
+Abinitio Halfset Pipeline - Experimental Data
+=============================================
+
+This demonstrates creating two half sets of experimental data,
+performing independent reconstructions, and computing resulting FSC.
+
+Specifically this pipeline uses the EMPIAR 10028 picked particles:
+
+https://www.ebi.ac.uk/empiar/EMPIAR-10028
+"""
+
+# %%
+# Imports
+# -------
+
+import logging
+from pathlib import Path
+
+from aspire.denoising import LegacyClassAvgSource
+from aspire.reconstruction import MeanEstimator
+from aspire.source import OrientedSource, RelionSource
+
+logger = logging.getLogger(__name__)
+
+
+# %%
+# Parameters
+# ---------------
+# Example configuration.
+#
+# Use of GPU is expected for a large configuration.
+# If running on a less capable machine, or simply experimenting, it is
+# strongly recommened to reduce ``img_size``, ``_n_imgs``, and
+# ``n_nbor``.
+
+img_size = 179  # Downsample the images/reconstruction to a desired resolution
+n_classes = 3000  # How many class averages to compute.
+n_nbor = 50  # How many neighbors to stack
+starfile_in = "10028/data/shiny_2sets_fixed9.star"
+data_folder = "."  # This depends on the specific starfile entries.
+pixel_size = 1.34
+fsc_cutoff = 0.143
+
+# %%
+# Load and split Source data
+# --------------------------
+#
+# ``RelionSource`` is used to access the experimental data via a `STAR` file.
+
+# Create a source object loading all the experimental images
+src = RelionSource(
+    starfile_in, pixel_size=pixel_size, max_rows=None, data_folder=data_folder
+)
+
+# Split the data into two sets.
+# This example uses evens and odds for simplicity, but random indices
+# can also be used with similar results.
+
+srcA = src[::2]
+srcB = src[1::2]
+
+# A dictionary can systematically organize our inputs and outputs for both sets.
+pipelines = {
+    "A": {
+        "input": srcA,
+    },
+    "B": {
+        "input": srcB,
+    },
+}
+
+# %%
+# Dual Pipelines
+# --------------
+# Preprocess by downsampling, correcting for CTF, and applying noise
+# correction.  After preprocessing, class averages are generated then
+# autmatically selected for use in reconstruction.
+#
+# Each of the above steps are performed totally independently
+# for each dataset, first for A, then for B.
+#
+# Caching is used throughout for speeding up large datasets on high memory machines.
+
+for src_id, pipeline in pipelines.items():
+
+    src = pipeline["input"]
+
+    # Downsample the images
+    logger.info(f"Set the resolution to {img_size} X {img_size}")
+    src = src.downsample(img_size).cache()
+
+    # Use phase_flip to attempt correcting for CTF.
+    logger.info(f"Perform phase flip to {len(src)} input images for set {src_id}.")
+    src = src.phase_flip().cache()
+
+    # Normalize the background of the images.
+    src = src.normalize_background().cache()
+
+    # Estimate the noise and whiten based on the estimated noise.
+    src = src.legacy_whiten().cache()
+
+    # Optionally invert image contrast.
+    logger.info("Invert the global density contrast")
+    src = src.invert_contrast().cache()
+
+    # Now perform classification and averaging for each class.
+    logger.info("Begin Class Averaging")
+
+    avgs = LegacyClassAvgSource(src, n_nbor=n_nbor)
+    avgs = avgs[:n_classes].cache()
+
+    # Common Line Estimation
+    logger.info("Begin Orientation Estimation")
+    oriented_src = OrientedSource(avgs)
+
+    # Volume Reconstruction
+    logger.info("Begin Volume reconstruction")
+
+    # Setup an estimator to perform the back projection.
+    estimator = MeanEstimator(oriented_src)
+
+    # Perform the estimation and save the volume.
+    pipeline["volume_output_filename"] = fn = (
+        f"10028_abinitio_c{n_classes}_m{n_nbor}_{img_size}px_{src_id}.mrc"
+    )
+    estimated_volume = estimator.estimate()
+    estimated_volume.save(fn, overwrite=True)
+    logger.info(f"Saved Volume to {str(Path(fn).resolve())}")
+
+    # Store volume result in pipeline dict.
+    pipeline["estimated_volume"] = estimated_volume
+
+# %%
+# Compute FSC Score
+# -----------------
+# At this point both pipelines have completed reconstructions and may be compared using FSC.
+
+# Recall our resulting volumes from the dictionary.
+vol_a = pipelines["A"]["estimated_volume"]
+vol_b = pipelines["B"]["estimated_volume"]
+
+# Compute the FSC
+# Save plot, in case display is not available.
+vol_a.fsc(vol_b, cutoff=fsc_cutoff, plot="fsc_plot.png")
+# Display plot and report
+fsc, _ = vol_a.fsc(vol_b, cutoff=fsc_cutoff, plot=True)
+logger.info(
+    f"Found FSC of {fsc} Angstrom at cutoff={fsc_cutoff} and pixel size {vol_a.pixel_size} Angstrom/pixel."
+)
diff --git a/gallery/experiments/experimental_abinitio_pipeline_10028.py b/gallery/experiments/experimental_abinitio_pipeline_10028.py
deleted file mode 100644
index 9e1570bb13..0000000000
--- a/gallery/experiments/experimental_abinitio_pipeline_10028.py
+++ /dev/null
@@ -1,182 +0,0 @@
-"""
-Abinitio Pipeline - Experimental Data Empiar 10028
-==================================================
-
-This notebook introduces a selection of
-components corresponding to loading real Relion picked
-particle cryo-EM data and running key ASPIRE-Python
-Abinitio model components as a pipeline.
-
-Specifically this pipeline uses the
-EMPIAR 10028 picked particles data, available here:
-
-https://www.ebi.ac.uk/empiar/EMPIAR-10028
-
-https://www.ebi.ac.uk/emdb/EMD-2660
-"""
-
-# %%
-# Imports
-# -------
-# First import some of the usual suspects.
-# In addition, import some classes from
-# the ASPIRE package that will be used throughout this experiment.
-
-import logging
-from pathlib import Path
-
-import matplotlib.pyplot as plt
-import numpy as np
-
-from aspire.abinitio import CLSync3N
-from aspire.denoising import DenoisedSource, DenoiserCov2D, LegacyClassAvgSource
-from aspire.reconstruction import MeanEstimator
-from aspire.source import OrientedSource, RelionSource
-
-logger = logging.getLogger(__name__)
-
-
-# %%
-# Parameters
-# ---------------
-# Example simulation configuration.
-
-interactive = False  # Draw blocking interactive plots?
-do_cov2d = False  # Use CWF coefficients
-n_imgs = None  # Set to None for all images in starfile, can set smaller for tests.
-img_size = 32  # Downsample the images/reconstruction to a desired resolution
-n_classes = 3000  # How many class averages to compute.
-n_nbor = 50  # How many neighbors to stack
-starfile_in = "10028/data/shiny_2sets_fixed9.star"
-data_folder = "."  # This depends on the specific starfile entries.
-volume_output_filename = f"10028_abinitio_c{n_classes}_m{n_nbor}_{img_size}.mrc"
-pixel_size = 1.34
-
-
-# %%
-# Source data and Preprocessing
-# -----------------------------
-#
-# `RelionSource` is used to access the experimental data via a `starfile`.
-# Begin by downsampling to our chosen resolution, then preprocess
-# to correct for CTF and noise.
-
-# Create a source object for the experimental images
-src = RelionSource(
-    starfile_in, pixel_size=pixel_size, max_rows=n_imgs, data_folder=data_folder
-)
-
-# Downsample the images
-# Caching is used for speeding up large datasets on high memory machines.
-logger.info(f"Set the resolution to {img_size} X {img_size}")
-src = src.downsample(img_size).cache()
-
-# Peek
-if interactive:
-    src.images[:10].show()
-
-# Use phase_flip to attempt correcting for CTF.
-logger.info("Perform phase flip to input images.")
-src = src.phase_flip().cache()
-
-# Estimate the noise and `Whiten` based on the estimated noise
-src = src.legacy_whiten().cache()
-
-# Peek, what do the whitened images look like...
-if interactive:
-    src.images[:10].show()
-
-# Optionally invert image contrast, depends on data conventions.
-logger.info("Invert the global density contrast")
-src = src.invert_contrast().cache()
-
-
-# %%
-# Optional: CWF Denoising
-# -----------------------
-#
-# Optionally generate an alternative source that is denoised with `cov2d`,
-# then configure a customized averager. This allows the use of CWF denoised
-# images for classification, but stacks the original images for averages
-# used in the remainder of the reconstruction pipeline.
-#
-# In this example, this behavior is controlled by the `do_cov2d` boolean variable.
-# When disabled, the original src and default averager is used.
-# If you will not be using cov2d,
-# you may remove this code block and associated variables.
-
-classification_src = src
-if do_cov2d:
-    # Use CWF denoising
-    cwf_denoiser = DenoiserCov2D(src)
-    # Use denoised src for classification
-    classification_src = DenoisedSource(src, cwf_denoiser)
-    # Cache for speedup.  Avoids recomputing.
-    classification_src = classification_src.cache()
-    # Peek, what do the denoised images look like...
-    if interactive:
-        classification_src.images[:10].show()
-
-# %%
-# Class Averaging
-# ----------------------
-#
-# Now perform classification and averaging for each class.
-
-logger.info("Begin Class Averaging")
-
-# Now perform classification and averaging for each class.
-# This also demonstrates the potential to use a different source for classification and averaging.
-
-avgs = LegacyClassAvgSource(
-    classification_src,
-    n_nbor=n_nbor,
-    averager_src=src,
-)
-# We'll continue our pipeline with the first `n_classes` from `avgs`.
-avgs = avgs[:n_classes].cache()
-
-# Save off the set of class average images.
-avgs.save("experimental_10028_class_averages.star")
-
-if interactive:
-    avgs.images[:10].show()
-
-
-# %%
-# Common Line Estimation
-# ----------------------
-#
-# Next create a CL instance for estimating orientation of projections
-# using the Common Line with Synchronization Voting method.
-
-logger.info("Begin Orientation Estimation")
-
-# Create a custom orientation estimation object for ``avgs``.
-# This is done to customize the ``n_theta`` value.
-orient_est = CLSync3N(avgs, n_theta=360)
-
-# Create an ``OrientedSource`` class instance that performs orientation
-# estimation in a lazy fashion upon request of images or rotations.
-oriented_src = OrientedSource(avgs, orient_est)
-
-# %%
-# Volume Reconstruction
-# ----------------------
-#
-# Using the oriented source, attempt to reconstruct a volume.
-
-logger.info("Begin Volume reconstruction")
-
-# Setup an estimator to perform the back projection.
-estimator = MeanEstimator(oriented_src)
-
-# Perform the estimation and save the volume.
-estimated_volume = estimator.estimate()
-estimated_volume.save(volume_output_filename, overwrite=True)
-logger.info(f"Saved Volume to {str(Path(volume_output_filename).resolve())}")
-
-# Peek at result
-if interactive:
-    plt.imshow(np.sum(estimated_volume.asnumpy()[0], axis=-1))
-    plt.show()
diff --git a/gallery/experiments/experimental_abinitio_pipeline_10028_jsb.py b/gallery/experiments/experimental_abinitio_pipeline_10028_jsb.py
new file mode 100644
index 0000000000..7d30522bb7
--- /dev/null
+++ b/gallery/experiments/experimental_abinitio_pipeline_10028_jsb.py
@@ -0,0 +1,154 @@
+"""
+Abinitio Pipeline - Experimental Data EMPIAR 10028
+==================================================
+
+This notebook introduces a selection of
+components corresponding to loading real Relion picked
+particle cryo-EM data and running key ASPIRE-Python
+Abinitio model components as a pipeline.
+
+This demonstrates reproducing results similar to those found in::
+
+  Common lines modeling for reference free Ab-initio reconstruction in cryo-EM
+  Journal of Structural Biology 2017
+  https://doi.org/10.1016/j.jsb.2017.09.007
+
+Specifically this pipeline uses the
+EMPIAR 10028 picked particles data, available here:
+
+https://www.ebi.ac.uk/empiar/EMPIAR-10028
+"""
+
+# %%
+# Imports
+# -------
+# Import packages that will be used throughout this experiment.
+
+import logging
+from pathlib import Path
+
+from aspire.denoising import LegacyClassAvgSource
+from aspire.reconstruction import MeanEstimator
+from aspire.source import OrientedSource, RelionSource
+
+logger = logging.getLogger(__name__)
+
+
+# %%
+# Parameters
+# ---------------
+#
+# Use of GPU is expected for a large configuration.
+# If running on a less capable machine, or simply experimenting, it is
+# strongly recommened to reduce ``img_size``, ``_n_imgs``, and
+# ``n_nbor``.
+
+# Inputs
+starfile_in = "10028/data/shiny_2sets_fixed9.star"
+data_folder = "."  # This depends on the specific starfile entries.
+pixel_size = 1.34  # Defined with the dataset from EMPIAR
+
+# Config
+n_imgs = None  # Set to None for all images in starfile, can set smaller for tests.
+img_size = 179  # Downsample the images/reconstruction to a desired resolution
+n_classes = 3000  # How many class averages to compute.
+n_nbor = 50  # How many neighbors to stack
+
+# Outputs
+preprocessed_fn = f"10028_preprocessed_{img_size}px.star"
+oriented_fn = f"10028_oriented_class_averages_{img_size}px.star"
+volume_output_filename = f"10028_abinitio_c{n_classes}_m{n_nbor}_{img_size}px.mrc"
+
+
+# %%
+# Source data and Preprocessing
+# -----------------------------
+#
+# ``RelionSource`` is used to access the experimental data via a `STAR` file.
+# Begin by preprocessing to correct for CTF, then downsample to ``img_size``
+# and apply noise correction.
+#
+# ASPIRE-Python has the ability to automatically adjust CTF filters
+# for downsampling, and this can be employed simply by changing the
+# order of preprocessing steps, saving time by phase flipping lower
+# resolution images.  Howver,tThis script intentionally follows the order
+# described in the original publication.
+
+# Create a source object for the experimental images
+src = RelionSource(
+    starfile_in, pixel_size=pixel_size, max_rows=n_imgs, data_folder=data_folder
+)
+
+# Use phase_flip to attempt correcting for CTF.
+# Caching is used throughout for speeding up large datasets on high memory machines.
+logger.info("Perform phase flip to input images.")
+src = src.phase_flip().cache()
+
+# Downsample the images.
+logger.info(f"Set the resolution to {img_size} X {img_size}")
+src = src.downsample(img_size).cache()
+
+# Normalize the background of the images.
+src = src.normalize_background().cache()
+
+# Estimate the noise and whiten based on the estimated noise.
+src = src.legacy_whiten().cache()
+
+# Optionally invert image contrast.
+logger.info("Invert the global density contrast")
+src = src.invert_contrast().cache()
+
+# Save the preprocessed images.
+# These can be resused to experiment with later stages of the pipeline
+# without repeating the preprocessing computations.
+src.save(preprocessed_fn, save_mode="single", overwrite=True)
+
+# %%
+# Class Averaging
+# ----------------------
+#
+# Now perform classification and averaging for each class.
+
+logger.info("Begin Class Averaging")
+avgs = LegacyClassAvgSource(src, n_nbor=n_nbor)
+
+# We'll continue our pipeline with the first ``n_classes`` from
+# ``avgs``.  The classes will be selected by the ``class_selector`` of a
+# ``ClassAvgSource``, which in this case will be the class averages
+# having the largest variance.  Note global sorting requires computing
+# all class averages, which is computationally intensive.
+avgs = avgs[:n_classes].cache()
+
+
+# %%
+# Common Line Estimation
+# ----------------------
+#
+# Estimate orientation of projections and assign to source by
+# applying ``OrientedSource`` to the class averages from the prior
+# step. By default this applies the Common Line with Synchronization
+# Voting ``CLSync3N`` method.
+
+logger.info("Begin Orientation Estimation")
+oriented_src = OrientedSource(avgs)
+
+# Save off the selected set of class average images, along with the
+# estimated orientations and shifts.  These can be reused to
+# experiment with alternative volume reconstructions.
+oriented_src.save(oriented_fn, save_mode="single", overwrite=True)
+
+# %%
+# Volume Reconstruction
+# ----------------------
+#
+# Using the oriented source, attempt to reconstruct a volume.
+
+logger.info("Begin Volume reconstruction")
+
+# Setup an estimator to perform the back projection.
+estimator = MeanEstimator(oriented_src)
+
+# Perform the estimation and save the volume.
+estimated_volume = estimator.estimate()
+estimated_volume.save(volume_output_filename, overwrite=True)
+logger.info(f"Saved Volume to {str(Path(volume_output_filename).resolve())}")
diff --git a/gallery/experiments/experimental_abinitio_pipeline_10073.py b/gallery/experiments/experimental_abinitio_pipeline_10073.py
deleted file mode 100644
index 7e576abf83..0000000000
--- a/gallery/experiments/experimental_abinitio_pipeline_10073.py
+++ /dev/null
@@ -1,158 +0,0 @@
-"""
-Abinitio Pipeline - Experimental Data Empiar 10073
-==================================================
-
-This notebook introduces a selection of
-components corresponding to loading real Relion picked
-particle cryo-EM data and running key ASPIRE-Python
-Abinitio model components as a pipeline.
-
-This demonstrates using the Global BandedSNRImageQualityFunction
-approach starting with Relion polished picked particles.
-
-Specifically this pipeline uses the
-EMPIAR 10073 picked particles data, available here:
-
-https://www.ebi.ac.uk/empiar/EMPIAR-10073
-
-https://www.ebi.ac.uk/emdb/EMD-8012
-"""
-
-# %%
-# Imports
-# -------
-# First import some of the usual suspects.
-# In addition, import some classes from
-# the ASPIRE package that will be used throughout this experiment.
-
-import logging
-from pathlib import Path
-
-import numpy as np
-
-from aspire.abinitio import CLSync3N
-from aspire.basis import FFBBasis2D
-from aspire.classification import (
-    BandedSNRImageQualityFunction,
-    BFRAverager2D,
-    GlobalWithRepulsionClassSelector,
-    RIRClass2D,
-)
-from aspire.denoising import ClassAvgSource
-from aspire.reconstruction import MeanEstimator
-from aspire.source import OrientedSource, RelionSource
-
-logger = logging.getLogger(__name__)
-
-
-# %%
-# Parameters
-# ---------------
-# Example simulation configuration.
-
-n_imgs = None  # Set to None for all images in starfile, can set smaller for tests.
-img_size = 32  # Downsample the images/reconstruction to a desired resolution
-n_classes = 2000  # How many class averages to compute.
-n_nbor = 50  # How many neighbors to stack
-starfile_in = "10073/data/shiny_correctpaths_cleanedcorruptstacks.star"
-data_folder = "."  # This depends on the specific starfile entries.
-volume_output_filename = f"10073_abinitio_c{n_classes}_m{n_nbor}_{img_size}.mrc"
-pixel_size = 1.43
-
-
-# %%
-# Source data and Preprocessing
-# -----------------------------
-#
-# `RelionSource` is used to access the experimental data via a `starfile`.
-# Begin by downsampling to our chosen resolution, then preprocess
-# to correct for CTF and noise.
-
-# Create a source object for the experimental images
-src = RelionSource(
-    starfile_in, pixel_size=pixel_size, max_rows=n_imgs, data_folder=data_folder
-)
-
-# Downsample the images
-logger.info(f"Set the resolution to {img_size} X {img_size}")
-src = src.downsample(img_size)
-
-src = src.cache()
-
-# %%
-# Class Averaging
-# ----------------------
-#
-# Now perform classification and averaging for each class.
-
-logger.info("Begin Class Averaging")
-
-# Now perform classification and averaging for each class.
-# This also demonstrates customizing a ClassAvgSource, by using global
-# contrast selection.  This computes the entire set of class averages,
-# and sorts them by (highest) contrast.
-
-# Build up the customized components.
-basis = FFBBasis2D(src.L, dtype=src.dtype)
-classifier = RIRClass2D(src, n_nbor=n_nbor, nn_implementation="sklearn")
-averager = BFRAverager2D(basis, src)
-quality_function = BandedSNRImageQualityFunction()
-class_selector = GlobalWithRepulsionClassSelector(averager, quality_function)
-
-# Assemble the components into the Source.
-avgs = ClassAvgSource(
-    src, classifier=classifier, averager=averager, class_selector=class_selector
-)
-
-# Save out the resulting Nearest Neighbor networks arrays.
-np.savez(
-    "experimental_10073_class_averages_class_indices.npz",
-    class_indices=avgs.class_indices,
-    class_refl=avgs.class_refl,
-    class_distances=avgs.class_distances,
-)
-
-# Save the class selection rankings.
-np.save(
-    "experimental_10073_class_averages_selection_indices.npy",
-    avgs.selection_indices,
-)
-
-# We'll continue our pipeline with the first `n_classes` from `avgs`.
-avgs = avgs[:n_classes].cache()
-
-# Save off the set of class average images.
-avgs.save("experimental_10073_class_averages_global.star")
-
-
-# %%
-# Common Line Estimation
-# ----------------------
-#
-# Next create a CL instance for estimating orientation of projections
-# using the Common Line with Synchronization Voting method.
-
-logger.info("Begin Orientation Estimation")
-
-# Create a custom orientation estimation object for ``avgs``.
-orient_est = CLSync3N(avgs, n_theta=72)
-
-# Create an ``OrientedSource`` class instance that performs orientation
-# estimation in a lazy fashion upon request of images or rotations.
-oriented_src = OrientedSource(avgs, orient_est)
-
-# %%
-# Volume Reconstruction
-# ----------------------
-#
-# Using the oriented source, attempt to reconstruct a volume.
-
-logger.info("Begin Volume reconstruction")
-
-# Setup an estimator to perform the back projection.
-estimator = MeanEstimator(oriented_src)
-
-# Perform the estimation and save the volume.
-estimated_volume = estimator.estimate()
-estimated_volume.save(volume_output_filename, overwrite=True)
-logger.info(f"Saved Volume to {str(Path(volume_output_filename).resolve())}")
diff --git a/gallery/experiments/experimental_abinitio_pipeline_10073_jsb.py b/gallery/experiments/experimental_abinitio_pipeline_10073_jsb.py
new file mode 100644
index 0000000000..e0178f29d3
--- /dev/null
+++ b/gallery/experiments/experimental_abinitio_pipeline_10073_jsb.py
@@ -0,0 +1,168 @@
+"""
+Abinitio Pipeline - Experimental Data EMPIAR 10073
+==================================================
+
+This notebook introduces a selection of
+components corresponding to loading real Relion picked
+particle cryo-EM data and running key ASPIRE-Python
+Abinitio model components as a pipeline.
+
+This demonstrates reproducing results similar to those found in::
+
+  Common lines modeling for reference free Ab-initio reconstruction in cryo-EM
+  Journal of Structural Biology 2017
+  https://doi.org/10.1016/j.jsb.2017.09.007
+
+Specifically this pipeline uses the
+EMPIAR 10073 picked particles data, available here:
+
+https://www.ebi.ac.uk/empiar/EMPIAR-10073
+"""
+
+# %%
+# Imports
+# -------
+# Import packages that will be used throughout this experiment.
+
+import logging
+import os
+from pathlib import Path
+
+import numpy as np
+
+from aspire.abinitio import CLSync3N
+from aspire.denoising import LegacyClassAvgSource
+from aspire.image import Image
+from aspire.reconstruction import MeanEstimator
+from aspire.source import ArrayImageSource, OrientedSource, RelionSource
+from aspire.utils import fuzzy_mask
+
+logger = logging.getLogger(__name__)
+
+
+# %%
+# Parameters
+# ---------------
+#
+# Use of GPU is expected for a large configuration.
+# If running on a less capable machine, or simply experimenting, it is
+# strongly recommened to reduce ``img_size``, ``_n_imgs``, and
+# ``n_nbor``.
+
+# Inputs
+starfile_in = "10073/data/shiny_correctpaths_cleanedcorruptstacks.star"
+data_folder = "."  # This depends on the specific starfile entries.
+pixel_size = 1.43  # Defined with the dataset from EMPIAR
+
+# Config
+n_imgs = None  # Set to None for all images in starfile, can set smaller for tests.
+img_size = 129  # Downsample the images/reconstruction to a desired resolution
+n_classes = 3000  # How many class averages to compute.
+n_nbor = 50  # How many neighbors to stack
+
+# Outputs
+preprocessed_fn = f"10073_preprocessed_{img_size}px.star"
+class_avg_fn = f"10073_var_sorted_cls_avgs_m{n_nbor}_{img_size}px.star"
+oriented_fn = f"10073_oriented_class_averages_{img_size}px.star"
+volume_output_filename = f"10073_abinitio_c{n_classes}_m{n_nbor}_{img_size}px.mrc"
+
+
+# %%
+# Source data and Preprocessing
+# -----------------------------
+#
+# ``RelionSource`` is used to access the experimental data via a `STAR` file.
+
+# Create a source object for the experimental images.
+src = RelionSource(
+    starfile_in, pixel_size=pixel_size, max_rows=n_imgs, data_folder=data_folder
+)
+
+# Use ``phase_flip`` to attempt correcting for CTF.
+logger.info("Perform phase flip to input images.")
+src = src.phase_flip().cache()
+
+# Downsample the images.
+logger.info(f"Set the resolution to {img_size} X {img_size}")
+src = src.downsample(img_size).cache()
+
+# Normalize the background of the images.
+src = src.normalize_background().cache()
+
+# Estimate the noise and whiten based on the estimated noise.
+src = src.legacy_whiten().cache()
+
+# Optionally invert image contrast.
+logger.info("Invert the global density contrast")
+src = src.invert_contrast().cache()
+
+# Save the preprocessed image stack.
+src.save(preprocessed_fn, save_mode="single", overwrite=True)
+
+
+# %%
+# Class Averaging
+# ----------------------
+#
+# Now perform classification and averaging for each class.
+
+logger.info("Begin Class Averaging")
+
+avgs = LegacyClassAvgSource(
+    src,
+    n_nbor=n_nbor,
+).cache()
+
+# Save the entire set of class averages to disk so they can be re-used.
+avgs.save(class_avg_fn, save_mode="single", overwrite=True)
+
+# We'll continue our pipeline by selecting``n_classes`` from ``avgs``.
+# To capture a broader range of viewing angles, uniformly select every ``k`` image.
+k = (avgs.n - 1) // n_classes
+avgs = avgs[::k].cache()
+
+
+# %%
+# Common Line Estimation
+# ----------------------
+#
+# Estimating orientation of projections and assign to source by
+# applying ``OrientedSource`` to the class averages from the prior
+# step. By default this applies the Common Line with Synchronization
+# Voting ``CLSync3N`` method.  Here additional weighting techniques
+# are applied for common lines detection by customizing the
+# orientation estimator component.
+
+logger.info("Apply custom mask")
+# 10073 benefits from a masking procedure that is more aggressive than the default.
+# Note, since we've manually masked, the default masking is disabled below in ``CLSync3N``.
+# This also upcasts to double precision, which is helpful for this reconstruction.
+mask = fuzzy_mask((img_size, img_size), np.float64, r0=0.4 * img_size, risetime=2)
+avgs = ArrayImageSource(avgs.images[:] * mask)
+
+logger.info("Begin Orientation Estimation")
+# Configure the CLSync3N algorithm,
+#  customized by enabling weighting and disabling default mask.
+ori_est = CLSync3N(avgs, mask=False, S_weighting=True, J_weighting=True)
+# Handles calling code to find and assign orientations and shifts.
+oriented_src = OrientedSource(avgs, ori_est)
+
+# Save off the set of class average images, along with the estimated orientations and shifts.
+oriented_src.save(oriented_fn, save_mode="single", overwrite=True)
+
+
+# %%
+# Volume Reconstruction
+# ----------------------
+#
+# Using the oriented source, attempt to reconstruct a volume.
+
+logger.info("Begin Volume reconstruction")
+
+# Setup an estimator to perform the back projection.
+estimator = MeanEstimator(oriented_src)
+
+# Perform the estimation and save the volume.
+estimated_volume = estimator.estimate()
+estimated_volume.save(volume_output_filename, overwrite=True)
+logger.info(f"Saved Volume to {str(Path(volume_output_filename).resolve())}")

From bf5b8ed03acc49fc23252664dd9d540bab58ea50 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Fri, 28 Feb 2025 15:16:49 -0500
Subject: [PATCH 148/184] Fix sim pipeline index mapping

index_map will only work with debug/top ordering

[skip ci]
---
 gallery/experiments/simulated_abinitio_pipeline.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/gallery/experiments/simulated_abinitio_pipeline.py b/gallery/experiments/simulated_abinitio_pipeline.py
index b622970cb9..adf228499f 100644
--- a/gallery/experiments/simulated_abinitio_pipeline.py
+++ b/gallery/experiments/simulated_abinitio_pipeline.py
@@ -179,7 +179,7 @@ def noise_function(x, y):
 
 # Stash true rotations for later comparison.
 # Note class selection re-ordered our images, so we remap the indices back to the original source.
-indices = avgs.index_map  # Also available from avgs.src.selection_indices[:n_classes]
+indices = avgs.src.selection_indices[:n_classes]
 true_rotations = src.rotations[indices]
 
 # Create a custom orientation estimation object for ``avgs``.
@@ -217,3 +217,8 @@ def noise_function(x, y):
 if interactive:
     plt.imshow(np.sum(estimated_volume.asnumpy()[0], axis=-1))
     plt.show()
+
+# FSC with ground truth volume
+ds_v = og_v.downsample(img_size)
+fsc, _ = estimated_volume.fsc(ds_v, cutoff=0.5)
+logger.info(f"Estimated FSC {fsc} Angstroms at {ds_v.pixel_size} Angstrom per pixel.")

From a99ee799b6b8139f184579655a2998d20d06d541 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Mon, 3 Mar 2025 07:46:47 -0500
Subject: [PATCH 149/184] cleanup, mostly strings, add extra save point to
 10028.

[skip ci]
---
 gallery/experiments/example_halfset_pipeline_10028.py     | 8 ++++----
 .../experimental_abinitio_pipeline_10028_jsb.py           | 8 ++++++--
 .../experimental_abinitio_pipeline_10073_jsb.py           | 7 ++-----
 3 files changed, 12 insertions(+), 11 deletions(-)

diff --git a/gallery/experiments/example_halfset_pipeline_10028.py b/gallery/experiments/example_halfset_pipeline_10028.py
index 3e35a736b2..3f270d05e1 100644
--- a/gallery/experiments/example_halfset_pipeline_10028.py
+++ b/gallery/experiments/example_halfset_pipeline_10028.py
@@ -29,10 +29,10 @@
 # ---------------
 # Example configuration.
 #
-# Use of GPU is expected for a large configuration.
-# If running on a less capable machine, or simply experimenting, it is
-# strongly recommened to reduce ``img_size``, ``_n_imgs``, and
-# ``n_nbor``.
+# Use of GPU is expected for a large configuration.  If running on a
+# less capable machine, or simply experimenting, it is strongly
+# recommened to reduce the problem size by altering ``img_size``,
+# ``n_classes``, ``n_nbor``, ``max_rows`` etc.
 
 img_size = 179  # Downsample the images/reconstruction to a desired resolution
 n_classes = 3000  # How many class averages to compute.
diff --git a/gallery/experiments/experimental_abinitio_pipeline_10028_jsb.py b/gallery/experiments/experimental_abinitio_pipeline_10028_jsb.py
index 7d30522bb7..b6a05008ab 100644
--- a/gallery/experiments/experimental_abinitio_pipeline_10028_jsb.py
+++ b/gallery/experiments/experimental_abinitio_pipeline_10028_jsb.py
@@ -40,7 +40,7 @@
 #
 # Use of GPU is expected for a large configuration.
 # If running on a less capable machine, or simply experimenting, it is
-# strongly recommened to reduce ``img_size``, ``_n_imgs``, and
+# strongly recommened to reduce ``img_size``, ``n_imgs``, and
 # ``n_nbor``.
 
 # Inputs
@@ -56,6 +56,7 @@
 
 # Outputs
 preprocessed_fn = f"10028_preprocessed_{img_size}px.star"
+class_avg_fn = f"10028_var_sorted_cls_avgs_m{n_nbor}_{img_size}px.star"
 oriented_fn = f"10028_oriented_class_averages_{img_size}px.star"
 volume_output_filename = f"10028_abinitio_c{n_classes}_m{n_nbor}_{img_size}px.mrc"
 
@@ -110,7 +111,10 @@
 # Now perform classification and averaging for each class.
 
 logger.info("Begin Class Averaging")
-avgs = LegacyClassAvgSource(src, n_nbor=n_nbor)
+avgs = LegacyClassAvgSource(src, n_nbor=n_nbor).cache()
+
+# Save the entire set of class averages to disk so they can be re-used.
+avgs.save(class_avg_fn, save_mode="single", overwrite=True)
 
 # We'll continue our pipeline with the first ``n_classes`` from
 # ``avgs``.  The classes will be selected by the ``class_selector`` of a
diff --git a/gallery/experiments/experimental_abinitio_pipeline_10073_jsb.py b/gallery/experiments/experimental_abinitio_pipeline_10073_jsb.py
index e0178f29d3..80d3e4000a 100644
--- a/gallery/experiments/experimental_abinitio_pipeline_10073_jsb.py
+++ b/gallery/experiments/experimental_abinitio_pipeline_10073_jsb.py
@@ -46,7 +46,7 @@
 #
 # Use of GPU is expected for a large configuration.
 # If running on a less capable machine, or simply experimenting, it is
-# strongly recommened to reduce ``img_size``, ``_n_imgs``, and
+# strongly recommened to reduce ``img_size``, ``n_imgs``, and
 # ``n_nbor``.
 
 # Inputs
@@ -108,10 +108,7 @@
 
 logger.info("Begin Class Averaging")
 
-avgs = LegacyClassAvgSource(
-    src,
-    n_nbor=n_nbor,
-).cache()
+avgs = LegacyClassAvgSource(src, n_nbor=n_nbor).cache()
 
 # Save the entire set of class averages to disk so they can be re-used.
 avgs.save(class_avg_fn, save_mode="single", overwrite=True)

From 3869e9d16d7bcb78f345ea0e2663225ab3107439 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Mon, 3 Mar 2025 08:38:30 -0500
Subject: [PATCH 150/184] more cleanup, strings found in html version review.

[skip ci]
---
 .../experiments/example_halfset_pipeline_10028.py    |  4 ++--
 .../experimental_abinitio_pipeline_10028_jsb.py      | 10 +++++-----
 .../experimental_abinitio_pipeline_10073_jsb.py      | 12 ++++++------
 3 files changed, 13 insertions(+), 13 deletions(-)

diff --git a/gallery/experiments/example_halfset_pipeline_10028.py b/gallery/experiments/example_halfset_pipeline_10028.py
index 3f270d05e1..87ad721622 100644
--- a/gallery/experiments/example_halfset_pipeline_10028.py
+++ b/gallery/experiments/example_halfset_pipeline_10028.py
@@ -31,7 +31,7 @@
 #
 # Use of GPU is expected for a large configuration.  If running on a
 # less capable machine, or simply experimenting, it is strongly
-# recommened to reduce the problem size by altering ``img_size``,
+# recommended to reduce the problem size by altering ``img_size``,
 # ``n_classes``, ``n_nbor``, ``max_rows`` etc.
 
 img_size = 179  # Downsample the images/reconstruction to a desired resolution
@@ -75,7 +75,7 @@
 # --------------
 # Preprocess by downsampling, correcting for CTF, and applying noise
 # correction.  After preprocessing, class averages are generated then
-# autmatically selected for use in reconstruction.
+# automatically selected for use in reconstruction.
 #
 # Each of the above steps are performed totally independently
 # for each dataset, first for A, then for B.
diff --git a/gallery/experiments/experimental_abinitio_pipeline_10028_jsb.py b/gallery/experiments/experimental_abinitio_pipeline_10028_jsb.py
index b6a05008ab..7dec66b07e 100644
--- a/gallery/experiments/experimental_abinitio_pipeline_10028_jsb.py
+++ b/gallery/experiments/experimental_abinitio_pipeline_10028_jsb.py
@@ -40,7 +40,7 @@
 #
 # Use of GPU is expected for a large configuration.
 # If running on a less capable machine, or simply experimenting, it is
-# strongly recommened to reduce ``img_size``, ``n_imgs``, and
+# strongly recommended to reduce ``img_size``, ``n_imgs``, and
 # ``n_nbor``.
 
 # Inputs
@@ -72,8 +72,8 @@
 # ASPIRE-Python has the ability to automatically adjust CTF filters
 # for downsampling, and this can be employed simply by changing the
 # order of preprocessing steps, saving time by phase flipping lower
-# resolution images.  Howver,tThis script intentionally follows the order
-# described in the original publication.
+# resolution images.  However, this script intentionally follows the
+# order described in the original publication.
 
 # Create a source object for the experimental images
 src = RelionSource(
@@ -100,7 +100,7 @@
 src = src.invert_contrast().cache()
 
 # Save the preprocessed images.
-# These can be resused to experiment with later stages of the pipeline
+# These can be reused to experiment with later stages of the pipeline
 # without repeating the preprocessing computations.
 src.save(preprocessed_fn, save_mode="single", overwrite=True)
 
@@ -149,7 +149,7 @@
 
 logger.info("Begin Volume reconstruction")
 
-# Setup an estimator to perform the back projection.
+# Setup an estimator to perform the back-projection.
 estimator = MeanEstimator(oriented_src)
 
 # Perform the estimation and save the volume.
diff --git a/gallery/experiments/experimental_abinitio_pipeline_10073_jsb.py b/gallery/experiments/experimental_abinitio_pipeline_10073_jsb.py
index 80d3e4000a..55032d4a9c 100644
--- a/gallery/experiments/experimental_abinitio_pipeline_10073_jsb.py
+++ b/gallery/experiments/experimental_abinitio_pipeline_10073_jsb.py
@@ -46,7 +46,7 @@
 #
 # Use of GPU is expected for a large configuration.
 # If running on a less capable machine, or simply experimenting, it is
-# strongly recommened to reduce ``img_size``, ``n_imgs``, and
+# strongly recommended to reduce ``img_size``, ``n_imgs``, and
 # ``n_nbor``.
 
 # Inputs
@@ -110,7 +110,7 @@
 
 avgs = LegacyClassAvgSource(src, n_nbor=n_nbor).cache()
 
-# Save the entire set of class averages to disk so they can be re-used.
+# Save the entire set of class averages to disk so they can be reused.
 avgs.save(class_avg_fn, save_mode="single", overwrite=True)
 
 # We'll continue our pipeline by selecting``n_classes`` from ``avgs``.
@@ -123,7 +123,7 @@
 # Common Line Estimation
 # ----------------------
 #
-# Estimating orientation of projections and assign to source by
+# Estimate orientation of projections and assign to source by
 # applying ``OrientedSource`` to the class averages from the prior
 # step. By default this applies the Common Line with Synchronization
 # Voting ``CLSync3N`` method.  Here additional weighting techniques
@@ -132,13 +132,13 @@
 
 logger.info("Apply custom mask")
 # 10073 benefits from a masking procedure that is more aggressive than the default.
-# Note, since we've manually masked, the default masking is disabled below in ``CLSync3N``.
+# Note, since we've manually masked, the default masking is disabled below in `CLSync3N`.
 # This also upcasts to double precision, which is helpful for this reconstruction.
 mask = fuzzy_mask((img_size, img_size), np.float64, r0=0.4 * img_size, risetime=2)
 avgs = ArrayImageSource(avgs.images[:] * mask)
 
 logger.info("Begin Orientation Estimation")
-# Configure the CLSync3N algorithm,
+# Configure the `CLSync3N` algorithm,
 #  customized by enabling weighting and disabling default mask.
 ori_est = CLSync3N(avgs, mask=False, S_weighting=True, J_weighting=True)
 # Handles calling code to find and assign orientations and shifts.
@@ -156,7 +156,7 @@
 
 logger.info("Begin Volume reconstruction")
 
-# Setup an estimator to perform the back projection.
+# Setup an estimator to perform the back-projection.
 estimator = MeanEstimator(oriented_src)
 
 # Perform the estimation and save the volume.

From 9fb0292350773a8102564646036afd52e644bdd9 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Mon, 3 Mar 2025 13:47:19 -0500
Subject: [PATCH 151/184] make JSB reference pretty

[skip ci]
---
 .../experimental_abinitio_pipeline_10028_jsb.py        | 10 ++++++----
 .../experimental_abinitio_pipeline_10073_jsb.py        | 10 ++++++----
 2 files changed, 12 insertions(+), 8 deletions(-)

diff --git a/gallery/experiments/experimental_abinitio_pipeline_10028_jsb.py b/gallery/experiments/experimental_abinitio_pipeline_10028_jsb.py
index 7dec66b07e..1f550d99e9 100644
--- a/gallery/experiments/experimental_abinitio_pipeline_10028_jsb.py
+++ b/gallery/experiments/experimental_abinitio_pipeline_10028_jsb.py
@@ -7,11 +7,13 @@
 particle cryo-EM data and running key ASPIRE-Python
 Abinitio model components as a pipeline.
 
-This demonstrates reproducing results similar to those found in::
+This demonstrates reproducing results similar to those found in:
 
-  Common lines modeling for reference free Ab-initio reconstruction in cryo-EM
-  Journal of Structural Biology 2017
-  https://doi.org/10.1016/j.jsb.2017.09.007
+.. admonition:: Publication
+
+   | Common lines modeling for reference free Ab-initio reconstruction in cryo-EM
+   | Journal of Structural Biology 2017
+   | https://doi.org/10.1016/j.jsb.2017.09.007
 
 Specifically this pipeline uses the
 EMPIAR 10028 picked particles data, available here:
diff --git a/gallery/experiments/experimental_abinitio_pipeline_10073_jsb.py b/gallery/experiments/experimental_abinitio_pipeline_10073_jsb.py
index 55032d4a9c..4c289a0d97 100644
--- a/gallery/experiments/experimental_abinitio_pipeline_10073_jsb.py
+++ b/gallery/experiments/experimental_abinitio_pipeline_10073_jsb.py
@@ -7,11 +7,13 @@
 particle cryo-EM data and running key ASPIRE-Python
 Abinitio model components as a pipeline.
 
-This demonstrates reproducing results similar to those found in::
+This demonstrates reproducing results similar to those found in:
 
-  Common lines modeling for reference free Ab-initio reconstruction in cryo-EM
-  Journal of Structural Biology 2017
-  https://doi.org/10.1016/j.jsb.2017.09.007
+.. admonition:: Publication
+
+   | Common lines modeling for reference free Ab-initio reconstruction in cryo-EM
+   | Journal of Structural Biology 2017
+   | https://doi.org/10.1016/j.jsb.2017.09.007
 
 Specifically this pipeline uses the
 EMPIAR 10073 picked particles data, available here:

From 110792df5aa9384780aa1b1006a0bd0b4c7f4715 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Tue, 4 Mar 2025 09:34:20 -0500
Subject: [PATCH 152/184] String updates

[skip ci]
---
 .../experiments/experimental_abinitio_pipeline_10028_jsb.py | 5 +++--
 .../experiments/experimental_abinitio_pipeline_10073_jsb.py | 6 +++---
 2 files changed, 6 insertions(+), 5 deletions(-)

diff --git a/gallery/experiments/experimental_abinitio_pipeline_10028_jsb.py b/gallery/experiments/experimental_abinitio_pipeline_10028_jsb.py
index 1f550d99e9..da306d546c 100644
--- a/gallery/experiments/experimental_abinitio_pipeline_10028_jsb.py
+++ b/gallery/experiments/experimental_abinitio_pipeline_10028_jsb.py
@@ -5,7 +5,7 @@
 This notebook introduces a selection of
 components corresponding to loading real Relion picked
 particle cryo-EM data and running key ASPIRE-Python
-Abinitio model components as a pipeline.
+ab initio model components as a pipeline.
 
 This demonstrates reproducing results similar to those found in:
 
@@ -46,6 +46,7 @@
 # ``n_nbor``.
 
 # Inputs
+# Note the published ``shiny_2sets.star`` requires removal of a stray '9' character on line 5476.
 starfile_in = "10028/data/shiny_2sets_fixed9.star"
 data_folder = "."  # This depends on the specific starfile entries.
 pixel_size = 1.34  # Defined with the dataset from EMPIAR
@@ -151,7 +152,7 @@
 
 logger.info("Begin Volume reconstruction")
 
-# Setup an estimator to perform the back-projection.
+# Set up an estimator to perform the backprojection.
 estimator = MeanEstimator(oriented_src)
 
 # Perform the estimation and save the volume.
diff --git a/gallery/experiments/experimental_abinitio_pipeline_10073_jsb.py b/gallery/experiments/experimental_abinitio_pipeline_10073_jsb.py
index 4c289a0d97..0df7ae8927 100644
--- a/gallery/experiments/experimental_abinitio_pipeline_10073_jsb.py
+++ b/gallery/experiments/experimental_abinitio_pipeline_10073_jsb.py
@@ -5,7 +5,7 @@
 This notebook introduces a selection of
 components corresponding to loading real Relion picked
 particle cryo-EM data and running key ASPIRE-Python
-Abinitio model components as a pipeline.
+ab initio model components as a pipeline.
 
 This demonstrates reproducing results similar to those found in:
 
@@ -115,7 +115,7 @@
 # Save the entire set of class averages to disk so they can be reused.
 avgs.save(class_avg_fn, save_mode="single", overwrite=True)
 
-# We'll continue our pipeline by selecting``n_classes`` from ``avgs``.
+# We'll continue our pipeline by selecting ``n_classes`` from ``avgs``.
 # To capture a broader range of viewing angles, uniformly select every ``k`` image.
 k = (avgs.n - 1) // n_classes
 avgs = avgs[::k].cache()
@@ -158,7 +158,7 @@
 
 logger.info("Begin Volume reconstruction")
 
-# Setup an estimator to perform the back-projection.
+# Set up an estimator to perform the backprojection.
 estimator = MeanEstimator(oriented_src)
 
 # Perform the estimation and save the volume.

From 0ee6d3410c7e33a7ae10e4530b5d6e3304f9aa0e Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Wed, 5 Mar 2025 08:36:01 -0500
Subject: [PATCH 153/184] remove unused imports

---
 gallery/experiments/experimental_abinitio_pipeline_10073_jsb.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/gallery/experiments/experimental_abinitio_pipeline_10073_jsb.py b/gallery/experiments/experimental_abinitio_pipeline_10073_jsb.py
index 0df7ae8927..c1e9d673b0 100644
--- a/gallery/experiments/experimental_abinitio_pipeline_10073_jsb.py
+++ b/gallery/experiments/experimental_abinitio_pipeline_10073_jsb.py
@@ -27,14 +27,12 @@
 # Import packages that will be used throughout this experiment.
 
 import logging
-import os
 from pathlib import Path
 
 import numpy as np
 
 from aspire.abinitio import CLSync3N
 from aspire.denoising import LegacyClassAvgSource
-from aspire.image import Image
 from aspire.reconstruction import MeanEstimator
 from aspire.source import ArrayImageSource, OrientedSource, RelionSource
 from aspire.utils import fuzzy_mask

From 1c4aa3f6184f850d4f4b440c1c90faf2d11900f1 Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Fri, 14 Mar 2025 08:49:57 -0400
Subject: [PATCH 154/184] pass ASPIREDIR to make

---
 .github/workflows/workflow.yml | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/.github/workflows/workflow.yml b/.github/workflows/workflow.yml
index e515a0de46..2003b39b66 100644
--- a/.github/workflows/workflow.yml
+++ b/.github/workflows/workflow.yml
@@ -298,10 +298,9 @@ jobs:
           echo "    cache_dir: .github_cache/ASPIRE-data" >> config.yaml
       - name: Build Sphinx docs
         run: |
-          export ASPIREDIR=.
           make distclean
           sphinx-apidoc -f -o ./source ../src -H Modules
-          make html
+          ASPIREDIR=. make html
         working-directory: ./docs
       - name: Archive Sphinx docs
         uses: actions/upload-artifact@v4

From 3daeeeef905f88a33922aac5661d3adc8683affa Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Fri, 14 Mar 2025 09:28:30 -0400
Subject: [PATCH 155/184] debug lines

---
 .github/workflows/workflow.yml     | 2 +-
 gallery/tutorials/pipeline_demo.py | 2 ++
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/workflow.yml b/.github/workflows/workflow.yml
index 2003b39b66..4fa199a373 100644
--- a/.github/workflows/workflow.yml
+++ b/.github/workflows/workflow.yml
@@ -300,7 +300,7 @@ jobs:
         run: |
           make distclean
           sphinx-apidoc -f -o ./source ../src -H Modules
-          ASPIREDIR=. make html
+          ASPIREDIR=. make html SPHINXOPTS="-v"
         working-directory: ./docs
       - name: Archive Sphinx docs
         uses: actions/upload-artifact@v4
diff --git a/gallery/tutorials/pipeline_demo.py b/gallery/tutorials/pipeline_demo.py
index 77d304b156..76813db8c0 100644
--- a/gallery/tutorials/pipeline_demo.py
+++ b/gallery/tutorials/pipeline_demo.py
@@ -18,6 +18,7 @@
 # flake8: noqa
 # sphinx_gallery_end_ignore
 from aspire.downloader import emdb_2660
+from aspire.downloader.data_fetcher import _data_fetcher
 
 # Load 80s Ribosome as a ``Volume`` object.
 original_vol = emdb_2660()
@@ -25,6 +26,7 @@
 # Downsample the volume
 res = 41
 vol = original_vol.downsample(res)
+print(_data_fetcher.path)
 
 # %%
 # .. note::

From e74e169304a0a65d015ab8fde4985d366b0a1e8b Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Fri, 14 Mar 2025 10:00:43 -0400
Subject: [PATCH 156/184] try absolute path

---
 .github/workflows/workflow.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/workflow.yml b/.github/workflows/workflow.yml
index 4fa199a373..646b96d459 100644
--- a/.github/workflows/workflow.yml
+++ b/.github/workflows/workflow.yml
@@ -300,7 +300,7 @@ jobs:
         run: |
           make distclean
           sphinx-apidoc -f -o ./source ../src -H Modules
-          ASPIREDIR=. make html SPHINXOPTS="-v"
+          ASPIREDIR=${{ github.workspace }} make html SPHINXOPTS="-v"
         working-directory: ./docs
       - name: Archive Sphinx docs
         uses: actions/upload-artifact@v4

From 1761ba6c78db94d3f19e7ce82c06086426ce40bc Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Fri, 14 Mar 2025 10:18:38 -0400
Subject: [PATCH 157/184] remove debug lines

---
 .github/workflows/workflow.yml     | 2 +-
 gallery/tutorials/pipeline_demo.py | 2 --
 2 files changed, 1 insertion(+), 3 deletions(-)

diff --git a/.github/workflows/workflow.yml b/.github/workflows/workflow.yml
index 646b96d459..de1281e35b 100644
--- a/.github/workflows/workflow.yml
+++ b/.github/workflows/workflow.yml
@@ -300,7 +300,7 @@ jobs:
         run: |
           make distclean
           sphinx-apidoc -f -o ./source ../src -H Modules
-          ASPIREDIR=${{ github.workspace }} make html SPHINXOPTS="-v"
+          ASPIREDIR=${{ github.workspace }} make html
         working-directory: ./docs
       - name: Archive Sphinx docs
         uses: actions/upload-artifact@v4
diff --git a/gallery/tutorials/pipeline_demo.py b/gallery/tutorials/pipeline_demo.py
index 76813db8c0..77d304b156 100644
--- a/gallery/tutorials/pipeline_demo.py
+++ b/gallery/tutorials/pipeline_demo.py
@@ -18,7 +18,6 @@
 # flake8: noqa
 # sphinx_gallery_end_ignore
 from aspire.downloader import emdb_2660
-from aspire.downloader.data_fetcher import _data_fetcher
 
 # Load 80s Ribosome as a ``Volume`` object.
 original_vol = emdb_2660()
@@ -26,7 +25,6 @@
 # Downsample the volume
 res = 41
 vol = original_vol.downsample(res)
-print(_data_fetcher.path)
 
 # %%
 # .. note::

From 72dc8f2c7a66c44c16245179373acdf97cba3416 Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Fri, 14 Mar 2025 10:24:42 -0400
Subject: [PATCH 158/184] Add cache hash validation to docs job. Check with
 test_vol in tutorial.

---
 .github/workflows/workflow.yml     | 10 ++++++++++
 gallery/tutorials/pipeline_demo.py |  3 ++-
 2 files changed, 12 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/workflow.yml b/.github/workflows/workflow.yml
index de1281e35b..1272b60d4b 100644
--- a/.github/workflows/workflow.yml
+++ b/.github/workflows/workflow.yml
@@ -308,6 +308,16 @@ jobs:
           name: sphinx-docs
           path: docs/build
           retention-days: 7
+      - name: Validate Cache Directory Hash
+	run: |
+	  echo "Computing hash for .github_cache/ASPIRE-data..."
+	  new_hash=$(ls -1 .github_cache/ASPIRE-data/** | md5sum)
+	  echo "Hash from data-cache job: ${{ needs.data-cache.outputs.cache_hash }}"
+	  echo "Computed hash now: $new_hash"
+	  if [ "${{ needs.data-cache.outputs.cache_hash }}" != "$new_hash" ]; then
+	    echo "Error: Cache directory hash has changed!"
+	    exit 1
+	  fi
 
   osx_arm:
     needs: [check, data-cache]
diff --git a/gallery/tutorials/pipeline_demo.py b/gallery/tutorials/pipeline_demo.py
index 77d304b156..613058963e 100644
--- a/gallery/tutorials/pipeline_demo.py
+++ b/gallery/tutorials/pipeline_demo.py
@@ -17,10 +17,11 @@
 # sphinx_gallery_start_ignore
 # flake8: noqa
 # sphinx_gallery_end_ignore
-from aspire.downloader import emdb_2660
+from aspire.downloader import emdb_2660, emdb_8012
 
 # Load 80s Ribosome as a ``Volume`` object.
 original_vol = emdb_2660()
+test_vol = emdb_8012()
 
 # Downsample the volume
 res = 41

From e9a940f455423c1e645bf2b053f03451a41164fb Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Fri, 14 Mar 2025 10:29:44 -0400
Subject: [PATCH 159/184] fix yaml syntax

---
 .github/workflows/workflow.yml | 18 +++++++++---------
 1 file changed, 9 insertions(+), 9 deletions(-)

diff --git a/.github/workflows/workflow.yml b/.github/workflows/workflow.yml
index 1272b60d4b..706221454a 100644
--- a/.github/workflows/workflow.yml
+++ b/.github/workflows/workflow.yml
@@ -309,15 +309,15 @@ jobs:
           path: docs/build
           retention-days: 7
       - name: Validate Cache Directory Hash
-	run: |
-	  echo "Computing hash for .github_cache/ASPIRE-data..."
-	  new_hash=$(ls -1 .github_cache/ASPIRE-data/** | md5sum)
-	  echo "Hash from data-cache job: ${{ needs.data-cache.outputs.cache_hash }}"
-	  echo "Computed hash now: $new_hash"
-	  if [ "${{ needs.data-cache.outputs.cache_hash }}" != "$new_hash" ]; then
-	    echo "Error: Cache directory hash has changed!"
-	    exit 1
-	  fi
+        run: |
+          echo "Computing hash for .github_cache/ASPIRE-data..."
+          new_hash=$(ls -1 .github_cache/ASPIRE-data/** | md5sum)
+          echo "Hash from data-cache job: ${{ needs.data-cache.outputs.cache_hash }}"
+          echo "Computed hash now: $new_hash"
+          if [ "${{ needs.data-cache.outputs.cache_hash }}" != "$new_hash" ]; then
+            echo "Error: Cache directory hash has changed!"
+            exit 1
+          fi
 
   osx_arm:
     needs: [check, data-cache]

From a5d5c8295b9d8abca3af4311ec41ea35d977d076 Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Fri, 14 Mar 2025 11:16:43 -0400
Subject: [PATCH 160/184] remove test vol download

---
 gallery/tutorials/pipeline_demo.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/gallery/tutorials/pipeline_demo.py b/gallery/tutorials/pipeline_demo.py
index 613058963e..77d304b156 100644
--- a/gallery/tutorials/pipeline_demo.py
+++ b/gallery/tutorials/pipeline_demo.py
@@ -17,11 +17,10 @@
 # sphinx_gallery_start_ignore
 # flake8: noqa
 # sphinx_gallery_end_ignore
-from aspire.downloader import emdb_2660, emdb_8012
+from aspire.downloader import emdb_2660
 
 # Load 80s Ribosome as a ``Volume`` object.
 original_vol = emdb_2660()
-test_vol = emdb_8012()
 
 # Downsample the volume
 res = 41

From 79ddf95295ab4b3df4925986b3558d19887035e8 Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Fri, 14 Mar 2025 11:48:39 -0400
Subject: [PATCH 161/184] revert to export but use correct directory

---
 .github/workflows/workflow.yml     | 3 ++-
 gallery/tutorials/pipeline_demo.py | 2 ++
 2 files changed, 4 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/workflow.yml b/.github/workflows/workflow.yml
index 706221454a..48f6ff914d 100644
--- a/.github/workflows/workflow.yml
+++ b/.github/workflows/workflow.yml
@@ -298,9 +298,10 @@ jobs:
           echo "    cache_dir: .github_cache/ASPIRE-data" >> config.yaml
       - name: Build Sphinx docs
         run: |
+          export ASPIREDIR=..
           make distclean
           sphinx-apidoc -f -o ./source ../src -H Modules
-          ASPIREDIR=${{ github.workspace }} make html
+          make html
         working-directory: ./docs
       - name: Archive Sphinx docs
         uses: actions/upload-artifact@v4
diff --git a/gallery/tutorials/pipeline_demo.py b/gallery/tutorials/pipeline_demo.py
index 77d304b156..76813db8c0 100644
--- a/gallery/tutorials/pipeline_demo.py
+++ b/gallery/tutorials/pipeline_demo.py
@@ -18,6 +18,7 @@
 # flake8: noqa
 # sphinx_gallery_end_ignore
 from aspire.downloader import emdb_2660
+from aspire.downloader.data_fetcher import _data_fetcher
 
 # Load 80s Ribosome as a ``Volume`` object.
 original_vol = emdb_2660()
@@ -25,6 +26,7 @@
 # Downsample the volume
 res = 41
 vol = original_vol.downsample(res)
+print(_data_fetcher.path)
 
 # %%
 # .. note::

From 224b6c3ec9309f9f9c5a546aa986b2b19c194df0 Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Fri, 14 Mar 2025 11:53:56 -0400
Subject: [PATCH 162/184] use workspace

---
 .github/workflows/workflow.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/workflow.yml b/.github/workflows/workflow.yml
index 48f6ff914d..a29105c790 100644
--- a/.github/workflows/workflow.yml
+++ b/.github/workflows/workflow.yml
@@ -298,7 +298,7 @@ jobs:
           echo "    cache_dir: .github_cache/ASPIRE-data" >> config.yaml
       - name: Build Sphinx docs
         run: |
-          export ASPIREDIR=..
+          export ASPIREDIR=${{ github.workspace }}
           make distclean
           sphinx-apidoc -f -o ./source ../src -H Modules
           make html

From c63670dbec7dc64505e1157781bd88b11754b968 Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Fri, 14 Mar 2025 12:53:52 -0400
Subject: [PATCH 163/184] remove debug

---
 gallery/tutorials/pipeline_demo.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/gallery/tutorials/pipeline_demo.py b/gallery/tutorials/pipeline_demo.py
index 76813db8c0..77d304b156 100644
--- a/gallery/tutorials/pipeline_demo.py
+++ b/gallery/tutorials/pipeline_demo.py
@@ -18,7 +18,6 @@
 # flake8: noqa
 # sphinx_gallery_end_ignore
 from aspire.downloader import emdb_2660
-from aspire.downloader.data_fetcher import _data_fetcher
 
 # Load 80s Ribosome as a ``Volume`` object.
 original_vol = emdb_2660()
@@ -26,7 +25,6 @@
 # Downsample the volume
 res = 41
 vol = original_vol.downsample(res)
-print(_data_fetcher.path)
 
 # %%
 # .. note::

From 65353d35e38a7de5ce66bf37ae140448749285b4 Mon Sep 17 00:00:00 2001
From: Garrett Wright <garrettwrong@gmail.com>
Date: Mon, 17 Mar 2025 09:13:42 -0400
Subject: [PATCH 164/184] =?UTF-8?q?Bump=20version:=200.13.1=20=E2=86=92=20?=
 =?UTF-8?q?0.13.2?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .bumpversion.cfg               | 2 +-
 README.md                      | 4 ++--
 docs/source/conf.py            | 2 +-
 docs/source/index.rst          | 2 +-
 pyproject.toml                 | 2 +-
 src/aspire/__init__.py         | 2 +-
 src/aspire/config_default.yaml | 2 +-
 7 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/.bumpversion.cfg b/.bumpversion.cfg
index 63a590f41a..7986615540 100644
--- a/.bumpversion.cfg
+++ b/.bumpversion.cfg
@@ -1,5 +1,5 @@
 [bumpversion]
-current_version = 0.13.1
+current_version = 0.13.2
 commit = True
 tag = True
 
diff --git a/README.md b/README.md
index afcf10e49b..56ab26d65f 100644
--- a/README.md
+++ b/README.md
@@ -5,7 +5,7 @@
 [![DOI](https://zenodo.org/badge/DOI/10.5281/zenodo.5657281.svg)](https://doi.org/10.5281/zenodo.5657281)
 [![Downloads](https://static.pepy.tech/badge/aspire/month)](https://pepy.tech/project/aspire)
 
-# ASPIRE - Algorithms for Single Particle Reconstruction - v0.13.1
+# ASPIRE - Algorithms for Single Particle Reconstruction - v0.13.2
 
 The ASPIRE-Python project supersedes [Matlab ASPIRE](https://github.com/PrincetonUniversity/aspire).
 
@@ -20,7 +20,7 @@ For more information about the project, algorithms, and related publications ple
 Please cite using the following DOI. This DOI represents all versions, and will always resolve to the latest one.
 
 ```
-ComputationalCryoEM/ASPIRE-Python: v0.13.1 https://doi.org/10.5281/zenodo.5657281
+ComputationalCryoEM/ASPIRE-Python: v0.13.2 https://doi.org/10.5281/zenodo.5657281
 
 ```
 
diff --git a/docs/source/conf.py b/docs/source/conf.py
index 40f2ec71fd..21e256e48f 100644
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -86,7 +86,7 @@
 # built documents.
 #
 # The full version, including alpha/beta/rc tags.
-release = version = "0.13.1"
+release = version = "0.13.2"
 
 # The language for content autogenerated by Sphinx. Refer to documentation
 # for a list of supported languages.
diff --git a/docs/source/index.rst b/docs/source/index.rst
index 261ed50e9b..88f8107752 100644
--- a/docs/source/index.rst
+++ b/docs/source/index.rst
@@ -1,4 +1,4 @@
-Aspire v0.13.1
+Aspire v0.13.2
 ==============
 
 Algorithms for Single Particle Reconstruction
diff --git a/pyproject.toml b/pyproject.toml
index 9db57652e0..d07c986643 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [project]
 name = "aspire"
-version = "0.13.1"
+version = "0.13.2"
 description = "Algorithms for Single Particle Reconstruction"
 readme = "README.md" # Optional
 requires-python = ">=3.9"
diff --git a/src/aspire/__init__.py b/src/aspire/__init__.py
index 960892ef7f..f74af4dd86 100644
--- a/src/aspire/__init__.py
+++ b/src/aspire/__init__.py
@@ -15,7 +15,7 @@
 from aspire.exceptions import handle_exception
 
 # version in maj.min.bld format
-__version__ = "0.13.1"
+__version__ = "0.13.2"
 
 
 # Setup `confuse` config
diff --git a/src/aspire/config_default.yaml b/src/aspire/config_default.yaml
index 591d8f2dc5..e1eadbb08b 100644
--- a/src/aspire/config_default.yaml
+++ b/src/aspire/config_default.yaml
@@ -1,4 +1,4 @@
-version: 0.13.1
+version: 0.13.2
 common:
     # numeric module to use - one of numpy/cupy
     numeric: numpy

From ae5201149ee21684583dc3da9d099cdf6e31b44e Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Tue, 4 Mar 2025 14:22:44 -0500
Subject: [PATCH 165/184] Change sym_group dtype logger warning to actual
 warning. Cleanup all warnings associated with this.

---
 src/aspire/reconstruction/mean.py     |  5 +++--
 src/aspire/source/image.py            | 11 +++++++++--
 src/aspire/source/simulation.py       |  2 +-
 src/aspire/volume/symmetry_groups.py  | 11 ++++++++++-
 src/aspire/volume/volume.py           |  2 +-
 tests/test_array_image_source.py      | 19 +++++++++++++++++++
 tests/test_bot_align.py               |  4 ++--
 tests/test_simulation.py              | 16 ++++++++++++----
 tests/test_symmetry_groups.py         | 21 +++++++++++++++++----
 tests/test_weighted_mean_estimator.py |  2 +-
 10 files changed, 75 insertions(+), 18 deletions(-)

diff --git a/src/aspire/reconstruction/mean.py b/src/aspire/reconstruction/mean.py
index e3888ee200..b7abe1bce8 100644
--- a/src/aspire/reconstruction/mean.py
+++ b/src/aspire/reconstruction/mean.py
@@ -305,9 +305,10 @@ class MeanEstimator(WeightedVolumesEstimator):
 
     def __init__(self, src, **kwargs):
         # Note, Handle boosting by adjusting weights based on symmetric order.
-        weights = np.ones((src.n, 1)) / np.sqrt(
-            src.n * len(src.symmetry_group.matrices)
+        weights = np.ones((src.n, 1), dtype=src.dtype) / np.sqrt(
+            src.n * len(src.symmetry_group.matrices), dtype=src.dtype
         )
+
         super().__init__(weights, src, **kwargs)
 
     def __getattr__(self, name):
diff --git a/src/aspire/source/image.py b/src/aspire/source/image.py
index 96dbaea30f..c0dd11ffd9 100644
--- a/src/aspire/source/image.py
+++ b/src/aspire/source/image.py
@@ -419,7 +419,10 @@ def filter_indices(self, indices):
     @property
     def offsets(self):
         return np.atleast_2d(
-            self.get_metadata(["_rlnOriginX", "_rlnOriginY"], default_value=0.0)
+            self.get_metadata(
+                ["_rlnOriginX", "_rlnOriginY"],
+                default_value=np.array(0.0, dtype=self.dtype),
+            )
         )
 
     @offsets.setter
@@ -430,7 +433,11 @@ def offsets(self, values):
 
     @property
     def amplitudes(self):
-        return np.atleast_1d(self.get_metadata("_rlnAmplitude", default_value=1.0))
+        return np.atleast_1d(
+            self.get_metadata(
+                "_rlnAmplitude", default_value=np.array(1.0, dtype=self.dtype)
+            )
+        )
 
     @amplitudes.setter
     def amplitudes(self, values):
diff --git a/src/aspire/source/simulation.py b/src/aspire/source/simulation.py
index de5f0957c7..09005a20c3 100644
--- a/src/aspire/source/simulation.py
+++ b/src/aspire/source/simulation.py
@@ -413,7 +413,7 @@ def eigs(self):
         C = self.C
         vols_c = self.vols - self.mean_true()
 
-        p = np.ones(C) / C
+        p = np.ones(C, dtype=self.dtype) / C.astype(self.dtype)
         Q, R = qr(vols_c.to_vec().T, mode="economic")
 
         # Rank is at most C-1, so remove last vector
diff --git a/src/aspire/volume/symmetry_groups.py b/src/aspire/volume/symmetry_groups.py
index 1384b951be..0029236c06 100644
--- a/src/aspire/volume/symmetry_groups.py
+++ b/src/aspire/volume/symmetry_groups.py
@@ -1,4 +1,5 @@
 import logging
+import warnings
 from abc import ABC, abstractmethod, abstractproperty
 
 import numpy as np
@@ -61,7 +62,11 @@ def parse(symmetry, dtype):
 
         if isinstance(symmetry, SymmetryGroup):
             if symmetry.dtype != dtype:
-                logger.warning(f"Recasting SymmetryGroup with dtype {dtype}.")
+                warnings.warn(
+                    f"Recasting SymmetryGroup with dtype {dtype}.",
+                    category=UserWarning,
+                    stacklevel=2,
+                )
                 group_kwargs = dict(dtype=dtype)
                 if getattr(symmetry, "order", False) and symmetry.order > 1:
                     group_kwargs["order"] = symmetry.order
@@ -98,6 +103,10 @@ def parse(symmetry, dtype):
 
         return symmetry_group(**group_kwargs)
 
+    def astype(self, dtype):
+        """Copy of the SymmetryGroup object, cast to a specified dtype."""
+        return SymmetryGroup.parse(self.to_string, dtype)
+
 
 class CnSymmetryGroup(SymmetryGroup):
     """
diff --git a/src/aspire/volume/volume.py b/src/aspire/volume/volume.py
index 04ae8429be..74034a7b47 100644
--- a/src/aspire/volume/volume.py
+++ b/src/aspire/volume/volume.py
@@ -150,7 +150,7 @@ def astype(self, dtype, copy=True):
         return self.__class__(
             self.asnumpy().astype(dtype, copy=copy),
             pixel_size=self.pixel_size,
-            symmetry_group=self.symmetry_group,
+            symmetry_group=self.symmetry_group.astype(dtype),
         )
 
     def _check_key_dims(self, key):
diff --git a/tests/test_array_image_source.py b/tests/test_array_image_source.py
index a43db1ee8e..ffab418c23 100644
--- a/tests/test_array_image_source.py
+++ b/tests/test_array_image_source.py
@@ -5,6 +5,7 @@
 from unittest import TestCase
 
 import numpy as np
+import pytest
 from pytest import raises
 
 from aspire.basis import FBBasis3D
@@ -242,3 +243,21 @@ def test_save_mrcs(self):
             # Test the content matched when loaded
             src2 = RelionSource(star_path)
             np.testing.assert_allclose(src.images[:], src2.images[:])
+
+
+@pytest.mark.parametrize("dtype", [np.float32, np.float64])
+def test_dtype_passthrough(dtype):
+    """
+    Test dtypes are passed appropriately to metadata.
+    """
+    n_ims = 10
+    res = 32
+    ims = np.ones((n_ims, res, res), dtype=dtype)
+
+    src = ArrayImageSource(ims)
+
+    # Check dtypes
+    np.testing.assert_equal(src.dtype, dtype)
+    np.testing.assert_equal(src.images[:].dtype, dtype)
+    np.testing.assert_equal(src.amplitudes.dtype, dtype)
+    np.testing.assert_equal(src.offsets.dtype, dtype)
diff --git a/tests/test_bot_align.py b/tests/test_bot_align.py
index dff3e573f7..a6758828ff 100644
--- a/tests/test_bot_align.py
+++ b/tests/test_bot_align.py
@@ -74,10 +74,10 @@ def vol_data_fixture(snr, dtype):
     L = v.resolution
     shape = (L, L, L)
     ns_std = np.sqrt(norm(v) ** 2 / (L**3 * snr)).astype(v.dtype)
-    reference_vol = v + normal(0, ns_std, shape)
+    reference_vol = v + normal(0, ns_std, shape).astype(dtype, copy=False)
     r = Rotation.generate_random_rotations(1, dtype=v.dtype, seed=1234)
     R_true = r.matrices[0]
-    test_vol = v.rotate(r) + normal(0, ns_std, shape)
+    test_vol = v.rotate(r) + normal(0, ns_std, shape).astype(dtype, copy=False)
 
     return reference_vol, test_vol, R_true
 
diff --git a/tests/test_simulation.py b/tests/test_simulation.py
index 944e2e7c06..9b4bb2e587 100644
--- a/tests/test_simulation.py
+++ b/tests/test_simulation.py
@@ -400,6 +400,7 @@ def testSimulationMean(self):
 
     def testSimulationVolCoords(self):
         coords, norms, inners = self.sim.vol_coords()
+
         np.testing.assert_allclose([4.72837704, -4.72837709], coords, atol=1e-4)
         np.testing.assert_allclose([8.20515764e-07, 1.17550184e-06], norms, atol=1e-4)
         np.testing.assert_allclose(
@@ -494,7 +495,9 @@ def testSimulationCovar(self):
         np.testing.assert_allclose(result, covar[:, :, 4, 4, 4, 4], atol=1e-4)
 
     def testSimulationEvalMean(self):
-        mean_est = Volume(np.load(os.path.join(DATA_DIR, "mean_8_8_8.npy")))
+        mean_est = Volume(
+            np.load(os.path.join(DATA_DIR, "mean_8_8_8.npy")), dtype=self.dtype
+        )
         result = self.sim.eval_mean(mean_est)
 
         np.testing.assert_allclose(result["err"], 2.664116055950763, atol=1e-4)
@@ -510,14 +513,17 @@ def testSimulationEvalCovar(self):
         np.testing.assert_allclose(result["corr"], 0.8405347287741631, atol=1e-4)
 
     def testSimulationEvalCoords(self):
-        mean_est = Volume(np.load(os.path.join(DATA_DIR, "mean_8_8_8.npy")))
+        mean_est = Volume(
+            np.load(os.path.join(DATA_DIR, "mean_8_8_8.npy")), dtype=self.dtype
+        )
         eigs_est = Volume(
-            np.load(os.path.join(DATA_DIR, "eigs_est_8_8_8_1.npy"))[..., 0]
+            np.load(os.path.join(DATA_DIR, "eigs_est_8_8_8_1.npy"))[..., 0],
+            dtype=self.dtype,
         )
 
         clustered_coords_est = np.load(
             os.path.join(DATA_DIR, "clustered_coords_est.npy")
-        )
+        ).astype(dtype=self.dtype)
 
         result = self.sim.eval_coords(mean_est, eigs_est, clustered_coords_est)
 
@@ -551,6 +557,8 @@ def testSimulationEvalCoords(self):
                 0.11048937,
                 0.11048937,
             ],
+            rtol=1e-05,
+            atol=1e-08,
         )
 
         np.testing.assert_allclose(
diff --git a/tests/test_symmetry_groups.py b/tests/test_symmetry_groups.py
index 236f1c60de..ee483ad33c 100644
--- a/tests/test_symmetry_groups.py
+++ b/tests/test_symmetry_groups.py
@@ -80,6 +80,20 @@ def test_group_rotations(group_fixture):
     assert isinstance(rotations, Rotation)
 
 
+def test_astype(group_fixture):
+    """Test `astype` returns correct SymmetryGroup with correct dtype."""
+    sym_group_singles = group_fixture.astype(np.float32)
+    sym_group_doubles = group_fixture.astype(np.float64)
+
+    # Check that astype returns the correct SymmetryGroup
+    np.testing.assert_equal(sym_group_singles, group_fixture)
+    np.testing.assert_equal(sym_group_doubles, group_fixture)
+
+    # Check that we have specified dtype
+    np.testing.assert_equal(sym_group_singles.dtype, np.float32)
+    np.testing.assert_equal(sym_group_doubles.dtype, np.float64)
+
+
 def test_parser_identity():
     result = SymmetryGroup.parse("C1", dtype=np.float32)
     assert isinstance(result, IdentitySymmetryGroup)
@@ -92,16 +106,15 @@ def test_parser_with_group(group_fixture):
     assert result.dtype == group_fixture.dtype
 
 
-def test_parser_dtype_casting(group_fixture, caplog):
+def test_parser_dtype_casting(group_fixture):
     """Test that dtype gets re-cast and warns."""
     dtype = np.float32
     if group_fixture.dtype == np.float32:
         dtype = np.float64
 
-    caplog.clear()
     msg = f"Recasting SymmetryGroup with dtype {dtype}."
-    _ = SymmetryGroup.parse(group_fixture, dtype)
-    assert msg in caplog.text
+    with pytest.warns(UserWarning, match=msg):
+        _ = SymmetryGroup.parse(group_fixture, dtype)
 
 
 def test_parser_error():
diff --git a/tests/test_weighted_mean_estimator.py b/tests/test_weighted_mean_estimator.py
index eabcd0574f..2af95b048f 100644
--- a/tests/test_weighted_mean_estimator.py
+++ b/tests/test_weighted_mean_estimator.py
@@ -71,7 +71,7 @@ def weights(sim):
     # Construct simple test weights;
     # one uniform positive and negative weighted volume respectively.
     r = 2  # Number of weighted volumes
-    weights = np.ones((sim.n, r)) / np.sqrt(sim.n)
+    weights = np.ones((sim.n, r), dtype=sim.dtype) / np.sqrt(sim.n, dtype=sim.dtype)
     weights[:, 1] *= -1  # negate second weight vector
 
     return weights

From 6825579f65c2a9e52f6959fc2f32af27326b0777 Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Wed, 5 Mar 2025 08:07:12 -0500
Subject: [PATCH 166/184] remove added blank lines

---
 src/aspire/reconstruction/mean.py | 1 -
 tests/test_simulation.py          | 1 -
 2 files changed, 2 deletions(-)

diff --git a/src/aspire/reconstruction/mean.py b/src/aspire/reconstruction/mean.py
index b7abe1bce8..083bb85f01 100644
--- a/src/aspire/reconstruction/mean.py
+++ b/src/aspire/reconstruction/mean.py
@@ -308,7 +308,6 @@ def __init__(self, src, **kwargs):
         weights = np.ones((src.n, 1), dtype=src.dtype) / np.sqrt(
             src.n * len(src.symmetry_group.matrices), dtype=src.dtype
         )
-
         super().__init__(weights, src, **kwargs)
 
     def __getattr__(self, name):
diff --git a/tests/test_simulation.py b/tests/test_simulation.py
index 9b4bb2e587..3909aa3e90 100644
--- a/tests/test_simulation.py
+++ b/tests/test_simulation.py
@@ -400,7 +400,6 @@ def testSimulationMean(self):
 
     def testSimulationVolCoords(self):
         coords, norms, inners = self.sim.vol_coords()
-
         np.testing.assert_allclose([4.72837704, -4.72837709], coords, atol=1e-4)
         np.testing.assert_allclose([8.20515764e-07, 1.17550184e-06], norms, atol=1e-4)
         np.testing.assert_allclose(

From e509eb6004ba16c95a94111bcf99f631d014cb88 Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Fri, 7 Mar 2025 08:50:45 -0500
Subject: [PATCH 167/184] directly create class in astype

---
 src/aspire/volume/symmetry_groups.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/src/aspire/volume/symmetry_groups.py b/src/aspire/volume/symmetry_groups.py
index 0029236c06..cf42ee7fea 100644
--- a/src/aspire/volume/symmetry_groups.py
+++ b/src/aspire/volume/symmetry_groups.py
@@ -105,7 +105,12 @@ def parse(symmetry, dtype):
 
     def astype(self, dtype):
         """Copy of the SymmetryGroup object, cast to a specified dtype."""
-        return SymmetryGroup.parse(self.to_string, dtype)
+        kwargs = {"dtype": dtype}
+
+        if hasattr(self, "order"):
+            kwargs["order"] = self.order
+
+        return self.__class__(**kwargs)
 
 
 class CnSymmetryGroup(SymmetryGroup):

From 41aa9d7d97f7258854b131d9de31132c82b189d6 Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Fri, 7 Mar 2025 10:13:25 -0500
Subject: [PATCH 168/184] handle IdentityGroup which has order as an attribbute
 but not as an argument.

---
 src/aspire/volume/symmetry_groups.py | 2 +-
 tests/test_symmetry_groups.py        | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/aspire/volume/symmetry_groups.py b/src/aspire/volume/symmetry_groups.py
index cf42ee7fea..b9bd4003bf 100644
--- a/src/aspire/volume/symmetry_groups.py
+++ b/src/aspire/volume/symmetry_groups.py
@@ -107,7 +107,7 @@ def astype(self, dtype):
         """Copy of the SymmetryGroup object, cast to a specified dtype."""
         kwargs = {"dtype": dtype}
 
-        if hasattr(self, "order"):
+        if hasattr(self, "order") and self.order > 1:
             kwargs["order"] = self.order
 
         return self.__class__(**kwargs)
diff --git a/tests/test_symmetry_groups.py b/tests/test_symmetry_groups.py
index ee483ad33c..954b7229ff 100644
--- a/tests/test_symmetry_groups.py
+++ b/tests/test_symmetry_groups.py
@@ -86,8 +86,8 @@ def test_astype(group_fixture):
     sym_group_doubles = group_fixture.astype(np.float64)
 
     # Check that astype returns the correct SymmetryGroup
-    np.testing.assert_equal(sym_group_singles, group_fixture)
-    np.testing.assert_equal(sym_group_doubles, group_fixture)
+    np.testing.assert_equal(str(sym_group_singles), str(group_fixture))
+    np.testing.assert_equal(str(sym_group_doubles), str(group_fixture))
 
     # Check that we have specified dtype
     np.testing.assert_equal(sym_group_singles.dtype, np.float32)

From bd93bb49ce3b6cb67dd7263b6b7536c800833ae9 Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Fri, 7 Mar 2025 10:58:35 -0500
Subject: [PATCH 169/184] dtype weights in parent class. Test with both weights
 dtypes

---
 src/aspire/reconstruction/mean.py     |  8 ++++----
 tests/test_weighted_mean_estimator.py | 11 +++++++++--
 2 files changed, 13 insertions(+), 6 deletions(-)

diff --git a/src/aspire/reconstruction/mean.py b/src/aspire/reconstruction/mean.py
index 083bb85f01..aa3bcc15fd 100644
--- a/src/aspire/reconstruction/mean.py
+++ b/src/aspire/reconstruction/mean.py
@@ -53,9 +53,9 @@ def __init__(self, weights, *args, **kwargs):
         :param weights: Matrix of weights, n x r.
         """
 
-        self.weights = weights
-        self.r = self.weights.shape[1]
         super().__init__(*args, **kwargs)
+        self.weights = weights.astype(self.src.dtype)
+        self.r = self.weights.shape[1]
         assert self.src.n == self.weights.shape[0]
 
     def __getattr__(self, name):
@@ -305,8 +305,8 @@ class MeanEstimator(WeightedVolumesEstimator):
 
     def __init__(self, src, **kwargs):
         # Note, Handle boosting by adjusting weights based on symmetric order.
-        weights = np.ones((src.n, 1), dtype=src.dtype) / np.sqrt(
-            src.n * len(src.symmetry_group.matrices), dtype=src.dtype
+        weights = np.ones((src.n, 1)) / np.sqrt(
+            src.n * len(src.symmetry_group.matrices)
         )
         super().__init__(weights, src, **kwargs)
 
diff --git a/tests/test_weighted_mean_estimator.py b/tests/test_weighted_mean_estimator.py
index 2af95b048f..b40979371d 100644
--- a/tests/test_weighted_mean_estimator.py
+++ b/tests/test_weighted_mean_estimator.py
@@ -43,6 +43,11 @@ def dtype(request):
     return request.param
 
 
+@pytest.fixture(params=DTYPE, ids=lambda x: f"weights_dtype={x}", scope="module")
+def weights_dtype(request):
+    return request.param
+
+
 @pytest.fixture(scope="module")
 def sim(L, dtype):
     sim = Simulation(
@@ -67,11 +72,13 @@ def basis(L, dtype):
 
 
 @pytest.fixture(scope="module")
-def weights(sim):
+def weights(sim, weights_dtype):
     # Construct simple test weights;
     # one uniform positive and negative weighted volume respectively.
     r = 2  # Number of weighted volumes
-    weights = np.ones((sim.n, r), dtype=sim.dtype) / np.sqrt(sim.n, dtype=sim.dtype)
+    weights = np.ones((sim.n, r), dtype=weights_dtype) / np.sqrt(
+        sim.n, dtype=weights_dtype
+    )
     weights[:, 1] *= -1  # negate second weight vector
 
     return weights

From 5d6e95eb57d31c81e73de69eb1312c03b76ed3f1 Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Fri, 7 Mar 2025 12:48:33 -0500
Subject: [PATCH 170/184] remove dtype from SymmetryGroup

---
 src/aspire/volume/symmetry_groups.py | 118 +++++++++------------------
 tests/test_symmetry_groups.py        |  63 ++++----------
 2 files changed, 58 insertions(+), 123 deletions(-)

diff --git a/src/aspire/volume/symmetry_groups.py b/src/aspire/volume/symmetry_groups.py
index b9bd4003bf..f4f33829ff 100644
--- a/src/aspire/volume/symmetry_groups.py
+++ b/src/aspire/volume/symmetry_groups.py
@@ -14,11 +14,10 @@ class SymmetryGroup(ABC):
     Base class for symmetry groups.
     """
 
-    def __init__(self, dtype):
+    def __init__(self):
         """
-        :param dtype: Numpy dtype to be used for rotation matrices.
+        Abstract class for symmetry groups.
         """
-        self.dtype = np.dtype(dtype)
         self.rotations = self.generate_rotations()
 
     @abstractmethod
@@ -47,30 +46,19 @@ def __str__(self):
         return f"{self.to_string}"
 
     @staticmethod
-    def parse(symmetry, dtype):
+    def parse(symmetry):
         """
-        Takes a SymmetryGroup instance or a string, ie. 'C1', 'C7', 'D3', 'T', 'O', and returns a concrete
-        SymmetryGroup object with the specified dtype.
+        Takes a SymmetryGroup instance or a string, ie. 'C1', 'C7', 'D3', 'T', 'O',
+        and returns a concrete SymmetryGroup object.
 
         :param symmetry: A string (or SymmetryGroup instance) indicating the symmetry of a molecule.
-        :param dtype: dtype for rotation matrices.
         :return: Concrete SymmetryGroup object.
         """
 
         if symmetry is None:
-            return IdentitySymmetryGroup(dtype=dtype)
+            return IdentitySymmetryGroup()
 
         if isinstance(symmetry, SymmetryGroup):
-            if symmetry.dtype != dtype:
-                warnings.warn(
-                    f"Recasting SymmetryGroup with dtype {dtype}.",
-                    category=UserWarning,
-                    stacklevel=2,
-                )
-                group_kwargs = dict(dtype=dtype)
-                if getattr(symmetry, "order", False) and symmetry.order > 1:
-                    group_kwargs["order"] = symmetry.order
-                symmetry = symmetry.__class__(**group_kwargs)
             return symmetry
 
         if not isinstance(symmetry, str):
@@ -78,9 +66,10 @@ def parse(symmetry, dtype):
                 f"`symmetry` must be a string or `SymmetryGroup` instance. Found {type(symmetry)}"
             )
 
+        # Parse symmetry provided as a string.
         symmetry = symmetry.upper()
         if symmetry == "C1":
-            return IdentitySymmetryGroup(dtype=dtype)
+            return IdentitySymmetryGroup()
 
         symmetry_type = symmetry[0]
         symmetric_order = symmetry[1:]
@@ -97,39 +86,29 @@ def parse(symmetry, dtype):
             )
 
         symmetry_group = map_to_sym_group[symmetry_type]
-        group_kwargs = dict(dtype=dtype)
+        group_kwargs = dict()
         if symmetric_order:
             group_kwargs["order"] = int(symmetric_order)
 
         return symmetry_group(**group_kwargs)
 
-    def astype(self, dtype):
-        """Copy of the SymmetryGroup object, cast to a specified dtype."""
-        kwargs = {"dtype": dtype}
-
-        if hasattr(self, "order") and self.order > 1:
-            kwargs["order"] = self.order
-
-        return self.__class__(**kwargs)
-
 
 class CnSymmetryGroup(SymmetryGroup):
     """
     Cyclic symmetry group.
     """
 
-    def __init__(self, order, dtype):
+    def __init__(self, order):
         """
         `CnSymmetryGroup` instance that serves up a `Rotation` object
         containing rotation matrices of the symmetry group (including
         the identity) accessed via the `matrices` attribute.
 
         :param order: The cyclic order for the symmetry group (int).
-        :param dtype: Numpy dtype to be used for rotation matrices.
         """
 
         self.order = int(order)
-        super().__init__(dtype=dtype)
+        super().__init__()
 
     @property
     def to_string(self):
@@ -146,7 +125,7 @@ def generate_rotations(self):
         :return: Rotation object containing the Cn symmetry group and the identity.
         """
         angles = 2 * np.pi * np.arange(self.order) / self.order
-        return Rotation.about_axis("z", angles, dtype=self.dtype)
+        return Rotation.about_axis("z", angles, dtype=np.float64)
 
 
 class IdentitySymmetryGroup(CnSymmetryGroup):
@@ -154,15 +133,13 @@ class IdentitySymmetryGroup(CnSymmetryGroup):
     The identity symmetry group.
     """
 
-    def __init__(self, dtype):
+    def __init__(self):
         """
         `IdentitySymmetryGroup` instance that serves up a `Rotation` object
         containing the identity matrix.
-
-        :param dtype: Numpy dtype to be used for rotation matrices.
         """
 
-        super().__init__(order=1, dtype=dtype)
+        super().__init__(order=1)
 
 
 class DnSymmetryGroup(SymmetryGroup):
@@ -170,7 +147,7 @@ class DnSymmetryGroup(SymmetryGroup):
     Dihedral symmetry group.
     """
 
-    def __init__(self, order, dtype):
+    def __init__(self, order):
         """
         `DnSymmetryGroup` instance that serves up a `Rotation` object
         containing rotation matrices of the symmetry group (including
@@ -178,11 +155,10 @@ def __init__(self, order, dtype):
         is the chiral dihedral symmetry group which does contain reflections.
 
         :param order: The cyclic order for the symmetry group (int).
-        :param dtype: Numpy dtype to be used for rotation matrices.
         """
 
         self.order = int(order)
-        super().__init__(dtype=dtype)
+        super().__init__()
 
     @property
     def to_string(self):
@@ -197,14 +173,14 @@ def generate_rotations(self):
         :return: Rotation object containing the Dn symmetry group and the identity.
         """
         # Rotations to induce cyclic symmetry
-        angles = 2 * np.pi * np.arange(self.order, dtype=self.dtype) / self.order
+        angles = 2 * np.pi * np.arange(self.order) / self.order
         rot_z = Rotation.about_axis("z", angles).matrices
 
         # Perpendicular rotation to induce dihedral symmetry
-        rot_perp = Rotation.about_axis("y", np.pi, dtype=self.dtype).matrices
+        rot_perp = Rotation.about_axis("y", np.pi).matrices
 
         # Full set of rotations.
-        rots = np.concatenate((rot_z, rot_z @ rot_perp[0]), dtype=self.dtype)
+        rots = np.concatenate((rot_z, rot_z @ rot_perp[0]))
 
         return Rotation(rots)
 
@@ -214,17 +190,15 @@ class TSymmetryGroup(SymmetryGroup):
     Tetrahedral symmetry group.
     """
 
-    def __init__(self, dtype):
+    def __init__(self):
         """
         `TSymmetryGroup` instance that serves up a `Rotation` object
         containing rotation matrices of the symmetry group (including the
         Identity) accessed via the `matrices` attribute. Note, this is the
         chiral tetrahedral symmetry group which does not contain reflections.
-
-        :param dtype: Numpy dtype to be used for rotation matrices.
         """
 
-        super().__init__(dtype=dtype)
+        super().__init__()
 
     @property
     def to_string(self):
@@ -241,27 +215,21 @@ def generate_rotations(self):
         :return: Rotation object containing the tetrahedral symmetry group and the identity.
         """
         # C3 rotation vectors, ie. angle * axis.
-        axes_C3 = np.array(
-            [[1, 1, 1], [-1, -1, 1], [1, -1, -1], [-1, 1, -1]], dtype=self.dtype
-        )
+        axes_C3 = np.array([[1, 1, 1], [-1, -1, 1], [1, -1, -1], [-1, 1, -1]], dtype=np.float64)
         axes_C3 /= np.linalg.norm(axes_C3, axis=-1)[..., np.newaxis]
-        angles_C3 = np.array([2 * np.pi / 3, 4 * np.pi / 3], dtype=self.dtype)
-        rot_vecs_C3 = np.concatenate(
-            [angle * axes_C3 for angle in angles_C3], dtype=self.dtype
-        )
+        angles_C3 = np.array([2 * np.pi / 3, 4 * np.pi / 3], dtype=np.float64)
+        rot_vecs_C3 = np.concatenate([angle * axes_C3 for angle in angles_C3])
 
         # C2 rotation vectors.
-        axes_C2 = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=self.dtype)
+        axes_C2 = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=np.float64)
         rot_vecs_C2 = np.pi * axes_C2
 
         # The full set of rotation vectors inducing tetrahedral symmetry.
-        rot_vec_I = np.zeros((1, 3), dtype=self.dtype)
-        rot_vecs = np.concatenate(
-            (rot_vec_I, rot_vecs_C3, rot_vecs_C2), dtype=self.dtype
-        )
+        rot_vec_I = np.zeros((1, 3))
+        rot_vecs = np.concatenate((rot_vec_I, rot_vecs_C3, rot_vecs_C2))
 
         # Return rotations.
-        return Rotation.from_rotvec(rot_vecs, dtype=self.dtype)
+        return Rotation.from_rotvec(rot_vecs)
 
 
 class OSymmetryGroup(SymmetryGroup):
@@ -269,16 +237,14 @@ class OSymmetryGroup(SymmetryGroup):
     Octahedral symmetry group.
     """
 
-    def __init__(self, dtype):
+    def __init__(self):
         """
         `OSymmetryGroup` instance that serves up a `Rotation` object
         containing rotation matrices of the symmetry group (including the
         Identity) accessed via the `matrices` attribute. Note, this is the
         chiral octahedral symmetry group which does not contain reflections.
-
-        :param dtype: Numpy dtype to be used for rotation matrices.
         """
-        super().__init__(dtype=dtype)
+        super().__init__()
 
         self._symmetry_group = self.generate_rotations()
 
@@ -298,35 +264,31 @@ def generate_rotations(self):
         """
 
         # C4 rotation vectors, ie angle * axis
-        axes_C4 = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=self.dtype)
-        angles_C4 = np.array([np.pi / 2, np.pi, 3 * np.pi / 2], dtype=self.dtype)
+        axes_C4 = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=np.float64)
+        angles_C4 = np.array([np.pi / 2, np.pi, 3 * np.pi / 2], dtype=np.float64)
         rot_vecs_C4 = np.array(
-            [angle * axes_C4 for angle in angles_C4], dtype=self.dtype
+            [angle * axes_C4 for angle in angles_C4], dtype=np.float64,
         ).reshape((9, 3))
 
         # C3 rotation vectors
-        axes_C3 = np.array(
-            [[1, 1, 1], [-1, 1, 1], [1, -1, 1], [1, 1, -1]], dtype=self.dtype
-        )
+        axes_C3 = np.array([[1, 1, 1], [-1, 1, 1], [1, -1, 1], [1, 1, -1]], dtype=np.float64)
         axes_C3 /= np.linalg.norm(axes_C3, axis=-1)[..., np.newaxis]
-        angles_C3 = np.array([2 * np.pi / 3, 4 * np.pi / 3], dtype=self.dtype)
+        angles_C3 = np.array([2 * np.pi / 3, 4 * np.pi / 3], dtype=np.float64)
         rot_vecs_C3 = np.array(
-            [angle * axes_C3 for angle in angles_C3], dtype=self.dtype
+            [angle * axes_C3 for angle in angles_C3], dtype=np.float64,
         ).reshape((8, 3))
 
         # C2 rotation vectors
         axes_C2 = np.array(
             [[1, 1, 0], [-1, 1, 0], [1, 0, 1], [-1, 0, 1], [0, 1, 1], [0, -1, 1]],
-            dtype=self.dtype,
+            dtype=np.float64,
         )
         axes_C2 /= np.linalg.norm(axes_C2, axis=-1)[..., np.newaxis]
         rot_vecs_C2 = np.pi * axes_C2
 
         # The full set of rotation vectors inducing octahedral symmetry.
-        rot_vec_I = np.zeros((1, 3), dtype=self.dtype)
-        rot_vecs = np.concatenate(
-            (rot_vec_I, rot_vecs_C4, rot_vecs_C3, rot_vecs_C2), dtype=self.dtype
-        )
+        rot_vec_I = np.zeros((1, 3))
+        rot_vecs = np.concatenate((rot_vec_I, rot_vecs_C4, rot_vecs_C3, rot_vecs_C2))
 
         # Return rotations.
-        return Rotation.from_rotvec(rot_vecs, dtype=self.dtype)
+        return Rotation.from_rotvec(rot_vecs)
diff --git a/tests/test_symmetry_groups.py b/tests/test_symmetry_groups.py
index 954b7229ff..2be56cbf31 100644
--- a/tests/test_symmetry_groups.py
+++ b/tests/test_symmetry_groups.py
@@ -21,36 +21,30 @@
     DnSymmetryGroup,
 ]
 GROUPS_WITHOUT_ORDER = [
-    TSymmetryGroup,
-    OSymmetryGroup,
+    (TSymmetryGroup,),
+    (OSymmetryGroup,),
 ]
 ORDERS = [2, 3, 4, 5]
-DTYPES = [np.float32, np.float64]
-PARAMS_ORDER = list(itertools.product(GROUPS_WITH_ORDER, DTYPES, ORDERS))
-PARAMS = list(itertools.product(GROUPS_WITHOUT_ORDER, DTYPES))
+PARAMS_ORDER = list(itertools.product(GROUPS_WITH_ORDER, ORDERS))
 
 
 def group_fixture_id(params):
     group_class = params[0]
-    dtype = params[1]
-    if len(params) > 2:
-        order = params[2]
-        return f"{group_class.__name__}, order={order}, dtype={dtype}"
+    if len(params) > 1:
+        order = params[1]
+        return f"{group_class.__name__}, order={order}"
     else:
-        return f"{group_class.__name__}, dtype={dtype}"
+        return f"{group_class.__name__}"
 
 
 # Create SymmetryGroup fixture for the set of parameters.
-@pytest.fixture(params=PARAMS + PARAMS_ORDER, ids=group_fixture_id)
+@pytest.fixture(params=GROUPS_WITHOUT_ORDER + PARAMS_ORDER, ids=group_fixture_id)
 def group_fixture(request):
     params = request.param
     group_class = params[0]
-    dtype = params[1]
-    group_kwargs = dict(
-        dtype=dtype,
-    )
-    if len(params) > 2:
-        group_kwargs["order"] = params[2]
+    group_kwargs = dict()
+    if len(params) > 1:
+        group_kwargs["order"] = params[1]
 
     return group_class(**group_kwargs)
 
@@ -68,7 +62,7 @@ def test_group_str(group_fixture):
 
 
 def test_group_equivalence(group_fixture):
-    C2_symmetry_group = CnSymmetryGroup(order=2, dtype=group_fixture.dtype)
+    C2_symmetry_group = CnSymmetryGroup(order=2)
     if str(group_fixture) == "C2":
         assert C2_symmetry_group == group_fixture
     else:
@@ -80,41 +74,20 @@ def test_group_rotations(group_fixture):
     assert isinstance(rotations, Rotation)
 
 
-def test_astype(group_fixture):
-    """Test `astype` returns correct SymmetryGroup with correct dtype."""
-    sym_group_singles = group_fixture.astype(np.float32)
-    sym_group_doubles = group_fixture.astype(np.float64)
-
-    # Check that astype returns the correct SymmetryGroup
-    np.testing.assert_equal(str(sym_group_singles), str(group_fixture))
-    np.testing.assert_equal(str(sym_group_doubles), str(group_fixture))
-
-    # Check that we have specified dtype
-    np.testing.assert_equal(sym_group_singles.dtype, np.float32)
-    np.testing.assert_equal(sym_group_doubles.dtype, np.float64)
+def test_dtype(group_fixture):
+    """Test SymmetryGroup matrices are always doubles."""
+    np.testing.assert_equal(group_fixture.matrices.dtype, np.float64)
 
 
 def test_parser_identity():
-    result = SymmetryGroup.parse("C1", dtype=np.float32)
+    result = SymmetryGroup.parse("C1")
     assert isinstance(result, IdentitySymmetryGroup)
 
 
 def test_parser_with_group(group_fixture):
     """Test SymmetryGroup instance are parsed correctly."""
-    result = SymmetryGroup.parse(group_fixture, group_fixture.dtype)
+    result = SymmetryGroup.parse(group_fixture)
     assert result == group_fixture
-    assert result.dtype == group_fixture.dtype
-
-
-def test_parser_dtype_casting(group_fixture):
-    """Test that dtype gets re-cast and warns."""
-    dtype = np.float32
-    if group_fixture.dtype == np.float32:
-        dtype = np.float64
-
-    msg = f"Recasting SymmetryGroup with dtype {dtype}."
-    with pytest.warns(UserWarning, match=msg):
-        _ = SymmetryGroup.parse(group_fixture, dtype)
 
 
 def test_parser_error():
@@ -122,4 +95,4 @@ def test_parser_error():
     with pytest.raises(
         ValueError, match=f"Symmetry type {junk_symmetry[0]} not supported.*"
     ):
-        _ = SymmetryGroup.parse(junk_symmetry, dtype=np.float32)
+        _ = SymmetryGroup.parse(junk_symmetry)

From 424f6081a58ef68cf97bf0655dbb5375d69d0b35 Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Tue, 11 Mar 2025 09:43:03 -0400
Subject: [PATCH 171/184] remove dtyping from SymmetryGroup. Fix tests

---
 src/aspire/abinitio/commonline_d2.py  |  2 +-
 src/aspire/image/image.py             |  2 +-
 src/aspire/reconstruction/mean.py     |  2 +-
 src/aspire/source/image.py            |  5 ++---
 src/aspire/volume/symmetry_groups.py  | 23 +++++++++++++++--------
 src/aspire/volume/volume.py           |  6 +++---
 src/aspire/volume/volume_synthesis.py | 10 +++++-----
 tests/test_image.py                   |  2 +-
 tests/test_mean_estimator_boosting.py |  2 +-
 tests/test_orient_d2.py               |  2 +-
 tests/test_synthetic_volume.py        | 10 ++++++++--
 tests/test_volume.py                  |  2 +-
 12 files changed, 40 insertions(+), 28 deletions(-)

diff --git a/src/aspire/abinitio/commonline_d2.py b/src/aspire/abinitio/commonline_d2.py
index a8e951c642..9c3106ea31 100644
--- a/src/aspire/abinitio/commonline_d2.py
+++ b/src/aspire/abinitio/commonline_d2.py
@@ -81,7 +81,7 @@ def __init__(
         # D2 symmetry group.
         # Rearrange in order Identity, about_x, about_y, about_z.
         # This ordering is necessary for reproducing MATLAB code results.
-        self.gs = DnSymmetryGroup(order=2, dtype=self.dtype).matrices[[0, 3, 2, 1]]
+        self.gs = DnSymmetryGroup(order=2).matrices.astype(self.dtype)[[0, 3, 2, 1]]
 
     def estimate_rotations(self):
         """
diff --git a/src/aspire/image/image.py b/src/aspire/image/image.py
index e03de5b237..a304b01964 100644
--- a/src/aspire/image/image.py
+++ b/src/aspire/image/image.py
@@ -732,7 +732,7 @@ def backproject(self, rot_matrices, symmetry_group=None, zero_nyquist=True):
         ), "Number of rotation matrices must match the number of images"
 
         # Get symmetry rotations from SymmetryGroup.
-        symmetry_rots = SymmetryGroup.parse(symmetry_group, dtype=self.dtype).matrices
+        symmetry_rots = SymmetryGroup.parse(symmetry_group).matrices.astype(self.dtype)
         if len(symmetry_rots) > 1:
             logger.info(f"Boosting with {len(symmetry_rots)} rotational symmetries.")
 
diff --git a/src/aspire/reconstruction/mean.py b/src/aspire/reconstruction/mean.py
index aa3bcc15fd..08b31d7457 100644
--- a/src/aspire/reconstruction/mean.py
+++ b/src/aspire/reconstruction/mean.py
@@ -96,7 +96,7 @@ def _compute_kernel(self):
         # Handle symmetry boosting.
         sym_rots = np.eye(3, dtype=self.dtype)[None]
         if self.boost:
-            sym_rots = self.src.symmetry_group.matrices
+            sym_rots = self.src.symmetry_group.matrices.astype(self.dtype)
 
         for i in range(0, self.src.n, self.batch_size):
             _range = np.arange(i, min(self.src.n, i + self.batch_size), dtype=int)
diff --git a/src/aspire/source/image.py b/src/aspire/source/image.py
index c0dd11ffd9..c6fe79aaea 100644
--- a/src/aspire/source/image.py
+++ b/src/aspire/source/image.py
@@ -233,7 +233,7 @@ def symmetry_group(self, value):
                 f"This source is no longer mutable. Try new_source = source.update(symmetry_group='{value}')."
             )
 
-        self._symmetry_group = SymmetryGroup.parse(value, dtype=self.dtype)
+        self._symmetry_group = SymmetryGroup.parse(value)
         self.set_metadata(["_rlnSymmetryGroup"], str(self.symmetry_group))
 
     def _populate_symmetry_group(self, symmetry_group):
@@ -249,10 +249,9 @@ def _populate_symmetry_group(self, symmetry_group):
             else:
                 symmetry_group = SymmetryGroup.parse(
                     symmetry=self.get_metadata(["_rlnSymmetryGroup"])[0],
-                    dtype=self.dtype,
                 )
 
-        self.symmetry_group = symmetry_group or IdentitySymmetryGroup(dtype=self.dtype)
+        self.symmetry_group = symmetry_group or IdentitySymmetryGroup()
 
     def __getitem__(self, indices):
         """
diff --git a/src/aspire/volume/symmetry_groups.py b/src/aspire/volume/symmetry_groups.py
index f4f33829ff..c76050be05 100644
--- a/src/aspire/volume/symmetry_groups.py
+++ b/src/aspire/volume/symmetry_groups.py
@@ -172,12 +172,13 @@ def generate_rotations(self):
 
         :return: Rotation object containing the Dn symmetry group and the identity.
         """
+
         # Rotations to induce cyclic symmetry
         angles = 2 * np.pi * np.arange(self.order) / self.order
-        rot_z = Rotation.about_axis("z", angles).matrices
+        rot_z = Rotation.about_axis("z", angles, dtype=np.float64).matrices
 
         # Perpendicular rotation to induce dihedral symmetry
-        rot_perp = Rotation.about_axis("y", np.pi).matrices
+        rot_perp = Rotation.about_axis("y", np.pi, dtype=np.float64).matrices
 
         # Full set of rotations.
         rots = np.concatenate((rot_z, rot_z @ rot_perp[0]))
@@ -215,7 +216,9 @@ def generate_rotations(self):
         :return: Rotation object containing the tetrahedral symmetry group and the identity.
         """
         # C3 rotation vectors, ie. angle * axis.
-        axes_C3 = np.array([[1, 1, 1], [-1, -1, 1], [1, -1, -1], [-1, 1, -1]], dtype=np.float64)
+        axes_C3 = np.array(
+            [[1, 1, 1], [-1, -1, 1], [1, -1, -1], [-1, 1, -1]], dtype=np.float64
+        )
         axes_C3 /= np.linalg.norm(axes_C3, axis=-1)[..., np.newaxis]
         angles_C3 = np.array([2 * np.pi / 3, 4 * np.pi / 3], dtype=np.float64)
         rot_vecs_C3 = np.concatenate([angle * axes_C3 for angle in angles_C3])
@@ -225,7 +228,7 @@ def generate_rotations(self):
         rot_vecs_C2 = np.pi * axes_C2
 
         # The full set of rotation vectors inducing tetrahedral symmetry.
-        rot_vec_I = np.zeros((1, 3))
+        rot_vec_I = np.zeros((1, 3), dtype=np.float64)
         rot_vecs = np.concatenate((rot_vec_I, rot_vecs_C3, rot_vecs_C2))
 
         # Return rotations.
@@ -267,15 +270,19 @@ def generate_rotations(self):
         axes_C4 = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=np.float64)
         angles_C4 = np.array([np.pi / 2, np.pi, 3 * np.pi / 2], dtype=np.float64)
         rot_vecs_C4 = np.array(
-            [angle * axes_C4 for angle in angles_C4], dtype=np.float64,
+            [angle * axes_C4 for angle in angles_C4],
+            dtype=np.float64,
         ).reshape((9, 3))
 
         # C3 rotation vectors
-        axes_C3 = np.array([[1, 1, 1], [-1, 1, 1], [1, -1, 1], [1, 1, -1]], dtype=np.float64)
+        axes_C3 = np.array(
+            [[1, 1, 1], [-1, 1, 1], [1, -1, 1], [1, 1, -1]], dtype=np.float64
+        )
         axes_C3 /= np.linalg.norm(axes_C3, axis=-1)[..., np.newaxis]
         angles_C3 = np.array([2 * np.pi / 3, 4 * np.pi / 3], dtype=np.float64)
         rot_vecs_C3 = np.array(
-            [angle * axes_C3 for angle in angles_C3], dtype=np.float64,
+            [angle * axes_C3 for angle in angles_C3],
+            dtype=np.float64,
         ).reshape((8, 3))
 
         # C2 rotation vectors
@@ -287,7 +294,7 @@ def generate_rotations(self):
         rot_vecs_C2 = np.pi * axes_C2
 
         # The full set of rotation vectors inducing octahedral symmetry.
-        rot_vec_I = np.zeros((1, 3))
+        rot_vec_I = np.zeros((1, 3), dtype=np.float64)
         rot_vecs = np.concatenate((rot_vec_I, rot_vecs_C4, rot_vecs_C3, rot_vecs_C2))
 
         # Return rotations.
diff --git a/src/aspire/volume/volume.py b/src/aspire/volume/volume.py
index 74034a7b47..013642bba3 100644
--- a/src/aspire/volume/volume.py
+++ b/src/aspire/volume/volume.py
@@ -150,7 +150,7 @@ def astype(self, dtype, copy=True):
         return self.__class__(
             self.asnumpy().astype(dtype, copy=copy),
             pixel_size=self.pixel_size,
-            symmetry_group=self.symmetry_group.astype(dtype),
+            symmetry_group=self.symmetry_group,
         )
 
     def _check_key_dims(self, key):
@@ -185,7 +185,7 @@ def _set_symmetry_group(self, value):
 
         :param value: A `SymmetryGroup` instance or string indicating symmetry, ie. "C5", "D7", "T", etc.
         """
-        self._symmetry_group = SymmetryGroup.parse(value, dtype=self.dtype)
+        self._symmetry_group = SymmetryGroup.parse(value)
 
     def _symmetry_group_warning(self, stacklevel):
         """
@@ -228,7 +228,7 @@ def _result_symmetry(self, other=None, stacklevel=4):
 
         if any([axes_altering_transformation, incompat_syms, arbitrary_array]):
             self._symmetry_group_warning(stacklevel=stacklevel)
-            result_symmetry = IdentitySymmetryGroup(dtype=self.dtype)
+            result_symmetry = IdentitySymmetryGroup()
 
         return result_symmetry
 
diff --git a/src/aspire/volume/volume_synthesis.py b/src/aspire/volume/volume_synthesis.py
index 43f794bfaf..ea2f50e677 100644
--- a/src/aspire/volume/volume_synthesis.py
+++ b/src/aspire/volume/volume_synthesis.py
@@ -207,7 +207,7 @@ def _check_order(self):
             )
 
     def _set_symmetry_group(self):
-        self._symmetry_group = CnSymmetryGroup(order=self.order, dtype=self.dtype)
+        self._symmetry_group = CnSymmetryGroup(order=self.order)
 
     @property
     def n_blobs(self):
@@ -220,7 +220,7 @@ class DnSymmetricVolume(CnSymmetricVolume):
     """
 
     def _set_symmetry_group(self):
-        self._symmetry_group = DnSymmetryGroup(order=self.order, dtype=self.dtype)
+        self._symmetry_group = DnSymmetryGroup(order=self.order)
 
     @property
     def n_blobs(self):
@@ -233,7 +233,7 @@ class TSymmetricVolume(GaussianBlobsVolume):
     """
 
     def _set_symmetry_group(self):
-        self._symmetry_group = TSymmetryGroup(dtype=self.dtype)
+        self._symmetry_group = TSymmetryGroup()
 
     @property
     def n_blobs(self):
@@ -246,7 +246,7 @@ class OSymmetricVolume(GaussianBlobsVolume):
     """
 
     def _set_symmetry_group(self):
-        self._symmetry_group = OSymmetryGroup(dtype=self.dtype)
+        self._symmetry_group = OSymmetryGroup()
 
     @property
     def n_blobs(self):
@@ -270,7 +270,7 @@ def _check_order(self):
             )
 
     def _set_symmetry_group(self):
-        self._symmetry_group = IdentitySymmetryGroup(dtype=self.dtype)
+        self._symmetry_group = IdentitySymmetryGroup()
 
     def _symmetrize_gaussians(self, Q, D, mu):
         return Q, D, mu
diff --git a/tests/test_image.py b/tests/test_image.py
index 4b76d96883..4e8513a2f5 100644
--- a/tests/test_image.py
+++ b/tests/test_image.py
@@ -338,7 +338,7 @@ def test_backproject_symmetry_group(dtype):
     assert isinstance(vol.symmetry_group, CnSymmetryGroup)
 
     # Symmetry from instance.
-    vol = im.backproject(rots, symmetry_group=CnSymmetryGroup(order=3, dtype=dtype))
+    vol = im.backproject(rots, symmetry_group=CnSymmetryGroup(order=3))
     assert isinstance(vol.symmetry_group, CnSymmetryGroup)
 
 
diff --git a/tests/test_mean_estimator_boosting.py b/tests/test_mean_estimator_boosting.py
index 6eac159115..568cad1e79 100644
--- a/tests/test_mean_estimator_boosting.py
+++ b/tests/test_mean_estimator_boosting.py
@@ -148,7 +148,7 @@ def test_boost_flag(source, estimated_volume):
     """Manually boost a source and reconstruct without boosting."""
     ims = source.projections[:]
     rots = source.rotations
-    sym_rots = source.symmetry_group.matrices
+    sym_rots = source.symmetry_group.matrices.astype(dtype=source.dtype)
     sym_order = len(sym_rots)
 
     # Manually boosted images and rotations.
diff --git a/tests/test_orient_d2.py b/tests/test_orient_d2.py
index ca972bcf7c..95446b871d 100644
--- a/tests/test_orient_d2.py
+++ b/tests/test_orient_d2.py
@@ -405,7 +405,7 @@ def g_sync_d2(rots, rots_gt):
     n_img = len(rots)
     dtype = rots.dtype
 
-    rots_symm = DnSymmetryGroup(2, dtype).matrices
+    rots_symm = DnSymmetryGroup(2).matrices.astype(dtype)
     order = len(rots_symm)
 
     A_g = np.zeros((n_img, n_img), dtype=complex)
diff --git a/tests/test_synthetic_volume.py b/tests/test_synthetic_volume.py
index fec7591764..019b99b4a5 100644
--- a/tests/test_synthetic_volume.py
+++ b/tests/test_synthetic_volume.py
@@ -114,6 +114,12 @@ def test_volume_generate(vol_fixture):
         np.testing.assert_approx_equal(v.pixel_size, PXSZ)
 
 
+def test_dtype_passthrough(vol_fixture, dtype_fixture):
+    """Test Volume are generated with correct dtype."""
+    vol = vol_fixture.generate()
+    np.testing.assert_equal(vol._data.dtype, dtype_fixture)
+
+
 def test_simulation_init(vol_fixture):
     """Test that a Simulation initializes provided a synthetic Volume."""
     _ = Simulation(L=vol_fixture.L, vols=vol_fixture.generate())
@@ -135,12 +141,12 @@ def test_compact_support(vol_fixture):
 
 # Supress expected warnings due to rotation of symmetric volume.
 @pytest.mark.filterwarnings("ignore:`symmetry_group` attribute is being set to `C1`")
-def test_volume_symmetry(vol_fixture):
+def test_volume_symmetry(vol_fixture, dtype_fixture):
     """Test that volumes have intended symmetry."""
     vol = vol_fixture.generate()
 
     # Rotations in symmetry group, excluding the Identity.
-    rots = vol_fixture.symmetry_group.matrices[1:]
+    rots = vol_fixture.symmetry_group.matrices[1:].astype(dtype_fixture)
 
     for rot in rots:
         # Rotate volume by an element of the symmetric group.
diff --git a/tests/test_volume.py b/tests/test_volume.py
index d220860784..42ed4675f7 100644
--- a/tests/test_volume.py
+++ b/tests/test_volume.py
@@ -821,7 +821,7 @@ def test_project_broadcast(dtype):
 # SYM_GROUP_PARAMS consists of (initializing method, string representation).
 # Testing just the basic cases of setting the symmetry group from
 # a SymmetryGroup instance, a string, and the default.
-SYM_GROUP_PARAMS = [(TSymmetryGroup(np.float32), "T"), ("D2", "D2"), (None, "C1")]
+SYM_GROUP_PARAMS = [(TSymmetryGroup(), "T"), ("D2", "D2"), (None, "C1")]
 
 
 @pytest.mark.parametrize("sym_group, sym_string", SYM_GROUP_PARAMS)

From a950cea4a81a8ae78e4ee8ca341b674ebbab48c4 Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Tue, 11 Mar 2025 09:45:20 -0400
Subject: [PATCH 172/184] remove unsused import

---
 src/aspire/volume/symmetry_groups.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/src/aspire/volume/symmetry_groups.py b/src/aspire/volume/symmetry_groups.py
index c76050be05..7e5c67a13b 100644
--- a/src/aspire/volume/symmetry_groups.py
+++ b/src/aspire/volume/symmetry_groups.py
@@ -1,5 +1,4 @@
 import logging
-import warnings
 from abc import ABC, abstractmethod, abstractproperty
 
 import numpy as np

From 1dbaeabf631d9f1bc1cf562dbae0450c2dab3ef8 Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Tue, 11 Mar 2025 09:51:06 -0400
Subject: [PATCH 173/184] missed astype

---
 src/aspire/volume/volume_synthesis.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/aspire/volume/volume_synthesis.py b/src/aspire/volume/volume_synthesis.py
index ea2f50e677..88dd7ec90a 100644
--- a/src/aspire/volume/volume_synthesis.py
+++ b/src/aspire/volume/volume_synthesis.py
@@ -126,7 +126,7 @@ def _symmetrize_gaussians(self, Q, D, mu):
         """
         Called to add symmetry to Volumes by generating for each Gaussian blob duplicates in symmetric positions.
         """
-        rots = self.symmetry_group.matrices
+        rots = self.symmetry_group.matrices.astype(self.dtype)
 
         Q_rot = np.zeros(shape=(self.n_blobs, 3, 3)).astype(self.dtype)
         D_sym = np.zeros(shape=(self.n_blobs, 3, 3)).astype(self.dtype)

From 2bab895c59511027ee5f3edc9bc585b2d15d9002 Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Tue, 11 Mar 2025 13:23:49 -0400
Subject: [PATCH 174/184] about_axis handle floats

---
 src/aspire/utils/rotation.py         |  2 ++
 src/aspire/volume/symmetry_groups.py |  6 +++---
 tests/test_rotation.py               | 12 ++++++++++++
 3 files changed, 17 insertions(+), 3 deletions(-)

diff --git a/src/aspire/utils/rotation.py b/src/aspire/utils/rotation.py
index 07a31df9df..028033ce95 100644
--- a/src/aspire/utils/rotation.py
+++ b/src/aspire/utils/rotation.py
@@ -274,6 +274,8 @@ def about_axis(axis, angles, dtype=None, gimble_lock_warnings=True):
 
         :return: Rotation object
         """
+        if isinstance(angles, float):
+            dtype = np.float64
         dtype = dtype or getattr(angles, "dtype", np.float32)
         axes = ["x", "y", "z"]
         if axis.lower() not in axes:
diff --git a/src/aspire/volume/symmetry_groups.py b/src/aspire/volume/symmetry_groups.py
index 7e5c67a13b..ed545da1be 100644
--- a/src/aspire/volume/symmetry_groups.py
+++ b/src/aspire/volume/symmetry_groups.py
@@ -124,7 +124,7 @@ def generate_rotations(self):
         :return: Rotation object containing the Cn symmetry group and the identity.
         """
         angles = 2 * np.pi * np.arange(self.order) / self.order
-        return Rotation.about_axis("z", angles, dtype=np.float64)
+        return Rotation.about_axis("z", angles)
 
 
 class IdentitySymmetryGroup(CnSymmetryGroup):
@@ -174,10 +174,10 @@ def generate_rotations(self):
 
         # Rotations to induce cyclic symmetry
         angles = 2 * np.pi * np.arange(self.order) / self.order
-        rot_z = Rotation.about_axis("z", angles, dtype=np.float64).matrices
+        rot_z = Rotation.about_axis("z", angles).matrices
 
         # Perpendicular rotation to induce dihedral symmetry
-        rot_perp = Rotation.about_axis("y", np.pi, dtype=np.float64).matrices
+        rot_perp = Rotation.about_axis("y", np.pi).matrices
 
         # Full set of rotations.
         rots = np.concatenate((rot_z, rot_z @ rot_perp[0]))
diff --git a/tests/test_rotation.py b/tests/test_rotation.py
index ee87b11d09..5ecf92841c 100644
--- a/tests/test_rotation.py
+++ b/tests/test_rotation.py
@@ -143,6 +143,18 @@ def test_dtype(dtype, rot_obj):
     assert dtype == rot_obj.dtype
 
 
+def test_about_axis_dtype(dtype):
+    angles = np.random.uniform(0, 2 * np.pi, 10).astype(dtype, copy=False)
+    rots = Rotation.about_axis("x", angles).matrices
+    np.testing.assert_equal(rots.dtype, dtype)
+
+
+def test_about_axis_float_dtype():
+    angle = np.pi
+    rot = Rotation.about_axis("y", angle).matrices
+    np.testing.assert_equal(rot.dtype, np.float64)
+
+
 def test_from_rotvec(rot_obj):
     # Build random rotation vectors.
     axis = np.array([1, 0, 0], dtype=rot_obj.dtype)

From d75568974b2ad6a5ed37699989a1fb02d5b3b594 Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Tue, 11 Mar 2025 13:45:58 -0400
Subject: [PATCH 175/184] revert to self.dtype, but set to doubles

---
 src/aspire/volume/symmetry_groups.py | 36 ++++++++++++++++------------
 1 file changed, 21 insertions(+), 15 deletions(-)

diff --git a/src/aspire/volume/symmetry_groups.py b/src/aspire/volume/symmetry_groups.py
index ed545da1be..970a82da2d 100644
--- a/src/aspire/volume/symmetry_groups.py
+++ b/src/aspire/volume/symmetry_groups.py
@@ -17,6 +17,7 @@ def __init__(self):
         """
         Abstract class for symmetry groups.
         """
+        self.dtype = np.float64
         self.rotations = self.generate_rotations()
 
     @abstractmethod
@@ -123,7 +124,9 @@ def generate_rotations(self):
 
         :return: Rotation object containing the Cn symmetry group and the identity.
         """
-        angles = 2 * np.pi * np.arange(self.order) / self.order
+        angles = (2 * np.pi * np.arange(self.order) / self.order).astype(
+            self.dtype, copy=False
+        )
         return Rotation.about_axis("z", angles)
 
 
@@ -173,11 +176,14 @@ def generate_rotations(self):
         """
 
         # Rotations to induce cyclic symmetry
-        angles = 2 * np.pi * np.arange(self.order) / self.order
+        angles = (2 * np.pi * np.arange(self.order) / self.order).astype(
+            self.dtype, copy=False
+        )
         rot_z = Rotation.about_axis("z", angles).matrices
 
         # Perpendicular rotation to induce dihedral symmetry
-        rot_perp = Rotation.about_axis("y", np.pi).matrices
+        pi = np.array(np.pi, dtype=self.dtype)
+        rot_perp = Rotation.about_axis("y", pi).matrices
 
         # Full set of rotations.
         rots = np.concatenate((rot_z, rot_z @ rot_perp[0]))
@@ -216,18 +222,18 @@ def generate_rotations(self):
         """
         # C3 rotation vectors, ie. angle * axis.
         axes_C3 = np.array(
-            [[1, 1, 1], [-1, -1, 1], [1, -1, -1], [-1, 1, -1]], dtype=np.float64
+            [[1, 1, 1], [-1, -1, 1], [1, -1, -1], [-1, 1, -1]], dtype=self.dtype
         )
         axes_C3 /= np.linalg.norm(axes_C3, axis=-1)[..., np.newaxis]
-        angles_C3 = np.array([2 * np.pi / 3, 4 * np.pi / 3], dtype=np.float64)
+        angles_C3 = np.array([2 * np.pi / 3, 4 * np.pi / 3], dtype=self.dtype)
         rot_vecs_C3 = np.concatenate([angle * axes_C3 for angle in angles_C3])
 
         # C2 rotation vectors.
-        axes_C2 = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=np.float64)
+        axes_C2 = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=self.dtype)
         rot_vecs_C2 = np.pi * axes_C2
 
         # The full set of rotation vectors inducing tetrahedral symmetry.
-        rot_vec_I = np.zeros((1, 3), dtype=np.float64)
+        rot_vec_I = np.zeros((1, 3), dtype=self.dtype)
         rot_vecs = np.concatenate((rot_vec_I, rot_vecs_C3, rot_vecs_C2))
 
         # Return rotations.
@@ -266,34 +272,34 @@ def generate_rotations(self):
         """
 
         # C4 rotation vectors, ie angle * axis
-        axes_C4 = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=np.float64)
-        angles_C4 = np.array([np.pi / 2, np.pi, 3 * np.pi / 2], dtype=np.float64)
+        axes_C4 = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=self.dtype)
+        angles_C4 = np.array([np.pi / 2, np.pi, 3 * np.pi / 2], dtype=self.dtype)
         rot_vecs_C4 = np.array(
             [angle * axes_C4 for angle in angles_C4],
-            dtype=np.float64,
+            dtype=self.dtype,
         ).reshape((9, 3))
 
         # C3 rotation vectors
         axes_C3 = np.array(
-            [[1, 1, 1], [-1, 1, 1], [1, -1, 1], [1, 1, -1]], dtype=np.float64
+            [[1, 1, 1], [-1, 1, 1], [1, -1, 1], [1, 1, -1]], dtype=self.dtype
         )
         axes_C3 /= np.linalg.norm(axes_C3, axis=-1)[..., np.newaxis]
-        angles_C3 = np.array([2 * np.pi / 3, 4 * np.pi / 3], dtype=np.float64)
+        angles_C3 = np.array([2 * np.pi / 3, 4 * np.pi / 3], dtype=self.dtype)
         rot_vecs_C3 = np.array(
             [angle * axes_C3 for angle in angles_C3],
-            dtype=np.float64,
+            dtype=self.dtype,
         ).reshape((8, 3))
 
         # C2 rotation vectors
         axes_C2 = np.array(
             [[1, 1, 0], [-1, 1, 0], [1, 0, 1], [-1, 0, 1], [0, 1, 1], [0, -1, 1]],
-            dtype=np.float64,
+            dtype=self.dtype,
         )
         axes_C2 /= np.linalg.norm(axes_C2, axis=-1)[..., np.newaxis]
         rot_vecs_C2 = np.pi * axes_C2
 
         # The full set of rotation vectors inducing octahedral symmetry.
-        rot_vec_I = np.zeros((1, 3), dtype=np.float64)
+        rot_vec_I = np.zeros((1, 3), dtype=self.dtype)
         rot_vecs = np.concatenate((rot_vec_I, rot_vecs_C4, rot_vecs_C3, rot_vecs_C2))
 
         # Return rotations.

From 46d6e39430e16e18ff1254d8c63cbfd38b012064 Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Tue, 11 Mar 2025 13:56:14 -0400
Subject: [PATCH 176/184] copy False

---
 src/aspire/abinitio/commonline_d2.py | 4 +++-
 src/aspire/image/image.py            | 4 +++-
 src/aspire/reconstruction/mean.py    | 2 +-
 3 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/src/aspire/abinitio/commonline_d2.py b/src/aspire/abinitio/commonline_d2.py
index 9c3106ea31..5d3d2c6b61 100644
--- a/src/aspire/abinitio/commonline_d2.py
+++ b/src/aspire/abinitio/commonline_d2.py
@@ -81,7 +81,9 @@ def __init__(
         # D2 symmetry group.
         # Rearrange in order Identity, about_x, about_y, about_z.
         # This ordering is necessary for reproducing MATLAB code results.
-        self.gs = DnSymmetryGroup(order=2).matrices.astype(self.dtype)[[0, 3, 2, 1]]
+        self.gs = DnSymmetryGroup(order=2).matrices.astype(self.dtype, copy=False)[
+            [0, 3, 2, 1]
+        ]
 
     def estimate_rotations(self):
         """
diff --git a/src/aspire/image/image.py b/src/aspire/image/image.py
index a304b01964..92f43caac6 100644
--- a/src/aspire/image/image.py
+++ b/src/aspire/image/image.py
@@ -732,7 +732,9 @@ def backproject(self, rot_matrices, symmetry_group=None, zero_nyquist=True):
         ), "Number of rotation matrices must match the number of images"
 
         # Get symmetry rotations from SymmetryGroup.
-        symmetry_rots = SymmetryGroup.parse(symmetry_group).matrices.astype(self.dtype)
+        symmetry_rots = SymmetryGroup.parse(symmetry_group).matrices.astype(
+            self.dtype, copy=False
+        )
         if len(symmetry_rots) > 1:
             logger.info(f"Boosting with {len(symmetry_rots)} rotational symmetries.")
 
diff --git a/src/aspire/reconstruction/mean.py b/src/aspire/reconstruction/mean.py
index 08b31d7457..58c894d224 100644
--- a/src/aspire/reconstruction/mean.py
+++ b/src/aspire/reconstruction/mean.py
@@ -54,7 +54,7 @@ def __init__(self, weights, *args, **kwargs):
         """
 
         super().__init__(*args, **kwargs)
-        self.weights = weights.astype(self.src.dtype)
+        self.weights = weights.astype(self.src.dtype, copy=False)
         self.r = self.weights.shape[1]
         assert self.src.n == self.weights.shape[0]
 

From b4c980f2ea4cd699bd8f4fe71943472aeaa7b533 Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Tue, 11 Mar 2025 14:11:40 -0400
Subject: [PATCH 177/184] use np.full in place of np.ones

---
 src/aspire/source/simulation.py       | 2 +-
 tests/test_weighted_mean_estimator.py | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/aspire/source/simulation.py b/src/aspire/source/simulation.py
index 09005a20c3..a2b9de712c 100644
--- a/src/aspire/source/simulation.py
+++ b/src/aspire/source/simulation.py
@@ -413,7 +413,7 @@ def eigs(self):
         C = self.C
         vols_c = self.vols - self.mean_true()
 
-        p = np.ones(C, dtype=self.dtype) / C.astype(self.dtype)
+        p = np.full(C, 1 / C, dtype=self.dtype)
         Q, R = qr(vols_c.to_vec().T, mode="economic")
 
         # Rank is at most C-1, so remove last vector
diff --git a/tests/test_weighted_mean_estimator.py b/tests/test_weighted_mean_estimator.py
index b40979371d..7961115e99 100644
--- a/tests/test_weighted_mean_estimator.py
+++ b/tests/test_weighted_mean_estimator.py
@@ -76,8 +76,8 @@ def weights(sim, weights_dtype):
     # Construct simple test weights;
     # one uniform positive and negative weighted volume respectively.
     r = 2  # Number of weighted volumes
-    weights = np.ones((sim.n, r), dtype=weights_dtype) / np.sqrt(
-        sim.n, dtype=weights_dtype
+    weights = np.full(
+        (sim.n, r), 1 / np.sqrt(sim.n, dtype=weights_dtype), dtype=weights_dtype
     )
     weights[:, 1] *= -1  # negate second weight vector
 

From 4972243537a43185ac3dca080f4e8618124de70f Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Tue, 11 Mar 2025 14:27:11 -0400
Subject: [PATCH 178/184] fix dtype logic in about_axis

---
 src/aspire/utils/rotation.py | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/src/aspire/utils/rotation.py b/src/aspire/utils/rotation.py
index 028033ce95..2a2e0f98a0 100644
--- a/src/aspire/utils/rotation.py
+++ b/src/aspire/utils/rotation.py
@@ -274,9 +274,11 @@ def about_axis(axis, angles, dtype=None, gimble_lock_warnings=True):
 
         :return: Rotation object
         """
-        if isinstance(angles, float):
-            dtype = np.float64
-        dtype = dtype or getattr(angles, "dtype", np.float32)
+        dtype = dtype or (
+            np.float64
+            if isinstance(angles, float)
+            else getattr(angles, "dtype", np.float32)
+        )
         axes = ["x", "y", "z"]
         if axis.lower() not in axes:
             raise ValueError("`axis` must be 'x', 'y', or 'z'.")

From 01ebaa5b42aef9fb1ab13f9c4a7107b1d6f9317d Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Tue, 11 Mar 2025 14:38:27 -0400
Subject: [PATCH 179/184] more copy False

---
 src/aspire/reconstruction/mean.py     | 2 +-
 src/aspire/volume/volume_synthesis.py | 2 +-
 tests/test_synthetic_volume.py        | 2 +-
 3 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/aspire/reconstruction/mean.py b/src/aspire/reconstruction/mean.py
index 58c894d224..167769305f 100644
--- a/src/aspire/reconstruction/mean.py
+++ b/src/aspire/reconstruction/mean.py
@@ -96,7 +96,7 @@ def _compute_kernel(self):
         # Handle symmetry boosting.
         sym_rots = np.eye(3, dtype=self.dtype)[None]
         if self.boost:
-            sym_rots = self.src.symmetry_group.matrices.astype(self.dtype)
+            sym_rots = self.src.symmetry_group.matrices.astype(self.dtype, copy=False)
 
         for i in range(0, self.src.n, self.batch_size):
             _range = np.arange(i, min(self.src.n, i + self.batch_size), dtype=int)
diff --git a/src/aspire/volume/volume_synthesis.py b/src/aspire/volume/volume_synthesis.py
index 88dd7ec90a..d0d9225f48 100644
--- a/src/aspire/volume/volume_synthesis.py
+++ b/src/aspire/volume/volume_synthesis.py
@@ -126,7 +126,7 @@ def _symmetrize_gaussians(self, Q, D, mu):
         """
         Called to add symmetry to Volumes by generating for each Gaussian blob duplicates in symmetric positions.
         """
-        rots = self.symmetry_group.matrices.astype(self.dtype)
+        rots = self.symmetry_group.matrices.astype(self.dtype, copy=False)
 
         Q_rot = np.zeros(shape=(self.n_blobs, 3, 3)).astype(self.dtype)
         D_sym = np.zeros(shape=(self.n_blobs, 3, 3)).astype(self.dtype)
diff --git a/tests/test_synthetic_volume.py b/tests/test_synthetic_volume.py
index 019b99b4a5..ad978aee38 100644
--- a/tests/test_synthetic_volume.py
+++ b/tests/test_synthetic_volume.py
@@ -146,7 +146,7 @@ def test_volume_symmetry(vol_fixture, dtype_fixture):
     vol = vol_fixture.generate()
 
     # Rotations in symmetry group, excluding the Identity.
-    rots = vol_fixture.symmetry_group.matrices[1:].astype(dtype_fixture)
+    rots = vol_fixture.symmetry_group.matrices[1:].astype(dtype_fixture, copy=False)
 
     for rot in rots:
         # Rotate volume by an element of the symmetric group.

From d090281396333b8ea84c6749face967999c136eb Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Tue, 11 Mar 2025 14:46:45 -0400
Subject: [PATCH 180/184] couple more

---
 tests/test_mean_estimator_boosting.py | 2 +-
 tests/test_orient_d2.py               | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/test_mean_estimator_boosting.py b/tests/test_mean_estimator_boosting.py
index 568cad1e79..8374acd58e 100644
--- a/tests/test_mean_estimator_boosting.py
+++ b/tests/test_mean_estimator_boosting.py
@@ -148,7 +148,7 @@ def test_boost_flag(source, estimated_volume):
     """Manually boost a source and reconstruct without boosting."""
     ims = source.projections[:]
     rots = source.rotations
-    sym_rots = source.symmetry_group.matrices.astype(dtype=source.dtype)
+    sym_rots = source.symmetry_group.matrices.astype(dtype=source.dtype, copy=False)
     sym_order = len(sym_rots)
 
     # Manually boosted images and rotations.
diff --git a/tests/test_orient_d2.py b/tests/test_orient_d2.py
index 95446b871d..968722e905 100644
--- a/tests/test_orient_d2.py
+++ b/tests/test_orient_d2.py
@@ -405,7 +405,7 @@ def g_sync_d2(rots, rots_gt):
     n_img = len(rots)
     dtype = rots.dtype
 
-    rots_symm = DnSymmetryGroup(2).matrices.astype(dtype)
+    rots_symm = DnSymmetryGroup(2).matrices.astype(dtype, copy=False)
     order = len(rots_symm)
 
     A_g = np.zeros((n_img, n_img), dtype=complex)

From 7147ae353b64b78bb5aa4369543a85e67b730ad7 Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Tue, 11 Mar 2025 15:05:36 -0400
Subject: [PATCH 181/184] copy True for test fixtures

---
 tests/test_mean_estimator_boosting.py | 2 +-
 tests/test_synthetic_volume.py        | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/test_mean_estimator_boosting.py b/tests/test_mean_estimator_boosting.py
index 8374acd58e..568cad1e79 100644
--- a/tests/test_mean_estimator_boosting.py
+++ b/tests/test_mean_estimator_boosting.py
@@ -148,7 +148,7 @@ def test_boost_flag(source, estimated_volume):
     """Manually boost a source and reconstruct without boosting."""
     ims = source.projections[:]
     rots = source.rotations
-    sym_rots = source.symmetry_group.matrices.astype(dtype=source.dtype, copy=False)
+    sym_rots = source.symmetry_group.matrices.astype(dtype=source.dtype)
     sym_order = len(sym_rots)
 
     # Manually boosted images and rotations.
diff --git a/tests/test_synthetic_volume.py b/tests/test_synthetic_volume.py
index ad978aee38..019b99b4a5 100644
--- a/tests/test_synthetic_volume.py
+++ b/tests/test_synthetic_volume.py
@@ -146,7 +146,7 @@ def test_volume_symmetry(vol_fixture, dtype_fixture):
     vol = vol_fixture.generate()
 
     # Rotations in symmetry group, excluding the Identity.
-    rots = vol_fixture.symmetry_group.matrices[1:].astype(dtype_fixture, copy=False)
+    rots = vol_fixture.symmetry_group.matrices[1:].astype(dtype_fixture)
 
     for rot in rots:
         # Rotate volume by an element of the symmetric group.

From 9052bf2b953bd1e39187a3ad4a04895c8e83374e Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Tue, 11 Mar 2025 15:17:37 -0400
Subject: [PATCH 182/184] use sqrt default type

---
 tests/test_weighted_mean_estimator.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/tests/test_weighted_mean_estimator.py b/tests/test_weighted_mean_estimator.py
index 7961115e99..3f58fcdd00 100644
--- a/tests/test_weighted_mean_estimator.py
+++ b/tests/test_weighted_mean_estimator.py
@@ -76,9 +76,7 @@ def weights(sim, weights_dtype):
     # Construct simple test weights;
     # one uniform positive and negative weighted volume respectively.
     r = 2  # Number of weighted volumes
-    weights = np.full(
-        (sim.n, r), 1 / np.sqrt(sim.n, dtype=weights_dtype), dtype=weights_dtype
-    )
+    weights = np.full((sim.n, r), 1 / np.sqrt(sim.n), dtype=weights_dtype)
     weights[:, 1] *= -1  # negate second weight vector
 
     return weights

From 90f275816476c252acabced6cda04fb898c6c7f6 Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Tue, 11 Mar 2025 15:21:04 -0400
Subject: [PATCH 183/184] linspace

---
 src/aspire/volume/symmetry_groups.py | 8 ++------
 1 file changed, 2 insertions(+), 6 deletions(-)

diff --git a/src/aspire/volume/symmetry_groups.py b/src/aspire/volume/symmetry_groups.py
index 970a82da2d..ac60c231e9 100644
--- a/src/aspire/volume/symmetry_groups.py
+++ b/src/aspire/volume/symmetry_groups.py
@@ -124,9 +124,7 @@ def generate_rotations(self):
 
         :return: Rotation object containing the Cn symmetry group and the identity.
         """
-        angles = (2 * np.pi * np.arange(self.order) / self.order).astype(
-            self.dtype, copy=False
-        )
+        angles = np.linspace(0, 2 * np.pi, self.order, endpoint=False, dtype=self.dtype)
         return Rotation.about_axis("z", angles)
 
 
@@ -176,9 +174,7 @@ def generate_rotations(self):
         """
 
         # Rotations to induce cyclic symmetry
-        angles = (2 * np.pi * np.arange(self.order) / self.order).astype(
-            self.dtype, copy=False
-        )
+        angles = np.linspace(0, 2 * np.pi, self.order, endpoint=False, dtype=self.dtype)
         rot_z = Rotation.about_axis("z", angles).matrices
 
         # Perpendicular rotation to induce dihedral symmetry

From c8a97562223b1dcc3fdef13158c2eae99628fc73 Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Tue, 11 Mar 2025 16:08:15 -0400
Subject: [PATCH 184/184] compute in double then cast

---
 src/aspire/volume/symmetry_groups.py | 45 +++++++++++++++-------------
 1 file changed, 24 insertions(+), 21 deletions(-)

diff --git a/src/aspire/volume/symmetry_groups.py b/src/aspire/volume/symmetry_groups.py
index ac60c231e9..e239144150 100644
--- a/src/aspire/volume/symmetry_groups.py
+++ b/src/aspire/volume/symmetry_groups.py
@@ -124,8 +124,8 @@ def generate_rotations(self):
 
         :return: Rotation object containing the Cn symmetry group and the identity.
         """
-        angles = np.linspace(0, 2 * np.pi, self.order, endpoint=False, dtype=self.dtype)
-        return Rotation.about_axis("z", angles)
+        angles = np.linspace(0, 2 * np.pi, self.order, endpoint=False)
+        return Rotation.about_axis("z", angles, dtype=self.dtype)
 
 
 class IdentitySymmetryGroup(CnSymmetryGroup):
@@ -174,12 +174,11 @@ def generate_rotations(self):
         """
 
         # Rotations to induce cyclic symmetry
-        angles = np.linspace(0, 2 * np.pi, self.order, endpoint=False, dtype=self.dtype)
-        rot_z = Rotation.about_axis("z", angles).matrices
+        angles = np.linspace(0, 2 * np.pi, self.order, endpoint=False)
+        rot_z = Rotation.about_axis("z", angles, dtype=self.dtype).matrices
 
         # Perpendicular rotation to induce dihedral symmetry
-        pi = np.array(np.pi, dtype=self.dtype)
-        rot_perp = Rotation.about_axis("y", pi).matrices
+        rot_perp = Rotation.about_axis("y", np.pi, dtype=self.dtype).matrices
 
         # Full set of rotations.
         rots = np.concatenate((rot_z, rot_z @ rot_perp[0]))
@@ -218,22 +217,22 @@ def generate_rotations(self):
         """
         # C3 rotation vectors, ie. angle * axis.
         axes_C3 = np.array(
-            [[1, 1, 1], [-1, -1, 1], [1, -1, -1], [-1, 1, -1]], dtype=self.dtype
+            [[1.0, 1.0, 1.0], [-1.0, -1.0, 1.0], [1.0, -1.0, -1.0], [-1.0, 1.0, -1.0]],
         )
         axes_C3 /= np.linalg.norm(axes_C3, axis=-1)[..., np.newaxis]
-        angles_C3 = np.array([2 * np.pi / 3, 4 * np.pi / 3], dtype=self.dtype)
+        angles_C3 = np.array([2 * np.pi / 3, 4 * np.pi / 3])
         rot_vecs_C3 = np.concatenate([angle * axes_C3 for angle in angles_C3])
 
         # C2 rotation vectors.
-        axes_C2 = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=self.dtype)
+        axes_C2 = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]])
         rot_vecs_C2 = np.pi * axes_C2
 
         # The full set of rotation vectors inducing tetrahedral symmetry.
-        rot_vec_I = np.zeros((1, 3), dtype=self.dtype)
+        rot_vec_I = np.zeros((1, 3))
         rot_vecs = np.concatenate((rot_vec_I, rot_vecs_C3, rot_vecs_C2))
 
         # Return rotations.
-        return Rotation.from_rotvec(rot_vecs)
+        return Rotation.from_rotvec(rot_vecs, dtype=self.dtype)
 
 
 class OSymmetryGroup(SymmetryGroup):
@@ -268,35 +267,39 @@ def generate_rotations(self):
         """
 
         # C4 rotation vectors, ie angle * axis
-        axes_C4 = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=self.dtype)
-        angles_C4 = np.array([np.pi / 2, np.pi, 3 * np.pi / 2], dtype=self.dtype)
+        axes_C4 = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]])
+        angles_C4 = np.array([np.pi / 2, np.pi, 3 * np.pi / 2])
         rot_vecs_C4 = np.array(
             [angle * axes_C4 for angle in angles_C4],
-            dtype=self.dtype,
         ).reshape((9, 3))
 
         # C3 rotation vectors
         axes_C3 = np.array(
-            [[1, 1, 1], [-1, 1, 1], [1, -1, 1], [1, 1, -1]], dtype=self.dtype
+            [[1.0, 1.0, 1.0], [-1.0, 1.0, 1.0], [1.0, -1.0, 1.0], [1.0, 1.0, -1.0]]
         )
         axes_C3 /= np.linalg.norm(axes_C3, axis=-1)[..., np.newaxis]
-        angles_C3 = np.array([2 * np.pi / 3, 4 * np.pi / 3], dtype=self.dtype)
+        angles_C3 = np.array([2 * np.pi / 3, 4 * np.pi / 3])
         rot_vecs_C3 = np.array(
             [angle * axes_C3 for angle in angles_C3],
-            dtype=self.dtype,
         ).reshape((8, 3))
 
         # C2 rotation vectors
         axes_C2 = np.array(
-            [[1, 1, 0], [-1, 1, 0], [1, 0, 1], [-1, 0, 1], [0, 1, 1], [0, -1, 1]],
-            dtype=self.dtype,
+            [
+                [1.0, 1.0, 0.0],
+                [-1.0, 1.0, 0.0],
+                [1.0, 0.0, 1.0],
+                [-1.0, 0.0, 1.0],
+                [0.0, 1.0, 1.0],
+                [0.0, -1.0, 1.0],
+            ],
         )
         axes_C2 /= np.linalg.norm(axes_C2, axis=-1)[..., np.newaxis]
         rot_vecs_C2 = np.pi * axes_C2
 
         # The full set of rotation vectors inducing octahedral symmetry.
-        rot_vec_I = np.zeros((1, 3), dtype=self.dtype)
+        rot_vec_I = np.zeros((1, 3))
         rot_vecs = np.concatenate((rot_vec_I, rot_vecs_C4, rot_vecs_C3, rot_vecs_C2))
 
         # Return rotations.
-        return Rotation.from_rotvec(rot_vecs)
+        return Rotation.from_rotvec(rot_vecs, dtype=self.dtype)