perf: apply the interferometer sparse operator with rfft2/irfft2

[Jammy2211]

Overview

InterferometerSparseOperator.apply_operator zero-pads a real batch to (2y, 2x) and applies the operator with the complex jnp.fft.fft2 / jnp.fft.ifft2 pair, taking the real part at the end; Khat = jnp.fft.fft2(nufft_precision_operator) is likewise the complex transform of a real array. Because both inputs are real, the real-transform pair rfft2 / irfft2 computes the same result exactly — not approximately — with roughly half the spectrum and half the butterflies. Measured 1.27–1.61× on the operator apply across every geometry and mesh tested, and the NumPy rfft2 route also beats the library's own JAX-CPU route by 1.10–1.30× at every cell.

Follow-up from autolens_profiling #226 (phase 2 of the retired epic numba-interferometer-revisit, ledger PyAutoMind/complete/archive/epics/numba_interferometer_likelihood_revisit.md). Numbers: autolens_profiling/results/notes/numba_interferometer_verdict.md and results/breakdown/interferometer/bakeoff_v2026.8.17.1.json.

Measured saving

Single thread, i9-10885H, autoarray 2026.8.17.1; same synthetic inputs, same batch size (128), scipy stack, so the comparison isolates the transform:

geometry (extent)	mesh	complex fft2	real rfft2	saving
sma (70²)	Delaunay S=1500	1.7479 s	1.2660 s	1.38×
alma (140²)	Delaunay S=1500	6.5003 s	4.3951 s	1.48×
alma_high (280²)	Delaunay S=1500	17.7362 s	11.0346 s	1.61×
sma (70²)	rect S=784	1.0124 s	0.7361 s	1.38×
alma (140²)	rect S=784	3.3786 s	2.6557 s	1.27×
alma_high (280²)	rect S=784	6.9360 s	4.8414 s	1.43×

Plan

• Cache Khat as the real transform of the preload: jnp.fft.rfft2(nufft_precision_operator), shape (2y, x+1) instead of (2y, 2x) complex.
• Apply the operator with rfft2 → multiply → irfft2(s=(2y, 2x)) → crop; the result is already real, so the trailing jnp.real goes away.
• Change nothing else: the algebra is untouched, and the four consumers of apply_operator only ever call it.
• Pin the change as exact — a new test compares apply_operator against an inline complex fft2/ifft2 reference at rtol=1e-10, on the small shared fixture and on a larger seeded case. A loosened tolerance would hide a bug.
• Record in the docstrings why the real transform is exact here, with the measured saving and its source.
• Confirm on the A100 before merging (see the /prm precondition below).

/prm precondition — A100 confirmation

The prompt requires: "GPU: the saving should hold or improve (half the FFT work and half the intermediate memory); confirm on the A100 before merging rather than assuming it."

There is no GPU in the development environment, so this PR ships on CPU parity plus the phase-2 measurements above. Before merge, run a short jax microbenchmark of apply_operator (280×280 preload, batch 128, old vs new) on RAL's A100 against a private merge-base checkout on PYTHONPATH (never the shared /mnt/ral/jnightin/PyAuto install), submitted through the project's hpc/sync push-submit gpu CLI, and paste the two timings into the PR before merging. The human decides at /prm whether to require it.

Detailed implementation plan

Work Classification

Library

Affected Repositories

• PyAutoArray (primary)

Branch Survey

Repository	Current Branch	Dirty?
./PyAutoArray	main (47a00e8c)	clean

Suggested branch: feature/interferometer-apply-operator-rfft2

Worktree root: ~/Code/PyAutoLabs-wt/interferometer-apply-operator-rfft2/

Implementation Steps

