Repository: ComputationalCryoEM/ASPIRE-Python

Issue:
Validate normalize_background against MATLAB

Hints from issue comments:
@j-c-c: Resolved in #1258 

Solution patch:
diff --git a/src/aspire/image/image.py b/src/aspire/image/image.py
index 27d5ced5..8fcf60fa 100644
--- a/src/aspire/image/image.py
+++ b/src/aspire/image/image.py
@@ -24,7 +24,7 @@ from aspire.volume import SymmetryGroup
 logger = logging.getLogger(__name__)
 
 
-def normalize_bg(imgs, bg_radius=1.0, do_ramp=True):
+def normalize_bg(imgs, bg_radius=1.0, do_ramp=True, legacy=False):
     """
     Normalize backgrounds and apply to a stack of images
 
@@ -33,16 +33,30 @@ def normalize_bg(imgs, bg_radius=1.0, do_ramp=True):
     :param do_ramp: When it is `True`, fit a ramping background to the data
             and subtract. Namely perform normalization based on values from each image.
             Otherwise, a constant background level from all images is used.
+    :param legacy: Option to match Matlab legacy normalize_background. Default, False,
+        uses ASPIRE-Python implementation. When True, ramping is disabled, a shifted
+        2d grid and alternative `bg_radius` is used to generate the background mask,
+        and standard deviation is computed using N - 1 degrees of freedom.
     :return: The modified images
     """
     if imgs.ndim > 3:
         raise NotImplementedError(
             "`normalize_bg` is currently limited to 1D image stacks."
         )
-
     L = imgs.shape[-1]
+
+    # Make adjustments for legacy mode
+    shifted = False
+    ddof = 0  # Degrees of freedom for standard deviation
+    if legacy:
+        do_ramp = False
+        shifted = True  # Shifts 2d grid by 1/2 pixel for even resolution
+        bg_radius = 2 * (L // 2) / L
+        ddof = 1
+
+    # Generate background mask
     input_dtype = imgs.dtype
-    grid = grid_2d(L, indexing="yx", dtype=input_dtype)
+    grid = grid_2d(L, shifted=shifted, indexing="yx", dtype=input_dtype)
     mask = grid["r"] > bg_radius
 
     if do_ramp:
@@ -71,15 +85,10 @@ def normalize_bg(imgs, bg_radius=1.0, do_ramp=True):
         imgs = imgs.reshape((-1, L, L))
 
     # Apply mask images and calculate mean and std values of background
-    imgs_masked = imgs * 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)
-    variance = second_moment.reshape(-1, 1, 1) - mean**2
-    std = np.sqrt(variance)
-
-    return (imgs - mean) / std
+    mean = np.mean(imgs[:, mask], axis=1)
+    std = np.std(imgs[:, mask], ddof=ddof, axis=1)
+
+    return (imgs - mean[:, None, None]) / std[:, None, None]
 
 
 def load_mrc(filepath):
diff --git a/src/aspire/source/image.py b/src/aspire/source/image.py
index f7bbfa4c..f529e153 100644
--- a/src/aspire/source/image.py
+++ b/src/aspire/source/image.py
@@ -945,7 +945,7 @@ class ImageSource(ABC):
         self.generation_pipeline.add_xform(Multiply(scale_factor))
 
     @_as_copy
-    def normalize_background(self, bg_radius=1.0, do_ramp=True):
+    def normalize_background(self, bg_radius=1.0, do_ramp=True, legacy=False):
         """
         Normalize the images by the noise background
 
@@ -958,6 +958,10 @@ class ImageSource(ABC):
         :param do_ramp: When it is `True`, fit a ramping background to the data
             and subtract. Namely perform normalization based on values from each image.
             Otherwise, a constant background level from all images is used.
+        :param legacy: Option to match Matlab legacy normalize_background. Default, False,
+            uses ASPIRE-Python implementation. When True, ramping is disable, a shifted
+            2d grid and alternative `bg_radius` is used to generate the background mask,
+            and standard deviation is computed using N - 1 degrees of freedom.
         :return: On return, the `ImageSource` object has been modified in place.
         """
 
