From d81fe20008229af7a0bc9a78694ada0d60062cb2 Mon Sep 17 00:00:00 2001
From: cedriclim1 <cedlim@stanford.edu>
Date: Thu, 11 Jun 2026 00:16:38 -0700
Subject: [PATCH] Restore real-space fourier_cropping in tomography.utils

The tomography utility fourier_cropping (real-space image in, fft -> center
crop -> ifft -> real out) was removed when the corner-centered variant moved
to core.utils.imaging_utils. The two have incompatible contracts: the core
version expects an already-FFT'd corner-centered array, so callers passing
real-space tilt images (tutorials 01/02) either hit an ImportError or, if
switched naively to the core function, get ~all-zero stacks that normalize
to NaN. Restore the original implementation with regression tests covering
import, legacy parity, and non-degenerate real-space output.
---
 src/quantem/tomography/utils.py | 23 ++++++++++++++
 tests/tomography/test_utils.py  | 53 +++++++++++++++++++++++++++++++++
 2 files changed, 76 insertions(+)

diff --git a/src/quantem/tomography/utils.py b/src/quantem/tomography/utils.py
index b1105d3c..96d0fe14 100644
--- a/src/quantem/tomography/utils.py
+++ b/src/quantem/tomography/utils.py
@@ -1,6 +1,29 @@
+import numpy as np
 import torch
 import torch.nn.functional as F
 
+
+def fourier_cropping(img, crop_size):
+    """
+    Crop a real-space image in Fourier space (band-limited downsampling).
+
+    Takes a real-space image, crops its centered FFT to ``crop_size``, and
+    returns the real part of the inverse transform. Distinct from
+    ``quantem.core.utils.imaging_utils.fourier_cropping``, which operates on an
+    already-FFT'd corner-centered array.
+    """
+    center = np.array(img.shape) // 2
+
+    fft_img = np.fft.fftshift(np.fft.fft2(img))
+
+    cropped_fft = fft_img[
+        center[0] - crop_size[0] // 2 : center[0] + crop_size[0] // 2,
+        center[1] - crop_size[1] // 2 : center[1] + crop_size[1] // 2,
+    ]
+    cropped_img = np.fft.ifft2(np.fft.ifftshift(cropped_fft)).real
+    return cropped_img
+
+
 # --- Projection Operator Utils ---
 
 
diff --git a/tests/tomography/test_utils.py b/tests/tomography/test_utils.py
index b8530491..21c27bbd 100644
--- a/tests/tomography/test_utils.py
+++ b/tests/tomography/test_utils.py
@@ -89,6 +89,7 @@ def test_grad_flows_with_mixed_float_and_tensor_angles(self):
         out.sum().backward()
         assert x.grad is not None
         assert torch.isfinite(x.grad)
+
     @pytest.mark.parametrize(
         "rot_fn", [differentiable_rotz_vectorized, differentiable_rotx_vectorized]
     )
@@ -111,3 +112,55 @@ def test_per_volume_angles(self, rot_fn):
         batched = rot_fn(vols, angles)
         per_volume = torch.cat([rot_fn(vols[i : i + 1], angles[i]) for i in range(3)])
         assert torch.allclose(batched, per_volume, atol=1e-6)
+
+
+class TestFourierCropping:
+    """Regression tests for the real-space fourier_cropping removed when the
+    corner-centered variant moved to core.utils.imaging_utils (tutorials 01/02
+    import it from quantem.tomography.utils and pass real-space images)."""
+
+    def test_importable_from_tomography_utils(self):
+        from quantem.tomography.utils import fourier_cropping  # noqa: F401
+
+    def test_matches_legacy_implementation(self):
+        import numpy as np
+
+        from quantem.tomography.utils import fourier_cropping
+
+        rng = np.random.default_rng(0)
+        img = rng.normal(size=(64, 64))
+        crop = (32, 32)
+
+        # legacy reference: fft -> shift -> center crop -> ifft -> real
+        center = np.array(img.shape) // 2
+        fft_img = np.fft.fftshift(np.fft.fft2(img))
+        ref = np.fft.ifft2(
+            np.fft.ifftshift(
+                fft_img[
+                    center[0] - crop[0] // 2 : center[0] + crop[0] // 2,
+                    center[1] - crop[1] // 2 : center[1] + crop[1] // 2,
+                ]
+            )
+        ).real
+
+        out = fourier_cropping(img, crop)
+        assert out.shape == crop
+        assert np.isrealobj(out)
+        assert np.allclose(out, ref)
+
+    def test_real_space_content_preserved(self):
+        # Band-limited downsample of a smooth blob must stay smooth and
+        # non-degenerate -- the corner-centered core variant fed a real-space
+        # image instead returns a ~all-zero array.
+        import numpy as np
+
+        from quantem.tomography.utils import fourier_cropping
+
+        y, x = np.mgrid[:128, :128]
+        img = np.exp(-(((y - 64) ** 2 + (x - 64) ** 2) / (2 * 20.0**2)))
+        out = fourier_cropping(img, (64, 64))
+
+        assert out.shape == (64, 64)
+        # peak stays near the center and the result is far from all-zero
+        assert np.unravel_index(np.argmax(out), out.shape) == (32, 32)
+        assert (np.abs(out) > 1e-3 * np.abs(out).max()).mean() > 0.05