1. autoarray/inversion/inversion/interferometer/inversion_interferometer_util.py, InterferometerSparseOperator.from_nufft_precision_operator (~line 621): Khat = jnp.fft.rfft2(nufft_precision_operator) → shape (2y, x+1). Update the field comment to Khat: "jax.Array" # (2y, x+1), rfft2 of the real preload.
2. Same file, apply_operator (~line 688): Fhat = jnp.fft.rfft2(F_pad); Ghat = Fhat * Khat[None]; G_pad = jnp.fft.irfft2(Ghat, s=(2 * y_shape, 2 * x_shape)) (already real — drop jnp.real); crop [:, :y_shape, :x_shape].
3. Leave the four consumers alone — curvature_matrix_diag_from, curvature_matrix_off_diag_from, curvature_matrix_func_list_from, operated_matrix_slim_from only call apply_operator.
4. Docstrings: one paragraph on why rfft2/irfft2 is exact for a real preload and a real padded batch, plus the measured saving, citing autolens_profiling docs: refactor docstrings for autoarray/mask package #226 / the verdict note.
5. Tests, test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py: keep the existing dense-oracle tests (_operator_dense, pytest.approx(..., 1e-8)). Add test__apply_operator__rfft2_matches_complex_fft2_reference building the operator from the shared 7×7 / K=5 fixture and a larger seeded case (e.g. 12×12 mask, K=64), comparing apply_operator(np.eye(M)) against an inline complex fft2/ifft2 reference at rtol=1e-10, atol=1e-10·max|G|.
6. Existing test_interferometer.py comparisons against the mapping path (1e-4) remain the end-to-end gate.

Compatibility notes

• The Khat shape change is invisible to serialisation: the class has no __getstate__ and the operator rides as pytree aux data via simulator.py:56 no_flatten.
• Workspace .npy caches store the raw preload, not Khat, so they are unaffected.

Verification

• New rfft2-vs-fft2 pin passes at rtol=1e-10.
• pytest test_autoarray/inversion/inversion/interferometer test_autoarray/dataset/interferometer test_autoarray/operators -q green.
• The full test_autoarray/inversion tree green once.

Key Files

• autoarray/inversion/inversion/interferometer/inversion_interferometer_util.py — InterferometerSparseOperator (~line 538): from_nufft_precision_operator (~621) and apply_operator (~688).
• test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py — parity pins.
• test_autoarray/inversion/inversion/interferometer/test_interferometer.py — end-to-end sparse-vs-mapping gate.

Related

• Epic ledger (retired 2026-09-08): PyAutoMind/complete/archive/epics/numba_interferometer_likelihood_revisit.md
• Phase 2: research: numba vs JAX-CPU interferometer curvature — levers and verdict autolens_profiling#226
• Follow-on task (stacked on this branch): perf: build nufft_precision_operator_from as a type-1 NUFFT

Original Prompt

Click to expand starting prompt

InterferometerSparseOperator.apply_operator pads real input to complex — use rfft2/irfft2

Type: feature
Target: autoarray
Repos:

• PyAutoArray
  Themes:
• interferometer
• jax-performance
• likelihood-profiling
  Difficulty: small
  Autonomy: supervised
  Priority: high
  Epic: numba-interferometer-revisit
  Filed: 2026-09-07

Follow-up from autolens_profiling#226 (phase 2 of numba-interferometer-revisit); the
numbers are in autolens_profiling/results/notes/numba_interferometer_verdict.md and
results/breakdown/interferometer/bakeoff_v2026.8.17.1.json.

What

autoarray/inversion/inversion/interferometer/inversion_interferometer_util.py
(InterferometerSparseOperator.apply_operator, ~689-728) zero-pads a real batch to
(2y, 2x) and applies the operator with complex jnp.fft.fft2 / jnp.fft.ifft2, taking
the real part at the end. Khat = jnp.fft.fft2(nufft_precision_operator) is likewise the
complex transform of a real array. Both halve-able: the real-transform pair
rfft2 / irfft2 computes the same result exactly (not approximately) with roughly half
the spectrum and half the butterflies.

Measured saving (single thread, i9-10885H, autoarray 2026.8.17.1)

Same synthetic inputs, same batch size (128), scipy stack, so the comparison isolates the
transform:

