research: numba positive-only solve vs JAX sparse-operator mode — same linear algebra?

[Jammy2211]

Overview

Research task (batch 2026-08-31-pm, member numba-vs-jax-sparse, --auto, effective level supervised = min(header safe, research cap)). The subhalo_validation end-to-end runs put the numba CPU path ~8x faster than the old code. The human asked whether the linear-algebra approach behind that speed is the same as the sparse-operator mode already implemented for JAX — and if not, whether the numba approach is amenable to GPU-JAX speed-up or rests on sparsity exploitation that suits it to CPU. Deliverable is a written verdict posted on this issue; no library changes.

Plan

• Characterise the numba CPU positive-only lineage as shipped: the SparseLinAlgImagingNumba PSF-precision preload, the sparse F contraction (perf: numba CPU curvature matrix F at HST resolution (phase 1: split + FFT mapper×func block) #505/perf(numba): HST curvature matrix phase 2 — mapper×mapper block + MGE operated matrix #507 two-stage accumulator), and the fnnls active-set solve with the cross-evaluation warm-start memo (perf: cut positive-only solve active-set iterations — warm-start memo + diagnostics (numba CPU likelihood, phase 3a) #498/perf: warm-start the positive-only NNLS solve from the previous evaluation's passive set (#498) #501).
• Characterise the JAX sparse-operator mode: the FFT-backed matrix-free ImagingSparseOperator, batched F assembly, and the jaxnnls PDIP solve with its relaxed-KKT custom VJP.
• Compare formulation: operator storage, curvature assembly, solve strategy — same normal-equations problem or different mathematics?
• For each point of difference, assess GPU-JAX portability (batched/dense enough? data-dependent control flow? cross-evaluation state?) against the closed solver ledger (nnls_solver_ledger.md: PDIP/BPP/ADMM/warm-start already measured).
• Verdict + recommendation: unify, port, or keep two deliberately different paths, with evidence.

Detailed implementation plan

Affected Repositories

• PyAutoArray (primary — read-only; no changes)

Branch Survey

Repository	Current Branch	Dirty?
./PyAutoArray	main	clean (read-only)

Suggested branch: n/a — research verdict only, no code changes, no worktree.

Investigation Steps

1. Read the solver dispatch: autoarray/inversion/inversion/inversion_util.py::reconstruction_positive_only_from (xp dispatch: JAX → util/jax_nnls.py::solve_nnls_primal PDIP; NumPy → util/fnnls.py::fnnls_cholesky active-set + inversion/inversion/nnls_memo.py warm-start memo).
2. Read the two operator realisations: inversion/inversion/imaging/inversion_imaging_util.py::ImagingSparseOperator (JAX, FFT matrix-free W = Hᵀ N⁻¹ H) vs inversion/inversion/imaging_numba/inversion_imaging_numba_util.py::psf_precision_operator_sparse_from + SparseLinAlgImagingNumba (numba, explicit per-pixel-pair preload); the dataset factories apply_sparse_operator / apply_sparse_operator_cpu and the inversion factory dispatch.
3. Read the two F assemblies: ImagingSparseOperator.curvature_matrix_diag_from (COO scatter → batched FFT → segment_sum) vs curvature_matrix_via_sparse_operator_two_stage_from (perf(numba): HST curvature matrix phase 2 — mapper×mapper block + MGE operated matrix #507 sparse contraction over unique data→pix mappings).
4. Ground the numbers against autolens_profiling/results/notes/nnls_warm_start_memo.md, nnls_warm_start_memo_matrix.md, nnls_solver_ledger.md and the breakdown JSONs — no re-profiling.
5. Pull Memory context: wiki/methods/concepts/autodiff-implicit-diff.md (gradients through the positive-only solve), wiki/methods/sources/linear-algebra.md (FNNLS lineage).
6. Write the verdict per the prompt's four questions as an issue comment; park for human review (judge-tier task — no decide-and-flag).

Key Files

• autoarray/inversion/inversion/inversion_util.py — solver dispatch, both paths
• autoarray/util/fnnls.py, autoarray/util/cholesky_funcs.py, autoarray/inversion/inversion/nnls_memo.py — numba/NumPy active-set solve
• autoarray/util/jax_nnls.py — JAX PDIP driver with custom VJP
• autoarray/inversion/inversion/imaging/{sparse.py,inversion_imaging_util.py} — JAX sparse-operator mode
• autoarray/inversion/inversion/imaging_numba/{sparse.py,inversion_imaging_numba_util.py} — numba sparse-operator mode
• autoarray/inversion/inversion/factory.py, autoarray/dataset/imaging/dataset.py — mode selection

Original Prompt

Click to expand starting prompt

Is the numba positive-only solve the same linear algebra as the JAX sparse-operator mode?

Type: research
Target: autoarray
Repos:

• PyAutoArray
  Themes:
• numba-cpu
• jax
  Difficulty: medium
  Autonomy: safe
  Priority: normal
  Status: formalised
  Consequence: judge
  Review-minutes: 20
  Unattended: ready
  Filed: 2026-08-31

Intake from the 2026-08-31-am batch review, prompted by the subhalo_validation
end-to-end runs (numba CPU path ~x8 faster than the old code end to end). Human,
verbatim:

"""

1. End to end runs are good in general, with numba speed up run times roughly x8
   faster than old code, incredible. Do an intake to investigate if the linear algebra
   approach taken for this is the same as the sparse operator mode alreasdy implemented
   for JAX, and if it is not ask whether it is amenable to GPU JAX speed up or if there
   is some sparsity exploitation that means this is suited to CPU.
   """

Deliverable is a written verdict (no library changes):

1. Characterise the linear algebra of the numba CPU positive-only solve as shipped
   (the NNLS warm-start/memo lineage — PyAutoArray perf: warm-start the positive-only NNLS solve from the previous evaluation's passive set (#498) #501, epic numba-cpu-likelihood;
   curvature-matrix F work perf: numba CPU curvature matrix F at HST resolution (phase 1: split + FFT mapper×func block) #505/perf(numba): HST curvature matrix phase 2 — mapper×mapper block + MGE operated matrix #507; the 177 MB sparse operator per process).
2. Compare against the sparse-operator mode already implemented for JAX: same
   formulation (operator storage, curvature assembly, solve strategy) or different?
3. If different: is the numba approach amenable to a GPU-JAX port (batched/dense
   enough, no data-dependent control flow), or does its speed rest on sparsity
   exploitation / branchy active-set iteration that is intrinsically CPU-suited?
4. Verdict + recommendation: unify, port, or keep two deliberately different paths —
   with the evidence for whichever it is.

Ground against the installed stack and the numba-cpu epic's memory/notes; do not
re-profile from scratch — the breakdown JSONs and nnls_warm_start_memo.md in
autolens_profiling already carry the numbers.

[Jammy2211]

Verdict: same linear algebra, deliberately different machines — keep both, port nothing

TL;DR. The numba CPU positive-only solve and the JAX sparse-operator mode solve the identical mathematical problem — the same regularised normal equations Q·s = q, s ≥ 0 with the same precision operator W = Hᵀ N⁻¹ H — but realise it differently at every layer: operator storage, curvature assembly, and solver. Each difference is a deliberate fit to its hardware, and each numba lever is either already superseded on GPU by the JAX mode's own formulation or has been measured to fail when transplanted (the closed solver ledger). Recommendation: keep two deliberately different paths. No port, no unification.

1. The numba CPU positive-only solve as shipped

The lineage (epic numba-cpu-likelihood, phases 1–3a + curvature-F successors #505/#507) is four CPU-specific levers stacked on one formulation:

Formulation (shared with every inversion path): F = Aᵀ W A with W = Hᵀ N⁻¹ H (H = PSF convolution, N = diagonal noise covariance), data vector q = Aᵀ Hᵀ N⁻¹ d (psf_weighted_data), Q = F + reg, then the positive-only reconstruction minimises |Z·s − x|² s.t. s ≥ 0 where Q = ZᵀZ and q = Zᵀx (inversion_util.reconstruction_positive_only_from docstring; Warren & Dye 2003 / arXiv:astro-ph/0302587 eqs. 2 and 4).

The four levers:

1. Explicit sparse operator storage (SparseLinAlgImagingNumba, built by Imaging.apply_sparse_operator_cpu()): the nonzeros of W are precomputed once per process — for every data pixel, every partner pixel within the (2K−1)² kernel-overlap footprint, value = the PSF-pair convolution against N⁻¹ (psf_precision_operator_sparse_from; upper triangle only, diagonal halved). Memory scales as image_pixels × (2Ky−1)(2Kx−1) × 16 B (values + indexes) — the "177 MB per process" the subhalo_validation runs report is this preload at their mask/PSF geometry. It trades memory for never doing a PSF convolution during curvature assembly.
2. Sparse gather/scatter curvature contraction (curvature_matrix_via_sparse_operator_two_stage_from, perf(numba): HST curvature matrix phase 2 — mapper×mapper block + MGE operated matrix #507/perf(numba): the mapper x mapper block and the MGE operated matrix (#507) #508): F is contracted directly from the preload and the unique data→pixelization mappings — data-dependent loop bounds (pix_lengths, psf_precision_lengths), irregular indexed accumulation, restructured in perf(numba): HST curvature matrix phase 2 — mapper×mapper block + MGE operated matrix #507 into a per-data-pixel L1-resident source-space accumulator + contiguous AXPYs (2.90× on the block). perf: numba CPU curvature matrix F at HST resolution (phase 1: split + FFT mapper×func block) #505/perf(numba): speed up the curvature matrix F at HST resolution (#505) #506 routed the mapper×linear-func block through the batched FFT Convolver. Together: HST rectangular eval 1.56 → 0.33 s (4.7×).
3. Active-set NNLS solve (fnnls_cholesky, Bro & De Jong 1997): takes (Q, q) directly; nested data-dependent while loops adding/removing passive-set entries; each subproblem solved by in-place rank-1 Cholesky up/downdates of the passive corner (perf: in-place Cholesky buffer + copy-free numba solves for fnnls_cholesky #453/fix: coerce fnnls_cholesky inputs to numpy — JAX ZTx broke the numba buffer kernels #463 preallocated buffer). Exact KKT optimum of a strictly convex problem — unique, so seed quality affects cost, never correctness.
4. Cross-evaluation warm-start memo (perf: cut positive-only solve active-set iterations — warm-start memo + diagnostics (numba CPU likelihood, phase 3a) #498 → perf: warm-start the positive-only NNLS solve from the previous evaluation's passive set (#498) #501, nnls_memo.py): a module-level dict keyed by the index space (never matrix values) seeds each solve from the previous evaluation's final passive set, with a self-calibrating relative-error guard (drop at 1.5× the dense-sign reference). Measured (nnls_warm_start_memo.md + 32-cell robustness matrix): 9.9× (euclid) / 4.0× (hst) median iteration reduction in the sampler-like random-walk regime — 70 → 7 iterations — parity ≤ 1e-14; iid (uncorrelated) regime 0.52–1.09×, still never slower in wall-time because the memo also skips the dense sign-start solve. A sibling memo caches the MGE operated mapping matrix across evaluations (AUTOARRAY_NUMBA_OPERATED_MEMO).

2. Comparison with the JAX sparse-operator mode

Same problem, three deliberate divergences. The factory dispatches purely on which operator the dataset carries (factory.py: SparseLinAlgImagingNumba → InversionImagingSparseNumba; ImagingSparseOperator → InversionImagingSparse), and both classes assemble the same Q and q by construction; since the NNLS optimum is unique, both solvers converge to the same reconstruction to solver tolerance.

layer	numba CPU (imaging_numba/)	JAX sparse-operator (imaging/sparse.py)
operator storage	explicit: precomputed nonzeros of W, O(N_pix · K²) memory (~177 MB/process class), computed once per process	matrix-free: stores only rFFT of PSF + flipped PSF and 1/σ² map (few MB); W applied on demand as FFT → weight → FFT
curvature assembly	sparse contraction over data-dependent index lists; sequential, cache-tuned, branchy	dense batched: scatter COO columns of A to the rectangular grid, apply W to 128-column blocks via three FFTs, segment_sum back — fixed shapes, no data-dependent control flow, jit/vmap-safe
solve	active-set (combinatorial; variable iteration count 7–95; Cholesky up/downdates; warm-startable across evaluations)	jaxnnls PDIP interior point (~20 fixed-shape iterations, each a fresh dense Cholesky of the KKT system, lax.while_loop, capped at 50)
differentiability	none (and none needed)	differentiable — relaxed-KKT custom VJP; the 2026-07 gradient audit FD-validates jax.grad through the positive-only solve
cross-eval state	two mutable memos (passive set, MGE matrix)	none — pure functions, jit-cacheable

So the answer to the human's first question is: the formulation is the same; the "approach" — everything that makes the numba path 8× faster — is not.

3. Is the numba approach amenable to a GPU-JAX port?

No — each lever, examined:

1. The explicit W preload + gather/scatter contraction is a cache-locality optimisation. On GPU the same F is the step that accelerates best under the existing JAX formulation — the solver ledger measured F at 3.01 s laptop-CPU vs 0.397 s on a consumer RTX 2060 ("F is one big GEMM and melts on GPU"). The GPU-appropriate realisation of this operator already exists and is the JAX mode. Porting the numba contraction would move a fix for a bottleneck the GPU does not have onto hardware that punishes its data-dependent loop bounds and irregular scatter.
2. The active-set solver is branchy by construction: nested while loops with data-dependent trip counts and passive-set-sized factor updates. The GPU-viable active-set analog was already tried and closed (nnls_solver_ledger.md, PyAutoArray#369/perf: BPP/ADMM NNLS solver experiment (PDIP alternative, #369 follow-up) #370): block principal pivoting needed 36–53 factorizations vs PDIP's 19–21 (strict Kim–Park: 151) on real extracted production systems — slower; ADMM fails to reach evidence tolerance; the ledger's conclusion 1: "PDIP at ~20 iterations × fresh Cholesky is near-optimal for this problem class among GPU-viable algorithms. Do not re-litigate without new structure."
3. The warm-start memo does not transfer, on three measured/structural grounds: (a) the ledger measured warm-starting the PDIP solver from the unconstrained solve and it hurt — 17 → 38 iterations, because it un-centers the interior point (the memo's own "Relation to the JAX solver ledger" section records exactly this non-transfer); a warm start is only a discrete passive-set guess in an active-set method. (b) It is cross-evaluation mutable state — incompatible with jit purity; it would have to become explicitly threaded carried state, a different API. (c) Its win requires sequentially correlated evaluations (random-walk 9.9×); the GPU sampler execution model is vmap-batched evaluation, where lax.while_loop runs every lane to the slowest lane and batch members are not each other's predecessors — the memo's iid rows (0.52–0.91×) are the measured preview of that regime.
4. The differentiability asymmetry seals it: the JAX path's solver must support jax.grad (relaxed-KKT VJP, FD-validated); an active-set solve with discrete set changes has no such machinery. Replacing PDIP with a ported active-set solver would cost the gradient path that JAX inference (HMC/NUTS, gradient optimisers) exists for.

So the human's dichotomy resolves cleanly to the second horn: the numba speed rests on sparsity exploitation, cache-resident branchy iteration, and cross-evaluation state — all intrinsically CPU-suited — while the GPU already holds the right formulation for its own strengths (dense batched FFT + fixed-shape PDIP). One further CPU-only fact: the numba path also serves Delaunay meshes, whose construction is a structural JAX hole (pure_callback, no JVP).

4. Recommendation

Keep two deliberately different paths. Do not unify, do not port.

• The paths already share what should be shared: the formulation (Q, q, regularisation, the single dispatch point in reconstruction_positive_only_from) and the parity bar (pinned likelihoods, rtol 1e-6).
• Every remaining difference is load-bearing for its hardware, and the expensive transplant experiments have already been run and closed (warm-start→PDIP: hurts; BPP/ADMM on GPU: slower/fails).
• Future speed-ups that help both paths live at the problem level, exactly as the ledger's conclusion 3 states: fewer source pixels (n³), regularisation design that improves conditioning, relaxing the positivity model — not solver cross-pollination.

Scholar-mode candidates (not filed — for review; /intake if wanted)

• [from: research numba-vs-jax-sparse · PyAutoArray#513] A short design note (PyAutoMemory wiki/methods/ or a PyAutoArray doc page) recording "two sparse-operator modes: same math, different machines" with the decision table above — the closed solver ledger lives in autolens_profiling where library sessions don't look, and this question will be re-asked.
• [from: research numba-vs-jax-sparse · PyAutoArray#513] The solver ledger's A100 rows (job 330046, submit_probe_nnls_a100) were still pending when the ledger closed; if they never landed, a small follow-up to land or formally retire them would complete the GPU half of the evidence base this verdict cites.

Evidence base: PyAutoArray main @ 302d5df3 (inversion_util.py, util/fnnls.py, util/jax_nnls.py, nnls_memo.py, imaging/sparse.py + inversion_imaging_util.py, imaging_numba/sparse.py + inversion_imaging_numba_util.py, factory.py, dataset/imaging/dataset.py); autolens_profiling results/notes/{nnls_warm_start_memo,nnls_warm_start_memo_matrix,nnls_solver_ledger}.md; Mind complete/2026/08/{numba-cpu-nnls-iteration-reduction,numba-hst-curvature-matrix-speedup,numba-hst-curvature-matrix-phase2}.md + epics.md; PyAutoMemory wiki/methods/concepts/autodiff-implicit-diff.md, wiki/methods/sources/linear-algebra.md. No re-profiling performed, per the prompt.

---

Task is Consequence: judge — this verdict parks for human review (supervised checkpoint-and-continue; no decide-and-flag). Sign-off retires the prompt via the usual close-out; if the verdict doesn't pass, the follow-up is an ordinary /intake.

---

Generated by Claude Code

[Jammy2211]

Batch review ruling (2026-08-31-pm, reviewed 2026-09-01 15:13): verdict ACCEPTED.

"Same linear algebra, deliberately different machines — keep both paths, port nothing" stands as the ruling of record. No library changes follow from it. The two scholar-mode candidates (a design note recording the two-modes decision; landing or retiring the solver ledger's pending A100 rows) remain unfiled — they can be /intake'd on demand.

Closing; the PyAutoMind prompt retires to complete/ via the batch close-out.

---

Generated by Claude Code