perf: cut positive-only solve active-set iterations — warm-start memo + diagnostics (numba CPU likelihood, phase 3a)

[Jammy2211]

Overview

Phase 3a of the numba CPU sparse-operator likelihood speed restoration (epic numba-cpu-likelihood). On the Delaunay-1250 + Hilbert AdaptImage campaign fiducial (apply_sparse_operator_cpu(), use_jax=False) the positive-only reconstruction solve is ~72% of a euclid evaluation and is iteration-bound (autolens_profiling#151 comments 5-6): the production warm start (P_initial = np.linalg.solve(curvature_reg, data_vector) > 0, inversion_util.py:378) gets ~1411/1560 entries right, and fnnls_cholesky adds one index per outer iteration and drops violators one at a time, so ~150 wrong entries ⇒ ~150 outer + ~100 inner iterations. PR #453/#463 already removed the per-iteration copy overhead (record complete/2026/08/numba-fnnls-inplace-cholesky-buffer.md); the iteration COUNT is untouched. The "restore deleted numba fnnls" idea is retired — the solver was never missing.

This task instruments the solver, measures on a realistic instance sequence which warm start beats the dense sign pattern, then ships a cross-evaluation warm start behind a settings knob + env kill-switch. Solution unchanged (unique NNLS optimum), pinned log-likelihoods must hold at rtol 1e-6. Phase 3b (multi-index / batched active-set moves, with a single-index fallback — note block principal pivoting already lost against the JAX PDIP solver, autolens_profiling/results/notes/nnls_solver_ledger.md) is a separate prompt driven by 3a's diagnostic.

Plan

• Expose solver diagnostics from fnnls_cholesky (final passive set, outer/inner iteration counts, warm-start error count) via an optional stats dict — default return unchanged.
• Perturbed-instance diagnostic on the euclid/hst Delaunay-1250 fiducial: per solve record iterations, |P_final ⊕ P_dense-sign|, |P_final(t) ⊕ P_final(t−1)| over a random-walk sequence AND an i.i.d.-draw sequence (nested-sampling order), solve seconds. Run with AUTOARRAY_NUMBA_OPERATED_MEMO=0.
• Cross-evaluation warm start: module-level FIFO memo (template: imaging_numba/sparse.py:20-52), keyed on problem size + mesh fingerprint, Settings knob + AUTOARRAY_NNLS_WARM_START=0 kill-switch, dense-sign fallback on miss. Fork-based pools give each worker its own memo automatically.
• Cheap exact wins alongside: seed the Cholesky buffer from the warm-start factorisation instead of a second from-scratch slg.cholesky; skip the separate dense np.linalg.solve when a warm passive set is available.
• Tests + measurement; gate for on-by-default = median iteration reduction ≥2× on the random-walk sequence with zero pin drift, else ship default-off and let the numbers drive 3b.
• autolens_profiling leg (diagnostic script + results + notes) ships after harvest-0827-gate-b-pt2 (autolens_profiling#183) merges — it holds the repo claim; files are disjoint.

Detailed implementation plan

Work Classification

Library (PyAutoArray) first; workspace leg (autolens_profiling) follows.

Affected Repositories

• PyAutoArray (primary)
• autolens_profiling (measurement/results leg, queued behind autolens_profiling#183)

Branch Survey

Repository	Current Branch	Dirty?
./PyAutoArray	main	clean
./autolens_profiling	main	clean — CLAIMED by harvest-0827-gate-b-pt2 (awaiting-merge, touches only scripts/misc/searches + inference results)

Suggested branch: feature/numba-cpu-nnls-iteration-reduction
Worktree root: ~/Code/PyAutoLabs-wt/numba-cpu-nnls-iteration-reduction/ (created by /start_library)

Implementation Steps

PyAutoArray — autoarray/util/fnnls.py (fnnls_cholesky, lines 25-167; fix_constraint_cholesky 170-207)

1. Add stats: Optional[dict] = None; when given, fill outer_iterations, inner_iterations, passive_set (int index array, final P_inorder), n_passive, warm_start_errors (= |P_initial ⊕ P_final|). No change to the returned d.
2. Accept P_initial as a bool mask or an index array (production passes a mask, test_cholesky_inplace.py:155 passes indices) — normalise once at line ~47.
3. Warm-start factorisation reuse: today the P_initial branch (lines 83-89) does a dense scipy solve for d and then the first outer iteration still rebuilds U_buffer with a full slg.cholesky (106-110). Seed U_buffer[:k,:k] from one slg.cholesky(ZTZ[P_inorder][:,P_inorder]) and derive d from it via _cho_solve_buffer — one factorisation instead of a solve + a factorisation. Bit-level: solution may differ at ulp level vs the scipy solve; pins at rtol 1e-6 are the guard.

PyAutoArray — autoarray/settings.py (fields lines 10-21; assignments 152-154)
4. Add nnls_warm_start_memo: Optional[bool] = None (resolve None from conf.instance["general"]["inversion"]["nnls_warm_start_memo"], default false in autoarray/config/general.yaml until measured). Docstring: numpy/numba path only.

PyAutoArray — new autoarray/inversion/inversion/nnls_memo.py
5. _nnls_passive_set_memo: Dict[str, np.ndarray], cap 8, FIFO; memo_key(n, fingerprint); env AUTOARRAY_NNLS_WARM_START=0 disables; stored arrays copied + setflags(write=False). Mirror the contract of imaging_numba/sparse.py:20-52 (misses, never stale hits).

PyAutoArray — autoarray/inversion/inversion/inversion_util.py:366-386 (numpy branch of reconstruction_positive_only_from)
6. If settings.nnls_warm_start_memo and memo hit for (n, fingerprint): P_initial = memo passive set (validate all indices < n, else miss). Else today's dense-sign P_initial. Pass stats dict; on success store stats["passive_set"]. Fingerprint: the caller (abstract.py:554-601) passes fingerprint= derived from the mapper's source_plane_mesh_grid.shape + data_vector.shape (+ ids_to_keep when the edge-zeroed subset branch is taken — the passive set must be in the SUBSET index space in that branch).
7. Keep the existing except (RuntimeError, LinAlgError, ValueError) → InversionException contract; a memo-seeded solve that raises must retry once from the dense-sign start before raising (a bad warm start must not turn into a resample).

Tests
8. test_autoarray/util/test_cholesky_inplace.py: stats populated and consistent (n_passive == len(passive_set), errors == 0 when warm-started from the true solution's support); warm start from a deliberately WRONG passive set converges to the cold-start solution (rel 1e-8); mask vs index P_initial equivalence; factorisation-seeded warm start matches previous behaviour to 1e-10.
9. test_autoarray/util/test_jax_nnls.py:39 (numpy_path_ignores_knobs): keep asserting nnls_solver_tol/nnls_max_iter are ignored; add that nnls_warm_start_memo is honoured (memo populated after one solve, second solve reports warm_start_errors == 0), and that AUTOARRAY_NNLS_WARM_START=0 disables.
10. test_autoarray/inversion/inversion/: memo-seeded inversion reconstruction equals un-memoed reconstruction (1e-10) on a small mapper fixture; subset (solve_ids_to_keep) branch covered.

Diagnostic + measurement (autolens_profiling, uncommitted until #183 merges)
11. Script scripts/imaging/likelihood_breakdown/delaunay_numba_nnls_iterations.py built from delaunay_numba.py's setup (lines 75-183; instance at 167 via model.instance_from_vector; rebuild adapt_images per instance — it is keyed on the instance's source galaxy). Sequences: (a) random walk, 30 steps, σ = 1% of prior width per parameter; (b) 30 i.i.d. draws from the priors' central 20%; run each with memo off / memo on; record per solve: outer/inner iterations, warm-start errors (dense-sign vs previous), solve seconds, log-likelihood. euclid + hst, AUTOARRAY_NUMBA_OPERATED_MEMO=0, OMP=1.
12. Decision gate on the numbers (≥2× median iteration reduction on (a) with zero pin drift → default true); post the table on this issue; then /start_workspace for the autolens_profiling leg once #183 lands (script + results JSON + note in results/notes/).

Key Files

• PyAutoArray/autoarray/util/fnnls.py — fnnls_cholesky, warm-start seeding, stats
• PyAutoArray/autoarray/util/cholesky_funcs.py — buffer kernels (cholinsertlast_inplace 192, choldeleteindexes_inplace 243, _cho_solve_buffer 178)
• PyAutoArray/autoarray/inversion/inversion/inversion_util.py:256-386 — reconstruction_positive_only_from
• PyAutoArray/autoarray/inversion/inversion/abstract.py:537-607 — reconstruction (subset branch 564-593)
• PyAutoArray/autoarray/inversion/inversion/imaging_numba/sparse.py:20-52 — memo template
• PyAutoArray/autoarray/settings.py — Settings
• autolens_profiling/scripts/imaging/likelihood_breakdown/delaunay_numba.py — fiducial setup; solve accessor line 231
• autolens_profiling/results/notes/nnls_solver_ledger.md — prior BPP/warm-start findings (JAX PDIP path — do not transfer blindly)

Epic

numba-cpu-likelihood: profiling ✔, first-call bug ✔, phase 1 ✔ (#497/#588), phase 2b ✔ (#453/#463), phase 3a = this, phase 3b = batched active-set moves (to file after 3a's diagnostic), phase 2a kernel-CDF deferred. Measurement prerequisite: draft/feature/autolens_profiling/numba_breakdown_harness_memo_blind.md.

Original Prompt

Click to expand starting prompt

Numba CPU likelihood phase 3: cut the positive-only solve's active-set iterations (warm start across evaluations + block pivoting)

Type: feature
Epic: numba-cpu-likelihood
Phase: 3
Target: autoarray
Repos:

• @PyAutoArray
• @autolens_profiling
  Difficulty: large
  Autonomy: supervised
  Priority: high
  Status: formalised
  Filed: 2026-08-27

Phase 3 of the CPU-likelihood speed restoration. Phase 1 (MGE batching +
caching, PyAutoArray#497 / PyAutoGalaxy#588) and phase 2b (fnnls in-place
Cholesky buffer, PyAutoArray#453 / #463, backfilled record
complete/2026/08/numba-fnnls-inplace-cholesky-buffer.md) are shipped.
The "restore the deleted numba fnnls" idea is RETIRED (autolens_profiling#151
comment 5) — the solver was never missing. Do not re-file it.

Context (autolens_profiling#151 comments 5-6, PyAutoArray PR#453 text)

On the campaign fiducial (Delaunay + Hilbert AdaptImage + ConstantSplit,
apply_sparse_operator_cpu(), use_jax=False) the positive-only
reconstruction solve is the dominant term: euclid 3.6 s of 5.0 s (~72%) at
1250 source pixels, hst 1.4 s (~35%) — pre-#453 numbers; #453 reports the
solve at 1.40 s -> 1.16 s on its 1310-param system, so a re-profile on current
main is step 0.

The instrumented probe (comment 5) decomposed a 5.06 s solve at n=1560 as: dense
warm-start solve 0.10 s (1411/1560 positives correct) + initial Cholesky 0.06 s +
up/down-dates 3.58 s across 154 outer + 102 constraint-fix iterations +
cho_solve 0.99 s + w matvecs 0.26 s. One from-scratch Cholesky at n=1560 is
0.08 s. Comment 6: "the solve is iteration-bound, not resolution- or
param-bound" — cost tracks how many entries the dense warm start gets wrong
(~150 at euclid, far fewer at hst). #453 removed the copy overhead per
iteration; the iteration COUNT is untouched.

Live solver: autoarray/util/fnnls.py::fnnls_cholesky(ZTZ, ZTx, P_initial)
(Bro & De Jong 1997 active set, Cholesky up/down-dating via
autoarray/util/cholesky_funcs.py numba kernels), called from
inversion_util.reconstruction_positive_only_from (numpy branch, ~line 371)
with P_initial = np.linalg.solve(curvature_reg_matrix, data_vector) > 0.
Tests: test_autoarray/util/test_cholesky_inplace.py,
test_cholesky_degenerate.py.

Goal

Reduce the number of active-set iterations per solve on the numba CPU path,
keeping the solution the unique NNLS optimum (curvature_reg is PD, so the
solution is unique; pinned log-likelihoods must hold at rtol 1e-6):

1. Step 0 — re-profile on current main with the delaunay_numba
   breakdown (euclid + hst, 1250) so the baseline post-perf: in-place Cholesky buffer + copy-free numba solves for fnnls_cholesky #453 is recorded; add an
   iteration counter / per-solve diagnostic (outer + inner iterations, passive
   set size, entries the warm start got wrong) to the breakdown so the win is
   measured in iterations, not just seconds.
2. Cross-evaluation warm start: seed P_initial from the previous
   evaluation's final passive set (same process; nearby parameter points share
   most of the active set) instead of the dense unconstrained solve's sign, with
   the dense solve as fallback when shapes change. Design decision in the plan:
   where the state lives (module-level like the operated-matrix memo, keyed on
   mapper shape; or a Settings field / preload passed through the analysis)
   and how multiprocessing workers behave (each worker keeps its own).
3. Block pivoting (Portugal-Judice-Vicente style) or a batched
   constraint-fix step so each outer iteration moves many indices at once
   rather than one, bounded by a fallback to the current single-index rule to
   preserve convergence guarantees.
4. Measure: iterations and solve seconds before/after at euclid + hst; pinned
   log-likelihood parity; test_autoarray green. Record in autolens_profiling
   (results + notes) and on the issue. Measure with
   draft/feature/autolens_profiling/numba_breakdown_harness_memo_blind.md
   landed or with AUTOARRAY_NUMBA_OPERATED_MEMO=0 (the harness's fixed instance
   would otherwise hide the MGE term and, for a cross-eval warm start, would
   fake a 100%-correct warm start — perturb the instance between repeats).

Out of scope: the JAX PDIP solver (jax_nnls.py), the kernel-CDF phase 2a.

[Jammy2211]

Paused 2026-08-27 (end of day) before implementation. No branch, worktree or source edits exist. Resume: /start_library numba-cpu-nnls-iteration-reduction PyAutoArray (PyAutoArray only — autolens_profiling is claimed by #183 until it merges; that leg follows via /start_workspace). Plan above stands; measure with AUTOARRAY_NUMBA_OPERATED_MEMO=0.

[Jammy2211]

Resumed 2026-08-28. Worktree numba-cpu-nnls-iteration-reduction created with PyAutoArray + autolens_profiling (autolens_profiling attached since #183 merged).

Library leg implemented (uncommitted, test_autoarray: 1270 passed):

• fnnls_cholesky(..., stats=None) — diagnostics (outer_iterations, inner_iterations, passive_set, n_passive, warm_start_errors); accepts mask or index P_initial; the warm-start factor is now computed once and extended in place (previously a dense solve + a from-scratch Cholesky including the first inserted column); a warm passive set with non-positive entries is repaired before the outer loop (the old clip-only path returned a wrong vector when P was all-True — unreachable from the dense-sign start, reachable from a memo seed). fix_constraint_cholesky cannot do that repair (its alpha is 0/0 with no prior feasible iterate), so it is an explicit drop-and-refactor.
• Settings.nnls_warm_start_memo (config general.yaml inversion.nnls_warm_start_memo: false until measured; KeyError → False for shadowing workspace configs).
• New autoarray/inversion/inversion/nnls_memo.py: process-local FIFO memo (8 entries), keyed on n + shape fingerprint (+ ids_to_keep hash in the edge-zeroed subset branch), AUTOARRAY_NNLS_WARM_START=0 kill-switch. A memo-seeded solve that raises retries once from the dense-sign start before the InversionException path.
• Synthetic check (n=600, 0.2% perturbation): dense-sign start 271–286 wrong entries / 1 outer + 4–5 inner iterations / ~20 ms → memo start 0 wrong / 0 iterations / ~1 ms, solutions equal to 1e-14.

Now running: the plan's diagnostic on the euclid + hst Delaunay-1250 fiducial (random-walk + i.i.d. sequences, memo off/on, AUTOARRAY_NUMBA_OPERATED_MEMO=0) plus the step-0 baseline re-profile. The ≥2× median-iteration gate on the random-walk sequence decides the default.

[Jammy2211]

Measured — gate met, default flipped to true; ship blocked on Heart RED.

Diagnostic (autolens_profiling/scripts/imaging/likelihood_breakdown/delaunay_numba_nnls_iterations.py, Delaunay-1250 + MGE-60, AUTOARRAY_NUMBA_OPERATED_MEMO=0, OMP_NUM_THREADS=1, 30-instance sequences, memo off → on, medians):

instrument	sequence	total iterations	warm-start errors	solve s	eval s
euclid	random walk (σ = 1% prior width)	69.5 → 7.0 (9.9×)	73.5 → 9.5	0.544 → 0.051	1.510 → 0.564
euclid	i.i.d. (central 20% of priors)	86 → 94	94.5 → 98	0.347 → 0.309	0.906 → 0.838
hst	random walk	32 → 8 (4.0×)	31 → 9.5	0.245 → 0.067	1.733 → 1.594
hst	i.i.d.	42 → 69.5	48 → 68	0.268 → 0.260	1.687 → 1.867

• Parity: max relative Δ log-likelihood 3e-14, max |Δreconstruction| 1.6e-8; 0 memo-seed failures / retries in 240 evaluations; pinned evidence values PASS.
• Solve seconds are ≤ off in every cell (the memo also skips the dense unconstrained solve), so uncorrelated jumps cost iterations, not time. Default shipped as true — in general.yaml and in the KeyError fallback, because workspace general.yamls shadow autoarray's and lack the key (the test suite's own config does too, so the tests exercise that fallback).
• Step-0 baseline after perf: in-place Cholesky buffer + copy-free numba solves for fnnls_cholesky #453/perf: reuse Convolver state, cache operated-matrix dict, hoist linear-func pair loop (#496 phase 1) #497: the solve is now 0.47 s of 1.19 s (40%) at euclid and 0.54 s of 3.22 s (17%) at hst — not the ~72% this issue was filed on; at hst the curvature matrix F (1.77 s) now dominates.
• Phase 3b: the warm path is down to 7–8 iterations, so batched active-set moves have little left there; its case is the cold / uncorrelated path (30–95 outer iterations: first evaluation and i.i.d.-like proposals).

Committed locally on feature/numba-cpu-nnls-iteration-reduction (PyAutoArray 105f6ea4, test_autoarray 1271 passed; autolens_profiling 26cbfef). Not pushed / no PR yet: pyauto-heart readiness is RED — release validation FAILED (stage integrate) — the unrepinned autolens_workspace_test rectangular_mge.py / rectangular_mge_rtu.py pair, unrelated to this change. Awaiting a human ack of that RED reason (or Heart going GREEN) before /ship_library + /ship_workspace.

[Jammy2211]

Robustness across lens models + a fallback guard (both committed locally: PyAutoArray cfd7f802, autolens_profiling 2ebf80b; still unpushed behind Heart RED).

32-cell matrix (results/notes/nnls_warm_start_memo_matrix.md; 11 variants: fiducial, PowerLaw, NFW subhalo, no lens light, Sersic light, rectangular mesh + Constant reg, AdaptSplit reg, no edge-zeroing, complex source, Hilbert-600, Hilbert-2000; euclid all, hst 6; random-walk + i.i.d. sequences, memo off/on): parity ≤ 4.3e-14 in every cell, 0 seeded-solve failures, median solve time never worse than memo-off by >2%. Every random-walk cell gains (3.2×–118×; rectangular meshes most because the dense-sign start is 44–70% wrong there). All 11 losing cells are i.i.d. (worst euclid/no_lens_light/iid 0.52×) and lose iterations only, not time.

Calibration finding: the seed's absolute error fraction does not separate helpful from harmful seeds (helpful up to 0.138, harmful from 0.048), but the seed/dense-sign ratio does (helpful ≤ 0.89, worst 1.42). Hence the guard is relative and self-calibrating: Settings.nnls_warm_start_error_tolerance (config default 1.5) — each memo entry keeps the error fraction of the last dense-sign-started solve for its key; a seed worse than 1.5 × that is dropped and the next solve restarts dense, refreshing the reference. Non-finite / ≤ 0 disables. Re-run with the guard on: fiducial euclid 10.0×/0.91×, hst 4.27×/0.62×, complex source 3.39×/0.73× (rw/iid), 3 fallbacks in 240 evaluations, all in hst/fiducial/iid (the 1.42 cell). test_autoarray: 1279 passed.

[Jammy2211]

Human decision 2026-08-28: Heart RED release validation FAILED (stage integrate) (unrelated autolens_workspace_test MGE pin drift, being fixed separately) overridden for this task; shipping PyAutoArray + autolens_profiling PRs, then /prm (CI green → merge library-first → close-out).

[Jammy2211]

Shipped

• Library: perf: warm-start the positive-only NNLS solve from the previous evaluation's passive set (#498) #501 (MERGED 1f5c636) — fnnls_cholesky diagnostics + single-factorisation warm start + infeasible-seed repair; cross-evaluation passive-set memo (nnls_memo.py, Settings.nnls_warm_start_memo default true, AUTOARRAY_NNLS_WARM_START=0 kill-switch); relative fallback guard Settings.nnls_warm_start_error_tolerance = 1.5.
• Measurement: results: NNLS warm-start memo diagnostic, 32-cell lens-model robustness matrix, post-#453/#497 baseline (PyAutoArray#498) autolens_profiling#184 (MERGED 9a04e5ff) — diagnostic script, 32-cell lens-model robustness matrix, guard re-check, post-perf: in-place Cholesky buffer + copy-free numba solves for fnnls_cholesky #453/perf: reuse Convolver state, cache operated-matrix dict, hoist linear-func pair loop (#496 phase 1) #497 baseline; notes results/notes/nnls_warm_start_memo{,_matrix}.md.
• Random walk: median active-set iterations 70→7 (euclid, 9.9×), 32→8 (hst, 4.0×); solve 0.54→0.05 s / 0.25→0.07 s; parity ≤ 4e-14 in all 32 cells; 0 seeded failures; solve time never worse than the dense-sign start by >2%.
• Gate notes: shipped over Heart RED release validation FAILED (stage integrate) and Feature/jax simplify visualization #184's lint red (test_hazards_prior_exit.py, pre-existing on main at 451b85f) on explicit human authorisation; the hazards anchor drift is being fixed separately.
• Phase 3b: the warm path is at 7–8 iterations; only the cold / i.i.d. path (30–95 outer) motivates batched moves. At hst the curvature matrix F now dominates the evaluation.