Merge pull request #540 from PyAutoLabs/feature/interferometer-apply-operator-rfft2

perf: apply the interferometer sparse operator with rfft2/irfft2

Author: Jammy2211

autoarray/inversion/inversion/interferometer/inversion_interferometer_util.py

@@ -563,7 +563,7 @@ class InterferometerSparseOperator:
     M: int
     batch_size: int
     w_dtype: "jax.numpy.dtype"
-    Khat: "jax.Array"  # (2y, 2x), complex
+    Khat: "jax.Array"  # (2y, x+1), rfft2 of the real preload
     col_offsets: "jax.Array"  # (batch_size,) int32
     """
     Cached FFT operator state for fast interferometer curvature-matrix assembly.
@@ -578,9 +578,10 @@ class InterferometerSparseOperator:
     By taking an FFT of this preload, the operator can be applied to batches of
     images via elementwise multiplication in Fourier space:
 
-        apply_W(F) = IFFT( FFT(F_pad) * Khat )
+        apply_W(F) = IRFFT( RFFT(F_pad) * Khat )
 
-    where `F_pad` is a (2y, 2x) padded version of `F` and `Khat = FFT(nufft_precision_operator)`.
+    where `F_pad` is a (2y, 2x) padded version of `F` and
+    `Khat = rfft2(nufft_precision_operator)`.
 
     The curvature matrix for a pixelization (mapper) is then assembled from sparse
     mapping triplets without forming dense mapping matrices:
@@ -614,7 +615,7 @@ class InterferometerSparseOperator:
     w_dtype
         Floating-point dtype for weights and accumulations (e.g. float64).
     Khat
-        FFT of the curvature preload, shape (2y_shape, 2x_shape), complex.
+        Real FFT of the curvature preload, shape (2y_shape, x_shape + 1), complex.
         This is the frequency-domain representation of the W~ operator kernel.
     """
 
