From 36a3dbc9b020902c08a31173a3886b00cb033adb Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Thu, 11 Jan 2024 14:32:11 -0500
Subject: [PATCH 01/29] Initial boosting implementation.

---
 src/aspire/image/image.py         | 20 +++++++++++++++++---
 src/aspire/reconstruction/mean.py |  8 +++++++-
 src/aspire/source/image.py        |  6 ++++--
 3 files changed, 28 insertions(+), 6 deletions(-)

diff --git a/src/aspire/image/image.py b/src/aspire/image/image.py
index 787725df83..6bba7dea94 100644
--- a/src/aspire/image/image.py
+++ b/src/aspire/image/image.py
@@ -489,7 +489,7 @@ def size(self):
         # probably not needed, transition
         return np.size(self._data)
 
-    def backproject(self, rot_matrices):
+    def backproject(self, rot_matrices, symmetry_group=None):
         """
         Backproject images along rotation
 
@@ -511,8 +511,21 @@ def backproject(self, rot_matrices):
             self.n_images == rot_matrices.shape[0]
         ), "Number of rotation matrices must match the number of images"
 
+        # Apply symmetry boosting to rotations.
+        if symmetry_group is None:
+            symmetry_rots = np.eye(3, dtype=self.dtype)[None]
+        else:
+            symmetry_rots = symmetry_group.matrices
+
+        n_sym = len(symmetry_rots)
+        boosted_rot_mats = np.zeros((n_sym * self.shape[0], 3, 3), dtype=self.dtype)
+        for i, sym_rot in enumerate(symmetry_rots):
+            boosted_rot_mats[i * self.shape[0] : (i + 1) * self.shape[0]] = (
+                sym_rot @ rot_matrices
+            )
+
         # TODO: rotated_grids might as well give us correctly shaped array in the first place
-        pts_rot = aspire.volume.rotated_grids(L, rot_matrices).astype(
+        pts_rot = aspire.volume.rotated_grids(L, boosted_rot_mats).astype(
             self.dtype, copy=False
         )
         pts_rot = pts_rot.reshape((3, -1))
@@ -522,7 +535,8 @@ def backproject(self, rot_matrices):
             im_f[:, 0, :] = 0
             im_f[:, :, 0] = 0
 
-        im_f = im_f.flatten()
+        # Apply boosting to images.
+        im_f = np.concatenate((im_f.flatten(),) * n_sym)
 
         vol = anufft(im_f, pts_rot[::-1], (L, L, L), real=True) / L
 
diff --git a/src/aspire/reconstruction/mean.py b/src/aspire/reconstruction/mean.py
index d25d914276..33f6530aba 100644
--- a/src/aspire/reconstruction/mean.py
+++ b/src/aspire/reconstruction/mean.py
@@ -164,7 +164,13 @@ def src_backward(self):
                 im = self.src.images[i : i + self.batch_size]
 
                 batch_vol_rhs = (
-                    self.src.im_backward(im, i, self.weights[:, k]) / self.src.n
+                    self.src.im_backward(
+                        im,
+                        i,
+                        self.weights[:, k],
+                        symmetry_group=self.src.symmetry_group,
+                    )
+                    / self.src.n
                 )
                 vol_rhs[k] += batch_vol_rhs.astype(self.dtype)
 
diff --git a/src/aspire/source/image.py b/src/aspire/source/image.py
index aa127bfb66..d7677d07a4 100644
--- a/src/aspire/source/image.py
+++ b/src/aspire/source/image.py
@@ -940,7 +940,7 @@ def normalize_background(self, bg_radius=1.0, do_ramp=True):
             LambdaXform(normalize_bg, bg_radius=bg_radius, do_ramp=do_ramp)
         )
 
-    def im_backward(self, im, start, weights=None):
+    def im_backward(self, im, start, weights=None, symmetry_group=None):
         """
         Apply adjoint mapping to set of images
 
@@ -960,7 +960,9 @@ def im_backward(self, im, start, weights=None):
         if weights is not None:
             im *= weights[all_idx, np.newaxis, np.newaxis]
 
-        vol = im.backproject(self.rotations[start : start + num, :, :])[0]
+        vol = im.backproject(
+            self.rotations[start : start + num, :, :], symmetry_group=symmetry_group
+        )[0]
 
         return vol
 

From 0a0b86ca3ac72acac780f2117b6ae7d79f7aa197 Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Fri, 12 Jan 2024 10:15:19 -0500
Subject: [PATCH 02/29] Add boosting flag

---
 src/aspire/image/image.py              | 3 +++
 src/aspire/reconstruction/estimator.py | 2 ++
 src/aspire/reconstruction/mean.py      | 5 ++++-
 3 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/src/aspire/image/image.py b/src/aspire/image/image.py
index 6bba7dea94..e68d4c25be 100644
--- a/src/aspire/image/image.py
+++ b/src/aspire/image/image.py
@@ -519,6 +519,9 @@ def backproject(self, rot_matrices, symmetry_group=None):
 
         n_sym = len(symmetry_rots)
         boosted_rot_mats = np.zeros((n_sym * self.shape[0], 3, 3), dtype=self.dtype)
+        import pdb
+
+        pdb.set_trace()
         for i, sym_rot in enumerate(symmetry_rots):
             boosted_rot_mats[i * self.shape[0] : (i + 1) * self.shape[0]] = (
                 sym_rot @ rot_matrices
diff --git a/src/aspire/reconstruction/estimator.py b/src/aspire/reconstruction/estimator.py
index 9e0e561229..05f60b69d4 100644
--- a/src/aspire/reconstruction/estimator.py
+++ b/src/aspire/reconstruction/estimator.py
@@ -18,6 +18,7 @@ def __init__(
         checkpoint_iterations=10,
         checkpoint_prefix="volume_checkpoint",
         maxiter=100,
+        boost=True,
     ):
         """
         An object representing a 2*L-by-2*L-by-2*L array containing the non-centered Fourier transform of the mean
@@ -51,6 +52,7 @@ def __init__(
         self.dtype = self.src.dtype
         self.batch_size = batch_size
         self.preconditioner = preconditioner
+        self.boost = boost
 
         # dtype configuration
         if not self.dtype == self.basis.dtype:
diff --git a/src/aspire/reconstruction/mean.py b/src/aspire/reconstruction/mean.py
index 33f6530aba..6791f55d62 100644
--- a/src/aspire/reconstruction/mean.py
+++ b/src/aspire/reconstruction/mean.py
@@ -153,6 +153,9 @@ def src_backward(self):
         :return: The adjoint mapping applied to the images, averaged over the whole dataset and expressed
             as coefficients of `basis`.
         """
+        symmetry_group = None
+        if self.boost:
+            symmetry_group = self.src.symmetry_group
 
         # src_vols_wt_backward
         vol_rhs = Volume(
@@ -168,7 +171,7 @@ def src_backward(self):
                         im,
                         i,
                         self.weights[:, k],
-                        symmetry_group=self.src.symmetry_group,
+                        symmetry_group=symmetry_group,
                     )
                     / self.src.n
                 )

From 5f1d1591a2cd8c03dd58470ac5881592c997feb1 Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Fri, 12 Jan 2024 13:20:30 -0500
Subject: [PATCH 03/29] Boost the kernel.

---
 src/aspire/image/image.py         |  9 +++----
 src/aspire/reconstruction/mean.py | 39 ++++++++++++++++++++++---------
 2 files changed, 31 insertions(+), 17 deletions(-)

