Work Classification

Library

Affected Repositories

• PyAutoArray (primary)

Branch Survey

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.

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:

(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).