perf: NumPy/scipy application path for InterferometerSparseOperator (xp=np never imports JAX)

[Jammy2211]

Overview

InterferometerSparseOperator (autoarray/inversion/inversion/interferometer/inversion_interferometer_util.py:857) imports jax.numpy inside every method body and has no xp argument. A CPU user without JAX still runs the whole interferometer curvature build through JAX — an unwanted dependency, and slower: a NumPy/scipy rfft2 convolution with a scipy.sparse projection beats the JAX-CPU route at every geometry measured.

This task gives the operator an xp-aware application path. The application backend follows the inversion's xp (what the fit passes), not the dataset's use_jax kwarg — which after PyAutoArray#541 only picks a brute-force builder and is otherwise irrelevant. That is the repo's own convention (transformer.image_from(xp=)) and it removes the "surprising split" the prompt names without changing any dataset API. The JAX branch stays exactly as it is.

Phase 3 of the retired numba-interferometer-revisit epic. It is also a prerequisite for the numba direct_conv follow-up: change 1 below keeps the raw real-space preload that the numba kernel indexes.

Measured (from the prompt)

geometry	mesh	JAX-CPU (fft2)	NumPy (rfft2)	saving
sma	Delaunay S=1500	1.5239 s	1.2660 s	1.20x
alma	Delaunay S=1500	5.5156 s	4.3951 s	1.26x
alma_high	Delaunay S=1500	12.9655 s	11.0346 s	1.18x
sma	rect S=784	0.8510 s	0.7361 s	1.16x
alma	rect S=784	2.9317 s	2.6557 s	1.10x
alma_high	rect S=784	6.2817 s	4.8414 s	1.30x

