feat: numba CPU interferometer curvature path — direct_conv, geometry-gated

[Jammy2211]

Overview

Reinstate a numba CPU curvature path for the interferometer inversion: a new package autoarray/inversion/inversion/interferometer_numba/ mirroring the live imaging_numba/, holding one kernel — the extent-grid direct convolution prototyped as direct_conv in autolens_profiling/scripts/misc/numba_interferometer/kernels.py. For each source column of A it convolves over the (Ny, Nx) unmasked-extent rectangle with contiguous W~ preload rows, then projects with A^T; cost O(nnz*M + S*nnz).

Do not reinstate the deleted kernel. The recovered O(N^2 P^2) pair loop beats a NumPy rfft2 convolution at only one of six measured cells; direct_conv beats it by 2.9-4.8x on the real likelihood's F row.

The path is selected, not defaulted to: it is gated on geometry (non-zeros per source column) and only ever runs when xp is np, numba imports, and the preconditions hold. Above the crossover the FFT path wins, so a blanket default would be a regression.

Scope of this issue: the PyAutoArray library phase only. The prompt also lists autolens_profiling; that is the follow-on in-situ phase — re-measuring delaunay_numba.py armed through the library dispatch rather than through the profiling pack — and it is filed separately via /intake at ship time as draft/research/autolens_profiling/interferometer_numba_library_dispatch_insitu.md. It is not done here.

Stacked on

This is stacked on #542 (perf: NumPy/scipy application path for InterferometerSparseOperator), which stores the raw real-space (2y, 2x) preload that this kernel indexes and establishes the xp=np application path this class specialises.

• Branch: feature/interferometer-numba-cpu-direct-conv, cut from feature/interferometer-sparse-operator-numpy-cpu-path (not from main).
• PR base = feature/interferometer-sparse-operator-numpy-cpu-path until perf: NumPy/scipy application path for InterferometerSparseOperator (xp=np never imports JAX) #542 merges; GitHub retargets the PR to main on that merge.
• Merge order: perf: NumPy/scipy application path for InterferometerSparseOperator (xp=np never imports JAX) #542 first, then this one.

Measured (from the prompt)

Single thread, i9-10885H, autolens 2026.8.17.1. In-situ F: mapper x mapper row, real likelihood, arms interleaved with dgemm controls:

instrument	mesh	recovered kernel	direct_conv	JAX/FFT (jit-warm)	direct_conv vs JAX
sma	Delaunay 1500	0.2475 s	0.0844 s	0.5917 s	7.01x
sma	rect 32^2	0.3967 s	0.0998 s	0.4101 s	4.11x
alma	Delaunay 1500	4.0710 s	1.2250 s	2.4297 s	1.98x
alma	rect 32^2	7.0205 s	1.4631 s	1.6110 s	1.10x

Whole evaluation, with the jit-warm F substituted into the measured non-F cost: 2.43x (sma Delaunay), 2.76x (sma rect), 1.77x (alma Delaunay), 1.08x (alma rect) faster than JAX-CPU.

The gate

The controlling variable is non-zeros per source column, nnz/S = N_pix*P/S, with nnz = mapper.pix_sizes_for_sub_slim_index.sum() and S = mapper.params. Measured crossover against a NumPy rfft2 convolution: ~60 (Delaunay) to ~77 (rectangular). Below it the numba kernel wins (up to 5.93x); above it the FFT wins — by 1.6x at alma_high Delaunay and 3.7x at alma_high rectangular. The default cap is therefore 60.0, the conservative of the two, and the constant is machine-dependent and documented as such.

Links

• Epic (retired ledger): PyAutoMind/complete/archive/epics/numba_interferometer_likelihood_revisit.md
• Phase 2 verdict, bake-off and in-situ arms: research: numba vs JAX-CPU interferometer curvature — levers and verdict autolens_profiling#226 (autolens_profiling/results/notes/numba_interferometer_verdict.md)
• Prototype pack (kernels + parity tests) landed by research: interferometer preload on CPU — builder timings + type-1 NUFFT construction autolens_profiling#229 / PR research: interferometer preload on CPU — builder timings + type-1 NUFFT construction autolens_profiling#234
• Merged today, same epic: PyAutoArray#540 (rfft2 in apply_operator), PyAutoArray#541 (NUFFT type-1 preload)
• Prerequisite: PyAutoArray#542 (NumPy/scipy application path)

