perf: apply the interferometer sparse operator with rfft2/irfft2

[Jammy2211]

Summary

InterferometerSparseOperator zero-padded a real batch to (2y, 2x) and applied the
operator with the complex fft2 / ifft2 pair, taking Re(...) of an already-real
result. Khat was likewise the complex transform of a real preload. Because both the
preload and the batch are real, the real-transform pair rfft2 / irfft2 computes the
same result — exactly, not approximately: the discarded half of the spectrum is the
conjugate mirror of the half that is kept, and irfft2 reconstructs it.

• from_nufft_precision_operator caches Khat = jnp.fft.rfft2(preload) — shape
  (2y, x + 1) instead of (2y, 2x).
• apply_operator does rfft2 → multiply → irfft2(s=(2y, 2x)) → crop, dropping the
  now-redundant jnp.real.

Closes #538. Follow-up from autolens_profiling #226 (phase 2 of
numba-interferometer-revisit); numbers in
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, 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.

Exactness pin

test__interferometer_sparse_operator__apply_operator__rfft2_matches_complex_fft2_reference
builds the operator from the shared 7×7 / K=5 fixture and a seeded 12×12 / K=64 case,
and compares apply_operator(np.eye(M)) against an inline NumPy complex fft2 / ifft2
reference at rtol=1e-10, atol=1e-10·max|G| — max absolute error 2.8e-14. The change
is exact, so a loosened tolerance would be hiding a bug; the pin stays at round-off. It also
asserts the new Khat.shape == (2y, x + 1).

API Changes

None — internal changes only. The Khat field's shape and meaning change
((2y, 2x) complex fft2 → (2y, x + 1) rfft2), but it is an internal cache field:
the four consumers (curvature_matrix_diag_from, _off_diag_from, _func_list_from,
operated_matrix_slim_from) only call apply_operator, and no public signature moves.

Downstream / workspace impact

None. No API change and no numerical result changes (round-off only), so no workspace
migration or new demo is needed:

• grep across autofit_workspace/, autogalaxy_workspace/, autolens_workspace/,
  autolens_workspace_test/, euclid_strong_lens_modeling_pipeline/ and HowToLens/
  finds exactly one consumer,
  autolens_workspace_test/scripts/misc/jax_assertions/sparse_operators.py, which calls
  InterferometerSparseOperator.from_nufft_precision_operator(...) with an unchanged
  signature and pins the result at rtol=1e-4 — well outside the 2.8e-14 shift.
• Khat is not serialised: the class has no __getstate__ and rides as pytree aux
  data (simulator.py:56 no_flatten). The workspace .npy preload caches store the raw
  preload, not Khat, so no cache is invalidated by the shape change.
• PyAutoHands/autohands/config/no_run.yaml is not present in this checkout, so no
  disabled-script hidden risk could be checked there.

/prm precondition — A100 confirmation

The prompt asks: "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."
No GPU was available in the session that opened this PR, so this is shipped on CPU
parity plus the phase-2 CPU measurements above, and the GPU check is left as a /prm
precondition for the human to require or waive.

Recipe, if required:

1. On RAL, make a private merge-base checkout of PyAutoArray (never touch the shared
   /mnt/ral/jnightin/PyAuto install — it lags main and carries live jobs) and put it
   first on PYTHONPATH; check out main for the "old" leg and this branch for the "new"
   leg.
2. Write a ~10-line jax microbenchmark: build a random real (560, 560) preload
   (a 280×280 extent → (2y, 2x)), InterferometerSparseOperator.from_nufft_precision_operator,
   then time apply_operator on a (M, 128) batch — block_until_ready(), one warm-up
   call for the jit, then ~10 timed repeats, x64 on.
3. Submit it with the project's hpc/sync push-submit gpu <script> (SLURM gpu partition,
   JAX picks up the GPU automatically), then hpc/sync jobs / tail gpu / pull.
4. Merge criterion: new ≥ old wall-clock; a regression on the A100 would mean irfft2's
   XLA lowering is worse than the complex pair there, and is the one outcome the CPU
   numbers cannot rule out.

Test Plan

• pytest test_autoarray/inversion test_autoarray/dataset/interferometer test_autoarray/operators -q → 604 passed (61 s), CPU JAX x64.
• New exactness pin passes on both fixtures at rtol=1e-10 (max abs err 2.8e-14).
• Existing test_interferometer.py sparse-vs-mapping end-to-end comparisons (1e-4) unchanged.
• Optional pre-merge: the A100 confirmation above.

Full API Changes (for automation & release notes)

Changed Behaviour

• autoarray.InterferometerSparseOperator.from_nufft_precision_operator(...) — the cached
  Khat field is now rfft2(preload) of shape (2y, x + 1) rather than fft2(preload)
  of shape (2y, 2x). Signature and return type unchanged.
• autoarray.InterferometerSparseOperator.apply_operator(Fbatch_flat) — computed with
  rfft2 / irfft2; results identical to the previous complex route to round-off
  (max abs error 2.8e-14 on the test fixtures). Signature and return shape unchanged.

Removed / Added / Renamed / Changed Signature

• None.

Migration

• None required.

Generated by the PyAutoLabs agent workflow.

🤖 Generated with Claude Code

https://claude.ai/code/session_018hLF3ZAcz5MmaSJBEcLkvF