Repository navigation
perf: cut positive-only solve active-set iterations — warm-start memo + diagnostics (numba CPU likelihood, phase 3a) #498
Description
Activity
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 withAUTOARRAY_NUMBA_OPERATED_MEMO=0.Resumed 2026-08-28. Worktree
numba-cpu-nnls-iteration-reductioncreated 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 indexP_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 whenPwas all-True — unreachable from the dense-sign start, reachable from a memo seed).fix_constraint_choleskycannot do that repair (itsalphais 0/0 with no prior feasible iterate), so it is an explicit drop-and-refactor.Settings.nnls_warm_start_memo(configgeneral.yamlinversion.nnls_warm_start_memo: falseuntil measured;KeyError→Falsefor shadowing workspace configs).- New
autoarray/inversion/inversion/nnls_memo.py: process-local FIFO memo (8 entries), keyed onn+ shape fingerprint (+ids_to_keephash in the edge-zeroed subset branch),AUTOARRAY_NNLS_WARM_START=0kill-switch. A memo-seeded solve that raises retries once from the dense-sign start before theInversionExceptionpath. - 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.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— ingeneral.yamland in theKeyErrorfallback, because workspacegeneral.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(PyAutoArray105f6ea4,test_autoarray1271 passed; autolens_profiling26cbfef). Not pushed / no PR yet:pyauto-heart readinessis RED —release validation FAILED (stage integrate)— the unrepinnedautolens_workspace_testrectangular_mge.py/rectangular_mge_rtu.pypair, unrelated to this change. Awaiting a human ack of that RED reason (or Heart going GREEN) before/ship_library+/ship_workspace.Robustness across lens models + a fallback guard (both committed locally: PyAutoArray
cfd7f802, autolens_profiling2ebf80b; 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. (worsteuclid/no_lens_light/iid0.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 than1.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 inhst/fiducial/iid(the 1.42 cell).test_autoarray: 1279 passed.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).- added a commit that references this issue
on Aug 28, 2026 Shipped
- Library: perf: warm-start the positive-only NNLS solve from the previous evaluation's passive set (#498) #501 (MERGED 1f5c636) —
fnnls_choleskydiagnostics + single-factorisation warm start + infeasible-seed repair; cross-evaluation passive-set memo (nnls_memo.py,Settings.nnls_warm_start_memodefault true,AUTOARRAY_NNLS_WARM_START=0kill-switch); relative fallback guardSettings.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 onmainat 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.
- Library: perf: warm-start the positive-only NNLS solve from the previous evaluation's passive set (#498) #501 (MERGED 1f5c636) —
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, andfnnls_choleskyadds 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 (recordcomplete/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
fnnls_cholesky(final passive set, outer/inner iteration counts, warm-start error count) via an optionalstatsdict — default return unchanged.AUTOARRAY_NUMBA_OPERATED_MEMO=0.imaging_numba/sparse.py:20-52), keyed on problem size + mesh fingerprint,Settingsknob +AUTOARRAY_NNLS_WARM_START=0kill-switch, dense-sign fallback on miss. Fork-based pools give each worker its own memo automatically.slg.cholesky; skip the separate densenp.linalg.solvewhen a warm passive set is available.autolens_profilingleg (diagnostic script + results + notes) ships afterharvest-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
Branch Survey
Suggested branch:
feature/numba-cpu-nnls-iteration-reductionWorktree 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_cholesky170-207)stats: Optional[dict] = None; when given, fillouter_iterations,inner_iterations,passive_set(int index array, finalP_inorder),n_passive,warm_start_errors(= |P_initial ⊕ P_final|). No change to the returnedd.P_initialas a bool mask or an index array (production passes a mask,test_cholesky_inplace.py:155passes indices) — normalise once at line ~47.P_initialbranch (lines 83-89) does a dense scipy solve fordand then the first outer iteration still rebuildsU_bufferwith a fullslg.cholesky(106-110). SeedU_buffer[:k,:k]from oneslg.cholesky(ZTZ[P_inorder][:,P_inorder])and derivedfrom it via_cho_solve_buffer— one factorisation instead of a solve + a factorisation. Bit-level: solution may differ at ulp level vs the scipysolve; 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(resolveNonefromconf.instance["general"]["inversion"]["nnls_warm_start_memo"], defaultfalseinautoarray/config/general.yamluntil measured). Docstring: numpy/numba path only.PyAutoArray — new
autoarray/inversion/inversion/nnls_memo.py5.
_nnls_passive_set_memo: Dict[str, np.ndarray], cap 8, FIFO;memo_key(n, fingerprint); envAUTOARRAY_NNLS_WARM_START=0disables; stored arrays copied +setflags(write=False). Mirror the contract ofimaging_numba/sparse.py:20-52(misses, never stale hits).PyAutoArray —
autoarray/inversion/inversion/inversion_util.py:366-386(numpy branch ofreconstruction_positive_only_from)6. If
settings.nnls_warm_start_memoand memo hit for(n, fingerprint):P_initial = memo passive set(validate all indices < n, else miss). Else today's dense-signP_initial. Passstatsdict; on success storestats["passive_set"]. Fingerprint: the caller (abstract.py:554-601) passesfingerprint=derived from the mapper'ssource_plane_mesh_grid.shape+data_vector.shape(+ids_to_keepwhen 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)→InversionExceptioncontract; 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 indexP_initialequivalence; factorisation-seeded warm start matches previous behaviour to 1e-10.9.
test_autoarray/util/test_jax_nnls.py:39(numpy_path_ignores_knobs): keep assertingnnls_solver_tol/nnls_max_iterare ignored; add thatnnls_warm_start_memois honoured (memo populated after one solve, second solve reportswarm_start_errors == 0), and thatAUTOARRAY_NNLS_WARM_START=0disables.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.pybuilt fromdelaunay_numba.py's setup (lines 75-183; instance at 167 viamodel.instance_from_vector; rebuildadapt_imagesper 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_workspacefor the autolens_profiling leg once #183 lands (script + results JSON + note inresults/notes/).Key Files
PyAutoArray/autoarray/util/fnnls.py—fnnls_cholesky, warm-start seeding, statsPyAutoArray/autoarray/util/cholesky_funcs.py— buffer kernels (cholinsertlast_inplace192,choldeleteindexes_inplace243,_cho_solve_buffer178)PyAutoArray/autoarray/inversion/inversion/inversion_util.py:256-386—reconstruction_positive_only_fromPyAutoArray/autoarray/inversion/inversion/abstract.py:537-607—reconstruction(subset branch 564-593)PyAutoArray/autoarray/inversion/inversion/imaging_numba/sparse.py:20-52— memo templatePyAutoArray/autoarray/settings.py—Settingsautolens_profiling/scripts/imaging/likelihood_breakdown/delaunay_numba.py— fiducial setup; solve accessor line 231autolens_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:
Difficulty: large
Autonomy: supervised
Priority: high
Status: formalised
Filed: 2026-08-27
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-onlyreconstruction 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
mainis 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_solve0.99 s +wmatvecs 0.26 s. One from-scratch Cholesky at n=1560 is0.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.pynumba kernels), called frominversion_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):
delaunay_numbabreakdown (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.
P_initialfrom the previousevaluation'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
Settingsfield / preload passed through the analysis)and how multiprocessing workers behave (each worker keeps its own).
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.
log-likelihood parity;
test_autoarraygreen. Record in autolens_profiling(results + notes) and on the issue. Measure with
draft/feature/autolens_profiling/numba_breakdown_harness_memo_blind.mdlanded or with
AUTOARRAY_NUMBA_OPERATED_MEMO=0(the harness's fixed instancewould 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.