diff --git a/autoarray/inversion/inversion/interferometer/inversion_interferometer_util.py b/autoarray/inversion/inversion/interferometer/inversion_interferometer_util.py index 9c4345153..c19a7401b 100644 --- a/autoarray/inversion/inversion/interferometer/inversion_interferometer_util.py +++ b/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): diff --git a/test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py b/test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py index ec3fe4770..600a6250b 100644 --- a/test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py +++ b/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(), + )