Repository: ComputationalCryoEM/ASPIRE-Python

Issue:
Validate downsampling against MATLAB.

Discussion:
@j-c-c: Here are some initial results from downsample validation. 

Using `emdb_2660`, I generated two image sets (n=10) at resolutions 360 and 359. I downsampled each to resolution 64 and 65 in Matlab and Python and compared results by getting the relative RMSE for each corresponding set. 

I used relative RMSE so we could get a percent error that is invariant to the scale of the images. Initially I used MSE and the values were not too bad, on the order of e-06 to e-09:
```
MSE-360_64: 9.439877001223464e-09
MSE-360_65: 2.1544360936172384e-06
MSE-359_64: 8.93557758114028e-08
MSE-359_65: 1.5270138240818523e-06
```
Switching to percent error via relative RMSE we get:
```
RMSE/percent_error ds-360_64: 9.715902943743038e-05, 1.891307028424202%
RMSE/percent_error ds-360_65: 0.0014677997457477768, 28.556557734194126%
RMSE/percent_error ds-359_64: 0.0002989243647001743, 5.818899575795327%
RMSE/percent_error ds-359_65: 0.0012357240080543277, 24.04167111774978%
```
It looks like we're getting less correspondence when downsampling into odd resolution.

I've attached the notebook (saved as a txt) containing the script to produce this and a pdf version for quick viewing.