@@ -966,7 +970,9 @@ class ImageSource(ABC):
             f"size of {bg_radius} and do_ramp of {do_ramp}"
         )
         self.generation_pipeline.add_xform(
-            LambdaXform(normalize_bg, bg_radius=bg_radius, do_ramp=do_ramp)
+            LambdaXform(
+                normalize_bg, bg_radius=bg_radius, do_ramp=do_ramp, legacy=legacy
+            )
         )
 
     def im_backward(self, im, start, weights=None, symmetry_group=None):
diff --git a/src/aspire/utils/coor_trans.py b/src/aspire/utils/coor_trans.py
index ce5e5e8a..c7c4ac81 100644
--- a/src/aspire/utils/coor_trans.py
+++ b/src/aspire/utils/coor_trans.py
@@ -73,14 +73,8 @@ def _mgrid_slice(n, shifted, normalized):
         end -= 1
 
     if normalized:
-        # Compute the denominator for normalization
-        denom = n / 2
-        if shifted and n % 2 == 0:
-            denom -= 1 / 2
-
-        # Apply the normalization
-        start /= denom
-        end /= denom
+        start /= n / 2
+        end /= n / 2
 
     return slice(start, end, num_points)
 


Test patch:
diff --git a/tests/test_grids.py b/tests/test_grids.py
index c2a4b657..1b156d11 100644
--- a/tests/test_grids.py
+++ b/tests/test_grids.py
@@ -16,6 +16,8 @@ class LegacyGridsTestCase(TestCase):
         self.shifts = [False, True]
         self.norms = [False, True]
 