diff --git a/src/aspire/image/image.py b/src/aspire/image/image.py
index e68d4c25be..ca1cadd48a 100644
--- a/src/aspire/image/image.py
+++ b/src/aspire/image/image.py
@@ -517,11 +517,8 @@ def backproject(self, rot_matrices, symmetry_group=None):
         else:
             symmetry_rots = symmetry_group.matrices
 
-        n_sym = len(symmetry_rots)
-        boosted_rot_mats = np.zeros((n_sym * self.shape[0], 3, 3), dtype=self.dtype)
-        import pdb
-
-        pdb.set_trace()
+        sym_order = len(symmetry_rots)
+        boosted_rot_mats = np.zeros((sym_order * self.shape[0], 3, 3), dtype=self.dtype)
         for i, sym_rot in enumerate(symmetry_rots):
             boosted_rot_mats[i * self.shape[0] : (i + 1) * self.shape[0]] = (
                 sym_rot @ rot_matrices
@@ -539,7 +536,7 @@ def backproject(self, rot_matrices, symmetry_group=None):
             im_f[:, :, 0] = 0
 
         # Apply boosting to images.
-        im_f = np.concatenate((im_f.flatten(),) * n_sym)
+        im_f = np.concatenate((im_f.flatten(),) * sym_order)
 
         vol = anufft(im_f, pts_rot[::-1], (L, L, L), real=True) / L
 
diff --git a/src/aspire/reconstruction/mean.py b/src/aspire/reconstruction/mean.py
index 6791f55d62..f3ecddb19e 100644
--- a/src/aspire/reconstruction/mean.py
+++ b/src/aspire/reconstruction/mean.py
@@ -90,10 +90,24 @@ def _compute_kernel(self):
         # Note, because we're iteratively summing it is critical we zero this array.
         kernel = np.zeros((self.r, self.r, _2L, _2L, _2L), dtype=self.dtype)
 
+        # Handle symmetry boosting.
+        sym_rots = np.eye(3, dtype=self.dtype)[None]
+        if self.boost:
+            sym_rots = self.src.symmetry_group.matrices
+        sym_order = len(sym_rots)
+
         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)
 
-            pts_rot = rotated_grids(self.src.L, self.src.rotations[_range, :, :])
+            # Apply symmetry to rotations.
+            n_rots_batch = len(_range)
+            rotations = np.zeros((sym_order * n_rots_batch, 3, 3), dtype=self.dtype)
+            for i, sym_rot in enumerate(sym_rots):
+                rotations[i * n_rots_batch : (i + 1) * n_rots_batch] = (
+                    sym_rot @ self.src.rotations[_range]
+                )
+
+            pts_rot = rotated_grids(self.src.L, rotations)
             pts_rot = pts_rot.reshape((3, -1))
             assert pts_rot.dtype == self.dtype
 
@@ -117,6 +131,9 @@ def _compute_kernel(self):
 
                     weights = np.transpose(weights, (2, 0, 1)).flatten()
 
+                    # Apply boosting to weights.
+                    weights = np.concatenate((weights,) * sym_order)
+
                     batch_kernel = (
                         1
                         / (self.r * self.src.L**4)
@@ -153,9 +170,12 @@ def src_backward(self):
         :return: The adjoint mapping applied to the images, averaged over the whole dataset and expressed
             as coefficients of `basis`.
         """
+        # Handle symmetry boosting.
         symmetry_group = None
+        sym_order = 1
         if self.boost:
             symmetry_group = self.src.symmetry_group
+            sym_order = len(symmetry_group.matrices)
 
         # src_vols_wt_backward
         vol_rhs = Volume(
@@ -166,18 +186,15 @@ def src_backward(self):
             for k in range(self.r):
                 im = self.src.images[i : i + self.batch_size]
 
-                batch_vol_rhs = (
-                    self.src.im_backward(
-                        im,
-                        i,
-                        self.weights[:, k],
-                        symmetry_group=symmetry_group,
-                    )
-                    / self.src.n
-                )
+                batch_vol_rhs = self.src.im_backward(
+                    im,
+                    i,
+                    self.weights[:, k],
+                    symmetry_group=symmetry_group,
+                ) / (self.src.n * sym_order)
                 vol_rhs[k] += batch_vol_rhs.astype(self.dtype)
 
-        res = np.sqrt(self.src.n) * self.basis.evaluate_t(vol_rhs)
+        res = np.sqrt(self.src.n * sym_order) * self.basis.evaluate_t(vol_rhs)
         logger.info(f"Determined weighted adjoint mappings. Shape = {res.shape}")
 
         return res

From 27c98cb8569171ec960204b3a9f21770100892a0 Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Wed, 17 Jan 2024 11:11:12 -0500
Subject: [PATCH 04/29] Add IdentityFilter for Simulation. Fixes issue 978.

---
 src/aspire/source/simulation.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/src/aspire/source/simulation.py b/src/aspire/source/simulation.py
index e2ef10da12..90facc6b67 100644
--- a/src/aspire/source/simulation.py
+++ b/src/aspire/source/simulation.py
@@ -7,6 +7,7 @@
 
 from aspire.image import Image
 from aspire.noise import NoiseAdder
+from aspire.operators import IdentityFilter
 from aspire.source import ImageSource
 from aspire.source.image import _ImageAccessor
 from aspire.utils import (
@@ -151,7 +152,9 @@ def __init__(
         self.angles = self._init_angles(angles)
 
         if unique_filters is None:
-            unique_filters = []
+            # Use IdentityFilter to pass unharmed through filter eval code
+            # that is potentially called by other methods later.
+            unique_filters = [IdentityFilter()]
         self.unique_filters = unique_filters
         # sim_filters must be a deep copy so that it is not changed
         # when unique_filters is changed

From 9059bc881f76d86ed671dc05be065056a276c4be Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Wed, 17 Jan 2024 12:05:52 -0500
Subject: [PATCH 05/29] Adjust weights in MeanEstimator.

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

diff --git a/src/aspire/reconstruction/mean.py b/src/aspire/reconstruction/mean.py
index f3ecddb19e..5dc43d8693 100644
--- a/src/aspire/reconstruction/mean.py
+++ b/src/aspire/reconstruction/mean.py
@@ -307,7 +307,8 @@ class MeanEstimator(WeightedVolumesEstimator):
     """
 
     def __init__(self, src, basis, **kwargs):
-        weights = np.ones((src.n, 1)) / np.sqrt(src.n)
+        # 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))
         super().__init__(weights, src, basis, **kwargs)
 
     def __getattr__(self, name):

From 93907a9e6b3f08f929a5c1bce10fa98da3ddfd8c Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Wed, 17 Jan 2024 14:31:18 -0500
Subject: [PATCH 06/29] Add testing for boosted MeanEstimator.

---
 src/aspire/reconstruction/mean.py     |   4 +-
 tests/test_mean_estimator_boosting.py | 104 ++++++++++++++++++++++++++
 2 files changed, 107 insertions(+), 1 deletion(-)
 create mode 100644 tests/test_mean_estimator_boosting.py

diff --git a/src/aspire/reconstruction/mean.py b/src/aspire/reconstruction/mean.py
index 5dc43d8693..bc47e0668d 100644
--- a/src/aspire/reconstruction/mean.py
+++ b/src/aspire/reconstruction/mean.py
@@ -308,7 +308,9 @@ class MeanEstimator(WeightedVolumesEstimator):
 
     def __init__(self, src, basis, **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)) / np.sqrt(
+            src.n * len(src.symmetry_group.matrices)
+        )
         super().__init__(weights, src, basis, **kwargs)
 
     def __getattr__(self, name):