geometry (extent)	mesh	complex fft2	real rfft2	saving
sma (70²)	Delaunay S=1500	1.7479 s	1.2660 s	1.38×
alma (140²)	Delaunay S=1500	6.5003 s	4.3951 s	1.48×
alma_high (280²)	Delaunay S=1500	17.7362 s	11.0346 s	1.61×
sma (70²)	rect S=784	1.0124 s	0.7361 s	1.38×
alma (140²)	rect S=784	3.3786 s	2.6557 s	1.27×
alma_high (280²)	rect S=784	6.9360 s	4.8414 s	1.43×

rfft2 on the NumPy stack also beats the library's own JAX-CPU route by 1.10-1.30× at every
one of those cells, i.e. the transform change is worth more than the backend.

Scope

• Change from_nufft_precision_operator to cache Khat = jnp.fft.rfft2(...) and
  apply_operator to rfft2 → multiply → irfft2(s=(2y, 2x)) → crop.
• The parity gate is the existing interferometer inversion tests plus an explicit pin of
  curvature_matrix before/after at rtol=1e-10; the change is exact, so a loosened
  tolerance would be hiding a bug.
• GPU: the saving should hold or improve (half the FFT work and half the intermediate
  memory); confirm on the A100 before merging rather than assuming it.
• Nothing else in the interferometer path changes — the algebra is untouched.

[Jammy2211]

Library PR Created

• PR: perf: apply the interferometer sparse operator with rfft2/irfft2 #540 (pending-release, base main, MERGEABLE)
• Commit: a67828e3
• Tests: pytest test_autoarray/inversion test_autoarray/dataset/interferometer test_autoarray/operators -q → 604 passed (61 s)

API changes

None — internal only. Khat changes shape/meaning ((2y, 2x) complex fft2 →
(2y, x + 1) rfft2) but it is an internal cache field; no public signature moves and no
numerical result changes beyond round-off (new exactness pin: max abs error 2.8e-14 at
rtol=1e-10).

Downstream workspace impact — none (option iii)

grep across autofit_workspace/, autogalaxy_workspace/, autolens_workspace/,
autolens_workspace_test/, euclid_strong_lens_modeling_pipeline/ and HowToLens/ finds a
single consumer, autolens_workspace_test/scripts/misc/jax_assertions/sparse_operators.py,
which calls from_nufft_precision_operator(...) with an unchanged signature and pins at
rtol=1e-4. Khat is not serialised (pytree aux data, no __getstate__) and the workspace
.npy caches hold the raw preload, so nothing is invalidated. No migration, no new demo.

Open precondition

The prompt asks for the saving to be confirmed on the A100 before merging. No GPU was
available in this session; the PR body carries the exact RAL recipe as a /prm precondition
for the human to require or waive.

Next: /prm (merge gate). Task 2 (interferometer-preload-nufft-type1, #539) is stacked on
this branch.

[Jammy2211]

Shipped

Merged in #540 (merge commit 7a4cb700, head a67828e3).

InterferometerSparseOperator.apply_operator now uses the exact real-transform pair
jnp.fft.rfft2 / jnp.fft.irfft2 instead of padding the real batch to complex and
running fft2 / ifft2 with a .real at the end. Khat is likewise built with
rfft2, giving it shape (2y, x + 1) rather than (2y, 2x).

• Exact, not approximate: the new path is pinned against the previous complex fft2
  route at rtol=1e-10; measured agreement is 2.8e-14.
• Measured speedup 1.27-1.61x over the complex path (autolens_profiling#226).
• Tests: test_inversion_interferometer_util.py extended (+109/-29 lines overall diff),
  full PyAutoArray suite green on unittest 3.12, unittest 3.13 and unittest-nojax.

Open verification item: the A100 confirmation was not run before merge — the human
waived it at merge time, so it stays recorded as outstanding in the completion record.

Pending release: PyAutoArray@https://github.com/PyAutoLabs/PyAutoArray/pull/540 — merged
is not released; the obligation is carried into the PyAutoMind completion record.