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

[Jammy2211]

Summary

InterferometerSparseOperator used jax.numpy unconditionally in every method, so a CPU
user running xp=np still executed the whole curvature build through JAX — 1.1–1.3× slower
than a NumPy/scipy rfft2 route, and it forced JAX into the process of a fit that had asked
for NumPy.

The operator now keeps the raw (2y, 2x) real preload as nufft_precision_operator and
builds both transforms of it lazily: Khat (JAX rfft2) and khat_np
(scipy.fft.rfft2). All six public methods take an xp array module and branch on it:

method	JAX branch	NumPy branch
apply_operator	jnp.fft.rfft2 × Khat → irfft2	scipy.fft.rfft2 × khat_np → irfft2
curvature_matrix_diag_from	lax.fori_loop + segment_sum	scipy.sparse CSC A, Python block loop, Aᵀ @ G
curvature_matrix_off_diag_from	same, over A1's columns	same, over A1's columns
operated_matrix_slim_from	.at[].set scatter	fancy-index scatter
curvature_matrix_off_diag_func_list_from	dynamic_update_slice	dense slice assignment
curvature_matrix_func_list_from	segment_sum	Aᵀ @

Design: the backend follows the inversion's xp. InversionInterferometerSparse passes
xp=self._xp to every operator call, so what the fit asked for decides the backend. The
dataset's use_jax kwarg only picks a brute-force builder for the preload (since #541) and
never touches the application path. This is the repo's existing convention
(transformer.image_from(xp=)), and it removes the surprising split the issue names without
changing any dataset API. The JAX bodies are unchanged; the np.array(curvature_matrix)
workaround in sparse.py (needed only because the operator used to hand JAX arrays back to a
NumPy fit) is deleted.

No-JAX guarantee

A NumPy fit never imports JAX. The lazy Khat / col_offsets are the mechanism; a
subprocess test runs a full InversionInterferometerSparse(xp=np) fit and asserts
"jax" not in sys.modules afterwards. The new tests carry no importorskip("jax"), so the
unittest-nojax CI leg runs them.

Parity and speed

quantity	NumPy vs JAX
all six methods (7×7/K=5, 12×12/K=64 seeded, batch_size=4 block loop, Delaunay duplicate-COO)	≤ 5e-16 relative
end-to-end InversionInterferometerSparse curvature matrix	2.8e-14
curvature_matrix_diag_from, 40×40 grid / S=400	3.8× faster than JAX-CPU

Pins are the repo idiom assert_allclose(rtol=1e-10, atol=1e-10*max|ref|).

Two traps worth recording

• Tracer leak. Building Khat lazily inside a method that is itself traced (the JAX
  curvature bodies run under lax.fori_loop) turns the preload transform into a traced
  constant. jax.ensure_compile_time_eval() around the transform is what makes the laziness
  safe — without it the operator's static aux data leaks into the trace.
• col = -1 padding. Delaunay's triplets are padded with col = -1 entries. JAX's
  segment_sum drops out-of-range segments silently; scipy.sparse.csc_matrix does not — it
  raises. The NumPy builder filters 0 <= col < S explicitly.

API Changes

Additive only. Six InterferometerSparseOperator methods gain a keyword-only xp=np
argument, defaulting to the NumPy branch. The dataclass gains a
nufft_precision_operator field (the raw preload, previously discarded after the FFT) and
Khat / col_offsets become cached properties rather than eager fields, with a new
khat_np alongside. No signature is broken and no symbol is removed.

See full details below.

Test Plan

• pytest test_autoarray -q → 1486 passed
• python -m black --check clean on all five files
• Six-method NumPy-vs-JAX parity at the exact pin, on three fixtures plus a forced block loop
• End-to-end inversion parity (curvature_matrix, data_vector, reconstruction, log_evidence)
• Subprocess no-JAX test: "jax" not in sys.modules after a full NumPy fit
• Every existing interferometer test and tolerance unchanged

Downstream note

PyAutoLens/autolens/potential_correction/fit_interferometer.py and
iterative_interferometer.py call curvature_matrix_diag_from without xp, so they now
take the NumPy branch. Both already wrap the result in np.asarray, so they work unchanged
and get faster. One profiling script,
autolens_profiling/scripts/interferometer/likelihood_breakdown/datacube/delaunay.py:1011,
also calls it without xp but passes JAX arrays inside a jit — it needs xp=jnp added; a
profiling-repo follow-up, no library or shipped-workspace impact.

Full API Changes (for automation & release notes)

Added

• InterferometerSparseOperator.nufft_precision_operator: np.ndarray — the raw (2y, 2x)
  real preload, now kept rather than discarded after the FFT (the NumPy and forthcoming
  numba CPU paths index it directly).
• InterferometerSparseOperator.khat_np — cached property, scipy.fft.rfft2 of the preload.

Changed Signature

• InterferometerSparseOperator.apply_operator(Fbatch_flat, xp=np)
• InterferometerSparseOperator.curvature_matrix_diag_from(rows, cols, vals, *, S, xp=np)
• InterferometerSparseOperator.curvature_matrix_off_diag_from(rows0, cols0, vals0, rows1, cols1, vals1, *, S0, S1, xp=np)
• InterferometerSparseOperator.operated_matrix_slim_from(matrix_slim, extent_index_for_masked_pixel, xp=np)
• InterferometerSparseOperator.curvature_matrix_off_diag_func_list_from(..., xp=np)
• InterferometerSparseOperator.curvature_matrix_func_list_from(..., xp=np)

All are additive: xp defaults to numpy.

Changed Behaviour

• InterferometerSparseOperator.Khat and .col_offsets are cached properties, not eager
  dataclass fields; they are computed on first access under
  jax.ensure_compile_time_eval(). Reading them still returns JAX arrays.
• InversionInterferometerSparse methods return NumPy arrays under xp=np (they returned
  JAX arrays before), so callers no longer need np.asarray.

Migration

• Before: sparse_operator.curvature_matrix_diag_from(rows, cols, vals, S=S) — always JAX.
• After: the same call runs NumPy/scipy; pass xp=jnp explicitly for the JAX branch (needed
  only if the arrays are traced).

Closes #542

Generated by the PyAutoLabs agent workflow.

🤖 Generated with Claude Code

https://claude.ai/code/session_018hLF3ZAcz5MmaSJBEcLkvF