@@ -633,8 +634,9 @@ def from_nufft_precision_operator(
 
         The curvature preload is assumed to be defined on a (2y, 2x) rectangular
         grid of pixel offsets, where y and x correspond to the *unmasked extent*
-        of the real-space grid. The preload is FFT'd once to obtain `Khat`, which
-        is then reused for every subsequent curvature matrix build.
+        of the real-space grid. The preload is real, so it is transformed once with
+        a real FFT (`rfft2`) to obtain `Khat` of shape (2y, x + 1), which is then
+        reused for every subsequent curvature matrix build.
 
         Parameters
         ----------
@@ -654,7 +656,8 @@ def from_nufft_precision_operator(
         Returns
         -------
         InterferometerSparseOperator
-            Immutable cached state object containing shapes and FFT kernel `Khat`.
+            Immutable cached state object containing shapes and FFT kernel `Khat`,
+            of shape (2y, x + 1) and complex dtype.
 
         Raises
         ------
@@ -673,7 +676,7 @@ def from_nufft_precision_operator(
         x_shape = W2 // 2
         M = y_shape * x_shape
 
-        Khat = jnp.fft.fft2(nufft_precision_operator)
+        Khat = jnp.fft.rfft2(nufft_precision_operator)
 
         return InterferometerSparseOperator(
             dirty_image=dirty_image,
@@ -697,10 +700,19 @@ def apply_operator(self, Fbatch_flat):
 
         via FFT-based convolution with the cached `Khat` kernel:
 
-            apply_W(F) = Re( IFFT( FFT(F_pad) * Khat ) )[:y, :x]
+            apply_W(F) = IRFFT( RFFT(F_pad) * Khat )[:y, :x]
 
         where `F_pad` is the (2y, 2x) zero-padded version of `F`.
 
+        Both the preload and the batch are real-valued, so the real-transform pair
+        (`rfft2` / `irfft2`) is exact here rather than an approximation: the discarded
+        half of the spectrum is the conjugate mirror of the half that is kept, and the
+        inverse real transform reconstructs it, so the product is identical to the
+        complex `fft2` / `ifft2` route to floating-point round-off (the old code took
+        `Re(...)` of an already-real result). It halves the transform work and the size
+        of `Khat`, measured at 1.27-1.61x faster on every backend (autolens_profiling
+        #226, `results/notes/numba_interferometer_verdict.md`).
+
         Parameters
         ----------
         Fbatch_flat
@@ -720,10 +732,10 @@ def apply_operator(self, Fbatch_flat):
         B = Fbatch_flat.shape[1]
         F_img = Fbatch_flat.T.reshape((B, y_shape, x_shape))
         F_pad = jnp.pad(F_img, ((0, 0), (0, y_shape), (0, x_shape)))
-        Fhat = jnp.fft.fft2(F_pad)
+        Fhat = jnp.fft.rfft2(F_pad)
         Ghat = Fhat * Khat[None, :, :]
-        G_pad = jnp.fft.ifft2(Ghat)
-        G = jnp.real(G_pad[:, :y_shape, :x_shape])
+        G_pad = jnp.fft.irfft2(Ghat, s=(2 * y_shape, 2 * x_shape))
+        G = G_pad[:, :y_shape, :x_shape]
         return G.reshape((B, M)).T
 
     def curvature_matrix_diag_from(self, rows, cols, vals, *, S: int):

test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py

@@ -68,12 +68,30 @@ def test__data_vector_via_transformed_mapping_matrix_from():
     assert (data_vector_complex_via_blurred == data_vector_via_transformed).all()
 
 
-def _sparse_operator_and_mask():
+def _dataset_from(mask, n_visibilities, seed):
     """
-    Returns a real `InterferometerSparseOperator` (and the mask it is defined on) built from a
-    small 7x7 `TransformerDFT` interferometer dataset.
+    Returns a small `TransformerDFT` interferometer dataset on the input mask, with `n_visibilities`
+    seeded random visibilities and unit noise, alongside the random generator used to build it.
     """
-    mask = aa.Mask2D(
+    rng = np.random.default_rng(seed=seed)
+
+    dataset = aa.Interferometer(
+        data=aa.Visibilities(
+            visibilities=rng.normal(size=(n_visibilities, 2)).astype(np.float64)
+        ),
+        noise_map=aa.VisibilitiesNoiseMap(
+            visibilities=np.ones((n_visibilities, 2), dtype=np.float64)
+        ),
+        uv_wavelengths=rng.normal(size=(n_visibilities, 2)).astype(np.float64),
+        real_space_mask=mask,
+        transformer_class=aa.TransformerDFT,
+    )
+
+    return dataset, rng
+
+
+def _mask_7x7():
+    return aa.Mask2D(
         mask=[
             [True, True, True, True, True, True, True],
             [True, True, True, True, True, True, True],
@@ -86,20 +104,15 @@ def _sparse_operator_and_mask():
         pixel_scales=2.0,
     )
 
-    n_visibilities = 5
-    rng = np.random.default_rng(seed=3)
 
-    dataset = aa.Interferometer(
-        data=aa.Visibilities(
-            visibilities=rng.normal(size=(n_visibilities, 2)).astype(np.float64)
-        ),
-        noise_map=aa.VisibilitiesNoiseMap(
-            visibilities=np.ones((n_visibilities, 2), dtype=np.float64)
-        ),
-        uv_wavelengths=rng.normal(size=(n_visibilities, 2)).astype(np.float64),
-        real_space_mask=mask,
-        transformer_class=aa.TransformerDFT,
-    )
+def _sparse_operator_and_mask():
+    """
+    Returns a real `InterferometerSparseOperator` (and the mask it is defined on) built from a
+    small 7x7 `TransformerDFT` interferometer dataset.
+    """
+    mask = _mask_7x7()
+
+    dataset, rng = _dataset_from(mask=mask, n_visibilities=5, seed=3)
 
     return dataset.apply_sparse_operator(use_jax=False).sparse_operator, mask, rng
 
@@ -253,3 +266,65 @@ def test__interferometer_sparse_operator__operated_matrix_slim_from():
 
     assert operated.shape == (mask.pixels_in_mask, 2)
     assert operated == pytest.approx(operated_dense, 1.0e-8)
+
+
+def _apply_operator_via_complex_fft2(preload, operator):
+    """
+    Returns `W~ @ I` computed with the complex `fft2` / `ifft2` pair, written out in NumPy so that
+    it is an independent reference for the `rfft2` / `irfft2` implementation in
+    `InterferometerSparseOperator.apply_operator`.
+    """
+    y_shape, x_shape = operator.y_shape, operator.x_shape
+    M = operator.M
+
+    Fbatch_flat = np.eye(M)
+    B = Fbatch_flat.shape[1]
+
+    F_img = Fbatch_flat.T.reshape((B, y_shape, x_shape))
+    F_pad = np.pad(F_img, ((0, 0), (0, y_shape), (0, x_shape)))
+
+    Khat = np.fft.fft2(preload)
+    Ghat = np.fft.fft2(F_pad) * Khat[None, :, :]
+    G_pad = np.fft.ifft2(Ghat)
+    G = np.real(G_pad[:, :y_shape, :x_shape])
+
+    return G.reshape((B, M)).T
+
+
+def test__interferometer_sparse_operator__apply_operator__rfft2_matches_complex_fft2_reference():
+    pytest.importorskip("jax")
+
+    cases = [
+        (_mask_7x7(), 5, 3),
+        (
+            aa.Mask2D.circular(shape_native=(12, 12), pixel_scales=1.0, radius=4.0),
+            64,
+            11,
+        ),
+    ]
+
+    for mask, n_visibilities, seed in cases:
+        dataset, _ = _dataset_from(mask=mask, n_visibilities=n_visibilities, seed=seed)
+
+        preload = dataset.psf_precision_operator_from(use_jax=False)
+        operator = dataset.apply_sparse_operator(
+            nufft_precision_operator=preload
+        ).sparse_operator
+
+        # The preload and the batch are both real, so `rfft2` stores only the non-redundant half
+        # of the spectrum: (2y, x + 1) rather than (2y, 2x).
+        assert operator.Khat.shape == (2 * operator.y_shape, operator.x_shape + 1)
+
+        operated = np.array(operator.apply_operator(np.eye(operator.M)))
+        operated_via_complex_fft2 = _apply_operator_via_complex_fft2(
+            preload=preload, operator=operator
+        )
+
+        # The real transform pair is exact for a real preload and a real batch, so this pin is at
+        # round-off, not at an algorithmic tolerance.
+        np.testing.assert_allclose(
+            operated,
+            operated_via_complex_fft2,
+            rtol=1.0e-10,
+            atol=1.0e-10 * np.abs(operated_via_complex_fft2).max(),
+        )