draft extension of grid_2d and image.filter to rectangle

Author: garrettwrong

src/aspire/image/image.py

@@ -587,10 +587,6 @@ def filter(self, filter):
         :param filter: An object of type `Filter`.
         :return: A new filtered `Image` object.
         """
-        if not self._is_square:
-            raise NotImplementedError(
-                "`Image.filter` is not currently implemented for non-square images."
-            )
 
         original_stack_shape = self.stack_shape
 
@@ -602,9 +598,14 @@ def filter(self, filter):
         #
         # Second note, filter dtype may not match image dtype.
         filter_values = xp.asarray(
-            filter.evaluate_grid(self.resolution), dtype=self.dtype
+            filter.evaluate_grid(self.shape[-2:]), dtype=self.dtype
         )
 
+        # sanity check
+        assert (
+            filter_values.shape == im._data.shape[-2:]
+        ), f"{filter_values.shape} != {im._data.shape[:-2]}"
+
         # Convolve
         im_f = fft.centered_fft2(xp.asarray(im._data))
         im_f = filter_values * im_f

src/aspire/utils/coor_trans.py

@@ -101,15 +101,22 @@ def grid_2d(n, shifted=False, normalized=True, indexing="yx", dtype=np.float32):
     Generate two dimensional grid.
 
     :param n: the number of grid points in each dimension.
+        May be a single integer value or 2-tuple of integers.
     :param shifted: shifted by half of grid or not when n is even.
     :param normalized: normalize the grid in the range of (-1, 1) or not.
     :param indexing: 'yx' (C) or 'xy' (F), defaulting to 'yx'.
         See https://numpy.org/doc/stable/reference/generated/numpy.meshgrid.html
     :return: the rectangular and polar coordinates of all grid points.
     """
 
-    grid = _mgrid_slice(n, shifted, normalized)
-    y, x = np.mgrid[grid, grid].astype(dtype)
+    if isinstance(n, int):
+        n = (n, n)
+    rows, cols = n
+
+    rows = _mgrid_slice(rows, shifted, normalized)
+    cols = _mgrid_slice(cols, shifted, normalized)
+
+    y, x = np.mgrid[rows, cols].astype(dtype)
     if indexing == "xy":
         x, y = y, x
     elif indexing != "yx":

tests/test_filters.py

@@ -5,6 +5,7 @@
 import numpy as np
 import pytest
 
+from aspire.image import Image
 from aspire.operators import (
     ArrayFilter,
     CTFFilter,
@@ -16,7 +17,7 @@
     ScaledFilter,
     ZeroFilter,
 )
-from aspire.utils import utest_tolerance
+from aspire.utils import gaussian_2d, grid_2d, utest_tolerance
 
 DATA_DIR = os.path.join(os.path.dirname(__file__), "saved_test_data")
 
@@ -432,3 +433,40 @@ def test_ctf_reference():
     # Test we're within 1%.
     #   There are minor differences in the formulas for wavelength and grids.
     np.testing.assert_allclose(h, ref_h, rtol=0.01)
+
+
+def test_rectangular_ctf():
+    """
+    Compare a truncated rectangular CTF application with the
+    application of CTF to a full square image.
+    """
+    # Configure square and truncated rectangle size
+    L = 128
+    rows, cols = 96, L
+    assert rows <= L and cols <= L and min(rows, cols) < L
+
+    # Create a test image of a disk
+    # A = gaussian_2d(size=L, mu=(0,-L//10), sigma=L//8, dtype=np.float64)
+    A = gaussian_2d(size=L, mu=(0, -L // 10), sigma=L // 32, dtype=np.float64)
+
+    full_img = Image(A)
+    truncated_img = Image(A[:rows, :cols])
+
+    # Create a CTFFilter
+    ctf_filter = CTFFilter(pixel_size=2)
+
+    # Apply to both Image instances
+    full_img_with_ctf = full_img.filter(ctf_filter)
+    truncated_img_with_ctf = truncated_img.filter(ctf_filter)
+
+    # Truncate the full square result
+    full_img_with_ctf_truncated = full_img_with_ctf.asnumpy()[:, :rows, :cols]
+
+    # Create mask for circular convolution effects
+    mask = (grid_2d(L, normalized=True)["r"] < 0.5)[:rows, :cols]
+
+    # Compare, we should be the same up to masked off differences in
+    #   circular convolution wrap around.
+    np.testing.assert_allclose(
+        truncated_img_with_ctf * mask, full_img_with_ctf_truncated * mask, atol=1e-6
+    )