These numbers predate the rfft2 change (PyAutoArray#540, merged 2026-09-08). The JAX-CPU column is now faster than shown, so the remaining gap is smaller than the table's saving column — measure against the post-rfft2 baseline, not against these.

Links

• Epic (retired ledger): PyAutoMind/complete/archive/epics/numba_interferometer_likelihood_revisit.md
• Phase 2 verdict + measurements: research: numba vs JAX-CPU interferometer curvature — levers and verdict autolens_profiling#226 (autolens_profiling/results/notes/numba_interferometer_verdict.md)
• Merged today, same epic: PyAutoArray#540 (rfft2 in apply_operator), PyAutoArray#541 (NUFFT type-1 preload)
• Follow-on: feat: numba CPU interferometer curvature path — direct_conv, geometry-gated is stacked on this branch.

Plan

• Give InterferometerSparseOperator a NumPy/scipy application path: every public method gains xp=np and branches to a scipy.fft / scipy.sparse body when xp is not JAX; the existing JAX bodies are untouched.
• Keep the raw (2y, 2x) real-space preload on the operator instead of discarding it after the FFT, and make the JAX-typed fields (Khat, col_offsets) lazily cached so constructing the operator under xp=np never imports JAX.
• Have InversionInterferometerSparse pass self._xp to every operator call, and delete the if self._xp is np: np.array(...) workaround that only existed because the operator always returned JAX arrays.
• One docstring sentence on dataset.apply_sparse_operator recording that the application backend follows the inversion's xp; no dataset signature change.
• Pin the NumPy branch against the JAX branch at rtol=1e-10, atol=1e-10*max|ref| for all six methods, end-to-end through the inversion, and add a no-JAX guard test that runs on the unittest-nojax CI leg.

Detailed implementation plan

Work Classification

Library

Affected Repositories

• PyAutoArray (primary)

Branch Survey

Repository	Current Branch	Dirty?
./PyAutoArray	main (9bd76799)	clean

Suggested branch: feature/interferometer-sparse-operator-numpy-cpu-path

Worktree root: ~/Code/PyAutoLabs-wt/interferometer-sparse-operator-numpy-cpu-path/

Implementation Steps

All changes are in autoarray/inversion/inversion/interferometer/.

1. inversion_interferometer_util.py::InterferometerSparseOperator

• New field nufft_precision_operator: np.ndarray — the raw (2y, 2x) float64 preload, kept rather than discarded after Khat = rfft2(preload) (task 4 indexes it; it is small).
• Khat (JAX rfft2) and a new khat_np (scipy.fft.rfft2) become lazily cached (functools.cached_property works on a frozen dataclass; or store khat_np eagerly since scipy is a hard dep and compute the JAX Khat lazily) so constructing the operator under xp=np never imports JAX. Keep w_dtype and col_offsets for the JAX branch; col_offsets becomes lazy too (it is a jnp.arange).
• Verify the pytree registration (simulator.py:56, no_flatten) and the existing dataset/simulator tests still pass — the operator is static aux data.
• Every public method gains xp=np — apply_operator, curvature_matrix_diag_from, curvature_matrix_off_diag_from, operated_matrix_slim_from, curvature_matrix_off_diag_func_list_from, curvature_matrix_func_list_from — branching if xp.__name__.startswith("jax") to the existing bodies unchanged, else to the NumPy bodies:
  • apply_operator: scipy.fft.rfft2(F_pad) * khat_np -> irfft2(s=(2y, 2x)) -> crop.
  • curvature_matrix_diag_from: A = scipy.sparse.csc_matrix((vals, (rows, cols)), shape=(M, S)) (duplicates sum, as required), AT = A.T.tocsr(); for column blocks of batch_size: F = A[:, s:e].toarray(), G = apply_operator(F), C[:, s:e] = AT @ G; return 0.5 * (C + C.T) — the JAX branch symmetrises, so the NumPy branch must too for an exact pin. This is the prototype's curvature_fft_numpy with the symmetrisation added.
  • curvature_matrix_off_diag_from: A0, A1; blocks over A1; C01[:, s:e] = A0.T @ apply(A1[:, s:e]).
  • operated_matrix_slim_from: grid = zeros((M, n)); grid[extent_index] = matrix_slim (a set, not an add — indices are unique); apply; gather [extent_index].
  • the two func-list methods: the same scatter-set / apply / A.T @ pattern.
• No batching masks or padding are needed on the NumPy side (they exist only for static shapes under lax.fori_loop).

2. sparse.py::InversionInterferometerSparse

• Pass xp=self._xp to every operator call.
• Delete the if self._xp is np: curvature_matrix = np.array(...) workaround at :171-175 (its comment says why it existed). data_vector unchanged.

3. dataset.py

• No signature change. One docstring sentence on apply_sparse_operator saying the application backend follows the inversion's xp.

4. Tests (test_inversion_interferometer_util.py, test_interferometer.py)

• For each of the six methods: NumPy branch vs JAX branch at rtol=1e-10, atol=1e-10*max|ref| on the 7x7 / K=5 fixture and the 12x12 / K=64 seeded case added by PyAutoArray#540 (_dataset_from), with S larger than batch_size in at least one case (pass batch_size=4) so the block loop is exercised, plus a Delaunay mapper case with duplicate COO entries (the existing test_interferometer.py Delaunay fixture).
• InversionInterferometerSparse(xp=np) end-to-end: curvature_matrix, data_vector, reconstruction, log_evidence equal the xp=jnp inversion at the same pin.
• A no-JAX guard: with jax made unimportable (monkeypatch.setitem(sys.modules, "jax", None) and "jax.numpy"), constructing the operator and running the NumPy branch works. These tests must not carry importorskip("jax"), so the unittest-nojax CI leg runs them.
• Keep every existing test and tolerance.

Key Files

• autoarray/inversion/inversion/interferometer/inversion_interferometer_util.py — InterferometerSparseOperator (from :857); the six methods and the Khat / col_offsets fields.
• autoarray/inversion/inversion/interferometer/sparse.py — InversionInterferometerSparse; already resolves self._xp is np cleanly, triplets from mapper_util.sparse_triplets_from are pure NumPy; the :171-175 workaround.
• autoarray/inversion/inversion/interferometer/dataset.py — apply_sparse_operator docstring.
• autolens_profiling/scripts/misc/numba_interferometer/kernels.py::curvature_fft_numpy — the reference implementation to port (scipy rfft2/irfft2, scipy.sparse CSC/CSR projection, batch 128); pinned to the numba reference at rtol=1e-10 in that pack's test_parity.py.
• test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py — exact-pin idiom at :345.
• test_autoarray/inversion/inversion/interferometer/test_interferometer.py — Delaunay / rectangular DFT fixtures.

Verification

• Six-method parity NumPy vs JAX at the exact pin; end-to-end inversion parity.
• No-JAX guard passes with jax blocked; the unittest-nojax CI leg green.
• No existing test or tolerance changed.
• Operator construction under xp=np imports no JAX (assert "jax" not in sys.modules in a subprocess test).
• pytest test_autoarray/inversion test_autoarray/dataset test_autoarray/operators -q green; black clean.

Original Prompt

Click to expand starting prompt

InterferometerSparseOperator has no NumPy path — a CPU user is forced through JAX

Type: feature
Target: autoarray
Repos:

• PyAutoArray
  Themes:
• interferometer
• numpy-cpu
• likelihood-profiling
  Difficulty: medium
  Autonomy: supervised
  Priority: medium
  Epic: numba-interferometer-revisit
  Filed: 2026-09-07

Follow-up from autolens_profiling#226 (phase 2 of numba-interferometer-revisit);
numbers in autolens_profiling/results/notes/numba_interferometer_verdict.md.

What

InterferometerSparseOperator.apply_operator, curvature_matrix_diag_from and
curvature_matrix_off_diag_from
(autoarray/inversion/inversion/interferometer/inversion_interferometer_util.py) all
import jax.numpy inside the method body. There is no xp=np branch: a user without a GPU
still runs the interferometer curvature through JAX, and InversionInterferometerSparse
inherits that whatever xp the fit was built with.

That is not only an unwanted dependency — it is slower. On identical inputs, a NumPy/scipy
rfft2 convolution with a scipy.sparse projection (the NumPy-stack counterpart of the JAX
route's segment_sum) beats the library's JAX-CPU route at every geometry measured:

geometry	mesh	JAX-CPU (fft2)	NumPy (rfft2)	saving
sma	Delaunay S=1500	1.5239 s	1.2660 s	1.20×
alma	Delaunay S=1500	5.5156 s	4.3951 s	1.26×
alma_high	Delaunay S=1500	12.9655 s	11.0346 s	1.18×
sma	rect S=784	0.8510 s	0.7361 s	1.16×
alma	rect S=784	2.9317 s	2.6557 s	1.10×
alma_high	rect S=784	6.2817 s	4.8414 s	1.30×

(Some of that gap is the rfft2 change filed separately as
interferometer_apply_operator_rfft2.md, which shipped on 2026-09-08 — PyAutoArray#540,
record complete/2026/09/interferometer-apply-operator-rfft2.md. The real-FFT-first
sequencing constraint is therefore already satisfied; measure this one against the
post-rfft2 baseline, not the numbers in the table above.)

Scope

• An xp-aware application path so xp=np uses scipy.fft and scipy.sparse and never
  imports JAX; the JAX branch stays exactly as it is.
• The reference implementation to port is
  autolens_profiling/scripts/misc/numba_interferometer/kernels.py::curvature_fft_numpy,
  which is pinned to the numba reference kernel at rtol=1e-10 in that pack's
  test_parity.py.
• Parity gate: the NumPy and JAX branches must agree at rtol=1e-10 on F with atol
  scaled by max|F|.
• Worth checking at the same time whether apply_sparse_operator(use_jax=False) should
  imply the NumPy application path (today it only affects how the preload is built, not
  how the operator is applied, which is a surprising split).

[Jammy2211]

Library PR Created

#544 (pending-release, MERGEABLE) — commit c2469b9d, 5 files, +926/−37.

API changes: additive only. Six InterferometerSparseOperator methods gain xp=np; the dataclass keeps the raw preload as nufft_precision_operator; Khat/col_offsets become lazy cached properties with a new khat_np (scipy) alongside. Nothing removed, no signature broken.

Verification: 1486 tests pass; six-method NumPy-vs-JAX parity ≤ 5e-16 relative; end-to-end curvature parity 2.8e-14; a subprocess test asserts "jax" not in sys.modules after a full NumPy fit; NumPy diag 3.8× faster than JAX-CPU on a 40×40 / S=400 probe.

Downstream: PyAutoLens/autolens/potential_correction/{fit_interferometer,iterative_interferometer}.py call curvature_matrix_diag_from without xp and therefore now take the NumPy branch — they already np.asarray the result, so they work unchanged and get faster. No workspace script needs migration. One profiling script (autolens_profiling/.../datacube/delaunay.py:1011) calls it without xp while passing traced JAX arrays and needs xp=jnp; profiling-repo follow-up.

Next: task 4 (interferometer-numba-cpu-direct-conv, #543) is stacked on this branch. Merge order: #544 first.

Heart YELLOW at ship time, organism-scope only: "workspace validation not passing (5 failed, 2 timeout, cloud#34099198772)" and "release validation incomplete: no rehearsal for current source" — neither touches the interferometer sparse operator.

[Jammy2211]

Shipped

Merged in #544 (merge commit 39d3024c), CI green on all three legs (unittest 3.12, unittest 3.13, unittest-nojax).

What changed

• InterferometerSparseOperator keeps the raw preload and builds its transform lazily: Khat (JAX, under ensure_compile_time_eval) or khat_np (scipy), whichever branch the caller asks for. No JAX import on the NumPy path.
• Six methods on the operator now take an explicit xp; InversionInterferometerSparse passes xp=self._xp, so xp=np propagates from the fit through to the application path.
• A NumPy interferometer fit imports no JAX at all — verified in a subprocess that asserts jax is absent from sys.modules.
• The Delaunay col=-1 padding is now dropped explicitly on the NumPy branch rather than relying on JAX's out-of-bounds scatter semantics.

Numbers

• Parity JAX vs NumPy: <= 5e-16 relative on F; end-to-end curvature matrix agreement 2.8e-14.
• 3.8x faster than the JAX-CPU route on a 40x40 / S=400 probe.

Downstream

PyAutoLens potential_correction callers that build interferometer inversions with xp=np now take the NumPy branch — same results, no JAX import.

Task 4 (#543) is stacked on this branch; its PR opens against main next.