diff --git a/tests/test_mean_estimator_boosting.py b/tests/test_mean_estimator_boosting.py
new file mode 100644
index 0000000000..5e8973e190
--- /dev/null
+++ b/tests/test_mean_estimator_boosting.py
@@ -0,0 +1,104 @@
+import numpy as np
+import pytest
+
+from aspire.basis import FFBBasis3D
+from aspire.reconstruction import MeanEstimator
+from aspire.source import Simulation
+from aspire.volume import (
+    AsymmetricVolume,
+    CnSymmetricVolume,
+    DnSymmetricVolume,
+    OSymmetricVolume,
+    TSymmetricVolume,
+)
+
+SEED = 23
+
+RESOLUTION = [
+    32,
+    33,
+]
+
+DTYPE = [
+    np.float32,
+    pytest.param(np.float64, marks=pytest.mark.expensive),
+]
+
+# Symmetric volume parameters, (volume_type, symmetric_order).
+VOL_PARAMS = [
+    (AsymmetricVolume, None),
+    (CnSymmetricVolume, 4),
+    (CnSymmetricVolume, 5),
+    (DnSymmetricVolume, 2),
+    pytest.param((TSymmetricVolume, None), marks=pytest.mark.expensive),
+    pytest.param((OSymmetricVolume, None), marks=pytest.mark.expensive),
+]
+
+
+# Fixtures.
+@pytest.fixture(params=RESOLUTION, ids=lambda x: f"resolution={x}", scope="module")
+def resolution(request):
+    return request.param
+
+
+@pytest.fixture(params=DTYPE, ids=lambda x: f"dtype={x}", scope="module")
+def dtype(request):
+    return request.param
+
+
+@pytest.fixture(params=VOL_PARAMS, ids=lambda x: f"volume={x[0]}", scope="module")
+def volume(request, resolution, dtype):
+    Volume, order = request.param
+    vol_kwargs = dict(
+        L=resolution,
+        C=1,
+        seed=SEED,
+        dtype=dtype,
+    )
+    if order:
+        vol_kwargs["order"] = order
+
+    return Volume(**vol_kwargs).generate()
+
+
+@pytest.fixture(scope="module")
+def source(volume):
+    src = Simulation(
+        n=200,
+        vols=volume,
+        offsets=0,
+        amplitudes=1,
+        seed=SEED,
+        dtype=volume.dtype,
+    )
+
+    return src
+
+
+@pytest.fixture(scope="module")
+def estimated_volume(source):
+    basis = FFBBasis3D(source.L, dtype=source.dtype)
+    estimator = MeanEstimator(source, basis)
+    estimated_volume = estimator.estimate()
+
+    return estimated_volume
+
+
+# MeanEstimator Tests.
+def test_mean_estimator_boosting(source, estimated_volume):
+    """Test MeanEstimator with boosting."""
+    # Fourier Shell Correlation
+    fsc_resolution, fsc = source.vols.fsc(estimated_volume, pixel_size=1, cutoff=0.5)
+
+    # Check that resolution is less than 2.1 pixels.
+    np.testing.assert_array_less(fsc_resolution, 2.1)
+
+    # Check that second to last correlation value is high (>.90).
+    np.testing.assert_array_less(0.90, fsc[0, -2])
+
+
+def test_total_energy(source, estimated_volume):
+    """Test that energy is preserved in reconstructed volume."""
+    og_total_energy = np.sum(source.vols)
+    recon_total_energy = np.sum(estimated_volume)
+    np.testing.assert_allclose(og_total_energy, recon_total_energy, rtol=1e-3)

From 8ff66f0bbe4ec1d2055bd499ca24d83a5b33f01f Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Wed, 17 Jan 2024 15:03:43 -0500
Subject: [PATCH 07/29] Add to backproject docstring.

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

diff --git a/src/aspire/image/image.py b/src/aspire/image/image.py
index ca1cadd48a..419c75c543 100644
--- a/src/aspire/image/image.py
+++ b/src/aspire/image/image.py
@@ -491,11 +491,15 @@ def size(self):
 
     def backproject(self, rot_matrices, symmetry_group=None):
         """
-        Backproject images along rotation
+        Backproject images along rotation. If a symmetry group is provided, images
+        used in back-projection are duplicated (boosted) for symmetric viewing directions.
+        Note, it is assumed that a main axis of symmetry aligns with the z-axis.
 
         :param im: An Image (stack) to backproject.
-        :param rot_matrices: An n-by-3-by-3 array of rotation matrices \
-        corresponding to viewing directions.
+        :param rot_matrices: An n-by-3-by-3 array of rotation matrices
+            corresponding to viewing directions.
+        :param symmetry_group: A SymmetryGroup instance. If supplied,
+            uses symmetry to increase number of images used in back-projeciton.
 
         :return: Volume instance corresonding to the backprojected images.
         """

From 4bdd0eb85d2fe09b41f1c8bbc5bd2f905bfef034 Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Wed, 17 Jan 2024 15:28:33 -0500
Subject: [PATCH 08/29] SymmetryGroup Check in backproject. With test.

---
 src/aspire/image/image.py |  5 +++++
 tests/test_image.py       | 13 +++++++++++++
 2 files changed, 18 insertions(+)

diff --git a/src/aspire/image/image.py b/src/aspire/image/image.py
index 419c75c543..69d28cec15 100644
--- a/src/aspire/image/image.py
+++ b/src/aspire/image/image.py
@@ -12,6 +12,7 @@
 from aspire.nufft import anufft
 from aspire.numeric import fft, xp
 from aspire.utils import FourierRingCorrelation, anorm, crop_pad_2d, grid_2d
+from aspire.volume import SymmetryGroup
 
 logger = logging.getLogger(__name__)
 
@@ -519,6 +520,10 @@ def backproject(self, rot_matrices, symmetry_group=None):
         if symmetry_group is None:
             symmetry_rots = np.eye(3, dtype=self.dtype)[None]
         else:
+            if not isinstance(symmetry_group, SymmetryGroup):
+                raise TypeError(
+                    f"`symmetry_group` must be a `SymmetryGroup` instance. Found {type(symmetry_group)}."
+                )
             symmetry_rots = symmetry_group.matrices
 
         sym_order = len(symmetry_rots)
diff --git a/tests/test_image.py b/tests/test_image.py
index 75e511cb3e..341c8fb6f9 100644
--- a/tests/test_image.py
+++ b/tests/test_image.py
@@ -296,6 +296,19 @@ def testShow():
     im.show()
 
 