Plan

• New package autoarray/inversion/inversion/interferometer_numba/ mirroring imaging_numba/: a util module holding the direct_conv kernel (plus its lazily-built prange sibling) and its input builder, and a sparse.py with InversionInterferometerSparseNumba.
• Port the kernel verbatim from the profiling prototype, but take the mask's extent_index_for_masked_pixel for the extent geometry so nothing new has to be stored on the dataset.
• InversionInterferometerSparseNumba raises InversionException on every precondition it cannot honour — non-NumPy xp, missing numba, linear-function lists, more than one mapper, over-sampling — rather than working around them.
• Add a Settings key for the geometry gate (interferometer_numba_nnz_per_source_max, config-backed, default 60.0, 0 disables) and dispatch on it in inversion/factory.py; preconditions failing at dispatch fall through silently to the existing sparse path, which is a routing decision, not an error.
• New test package pinning F, D, reconstruction and log_evidence against InversionInterferometerSparse(xp=np) at rtol=1e-10, atol=1e-10*max|F|, with a control that a 1 % scale of F fails the pin, parallel == serial, builder parity with the prototype, factory routing, and each precondition's message.
• The in-situ re-measurement on autolens_profiling is a separate follow-up prompt filed at ship time, not part of this issue.

Detailed implementation plan

Work Classification

Library

Affected Repositories

• PyAutoArray (primary, and the only repo this issue touches)

Branch Survey

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