[validate_downsample (1).pdf](https://github.com/user-attachments/files/19451765/validate_downsample.1.pdf)

[validate_downsample.txt](https://github.com/user-attachments/files/19451772/validate_downsample.txt)@janden: Cool. I'm not following completely, though. The first numbers are absolute MSE while the latter are relative MSE? What are the two columns representing here?@j-c-c: > Cool. I'm not following completely, though. The first numbers are absolute MSE while the latter are relative MSE? What are the two columns representing here?

Ah, I think I had a typo in my script. For the second set of numbers just look at the second column (percent error) which is calculated using the relative RMSE:
```
relative_rmse = np.linalg.norm(A - B) / np.linalg.norm(B)
percent_error = relative_rmse * 100
```
Where `A` are the images downsampled in Python and `B` are the images downsampled in Matlab. 

I guess the takeaway is that when downsampling from, or into, an odd resolution, Python downsample is not doing the same thing as Matlab downsample. I'll explain why in a followup comment.@janden: Ok got it.@j-c-c: Resolved by #1252 

Solution patch:
diff --git a/src/aspire/image/image.py b/src/aspire/image/image.py
index 0395126b..27d5ced5 100644
--- a/src/aspire/image/image.py
+++ b/src/aspire/image/image.py
@@ -510,14 +510,16 @@ class Image:
 
         return Image(res)
 
-    def downsample(self, ds_res, zero_nyquist=True):
+    def downsample(self, ds_res, zero_nyquist=True, legacy=False):
         """
         Downsample Image to a specific resolution. This method returns a new Image.
 
         :param ds_res: int - new resolution, should be <= the current resolution
             of this Image
-        :param zero_nyquist: Option to keep or remove Nyquist frequency for even resolution.
-            Defaults to zero_nyquist=True, removing the Nyquist frequency.
+        :param zero_nyquist: Option to keep or remove Nyquist frequency for even
+            resolution (boolean). Defaults to zero_nyquist=True, removing the Nyquist frequency.
+        :param legacy: Option to match legacy Matlab downsample method (boolean).
+            Default of False uses `centered_fft` to maintain ASPIRE-Python centering conventions.
         :return: The downsampled Image object.
         """
 
@@ -528,19 +530,26 @@ class Image:
         # because all of the subsequent calls until `asnumpy` are GPU
         # when xp and fft in `cupy` mode.
 
-        # compute FT with centered 0-frequency
-        fx = fft.centered_fft2(xp.asarray(im._data))
+        if legacy:
+            fx = fft.fftshift(fft.fft2(xp.asarray(im._data)))
+        else:
+            # compute FT with centered 0-frequency
+            fx = fft.centered_fft2(xp.asarray(im._data))
+
         # crop 2D Fourier transform for each image
         crop_fx = crop_pad_2d(fx, ds_res)
 
         # If downsampled resolution is even, optionally zero out the nyquist frequency.
-        if ds_res % 2 == 0 and zero_nyquist is True:
+        if ds_res % 2 == 0 and zero_nyquist and not legacy:
             crop_fx[:, 0, :] = 0
             crop_fx[:, :, 0] = 0
 
         # take back to real space, discard complex part, and scale
-        out = fft.centered_ifft2(crop_fx).real * (ds_res**2 / self.resolution**2)
-        out = xp.asnumpy(out)
+        if legacy:
+            out = fft.ifft2(fft.ifftshift(crop_fx))
+        else:
+            out = fft.centered_ifft2(crop_fx)
+        out = xp.asnumpy(out.real * ds_res**2 / self.resolution**2)
 
         # Optionally scale pixel size
         ds_pixel_size = self.pixel_size
diff --git a/src/aspire/image/xform.py b/src/aspire/image/xform.py
index 4710f02d..10fddf1d 100644
--- a/src/aspire/image/xform.py
+++ b/src/aspire/image/xform.py
@@ -199,12 +199,26 @@ class Downsample(LinearXform):
     A Xform that downsamples an Image object to a resolution specified by this Xform's resolution.
     """
 
-    def __init__(self, resolution):
+    def __init__(self, resolution, zero_nyquist=True, legacy=False):
+        """
+        Initialize Xform to downsample Image to a specific resolution.
+
+        :param resolution: int - new resolution, should be <= the current resolution
+            of this Image
+        :param zero_nyquist: Option to keep or remove Nyquist frequency for even
+            resolution (boolean). Defaults to zero_nyquist=True, removing the Nyquist frequency.
+        :param legacy: Option to match legacy Matlab downsample method (boolean).
+            Default of False uses `centered_fft` to maintain ASPIRE-Python centering conventions.
+        """
         self.resolution = resolution
+        self.zero_nyquist = zero_nyquist
+        self.legacy = legacy
         super().__init__()
 
     def _forward(self, im, indices):
-        return im.downsample(self.resolution)
+        return im.downsample(
+            self.resolution, zero_nyquist=self.zero_nyquist, legacy=self.legacy
+        )
 
     def _adjoint(self, im, indices):
         # TODO: Implement up-sampling with zero-padding
diff --git a/src/aspire/source/image.py b/src/aspire/source/image.py
index c6fe79aa..f7bbfa4c 100644
--- a/src/aspire/source/image.py
+++ b/src/aspire/source/image.py
@@ -769,14 +769,16 @@ class ImageSource(ABC):
         """
 
     @_as_copy
-    def downsample(self, L):
+    def downsample(self, L, zero_nyquist=True, legacy=False):
         if L > self.L:
             raise ValueError(
                 "Max desired resolution {L} should be less than the current resolution {self.L}."
             )
         logger.info(f"Setting max. resolution of source = {L}")
 
-        self.generation_pipeline.add_xform(Downsample(resolution=L))
+        self.generation_pipeline.add_xform(
+            Downsample(resolution=L, zero_nyquist=zero_nyquist, legacy=legacy)
+        )
 
         ds_factor = self.L / L
         self.unique_filters = [f.scale(ds_factor) for f in self.unique_filters]
diff --git a/src/aspire/volume/volume.py b/src/aspire/volume/volume.py
index 013642bb..8339d8c2 100644
--- a/src/aspire/volume/volume.py
+++ b/src/aspire/volume/volume.py
@@ -517,27 +517,33 @@ class Volume:
             symmetry_group=symmetry,
         )
 
-    def downsample(self, ds_res, mask=None, zero_nyquist=True):
+    def downsample(self, ds_res, mask=None, zero_nyquist=True, legacy=False):
         """
         Downsample each volume to a desired resolution (only cubic supported).
 
         :param ds_res: Desired resolution.
-        :param zero_nyquist: Option to keep or remove Nyquist frequency for even resolution.
-            Defaults to zero_nyquist=True, removing the Nyquist frequency.
+        :param zero_nyquist: Option to keep or remove Nyquist frequency for even
+            resolution (boolean). Defaults to zero_nyquist=True, removing the Nyquist frequency.
         :param mask: Optional NumPy array mask to multiply in Fourier space.
+        :param legacy: Option to match legacy Matlab downsample method (boolean).
+            Default of False uses `centered_fft` to maintain ASPIRE-Python centering conventions.
+        :return: The downsampled Volume object.
         """
 
         original_stack_shape = self.stack_shape
         v = self.stack_reshape(-1)
 
         # take 3D Fourier transform of each volume in the stack
-        fx = fft.centered_fftn(xp.asarray(v._data))
+        if legacy:
+            fx = fft.fftshift(fft.fftn(xp.asarray(v._data)))
+        else:
+            fx = fft.centered_fftn(xp.asarray(v._data))
 
         # crop each volume to the desired resolution in frequency space
         fx = crop_pad_3d(fx, ds_res)
 
         # If downsample resolution is even, optionally zero out the nyquist frequency.
-        if ds_res % 2 == 0 and zero_nyquist is True:
+        if ds_res % 2 == 0 and zero_nyquist and not legacy:
             fx[:, 0, :, :] = 0
             fx[:, :, 0, :] = 0
             fx[:, :, :, 0] = 0
@@ -547,7 +553,10 @@ class Volume:
             fx = fx * xp.asarray(mask)
 
         # inverse Fourier transform of each volume
-        out = fft.centered_ifftn(fx)
+        if legacy:
+            out = fft.ifftn(fft.ifftshift(fx))
+        else:
+            out = fft.centered_ifftn(fx)
         out = out.real * (ds_res**3 / self.resolution**3)
 
         # Optionally scale pixel size


Test patch:
diff --git a/tests/test_downsample.py b/tests/test_downsample.py
index 6c1cac82..525dba1d 100644
--- a/tests/test_downsample.py
+++ b/tests/test_downsample.py
@@ -124,6 +124,7 @@ def test_integer_offsets():
 DTYPES = [np.float32, pytest.param(np.float64, marks=pytest.mark.expensive)]
 RES = [65, 66]
 RES_DS = [32, 33]
+LEGACY = [True, False]
 
 
 @pytest.fixture(params=DTYPES, ids=lambda x: f"dtype={x}", scope="module")
@@ -141,6 +142,11 @@ def res_ds(request):
     return request.param
 
 
+@pytest.fixture(params=LEGACY, ids=lambda x: f"legacy={x}", scope="module")
+def legacy(request):
+    return request.param
+
+
 @pytest.fixture(scope="module")
 def emdb_vol():
     return emdb_2660()
@@ -153,20 +159,59 @@ def volume(emdb_vol, res, dtype):
     return vol
 
 
-def test_downsample_project(volume, res_ds):
+def test_downsample_project(volume, res_ds, legacy):
     """
     Test that vol.downsample.project == vol.project.downsample.
     """
     rot = np.eye(3, dtype=volume.dtype)  # project along z-axis
-    im_ds_proj = volume.downsample(res_ds).project(rot)
-    im_proj_ds = volume.project(rot).downsample(res_ds)
+    im_ds_proj = volume.downsample(res_ds, legacy=legacy).project(rot)
+    im_proj_ds = volume.project(rot).downsample(res_ds, legacy=legacy)
+
+    tol = 1e-09
+    if volume.dtype == np.float32:
+        tol = 1e-07
+    if legacy:
+        # project does not enforce legacy centering convention,
+        # so this property will not hold up to allclose tolerance.
+        tol = 1e-03
 
-    tol = 1e-07
-    if volume.dtype == np.float64:
-        tol = 1e-09
     np.testing.assert_allclose(im_ds_proj, im_proj_ds, atol=tol)
 
 
+def test_downsample_legacy(volume, res_ds):
+    """
+    The legay Matlab downsample method differs from ASPIRE-Python
+    downsample in that is uses a different centering convention,
+    off by a half pixel for odd images, and does not zero out
+    the nyquist frequency. By making these alterations to the
+    ASPIRE-Python downsampled images we can match legacy downsample
+    upt to `allclose`.
+    """
+    n_img = 10
+    dtype = volume.dtype
+    src = Simulation(
+        n=n_img,
+        vols=volume,
+        amplitudes=1,
+        dtype=dtype,
+        seed=1980,
+    )
+    ims = src.images[:]
+
+    # Legacy downsampled images.
+    ims_ds_legacy = ims.downsample(res_ds, legacy=True)
+
+    # ASPIRE-Python downsample with centering adjustments for odd resolution images.
+    shifts = 0.5 * np.ones((n_img, 2), dtype=dtype)
+    if src.L % 2 == 1:
+        ims = ims.shift(shifts)
+    ims_ds_py = ims.downsample(res_ds, zero_nyquist=False)
+    if res_ds % 2 == 1:
+        ims_ds_py = ims_ds_py.shift(-shifts)
+
+    np.testing.assert_allclose(ims_ds_legacy, ims_ds_py, atol=1e-08)
+
+
 def test_simulation_relion_downsample():
     """
     Test that Simulation.downsample corresponds to RelionSource.downsample
diff --git a/tests/test_volume.py b/tests/test_volume.py
index 42ed4675..db6fae82 100644
--- a/tests/test_volume.py
+++ b/tests/test_volume.py
@@ -11,6 +11,7 @@ import pytest
 from numpy import pi
 from pytest import raises, skip
 
+from aspire.downloader import emdb_2660
 from aspire.source import _LegacySimulation
 from aspire.utils import Rotation, anorm, grid_2d, powerset, utest_tolerance
 from aspire.volume import (
@@ -32,6 +33,7 @@ def res_id(params):
 
 
 RES = [42, 43]
+RES_DS = [32, 33]
 TEST_PX_SZ = 4.56
 
 
@@ -40,6 +42,11 @@ def res(request):
     return request.param
 
 
+@pytest.fixture(params=RES_DS, ids=lambda x: f"resolution_ds={x}", scope="module")
+def res_ds(request):
+    return request.param
+
+
 def dtype_id(params):
     return f"dtype={params}"
 
@@ -122,6 +129,18 @@ def vols_hot_cold(res, dtype):
     return vols, hot_cold_locs, vol_center
 
 
+@pytest.fixture(scope="module")
+def emdb_vol():
+    return emdb_2660()
+
+
+@pytest.fixture(scope="module")
+def volume(emdb_vol, res, dtype):
+    vol = emdb_vol.astype(dtype, copy=False)
+    vol = vol.downsample(res)
+    return vol
+
+
 @pytest.fixture
 def random_data(res, dtype):
     return np.random.randn(res, res, res).astype(dtype)
@@ -673,6 +692,17 @@ def test_downsample(res):
     )
 
 
+def test_downsample_legacy(volume, res_ds):
+    """
+    Smoke test for the downsample legacy flag.
+    """
+    # Legacy downsampled images.
+    vol_ds = volume.downsample(res_ds, legacy=True)
+
+    # Check downsampled volume resolution.
+    np.testing.assert_equal(vol_ds.resolution, res_ds)
+
+
 def test_shape(vols_1, res):
     assert vols_1.shape == (N, res, res, res)
     assert vols_1.stack_shape == (N,)