+def test_backproject_symmetry_group_error():
+    """
+    Providing non-SymmetryGroup instance to backproject should raise an error.
+    """
+    ary = np.random.random((3, 8, 8))
+    im = Image(ary)
+    rots = np.random.random((3, 3, 3))
+
+    # Attempt backproject.
+    with raises(TypeError, match=r"`symmetry_group` must be a `SymmetryGroup`*"):
+        _ = im.backproject(rots, symmetry_group="Junk")
+
+
 def test_asnumpy_readonly():
     """
     Attempting assignment should raise an error.

From 19e139fffdc148ff8910b31127d3b6aa393dab0f Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Fri, 19 Jan 2024 09:15:24 -0500
Subject: [PATCH 09/29] revert Simulation unique_filters.

---
 src/aspire/source/simulation.py | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

diff --git a/src/aspire/source/simulation.py b/src/aspire/source/simulation.py
index 90facc6b67..e2ef10da12 100644
--- a/src/aspire/source/simulation.py
+++ b/src/aspire/source/simulation.py
@@ -7,7 +7,6 @@
 
 from aspire.image import Image
 from aspire.noise import NoiseAdder
-from aspire.operators import IdentityFilter
 from aspire.source import ImageSource
 from aspire.source.image import _ImageAccessor
 from aspire.utils import (
@@ -152,9 +151,7 @@ def __init__(
         self.angles = self._init_angles(angles)
 
         if unique_filters is None:
-            # Use IdentityFilter to pass unharmed through filter eval code
-            # that is potentially called by other methods later.
-            unique_filters = [IdentityFilter()]
+            unique_filters = []
         self.unique_filters = unique_filters
         # sim_filters must be a deep copy so that it is not changed
         # when unique_filters is changed

From f94aec717fe927b0eb575eaeef48db6d41a287a3 Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Fri, 19 Jan 2024 09:27:24 -0500
Subject: [PATCH 10/29] Temp fix for evaluate_src_filters_on_grid.

---
 tests/test_mean_estimator_boosting.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tests/test_mean_estimator_boosting.py b/tests/test_mean_estimator_boosting.py
index 5e8973e190..3a18774c30 100644
--- a/tests/test_mean_estimator_boosting.py
+++ b/tests/test_mean_estimator_boosting.py
@@ -2,6 +2,7 @@
 import pytest
 
 from aspire.basis import FFBBasis3D
+from aspire.operators import IdentityFilter
 from aspire.reconstruction import MeanEstimator
 from aspire.source import Simulation
 from aspire.volume import (
@@ -70,6 +71,7 @@ def source(volume):
         amplitudes=1,
         seed=SEED,
         dtype=volume.dtype,
+        unique_filters=[IdentityFilter()],  # Can remove after PR 1076
     )
 
     return src

From 15801e1615ef5f2cdca22d566c4720bf2c7540f8 Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Fri, 19 Jan 2024 10:19:02 -0500
Subject: [PATCH 11/29] accept symmetry_group strings in backproject.

---
 src/aspire/image/image.py | 11 +++++++----
 tests/test_image.py       |  3 ++-
 2 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/src/aspire/image/image.py b/src/aspire/image/image.py
index 69d28cec15..14ebbb1c84 100644
--- a/src/aspire/image/image.py
+++ b/src/aspire/image/image.py
@@ -492,15 +492,14 @@ def size(self):
 
     def backproject(self, rot_matrices, symmetry_group=None):
         """
-        Backproject images along rotation. If a symmetry group is provided, images
+        Backproject images along rotations. If a symmetry group is provided, images
         used in back-projection are duplicated (boosted) for symmetric viewing directions.
         Note, it is assumed that a main axis of symmetry aligns with the z-axis.
 
-        :param im: An Image (stack) to backproject.
         :param rot_matrices: An n-by-3-by-3 array of rotation matrices
             corresponding to viewing directions.
-        :param symmetry_group: A SymmetryGroup instance. If supplied,
-            uses symmetry to increase number of images used in back-projeciton.
+        :param symmetry_group: A SymmetryGroup instance or string indicating symmetry, ie. "C3".
+            If supplied, uses symmetry to increase number of images used in back-projeciton.
 
         :return: Volume instance corresonding to the backprojected images.
         """
@@ -520,6 +519,10 @@ def backproject(self, rot_matrices, symmetry_group=None):
         if symmetry_group is None:
             symmetry_rots = np.eye(3, dtype=self.dtype)[None]
         else:
+            if isinstance(symmetry_group, str):
+                symmetry_group = SymmetryGroup.from_string(
+                    symmetry_group, dtype=self.dtype
+                )
             if not isinstance(symmetry_group, SymmetryGroup):
                 raise TypeError(
                     f"`symmetry_group` must be a `SymmetryGroup` instance. Found {type(symmetry_group)}."
diff --git a/tests/test_image.py b/tests/test_image.py
index 341c8fb6f9..c090a0af60 100644
--- a/tests/test_image.py
+++ b/tests/test_image.py
@@ -303,10 +303,11 @@ def test_backproject_symmetry_group_error():
     ary = np.random.random((3, 8, 8))
     im = Image(ary)
     rots = np.random.random((3, 3, 3))
+    not_a_symmetry_group = []
 
     # Attempt backproject.
     with raises(TypeError, match=r"`symmetry_group` must be a `SymmetryGroup`*"):
-        _ = im.backproject(rots, symmetry_group="Junk")
+        _ = im.backproject(rots, symmetry_group=not_a_symmetry_group)
 
 
 def test_asnumpy_readonly():

From 6f54f61bd3945c52114f46b190423baa2efbf6e0 Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Fri, 19 Jan 2024 11:05:44 -0500
Subject: [PATCH 12/29] Add mse test. Add docstring param.

---
 src/aspire/reconstruction/estimator.py |  2 ++
 tests/test_mean_estimator_boosting.py  | 13 +++++++++++--
 2 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/src/aspire/reconstruction/estimator.py b/src/aspire/reconstruction/estimator.py
index 05f60b69d4..b0832be174 100644
--- a/src/aspire/reconstruction/estimator.py
+++ b/src/aspire/reconstruction/estimator.py
@@ -45,6 +45,8 @@ def __init__(
             before returning.  This should be used in conjunction with
             `checkpoint_iterations` to prevent excessive disk usage.
             `None` disables.
+        :param boost: Option to use `src` symmetry to boost number of images used for mean estimation (Boolean).
+            Default of `True` employs symmetry boosting.
         """
 
         self.src = src
diff --git a/tests/test_mean_estimator_boosting.py b/tests/test_mean_estimator_boosting.py
index 3a18774c30..69a2b99a91 100644
--- a/tests/test_mean_estimator_boosting.py
+++ b/tests/test_mean_estimator_boosting.py
@@ -87,8 +87,8 @@ def estimated_volume(source):
 
 
 # MeanEstimator Tests.
-def test_mean_estimator_boosting(source, estimated_volume):
-    """Test MeanEstimator with boosting."""
+def test_fsc(source, estimated_volume):
+    """Compare estimated volume to source volume with FSC."""
     # Fourier Shell Correlation
     fsc_resolution, fsc = source.vols.fsc(estimated_volume, pixel_size=1, cutoff=0.5)
 
@@ -99,6 +99,15 @@ def test_mean_estimator_boosting(source, estimated_volume):
     np.testing.assert_array_less(0.90, fsc[0, -2])
 
 
+def test_mse(source, estimated_volume):
+    """Check the mean-squared error between source and estimated volumes."""
+    mse = (
+        np.sum((source.vols.asnumpy() - estimated_volume.asnumpy()) ** 2)
+        / source.L**3
+    )
+    np.testing.assert_allclose(mse, 0, atol=1e-3)
+
+
 def test_total_energy(source, estimated_volume):
     """Test that energy is preserved in reconstructed volume."""
     og_total_energy = np.sum(source.vols)

From 4f915fea47ad5b8b59066bd5139db617f7c2719c Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Fri, 19 Jan 2024 11:32:47 -0500
Subject: [PATCH 13/29] refactor mse

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

diff --git a/tests/test_mean_estimator_boosting.py b/tests/test_mean_estimator_boosting.py
index 69a2b99a91..a456df5da9 100644
--- a/tests/test_mean_estimator_boosting.py
+++ b/tests/test_mean_estimator_boosting.py
@@ -101,10 +101,7 @@ def test_fsc(source, estimated_volume):
 
 def test_mse(source, estimated_volume):
     """Check the mean-squared error between source and estimated volumes."""
-    mse = (
-        np.sum((source.vols.asnumpy() - estimated_volume.asnumpy()) ** 2)
-        / source.L**3
-    )
+    mse = np.mean((source.vols.asnumpy() - estimated_volume.asnumpy()) ** 2)
     np.testing.assert_allclose(mse, 0, atol=1e-3)
 
 

