Repository navigation
perf: cut positive-only solve active-set iterations — warm-start memo + diagnostics (numba CPU likelihood, phase 3a) #498
Description
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 optionalstatsdict — 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,Settingsknob +AUTOARRAY_NNLS_WARM_START=0kill-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 densenp.linalg.solvewhen 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_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
- 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)
- Add
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. - Accept
P_initialas a bool mask or an index array (production passes a mask,test_cholesky_inplace.py:155passes indices) — normalise once at line ~47. - Warm-start factorisation reuse: today the
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 (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, 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:
- @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):
- 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. - Cross-evaluation warm start: seed
P_initialfrom 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 aSettingsfield / preload passed through the analysis)
and how multiprocessing workers behave (each worker keeps its own). - 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. - Measure: iterations and solve seconds before/after at euclid + hst; pinned
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.md
landed or withAUTOARRAY_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.
Activity
Jammy2211 commentedon Aug 28, 2026
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.Jammy2211 commentedon Aug 28, 2026
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).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.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 commentedon Aug 28, 2026
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):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).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.Jammy2211 commentedon Aug 28, 2026
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.Jammy2211 commentedon Aug 28, 2026
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).Merge pull request #501 from PyAutoLabs/feature/numba-cpu-nnls-iterat…
Jammy2211 commentedon Aug 28, 2026
Shipped
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.results/notes/nnls_warm_start_memo{,_matrix}.md.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.