+        # Note: Hard-coded references for (4, True, True) changed to reflect alteration
+        # in normalization factor for even/shifted/normalized grids (see PR #1258).
         self.legacy_references_2d = {
             (3, False, False): {
                 "x": np.array(
@@ -254,19 +256,19 @@ class LegacyGridsTestCase(TestCase):
             (4, True, True): {
                 "x": np.array(
                     [
-                        [-1.0, -1.0, -1.0, -1.0],
-                        [-0.33333334, -0.33333334, -0.33333334, -0.33333334],
-                        [0.33333334, 0.33333334, 0.33333334, 0.33333334],
-                        [1.0, 1.0, 1.0, 1.0],
+                        [-0.75, -0.75, -0.75, -0.75],
+                        [-0.25, -0.25, -0.25, -0.25],
+                        [0.25, 0.25, 0.25, 0.25],
+                        [0.75, 0.75, 0.75, 0.75],
                     ],
                     dtype=np.float32,
                 ),
                 "y": np.array(
                     [
-                        [-1.0, -0.33333334, 0.33333334, 1.0],
-                        [-1.0, -0.33333334, 0.33333334, 1.0],
-                        [-1.0, -0.33333334, 0.33333334, 1.0],
-                        [-1.0, -0.33333334, 0.33333334, 1.0],
+                        [-0.75, -0.25, 0.25, 0.75],
+                        [-0.75, -0.25, 0.25, 0.75],
+                        [-0.75, -0.25, 0.25, 0.75],
+                        [-0.75, -0.25, 0.25, 0.75],
                     ],
                     dtype=np.float32,
                 ),
@@ -281,10 +283,10 @@ class LegacyGridsTestCase(TestCase):
                 ),
                 "r": np.array(
                     [
-                        [1.4142135, 1.0540925, 1.0540925, 1.4142135],
-                        [1.0540925, 0.47140452, 0.47140452, 1.0540925],
-                        [1.0540925, 0.47140452, 0.47140452, 1.0540925],
-                        [1.4142135, 1.0540925, 1.0540925, 1.4142135],
+                        [1.0606601, 0.7905694, 0.7905694, 1.0606601],
+                        [0.7905694, 0.35355338, 0.35355338, 0.7905694],
+                        [0.7905694, 0.35355338, 0.35355338, 0.7905694],
+                        [1.0606601, 0.7905694, 0.7905694, 1.0606601],
                     ],
                     dtype=np.float32,
                 ),
@@ -1232,28 +1234,28 @@ class LegacyGridsTestCase(TestCase):
                 "x": np.array(
                     [
                         [
-                            [-1.0, -1.0, -1.0, -1.0],
-                            [-1.0, -1.0, -1.0, -1.0],
-                            [-1.0, -1.0, -1.0, -1.0],
-                            [-1.0, -1.0, -1.0, -1.0],
+                            [-0.75, -0.75, -0.75, -0.75],
+                            [-0.75, -0.75, -0.75, -0.75],
+                            [-0.75, -0.75, -0.75, -0.75],
+                            [-0.75, -0.75, -0.75, -0.75],
                         ],
                         [
-                            [-0.33333334, -0.33333334, -0.33333334, -0.33333334],
-                            [-0.33333334, -0.33333334, -0.33333334, -0.33333334],
-                            [-0.33333334, -0.33333334, -0.33333334, -0.33333334],
-                            [-0.33333334, -0.33333334, -0.33333334, -0.33333334],
+                            [-0.25, -0.25, -0.25, -0.25],
+                            [-0.25, -0.25, -0.25, -0.25],
+                            [-0.25, -0.25, -0.25, -0.25],
+                            [-0.25, -0.25, -0.25, -0.25],
                         ],
                         [
-                            [0.33333334, 0.33333334, 0.33333334, 0.33333334],
-                            [0.33333334, 0.33333334, 0.33333334, 0.33333334],
-                            [0.33333334, 0.33333334, 0.33333334, 0.33333334],
-                            [0.33333334, 0.33333334, 0.33333334, 0.33333334],
+                            [0.25, 0.25, 0.25, 0.25],
+                            [0.25, 0.25, 0.25, 0.25],
+                            [0.25, 0.25, 0.25, 0.25],
+                            [0.25, 0.25, 0.25, 0.25],
                         ],
                         [
-                            [1.0, 1.0, 1.0, 1.0],
-                            [1.0, 1.0, 1.0, 1.0],
-                            [1.0, 1.0, 1.0, 1.0],
-                            [1.0, 1.0, 1.0, 1.0],
+                            [0.75, 0.75, 0.75, 0.75],
+                            [0.75, 0.75, 0.75, 0.75],
+                            [0.75, 0.75, 0.75, 0.75],
+                            [0.75, 0.75, 0.75, 0.75],
                         ],
                     ],
                     dtype=np.float32,
@@ -1261,28 +1263,28 @@ class LegacyGridsTestCase(TestCase):
                 "y": np.array(
                     [
                         [
-                            [-1.0, -1.0, -1.0, -1.0],
-                            [-0.33333334, -0.33333334, -0.33333334, -0.33333334],
-                            [0.33333334, 0.33333334, 0.33333334, 0.33333334],
-                            [1.0, 1.0, 1.0, 1.0],
+                            [-0.75, -0.75, -0.75, -0.75],
+                            [-0.25, -0.25, -0.25, -0.25],
+                            [0.25, 0.25, 0.25, 0.25],
+                            [0.75, 0.75, 0.75, 0.75],
                         ],
                         [
-                            [-1.0, -1.0, -1.0, -1.0],
-                            [-0.33333334, -0.33333334, -0.33333334, -0.33333334],
-                            [0.33333334, 0.33333334, 0.33333334, 0.33333334],
-                            [1.0, 1.0, 1.0, 1.0],
+                            [-0.75, -0.75, -0.75, -0.75],
+                            [-0.25, -0.25, -0.25, -0.25],
+                            [0.25, 0.25, 0.25, 0.25],
+                            [0.75, 0.75, 0.75, 0.75],
                         ],
                         [
-                            [-1.0, -1.0, -1.0, -1.0],
-                            [-0.33333334, -0.33333334, -0.33333334, -0.33333334],
-                            [0.33333334, 0.33333334, 0.33333334, 0.33333334],
-                            [1.0, 1.0, 1.0, 1.0],
+                            [-0.75, -0.75, -0.75, -0.75],
+                            [-0.25, -0.25, -0.25, -0.25],
+                            [0.25, 0.25, 0.25, 0.25],
+                            [0.75, 0.75, 0.75, 0.75],
                         ],
                         [
-                            [-1.0, -1.0, -1.0, -1.0],
-                            [-0.33333334, -0.33333334, -0.33333334, -0.33333334],
-                            [0.33333334, 0.33333334, 0.33333334, 0.33333334],
-                            [1.0, 1.0, 1.0, 1.0],
+                            [-0.75, -0.75, -0.75, -0.75],
+                            [-0.25, -0.25, -0.25, -0.25],
+                            [0.25, 0.25, 0.25, 0.25],
+                            [0.75, 0.75, 0.75, 0.75],
                         ],
                     ],
                     dtype=np.float32,
@@ -1290,28 +1292,28 @@ class LegacyGridsTestCase(TestCase):
                 "z": np.array(
                     [
                         [
-                            [-1.0, -0.33333334, 0.33333334, 1.0],
-                            [-1.0, -0.33333334, 0.33333334, 1.0],
-                            [-1.0, -0.33333334, 0.33333334, 1.0],
-                            [-1.0, -0.33333334, 0.33333334, 1.0],
+                            [-0.75, -0.25, 0.25, 0.75],
+                            [-0.75, -0.25, 0.25, 0.75],
+                            [-0.75, -0.25, 0.25, 0.75],
+                            [-0.75, -0.25, 0.25, 0.75],
                         ],
                         [
-                            [-1.0, -0.33333334, 0.33333334, 1.0],
-                            [-1.0, -0.33333334, 0.33333334, 1.0],
-                            [-1.0, -0.33333334, 0.33333334, 1.0],
-                            [-1.0, -0.33333334, 0.33333334, 1.0],
+                            [-0.75, -0.25, 0.25, 0.75],
+                            [-0.75, -0.25, 0.25, 0.75],
+                            [-0.75, -0.25, 0.25, 0.75],
+                            [-0.75, -0.25, 0.25, 0.75],
                         ],
                         [
-                            [-1.0, -0.33333334, 0.33333334, 1.0],
-                            [-1.0, -0.33333334, 0.33333334, 1.0],
-                            [-1.0, -0.33333334, 0.33333334, 1.0],
-                            [-1.0, -0.33333334, 0.33333334, 1.0],
+                            [-0.75, -0.25, 0.25, 0.75],
+                            [-0.75, -0.25, 0.25, 0.75],
+                            [-0.75, -0.25, 0.25, 0.75],
+                            [-0.75, -0.25, 0.25, 0.75],
                         ],
                         [
-                            [-1.0, -0.33333334, 0.33333334, 1.0],
-                            [-1.0, -0.33333334, 0.33333334, 1.0],
-                            [-1.0, -0.33333334, 0.33333334, 1.0],
-                            [-1.0, -0.33333334, 0.33333334, 1.0],
+                            [-0.75, -0.25, 0.25, 0.75],
+                            [-0.75, -0.25, 0.25, 0.75],
+                            [-0.75, -0.25, 0.25, 0.75],
+                            [-0.75, -0.25, 0.25, 0.75],
                         ],
                     ],
                     dtype=np.float32,
@@ -1377,28 +1379,28 @@ class LegacyGridsTestCase(TestCase):
                 "r": np.array(
                     [
                         [
-                            [1.7320508, 1.4529663, 1.4529663, 1.7320508],
-                            [1.4529663, 1.1055416, 1.1055416, 1.4529663],
-                            [1.4529663, 1.1055416, 1.1055416, 1.4529663],
-                            [1.7320508, 1.4529663, 1.4529663, 1.7320508],
+                            [1.299038, 1.0897247, 1.0897247, 1.299038],
+                            [1.0897248, 0.8291562, 0.8291562, 1.0897248],
+                            [1.0897248, 0.8291562, 0.8291562, 1.0897248],
+                            [1.299038, 1.0897247, 1.0897247, 1.299038],
                         ],
                         [
-                            [1.4529663, 1.1055416, 1.1055416, 1.4529663],
-                            [1.1055416, 0.57735026, 0.57735026, 1.1055416],
-                            [1.1055416, 0.57735026, 0.57735026, 1.1055416],
-                            [1.4529663, 1.1055416, 1.1055416, 1.4529663],
+                            [1.0897248, 0.8291562, 0.8291562, 1.0897248],
+                            [0.8291562, 0.4330127, 0.4330127, 0.8291562],
+                            [0.8291562, 0.4330127, 0.4330127, 0.8291562],
+                            [1.0897248, 0.8291562, 0.8291562, 1.0897248],
                         ],
                         [
-                            [1.4529663, 1.1055416, 1.1055416, 1.4529663],
-                            [1.1055416, 0.57735026, 0.57735026, 1.1055416],
-                            [1.1055416, 0.57735026, 0.57735026, 1.1055416],
-                            [1.4529663, 1.1055416, 1.1055416, 1.4529663],
+                            [1.0897248, 0.8291562, 0.8291562, 1.0897248],
+                            [0.8291562, 0.4330127, 0.4330127, 0.8291562],
+                            [0.8291562, 0.4330127, 0.4330127, 0.8291562],
+                            [1.0897248, 0.8291562, 0.8291562, 1.0897248],
                         ],
                         [
-                            [1.7320508, 1.4529663, 1.4529663, 1.7320508],
-                            [1.4529663, 1.1055416, 1.1055416, 1.4529663],
-                            [1.4529663, 1.1055416, 1.1055416, 1.4529663],
-                            [1.7320508, 1.4529663, 1.4529663, 1.7320508],
+                            [1.299038, 1.0897247, 1.0897247, 1.299038],
+                            [1.0897248, 0.8291562, 0.8291562, 1.0897248],
+                            [1.0897248, 0.8291562, 0.8291562, 1.0897248],
+                            [1.299038, 1.0897247, 1.0897247, 1.299038],
                         ],
                     ],
                     dtype=np.float32,
@@ -1413,7 +1415,7 @@ class LegacyGridsTestCase(TestCase):
                 a = self.legacy_references_2d[k][j]
                 b = grid_2d(*k, indexing="xy")[j]
                 if j != "phi":
-                    assert np.allclose(a, b, atol=utest_tolerance(np.float32))
+                    np.testing.assert_allclose(a, b, atol=utest_tolerance(np.float32))
 
     def test_grid_3d(self):
         for k in product(self.ns, self.shifts, self.norms):
@@ -1425,6 +1427,6 @@ class LegacyGridsTestCase(TestCase):
                     continue
                 elif j == "theta":
                     a, b = np.mod(a, np.pi), np.mod(b, np.pi)
-                    assert np.allclose(a, b, atol=utest_tolerance(np.float32))
+                    np.testing.assert_allclose(a, b, atol=utest_tolerance(np.float32))
                 else:
-                    assert np.allclose(a, b, atol=utest_tolerance(np.float32))
+                    np.testing.assert_allclose(a, b, atol=utest_tolerance(np.float32))
diff --git a/tests/test_preprocess_pipeline.py b/tests/test_preprocess_pipeline.py
index e53322a2..94c5c1da 100644
--- a/tests/test_preprocess_pipeline.py
+++ b/tests/test_preprocess_pipeline.py
@@ -92,6 +92,28 @@ def testNormBackground(L, dtype):
     assert dtype == imgs_nb.dtype
 
 
+@pytest.mark.parametrize("L, dtype", params)
+def test_norm_background_legacy(L, dtype):
+    # Legacy normalize_background uses a shifted grid, a different
+    # mask radius, disabled ramping, and N - 1 degrees of freedom
+    # when computing standard deviation.
+    sim = get_sim_object(L, dtype)
+    bg_radius = 2 * (L // 2) / L
+    grid = grid_2d(sim.L, shifted=True, indexing="yx", dtype=dtype)
+    mask = grid["r"] > bg_radius
+    sim = sim.normalize_background(legacy=True)
+    imgs_nb = sim.images[:].asnumpy()
+    new_mean = np.mean(imgs_nb[:, mask])
+    new_variance = np.var(imgs_nb[:, mask], ddof=1)
+
+    # new mean of noise should be close to zero and variance should be close to 1
+    np.testing.assert_array_less(new_mean, utest_tolerance(dtype))
+    np.testing.assert_array_less(abs(new_variance - 1), 2e-3)
+
+    # dtype of returned images should be the same
+    np.testing.assert_equal(dtype, imgs_nb.dtype)
+
+
 @pytest.mark.parametrize("dtype", [np.float32, np.float64])
 def testWhiten(dtype):
     # Note this atol holds only for L even. Odd tested in testWhiten2.