From fdd14628638958c2fc6a2162d8f724c24c40ab69 Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Fri, 19 Jan 2024 14:31:14 -0500
Subject: [PATCH 14/29] backproject return vol with symmetry group.

---
 src/aspire/image/image.py |  2 +-
 tests/test_image.py       | 23 +++++++++++++++++------
 2 files changed, 18 insertions(+), 7 deletions(-)

diff --git a/src/aspire/image/image.py b/src/aspire/image/image.py
index 14ebbb1c84..650a377c85 100644
--- a/src/aspire/image/image.py
+++ b/src/aspire/image/image.py
@@ -552,7 +552,7 @@ def backproject(self, rot_matrices, symmetry_group=None):
 
         vol = anufft(im_f, pts_rot[::-1], (L, L, L), real=True) / L
 
-        return aspire.volume.Volume(vol)
+        return aspire.volume.Volume(vol, symmetry_group=symmetry_group)
 
     def show(self, columns=5, figsize=(20, 10), colorbar=True):
         """
diff --git a/tests/test_image.py b/tests/test_image.py
index c090a0af60..310c06630e 100644
--- a/tests/test_image.py
+++ b/tests/test_image.py
@@ -10,7 +10,8 @@
 from scipy.datasets import face
 
 from aspire.image import Image
-from aspire.utils import powerset, utest_tolerance
+from aspire.utils import Rotation, powerset, utest_tolerance
+from aspire.volume import CnSymmetryGroup
 
 from .test_utils import matplotlib_dry_run
 
@@ -296,19 +297,29 @@ def testShow():
     im.show()
 
 
-def test_backproject_symmetry_group_error():
+def test_backproject_symmetry_group():
     """
     Providing non-SymmetryGroup instance to backproject should raise an error.
     """
-    ary = np.random.random((3, 8, 8))
+    ary = np.random.random((5, 8, 8))
     im = Image(ary)
-    rots = np.random.random((3, 3, 3))
-    not_a_symmetry_group = []
+    rots = Rotation.generate_random_rotations(5).matrices
 
-    # Attempt backproject.
+    # Attempt backproject with bad symmetry group.
+    not_a_symmetry_group = []
     with raises(TypeError, match=r"`symmetry_group` must be a `SymmetryGroup`*"):
         _ = im.backproject(rots, symmetry_group=not_a_symmetry_group)
 
+    # Symmetry from string.
+    vol = im.backproject(rots, symmetry_group="C3")
+    assert isinstance(vol.symmetry_group, CnSymmetryGroup)
+
+    # Symmetry from instance.
+    vol = im.backproject(
+        rots, symmetry_group=CnSymmetryGroup(order=3, dtype=np.float32)
+    )
+    assert isinstance(vol.symmetry_group, CnSymmetryGroup)
+
 
 def test_asnumpy_readonly():
     """

From 78c8d8deb3cbc3fd0552bea820791be7d8f4bd4d Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Mon, 22 Jan 2024 14:38:19 -0500
Subject: [PATCH 15/29] Use tile instead of concatenate.

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

diff --git a/src/aspire/image/image.py b/src/aspire/image/image.py
index 650a377c85..b328336111 100644
--- a/src/aspire/image/image.py
+++ b/src/aspire/image/image.py
@@ -548,7 +548,7 @@ def backproject(self, rot_matrices, symmetry_group=None):
             im_f[:, :, 0] = 0
 
         # Apply boosting to images.
-        im_f = np.concatenate((im_f.flatten(),) * sym_order)
+        im_f = np.tile(im_f.flatten(), sym_order)
 
         vol = anufft(im_f, pts_rot[::-1], (L, L, L), real=True) / L
 
diff --git a/src/aspire/reconstruction/mean.py b/src/aspire/reconstruction/mean.py
index bc47e0668d..87df705b46 100644
--- a/src/aspire/reconstruction/mean.py
+++ b/src/aspire/reconstruction/mean.py
@@ -132,7 +132,7 @@ def _compute_kernel(self):
                     weights = np.transpose(weights, (2, 0, 1)).flatten()
 
                     # Apply boosting to weights.