Suggested branch: feature/interferometer-numba-cpu-direct-conv (cut from feature/interferometer-sparse-operator-numpy-cpu-path, PR base the same until #542 merges)

Worktree root: ~/Code/PyAutoLabs-wt/interferometer-numba-cpu-direct-conv/

Implementation Steps

1. interferometer_numba/inversion_interferometer_numba_util.py

• curvature_direct_conv(preload, iy, ix, flat, indptr, col, val, cscptr, csc_row, csc_val, ny, nx, pix_pixels) ported verbatim from the prototype under @numba_util.jit().
• direct_conv_parallel_kernel() lazily built with numba.njit(cache=True, parallel=True, nogil=True) exactly as the prototype does (numba_util.jit cannot express parallel=True), selected when conf["general"]["numba"]["parallel"] is true — the same flag the shared decorator reads.
• kernel_inputs_from(pix_indexes, pix_sizes, pix_weights, extent_index_for_masked_pixel, extent_shape, pix_pixels) — the prototype's builder, but taking the mask's extent_index_for_masked_pixel (iy = flat // nx, ix = flat % nx) instead of native_index_for_slim_index, so nothing new has to be stored on the dataset.
• nnz_per_source_column_from(mapper) = pix_sizes_for_sub_slim_index.sum() / params.

2. interferometer_numba/sparse.py::InversionInterferometerSparseNumba(InversionInterferometerSparse)

• __init__ raises InversionException (messages ported from the prototype) on: xp is not np, missing numba, any AbstractLinearObjFuncList, more than one mapper, or over-sampling (sub_fraction != 1). Preconditions are raised, not worked around.
• kernel_index_arrays cached per mapper.
• curvature_matrix_diag calls the kernel with self.dataset.sparse_operator.nufft_precision_operator (the preload stored by perf: NumPy/scipy application path for InterferometerSparseOperator (xp=np never imports JAX) #542; assert (2ny, 2nx) == preload.shape).
• data_vector inherited. curvature_matrix keeps the base class's regularization / diag handling — the kernel returns a complete symmetric F, so the mirroring pass is a no-op; verify that rather than skipping it.

3. autoarray/settings.py

• interferometer_numba_nnz_per_source_max: Optional[float] = None -> conf["general"]["inversion"]["interferometer_numba_nnz_per_source_max"] with a try/except KeyError fallback of 60.0 (the Delaunay crossover, the conservative of the two measured); 0 disables the path. Documented on the property with the measured crossovers and the note that the constant is machine-dependent.

4. autoarray/inversion/inversion/factory.py — interferometer branch

• After the sparse-operator check: if xp is np and settings.interferometer_numba_nnz_per_source_max > 0 and the preconditions hold and nnz_per_source_column <= max and numba imports -> InversionInterferometerSparseNumba; otherwise InversionInterferometerSparse as today.
• Preconditions failing at dispatch fall through silently to the sparse path (a routing decision, not an error); the numba class itself still raises if constructed directly with bad inputs.
• Mirrors how the imaging branch already selects by operator type (SparseLinAlgImagingNumba -> InversionImagingSparseNumba); the interferometer branch has no numba case today.

5. Tests — new test_autoarray/inversion/inversion/interferometer_numba/

• F, D, reconstruction, log_evidence vs InversionInterferometerSparse(xp=np) at rtol=1e-10, atol=1e-10*max|F| on the DFT Delaunay and rectangular fixtures from test_interferometer.py.
• A control that a 1 % scale of F fails the pin.
• Parallel kernel == serial kernel.
• kernel_inputs_from matches the prototype's on a hand-built mapper.
• Factory routing: gate above / below, xp=jnp never routes numba, func-list and multi-mapper fall through.
• Each precondition raises with its message.
• Run the existing interferometer test tree to prove the default sparse path is unchanged.
• Note: the conftest nufftax-skip is path-keyed and a new interferometer_numba/ test path is not on the list, so the DFT fixtures run unskipped.

6. Workspace / in-situ phase (not this issue)

• The in-situ re-measurement on autolens_profiling — delaunay_numba.py arming through the library dispatch instead of through the prototype pack — is filed as a follow-up prompt via /intake at ship time (draft/research/autolens_profiling/interferometer_numba_library_dispatch_insitu.md).

Key Files

• autoarray/inversion/inversion/interferometer_numba/ — new package (inversion_interferometer_numba_util.py, sparse.py), mirroring autoarray/inversion/inversion/imaging_numba/ (two modules, @numba_util.jit() kernels + sparse.py).
• autoarray/inversion/inversion/factory.py — interferometer dispatch branch.
• autoarray/settings.py — the new gate key (None -> conf["general"]["inversion"][key] with try/except KeyError).
• autoarray/inversion/inversion/interferometer/sparse.py — base class InversionInterferometerSparse.
• autoarray/inversion/inversion/interferometer/inversion_interferometer_util.py — nufft_precision_operator (stored by perf: NumPy/scipy application path for InterferometerSparseOperator (xp=np never imports JAX) #542).
• autolens_profiling/scripts/misc/numba_interferometer/kernels.py — curvature_direct_conv, _curvature_direct_conv_parallel, kernel_inputs_from; inversion.py for the precondition messages.
• test_autoarray/inversion/inversion/interferometer/test_interferometer.py — Delaunay / rectangular DFT fixtures to reuse.

Verification

• Parity + control + parallel == serial + routing + precondition tests green.
• The default interferometer path unchanged (existing interferometer tree green).
• pytest test_autoarray/inversion test_autoarray/dataset test_autoarray/operators -q green; black clean.
• PR base is feature/interferometer-sparse-operator-numpy-cpu-path until perf: NumPy/scipy application path for InterferometerSparseOperator (xp=np never imports JAX) #542 merges.

Original Prompt

Click to expand starting prompt

Reinstate a numba CPU interferometer curvature path — the extent-grid convolution, geometry-gated

Type: feature
Target: autoarray
Repos:

• PyAutoArray
• autolens_profiling
  Themes:
• numba-cpu
• interferometer
• likelihood-profiling
  Difficulty: large
  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
verdict, the bake-off and the in-situ arms are in
autolens_profiling/results/notes/numba_interferometer_verdict.md.

What

Add autoarray/inversion/inversion/interferometer_numba/ mirroring the live
imaging_numba/ package, with one kernel: the extent-grid direct convolution prototyped
as direct_conv in autolens_profiling/scripts/misc/numba_interferometer/kernels.py. For
each source column of A, convolve it over the (Ny, Nx) unmasked-extent rectangle with
contiguous W~ preload rows, then project with Aᵀ. Cost O(nnz·M + S·nnz).

Do not reinstate the deleted kernel. The recovered O(N² P²) pair loop beats a NumPy
rfft2 convolution at only one of six measured cells; direct_conv beats it by 2.9-4.8× on
the real likelihood's F row.

Measured (single thread, i9-10885H, autolens 2026.8.17.1)

In-situ F: mapper×mapper row, real likelihood, arms interleaved with dgemm controls:

instrument	mesh	recovered kernel	direct_conv	JAX/FFT (jit-warm)	direct_conv vs JAX
sma	Delaunay 1500	0.2475 s	0.0844 s	0.5917 s	7.01×
sma	rect 32²	0.3967 s	0.0998 s	0.4101 s	4.11×
alma	Delaunay 1500	4.0710 s	1.2250 s	2.4297 s	1.98×
alma	rect 32²	7.0205 s	1.4631 s	1.6110 s	1.10×

Whole evaluation, with the jit-warm F substituted into the measured non-F cost:
2.43× (sma Delaunay), 2.76× (sma rect), 1.77× (alma Delaunay), 1.08× (alma rect) faster than
JAX-CPU.

The gate this needs

The controlling variable is non-zeros per source column, nnz/S = N_pix·P/S. Measured
crossover against a NumPy rfft2 convolution: ≈60 (Delaunay) to ≈77 (rectangular). Below
it the numba kernel wins (up to 5.93×); above it the FFT wins (by 1.6× at alma_high Delaunay,
3.7× at alma_high rectangular). So the path must be selected, not defaulted to — a
dispatch rule on nnz/S with the constant measured on the target machine, or an explicit
user setting with the rule documented.

Scope

• interferometer_numba/inversion_interferometer_numba_util.py — the direct_conv kernel
  and its prange sibling (3.5× at sma, 4.6× at alma on 8 threads; see the pool caveat).
• interferometer_numba/sparse.py — InversionInterferometerSparseNumba, mirroring
  imaging_numba/sparse.py, dispatched from inversion/factory.py.
• The mapper must emit the CSR/CSC + extent-flat layout directly (as
  _sparse_triplets_curvature_from already emits COO), so the kernel does not marshal its
  own inputs per evaluation.
• Preconditions to raise on, not work around (the profiling pack already does): linear-
  function lists, multiple mappers, over_sample_size != 1, non-NumPy xp.
• Parity gate: F, D, reconstruction and log evidence against
  InversionInterferometerSparse at rtol=1e-10 on F with atol scaled by max|F|,
  plus a control that a 1 % scale fails the pin.
• The real-space W~ preload must be kept on the dataset: InterferometerSparseOperator
  currently stores only Khat = fft2(preload) and discards the array this kernel indexes.

[Jammy2211]

Shipped

Merged as #545 (merge commit 35aa681f).

What landed. A new package autoarray/inversion/inversion/interferometer_numba/:
the direct_conv numba CPU curvature kernel (ported from the autolens_profiling
investigation pack, autolens_profiling#226) as a serial @numba_util.jit() function plus
a prange variant compiled lazily on first use, and
InversionInterferometerSparseNumba — an InversionInterferometerSparse subclass that
overrides curvature_matrix_diag only and raises InversionException on every
configuration the kernel cannot represent (non-NumPy xp, more than one mapper, any
AbstractLinearObjFuncList, over-sampling).

The gate. Settings.interferometer_numba_nnz_per_source_max, default 60.0 from
autoarray/config/general.yaml (KeyError fallback to the same value; 0 disables the
path). inversion_interferometer_from routes a single-mapper, no-func-list,
no-over-sampling, xp=np interferometer inversion at or below the gate to the numba
class. Measured crossovers: ~60 non-zeros per source column on Delaunay, ~77 on
rectangular (autolens_profiling#226 verdict §2) — a machine-dependent constant.

Parity vs the sparse NumPy path: curvature matrix F max abs diff 1.8e-15, data
vector D exact, reconstruction 4.9e-17; parallel kernel exactly equal to serial; the
1 % perturbation control fails as it must.

Behaviour change. The default CPU interferometer route changes with this PR: CPU
xp=np single-mapper interferometer fits at or below the gate now go through numba by
default when numba is installed. Set
inversion: interferometer_numba_nnz_per_source_max: 0 in general.yaml (or pass it to
Settings) to restore the FFT route.

Tests: pytest test_autoarray 1501 passed locally; CI green on all three legs
(py3.12, py3.13, nojax). black --check clean.

Stacked on #544 (merged 2026-09-08), so this diffed cleanly against main.

Follow-on. The crossover numbers were measured on the standalone prototype pack, not
the library dispatch. The in-situ re-measurement is filed as
PyAutoMind/draft/research/autolens_profiling/interferometer_numba_library_dispatch_insitu.md.

This closes the last library follow-up of the retired numba-interferometer-revisit
epic.