-                    weights = np.concatenate((weights,) * sym_order)
+                    weights = np.tile(weights, sym_order)
 
                     batch_kernel = (
                         1

From 27a36084a7ba8e3daaba5883771335f881d948da Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Tue, 23 Jan 2024 10:12:07 -0500
Subject: [PATCH 16/29] Add boost flag test.

---
 tests/test_mean_estimator_boosting.py | 28 +++++++++++++++++++++++++--
 1 file changed, 26 insertions(+), 2 deletions(-)

diff --git a/tests/test_mean_estimator_boosting.py b/tests/test_mean_estimator_boosting.py
index a456df5da9..b9c896654c 100644
--- a/tests/test_mean_estimator_boosting.py
+++ b/tests/test_mean_estimator_boosting.py
@@ -4,7 +4,8 @@
 from aspire.basis import FFBBasis3D
 from aspire.operators import IdentityFilter
 from aspire.reconstruction import MeanEstimator
-from aspire.source import Simulation
+from aspire.source import ArrayImageSource, Simulation
+from aspire.utils import Rotation
 from aspire.volume import (
     AsymmetricVolume,
     CnSymmetricVolume,
@@ -47,7 +48,7 @@ def dtype(request):
     return request.param
 
 
-@pytest.fixture(params=VOL_PARAMS, ids=lambda x: f"volume={x[0]}", scope="module")
+@pytest.fixture(params=VOL_PARAMS, ids=lambda x: f"volume={x[0]}, order={x[1]}", scope="module")
 def volume(request, resolution, dtype):
     Volume, order = request.param
     vol_kwargs = dict(
@@ -110,3 +111,26 @@ def test_total_energy(source, estimated_volume):
     og_total_energy = np.sum(source.vols)
     recon_total_energy = np.sum(estimated_volume)
     np.testing.assert_allclose(og_total_energy, recon_total_energy, rtol=1e-3)
+
+
+def test_boost_flag(source, estimated_volume):
+    """Manually boost a source and reconstruct without boosting."""
+    ims = source.projections[:]
+    rots = source.rotations
+    sym_order = len(source.symmetry_group.matrices)
+
+    # Manually boosted images and rotations.
+    ims_boosted = np.tile(ims, (sym_order, 1, 1))
+    rots_boosted = Rotation(np.tile(rots, (sym_order, 1, 1)))
+
+    # Manually boosted source.
+    boosted_source = ArrayImageSource(ims_boosted, angles=rots_boosted.angles)
+
+    # Estimate volume with boosting OFF.
+    basis = FFBBasis3D(boosted_source.L, dtype=boosted_source.dtype)
+    estimator = MeanEstimator(boosted_source, basis, boost=False)
+    est_vol = estimator.estimate()
+
+    # Check reconstructions are close.
+    mse = np.mean((estimated_volume.asnumpy() - est_vol.asnumpy()) ** 2)
+    np.testing.assert_array_less(mse, 1e-4)

From a5c50b87528027474557d0620c7737fa08d8e1b8 Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Tue, 23 Jan 2024 14:08:53 -0500
Subject: [PATCH 17/29] Add WeightedVolumesEstimator boosting test.

---
 tests/test_mean_estimator_boosting.py | 66 +++++++++++++++++++++++++--
 1 file changed, 63 insertions(+), 3 deletions(-)

diff --git a/tests/test_mean_estimator_boosting.py b/tests/test_mean_estimator_boosting.py
index b9c896654c..1532f7bfa4 100644
--- a/tests/test_mean_estimator_boosting.py
+++ b/tests/test_mean_estimator_boosting.py
@@ -3,7 +3,7 @@
 
 from aspire.basis import FFBBasis3D
 from aspire.operators import IdentityFilter
-from aspire.reconstruction import MeanEstimator
+from aspire.reconstruction import MeanEstimator, WeightedVolumesEstimator
 from aspire.source import ArrayImageSource, Simulation
 from aspire.utils import Rotation
 from aspire.volume import (
@@ -18,7 +18,7 @@
 
 RESOLUTION = [
     32,
-    33,
+    pytest.param(33, marks=pytest.mark.expensive),
 ]
 
 DTYPE = [
@@ -48,7 +48,9 @@ def dtype(request):
     return request.param
 
 
-@pytest.fixture(params=VOL_PARAMS, ids=lambda x: f"volume={x[0]}, order={x[1]}", scope="module")
+@pytest.fixture(
+    params=VOL_PARAMS, ids=lambda x: f"volume={x[0]}, order={x[1]}", scope="module"
+)
 def volume(request, resolution, dtype):
     Volume, order = request.param
     vol_kwargs = dict(
@@ -87,6 +89,39 @@ def estimated_volume(source):
     return estimated_volume
 
 
+# Weighted volume fixture. Only tesing C1 and C4.
+@pytest.fixture(
+    params=VOL_PARAMS[:2], ids=lambda x: f"volume={x[0]}, order={x[1]}", scope="module"
+)
+def weighted_volume(request, resolution, dtype):
+    Volume, order = request.param
+    vol_kwargs = dict(
+        L=resolution,
+        C=2,
+        seed=SEED,
+        dtype=dtype,
+    )
+    if order:
+        vol_kwargs["order"] = order
+
+    return Volume(**vol_kwargs).generate()
+
+
+@pytest.fixture(scope="module")
+def weighted_source(weighted_volume):
+    src = Simulation(
+        n=400,
+        vols=weighted_volume,
+        offsets=0,
+        amplitudes=1,
+        seed=SEED,
+        dtype=weighted_volume.dtype,
+        unique_filters=[IdentityFilter(), IdentityFilter()],  # Can remove after PR 1076
+    )
+
+    return src
+
+
 # MeanEstimator Tests.
 def test_fsc(source, estimated_volume):
     """Compare estimated volume to source volume with FSC."""
@@ -134,3 +169,28 @@ def test_boost_flag(source, estimated_volume):
     # Check reconstructions are close.
     mse = np.mean((estimated_volume.asnumpy() - est_vol.asnumpy()) ** 2)
     np.testing.assert_array_less(mse, 1e-4)
+
+
+# WeightVolumesEstimator Tests.
+def test_weighted_volumes(weighted_source):
+    src = weighted_source
+
+    # Use source states to assign weights to volumes.
+    weights = np.zeros((src.n, src.C), dtype=src.dtype)
+    weights[:, 0] = abs(src.states - 1.99)  # sends states [1, 2] to weights [.99, .01]
+    weights[:, 1] = abs(-src.states + 1.01)  # sends states [1, 2] to weights [.01, .99]
+
+    # Scale weights
+    n0 = (-src.states + 2).sum()  # number of images from vol[0]
+    n1 = (src.states - 1).sum()  # number of images from vol[1]
+    weights[:, 0] = weights[:, 0] / weights[:, 0].sum() * np.sqrt(n0)
+    weights[:, 1] = weights[:, 1] / weights[:, 1].sum() * np.sqrt(n1)
+
+    # Initialize estimator.
+    basis = FFBBasis3D(src.L, dtype=src.dtype)
+    estimator = WeightedVolumesEstimator(src=src, basis=basis, weights=weights)
+    est_vols = estimator.estimate()
+
+    # Check FSC (scaling may not be close enough to match mse)
+    _, corr = src.vols.fsc(est_vols)
+    np.testing.assert_array_less(0.95, corr[:, -2])

From cbad44d4b53059e59dfe5020b9050a7b71a0e097 Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Fri, 26 Jan 2024 10:31:12 -0500
Subject: [PATCH 18/29] Apply backproject boosting in loop.

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

diff --git a/src/aspire/image/image.py b/src/aspire/image/image.py
index b328336111..075da20894 100644
--- a/src/aspire/image/image.py
+++ b/src/aspire/image/image.py
@@ -515,7 +515,7 @@ def backproject(self, rot_matrices, symmetry_group=None):
             self.n_images == rot_matrices.shape[0]
         ), "Number of rotation matrices must match the number of images"
 
-        # Apply symmetry boosting to rotations.
+        # Get symmetry rotations from SymmetryGroup.
         if symmetry_group is None:
             symmetry_rots = np.eye(3, dtype=self.dtype)[None]
         else:
@@ -529,28 +529,29 @@ def backproject(self, rot_matrices, symmetry_group=None):
                 )
             symmetry_rots = symmetry_group.matrices
 
-        sym_order = len(symmetry_rots)
-        boosted_rot_mats = np.zeros((sym_order * self.shape[0], 3, 3), dtype=self.dtype)
-        for i, sym_rot in enumerate(symmetry_rots):
-            boosted_rot_mats[i * self.shape[0] : (i + 1) * self.shape[0]] = (
-                sym_rot @ rot_matrices
-            )
-
-        # TODO: rotated_grids might as well give us correctly shaped array in the first place
-        pts_rot = aspire.volume.rotated_grids(L, boosted_rot_mats).astype(
-            self.dtype, copy=False
-        )
-        pts_rot = pts_rot.reshape((3, -1))
-
+        # Compute Fourier transform of images.
         im_f = xp.asnumpy(fft.centered_fft2(xp.asarray(self._data))) / (L**2)
         if L % 2 == 0:
             im_f[:, 0, :] = 0
             im_f[:, :, 0] = 0
 
-        # Apply boosting to images.
-        im_f = np.tile(im_f.flatten(), sym_order)
+        im_f = im_f.flatten()
+
+        # Backproject. Apply boosting by looping over symmetry rotations.
+        sym_order = len(symmetry_rots)
+        vol = np.zeros((L, L, L), dtype=self.dtype)
+        for sym_rot in symmetry_rots:
+            rotations = sym_rot @ rot_matrices
+
+            # 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
+            )
+            pts_rot = pts_rot.reshape((3, -1))
+
+            vol += anufft(im_f, pts_rot[::-1], (L, L, L), real=True)
 
-        vol = anufft(im_f, pts_rot[::-1], (L, L, L), real=True) / L
+        vol /= L
 
         return aspire.volume.Volume(vol, symmetry_group=symmetry_group)
 

From 2a7d9a77496441b7fca7e27f72987e4fd0b5a745 Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Fri, 26 Jan 2024 11:01:34 -0500
Subject: [PATCH 19/29] Fix manually boosted rots in test_boost_flag.

---
 tests/test_mean_estimator_boosting.py | 18 +++++++++++-------
 1 file changed, 11 insertions(+), 7 deletions(-)

diff --git a/tests/test_mean_estimator_boosting.py b/tests/test_mean_estimator_boosting.py
index 1532f7bfa4..b6d60d0168 100644
--- a/tests/test_mean_estimator_boosting.py
+++ b/tests/test_mean_estimator_boosting.py
@@ -5,7 +5,7 @@
 from aspire.operators import IdentityFilter
 from aspire.reconstruction import MeanEstimator, WeightedVolumesEstimator
 from aspire.source import ArrayImageSource, Simulation
-from aspire.utils import Rotation
+from aspire.utils import Rotation, utest_tolerance
 from aspire.volume import (
     AsymmetricVolume,
     CnSymmetricVolume,
@@ -89,9 +89,9 @@ def estimated_volume(source):
     return estimated_volume
 
 
-# Weighted volume fixture. Only tesing C1 and C4.
+# Weighted volume fixture.
 @pytest.fixture(
-    params=VOL_PARAMS[:2], ids=lambda x: f"volume={x[0]}, order={x[1]}", scope="module"
+    params=VOL_PARAMS, ids=lambda x: f"volume={x[0]}, order={x[1]}", scope="module"
 )
 def weighted_volume(request, resolution, dtype):
     Volume, order = request.param
@@ -152,11 +152,15 @@ def test_boost_flag(source, estimated_volume):
     """Manually boost a source and reconstruct without boosting."""
     ims = source.projections[:]
     rots = source.rotations
-    sym_order = len(source.symmetry_group.matrices)
+    sym_rots = source.symmetry_group.matrices
+    sym_order = len(sym_rots)
 
     # Manually boosted images and rotations.
     ims_boosted = np.tile(ims, (sym_order, 1, 1))
-    rots_boosted = Rotation(np.tile(rots, (sym_order, 1, 1)))
+    rots_boosted = np.zeros((sym_order * source.n, 3, 3), dtype=source.dtype)
+    for i, sym_rot in enumerate(sym_rots):
+        rots_boosted[i * source.n : (i + 1) * source.n] = sym_rot @ rots
+    rots_boosted = Rotation(rots_boosted)
 
     # Manually boosted source.
     boosted_source = ArrayImageSource(ims_boosted, angles=rots_boosted.angles)
@@ -166,9 +170,9 @@ def test_boost_flag(source, estimated_volume):
     estimator = MeanEstimator(boosted_source, basis, boost=False)
     est_vol = estimator.estimate()
 
-    # Check reconstructions are close.
+    # Check reconstructions are equal.
     mse = np.mean((estimated_volume.asnumpy() - est_vol.asnumpy()) ** 2)
-    np.testing.assert_array_less(mse, 1e-4)
+    np.testing.assert_allclose(mse, 0, atol=utest_tolerance(source.dtype))
 
 
 # WeightVolumesEstimator Tests.

From 93e64fb19fe8b26a16bd4c9458d6df597e18b066 Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Fri, 26 Jan 2024 11:46:42 -0500
Subject: [PATCH 20/29] Boost in loop in compute_kernel.

---
 src/aspire/image/image.py             |  1 -
 src/aspire/reconstruction/mean.py     | 34 ++++++++++-----------------
 tests/test_mean_estimator_boosting.py |  6 ++---
 3 files changed, 15 insertions(+), 26 deletions(-)

diff --git a/src/aspire/image/image.py b/src/aspire/image/image.py
index 075da20894..f49d8a1eb1 100644
--- a/src/aspire/image/image.py
+++ b/src/aspire/image/image.py
@@ -538,7 +538,6 @@ def backproject(self, rot_matrices, symmetry_group=None):
         im_f = im_f.flatten()
 
         # Backproject. Apply boosting by looping over symmetry rotations.
-        sym_order = len(symmetry_rots)
         vol = np.zeros((L, L, L), dtype=self.dtype)
         for sym_rot in symmetry_rots:
             rotations = sym_rot @ rot_matrices
diff --git a/src/aspire/reconstruction/mean.py b/src/aspire/reconstruction/mean.py
index 87df705b46..760bca3e3d 100644
--- a/src/aspire/reconstruction/mean.py
+++ b/src/aspire/reconstruction/mean.py
@@ -94,25 +94,10 @@ def _compute_kernel(self):
         sym_rots = np.eye(3, dtype=self.dtype)[None]
         if self.boost:
             sym_rots = self.src.symmetry_group.matrices
-        sym_order = len(sym_rots)
 
         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)
-
-            # Apply symmetry to rotations.
-            n_rots_batch = len(_range)
-            rotations = np.zeros((sym_order * n_rots_batch, 3, 3), dtype=self.dtype)
-            for i, sym_rot in enumerate(sym_rots):
-                rotations[i * n_rots_batch : (i + 1) * n_rots_batch] = (
-                    sym_rot @ self.src.rotations[_range]
-                )
-
-            pts_rot = rotated_grids(self.src.L, rotations)
-            pts_rot = pts_rot.reshape((3, -1))
-            assert pts_rot.dtype == self.dtype
-
             sq_filters_f = evaluate_src_filters_on_grid(self.src, _range) ** 2
-
             amplitudes_sq = (self.src.amplitudes[_range] ** 2).astype(
                 self.dtype, copy=False
             )
@@ -131,14 +116,19 @@ def _compute_kernel(self):
 
                     weights = np.transpose(weights, (2, 0, 1)).flatten()
 
-                    # Apply boosting to weights.
-                    weights = np.tile(weights, sym_order)
+                    # Apply boosting.
+                    batch_kernel = np.zeros((_2L, _2L, _2L), dtype=self.dtype)
+                    for sym_rot in sym_rots:
+                        rotations = sym_rot @ self.src.rotations[_range]
+                        pts_rot = rotated_grids(self.src.L, rotations)
+                        pts_rot = pts_rot.reshape((3, -1))
+
+                        batch_kernel += (
+                            1
+                            / (self.r * self.src.L**4)
+                            * anufft(weights, pts_rot[::-1], (_2L, _2L, _2L), real=True)
+                        )
 
-                    batch_kernel = (
-                        1
-                        / (self.r * self.src.L**4)
-                        * anufft(weights, pts_rot[::-1], (_2L, _2L, _2L), real=True)
-                    )
                     kernel[k, j] += batch_kernel
 
                     # r x r symmetric
diff --git a/tests/test_mean_estimator_boosting.py b/tests/test_mean_estimator_boosting.py
index b6d60d0168..ef7ef225ff 100644
--- a/tests/test_mean_estimator_boosting.py
+++ b/tests/test_mean_estimator_boosting.py
@@ -89,9 +89,9 @@ def estimated_volume(source):
     return estimated_volume
 
 
-# Weighted volume fixture.
+# Weighted volume fixture. Only tesing C1, C4, and C5.
 @pytest.fixture(
-    params=VOL_PARAMS, ids=lambda x: f"volume={x[0]}, order={x[1]}", scope="module"
+    params=VOL_PARAMS[:3], ids=lambda x: f"volume={x[0]}, order={x[1]}", scope="module"
 )
 def weighted_volume(request, resolution, dtype):
     Volume, order = request.param
@@ -197,4 +197,4 @@ def test_weighted_volumes(weighted_source):
 
     # Check FSC (scaling may not be close enough to match mse)
     _, corr = src.vols.fsc(est_vols)
-    np.testing.assert_array_less(0.95, corr[:, -2])
+    np.testing.assert_array_less(0.91, corr[:, -2])

From 865c4d781d7742f77503bb0cf7d614778c9d682c Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Fri, 26 Jan 2024 13:55:53 -0500
Subject: [PATCH 21/29] remove IdentityFilter in boosting test.

---
 tests/test_mean_estimator_boosting.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/tests/test_mean_estimator_boosting.py b/tests/test_mean_estimator_boosting.py
index ef7ef225ff..91a400a1e1 100644
--- a/tests/test_mean_estimator_boosting.py
+++ b/tests/test_mean_estimator_boosting.py
@@ -74,7 +74,6 @@ def source(volume):
         amplitudes=1,
         seed=SEED,
         dtype=volume.dtype,
-        unique_filters=[IdentityFilter()],  # Can remove after PR 1076
     )
 
     return src
@@ -116,7 +115,6 @@ def weighted_source(weighted_volume):
         amplitudes=1,
         seed=SEED,
         dtype=weighted_volume.dtype,
-        unique_filters=[IdentityFilter(), IdentityFilter()],  # Can remove after PR 1076
     )
 
     return src

From c4a8c351224e0903364cf4961b67aa38e989ab79 Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Fri, 26 Jan 2024 14:00:18 -0500
Subject: [PATCH 22/29] remove import.

---
 tests/test_mean_estimator_boosting.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tests/test_mean_estimator_boosting.py b/tests/test_mean_estimator_boosting.py
index 91a400a1e1..da748cd6ac 100644
--- a/tests/test_mean_estimator_boosting.py
+++ b/tests/test_mean_estimator_boosting.py
@@ -2,7 +2,6 @@
 import pytest
 
 from aspire.basis import FFBBasis3D
-from aspire.operators import IdentityFilter
 from aspire.reconstruction import MeanEstimator, WeightedVolumesEstimator
 from aspire.source import ArrayImageSource, Simulation
 from aspire.utils import Rotation, utest_tolerance

From fadf481500f671699e8d06ee8bc9bd263bb53fa8 Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Mon, 29 Jan 2024 11:31:03 -0500
Subject: [PATCH 23/29] reword docstring

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

diff --git a/tests/test_image.py b/tests/test_image.py
index 310c06630e..0bd217c7c9 100644
--- a/tests/test_image.py
+++ b/tests/test_image.py
@@ -299,7 +299,7 @@ def testShow():
 
 def test_backproject_symmetry_group():
     """
-    Providing non-SymmetryGroup instance to backproject should raise an error.
+    Test backproject SymmetryGroup pass through and error message.
     """
     ary = np.random.random((5, 8, 8))
     im = Image(ary)

From 783534783832e55905fe2ade419ece78835c44f2 Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Mon, 29 Jan 2024 11:33:25 -0500
Subject: [PATCH 24/29] Add testing docstring

---
 tests/test_mean_estimator_boosting.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/tests/test_mean_estimator_boosting.py b/tests/test_mean_estimator_boosting.py
index da748cd6ac..278a6890d0 100644
--- a/tests/test_mean_estimator_boosting.py
+++ b/tests/test_mean_estimator_boosting.py
@@ -174,6 +174,9 @@ def test_boost_flag(source, estimated_volume):
 
 # WeightVolumesEstimator Tests.
 def test_weighted_volumes(weighted_source):
+    """
+    Test WeightedVolumeEstimator reconstructs multiple volumes using symmetry boosting.
+    """
     src = weighted_source
 
     # Use source states to assign weights to volumes.

From 3878cc5e4d3e0fd285f75bdd7b2af82b88aa30a5 Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Thu, 1 Feb 2024 09:00:50 -0500
Subject: [PATCH 25/29] Add missing docstring.

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

diff --git a/src/aspire/source/image.py b/src/aspire/source/image.py
index d7677d07a4..dfed0825f7 100644
--- a/src/aspire/source/image.py
+++ b/src/aspire/source/image.py
@@ -947,7 +947,9 @@ def im_backward(self, im, start, weights=None, symmetry_group=None):
         :param im: An Image instance to which we wish to apply the adjoint of the forward model.
         :param start: Start index of image to consider
         :param weights: Optional vector of weights to apply to images.
-        Weights should be length `self.n`.
+            Weights should be length `self.n`.
+        :param symmetry_group: A SymmetryGroup instance. If supplied, uses symmetry to increase
+             number of images used in back-projectioon.
         :return: An L-by-L-by-L volume containing the sum of the adjoint mappings applied to the start+num-1 images.
         """
         num = im.n_images

From 96954187cfe863ebe4a64c7df66b8303da9fc8a9 Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Thu, 1 Feb 2024 09:05:11 -0500
Subject: [PATCH 26/29] remove unnecesary wildcard

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

diff --git a/tests/test_image.py b/tests/test_image.py
index 0bd217c7c9..ab616d04e6 100644
--- a/tests/test_image.py
+++ b/tests/test_image.py
@@ -307,7 +307,7 @@ def test_backproject_symmetry_group():
 
     # Attempt backproject with bad symmetry group.
     not_a_symmetry_group = []
-    with raises(TypeError, match=r"`symmetry_group` must be a `SymmetryGroup`*"):
+    with raises(TypeError, match=r"`symmetry_group` must be a `SymmetryGroup`"):
         _ = im.backproject(rots, symmetry_group=not_a_symmetry_group)
 
     # Symmetry from string.

From 6a70ee0d02bd7a71080470d1444a93c4f7470035 Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Thu, 1 Feb 2024 09:14:14 -0500
Subject: [PATCH 27/29] Cache testing source.

---
 tests/test_mean_estimator_boosting.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/test_mean_estimator_boosting.py b/tests/test_mean_estimator_boosting.py
index 278a6890d0..8220e1bd5e 100644
--- a/tests/test_mean_estimator_boosting.py
+++ b/tests/test_mean_estimator_boosting.py
@@ -74,6 +74,7 @@ def source(volume):
         seed=SEED,
         dtype=volume.dtype,
     )
+    src = src.cache()  # precompute images
 
     return src
 

From 9f421fe0270c55d5a51c3807e7046548ad508ac6 Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Thu, 1 Feb 2024 09:38:34 -0500
Subject: [PATCH 28/29] Use count instead of confusing math

---
 tests/test_mean_estimator_boosting.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tests/test_mean_estimator_boosting.py b/tests/test_mean_estimator_boosting.py
index 8220e1bd5e..005c0597af 100644
--- a/tests/test_mean_estimator_boosting.py
+++ b/tests/test_mean_estimator_boosting.py
@@ -183,11 +183,11 @@ def test_weighted_volumes(weighted_source):
     # Use source states to assign weights to volumes.
     weights = np.zeros((src.n, src.C), dtype=src.dtype)
     weights[:, 0] = abs(src.states - 1.99)  # sends states [1, 2] to weights [.99, .01]
-    weights[:, 1] = abs(-src.states + 1.01)  # sends states [1, 2] to weights [.01, .99]
+    weights[:, 1] = 1 - weights[:, 0]  # sets weights for states [1, 2] as [.01, .99]
 
     # Scale weights
-    n0 = (-src.states + 2).sum()  # number of images from vol[0]
-    n1 = (src.states - 1).sum()  # number of images from vol[1]
+    n0 = np.count_nonzero(src.states == 1)  # number of images from vol[0]
+    n1 = np.count_nonzero(src.states == 2)  # number of images from vol[1]
     weights[:, 0] = weights[:, 0] / weights[:, 0].sum() * np.sqrt(n0)
     weights[:, 1] = weights[:, 1] / weights[:, 1].sum() * np.sqrt(n1)
 

From 3a7e10391174ceb760fd92829b602c1bd372b236 Mon Sep 17 00:00:00 2001
From: Josh Carmichael <carmichael@princeton.edu>
Date: Thu, 1 Feb 2024 09:40:43 -0500
Subject: [PATCH 29/29] typo

---
 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 f49d8a1eb1..81589b1472 100644
--- a/src/aspire/image/image.py
+++ b/src/aspire/image/image.py
@@ -499,7 +499,7 @@ def backproject(self, rot_matrices, symmetry_group=None):
         :param rot_matrices: An n-by-3-by-3 array of rotation matrices
             corresponding to viewing directions.
         :param symmetry_group: A SymmetryGroup instance or string indicating symmetry, ie. "C3".
-            If supplied, uses symmetry to increase number of images used in back-projeciton.
+            If supplied, uses symmetry to increase number of images used in back-projection.
 
         :return: Volume instance corresonding to the backprojected images.
         """