perf: NNLS solver controls for rectangular + Delaunay JAX pipelines

[Jammy2211]

Overview

Optimize the non-negative least squares (NNLS) solve used by the pixelization inversion pipelines. The JAX path solves via jaxnnls.solve_nnls_primal (PDIP); the 2026-07-09 Delaunay review measured that the solver is a modest but real lever (tolerance loosening ≈15–20%, hard-coded MAX_ITER=50 under vmap pays worst-case-lane, fresh ~n² Cholesky every PDIP iteration) while refuting warm-starts (cond(Q)~1e16). Priority is the rectangular pipeline — upcoming work concentrates there and it has no host callback, so the linear-algebra solve dominates its 52.7ms/32.9ms (A100 single/vmap) per-eval cost. All profiling uses the production pixelization + MGE lens light configuration.

Two deliverables beyond the speedup itself:

1. A go/no-go checkpoint with measured NNLS shares per pipeline before any source edit.
2. An evidence-based answer to: what is the right next optimization target for Delaunay post-refactor: qhull-only Delaunay callback, exact JAX visibility-walk point location #368 — residual scipy/qhull callback (still host-side and vmap_method="sequential"), the NNLS solve, or something else in the per-eval budget.

Plan

• Re-baseline Delaunay + rectangular per-eval costs post PyAutoArray#368 (qhull-only callback), pixelization+MGE configs; local CPU now, A100 when RAL returns.
• Measure the NNLS PDIP share and iteration counts per pipeline; instrument per-iteration Cholesky cost and vmap worst-lane amplification.
• Checkpoint: report go/no-go + chosen lever set on this issue, including the Delaunay next-target verdict.
• If go: expose solver controls (solver_tol, max_iter) — vendor the PDIP loop into autoarray or land params upstream in jaxnnls (decided at checkpoint with numbers).
• New config keys default to current behavior; mirror into workspace configs.
• Stability gates: Δlog_evidence parity on production rect + Delaunay fits, JIT-vs-eager drift, vmap batch checks, noise-scale invariance; before/after benchmarks.
• Design note written for later PyAutoBrain ingestion; ship via ship_library.

Detailed implementation plan

Affected Repositories

• PyAutoArray (primary — all source edits)
• PyAutoLens (context/validation only; no edits expected)
• autolens_profiling (untracked scratch/ only — repo is claimed by profiling-preopt-campaign, so no branch there)

Branch Survey

Repository	Current Branch	Dirty?
./PyAutoArray	main	clean
./PyAutoLens	main	clean
./autolens_profiling	main	untracked results only (claimed by profiling-preopt-campaign)

Suggested branch: feature/nnls-solver-optimization
Work Classification: Library
Worktree root: ~/Code/PyAutoLabs-wt/nnls-solver-optimization/ (created by /start_library)

Implementation Steps

Phase A — assessment (standalone, no source edits)

1. Build autolens_profiling/scratch/nnls_speedup/ PoC scripts (extend the prior scratch/delaunay_speedup/demo_2_nnls_speedup.py), HST pixelization+MGE configs, rectangular first, Delaunay second.
2. Re-baseline both pipelines against the merged refactor: qhull-only Delaunay callback, exact JAX visibility-walk point location #368 code path (qhull-only callback + visibility walk) so Delaunay numbers reflect current main.
3. Instrument jaxnnls.pdip: iteration counts, per-iteration cho_factor cost (n≈1560), share of total per-eval; measure lax.while_loop-under-vmap worst-lane amplification on realistic batches.
4. Decompose the Delaunay per-eval budget post-refactor: qhull-only Delaunay callback, exact JAX visibility-walk point location #368: residual host callback (qhull, ~8.8ms local, still vmap_method="sequential") vs NNLS vs mapping/curvature ops → the "right next Delaunay target" verdict.
5. Checkpoint on this issue: measured shares, go/no-go, vendor-vs-upstream decision.

Phase B — implementation (PyAutoArray, only if go)
6. Expose solver_tol / max_iter controls. Preferred route decided at checkpoint:

• vendor: autoarray/util/jax_nnls.py PDIP loop with configurable tol/max_iter (and any measured Cholesky-reuse win), dispatch in autoarray/inversion/inversion/inversion_util.py (~lines 275–321);
• upstream: params added to jaxnnls.solve_nnls_primal, autoarray passes them through.

7. Config keys in autoarray/config/general.yaml → inversion (alongside nnls_jacobi_preconditioning, nnls_target_kappa), defaults = current behavior; mirror into workspace configs.
8. NumPy-only unit tests for config plumbing/dispatch (no JAX in library unit tests).

Phase C — validation + ship
9. JAX validation via autolens_workspace_test parity scripts + fitness._vmap; Δlog_evidence parity on production fits (caustic-geometry lesson from #367: production grids, not toy grids).
10. Before/after benchmark table; design note for PyAutoBrain ingestion; ship_library + smoke tests.

Key Files

• autoarray/inversion/inversion/inversion_util.py — NNLS dispatch (JAX → jaxnnls PDIP, NumPy → fnnls_cholesky)
• autoarray/util/fnnls.py — NumPy Bro & Jong solver (not the lever; unchanged unless profiling says otherwise)
• autoarray/config/general.yaml — inversion solver config keys
• site-packages jaxnnls/pdip.py — hard-coded MAX_ITER=50, solver_tol = n*EPSILON, per-iter cho_factor (third-party, already PyAuto-modified)
• autolens_profiling/scratch/delaunay_speedup/ — prior PoC + design notes this work builds on

Prior evidence (2026-07-09 Delaunay review)

• Warm-start REFUTED: exact unconstrained init worsens 17→38 iters; jaxnnls (Q+I) init already optimal.
• tol 1e-6 saves 1–2 of 17 iters (~15–20%), objective drift 5e-16 — config-key material.
• A100 baselines (pre-refactor: qhull-only Delaunay callback, exact JAX visibility-walk point location #368): MGE 5.9ms/0.4–0.8ms; Rect 52.7ms/32.9ms; Delaunay 80ms/106.8ms (single/vmap). Delaunay callback shrink merged as PyAutoArray#368; A100 re-profile outstanding.

Original Prompt

Click to expand starting prompt

Optimize the fast non-negative least squares solver for the Delaunay and rectangular mesh pipelines

Type: feature
Target: autolens_profiling
Repos:

• PyAutoArray
• PyAutoLens
  Difficulty: medium
  Autonomy: supervised
  Priority: normal
  Status: formalised

You are working in the PyAutoLens codebase. First, review the recent Delaunay mesh speed-up work and understand what was changed and why. Then assess whether optimizing the fast non-negative least squares solver is the right next target, considering both Delaunay and rectangular mesh pipelines. If it is, propose and implement a careful optimization with benchmarks and numerical stability checks. Keep the design clean and document it so that it can be ingested into PyAutoBrain later. Before coding, summarize your understanding, confirm the optimization target, and outline a step-by-step plan.

(formalised by the Intake (Conception) Agent on 2026-07-09; user refinements at issue time: rectangular pipeline is the priority; profile pixelization+MGE configs only; re-baseline Delaunay post-#368; work standalone; report the correct next Delaunay target before implementation.)

[Jammy2211]

Phase A checkpoint — measured shares, go/no-go, and the Delaunay verdict

All numbers are post-#368 main (89a5fbd), local CPU fp64, HST production configs (MGE-60 lens light; rect 39×39 RectangularAdaptImage+Constant, n=1581; Delaunay 1500-vertex Hilbert+ConstantSplit, n=1560). Probes ran on the real extracted inversion matrices, not synthetic systems. Working dir: autolens_profiling/scratch/nnls_speedup/ (standalone, untracked).

Per-eval breakdown (likelihood_breakdown, CPU)

Step	Rect	Delaunay
Curvature matrix (F)	3.01s (50.6%)	2.99s (51.2%)
Inversion setup (relocate→map→PSF-convolve)	1.58s (26.5%)	1.76s (30.2%)
NNLS reconstruction	1.21s (20.3%)	0.97s (16.6%)
Everything else	~2.5%	~2%

F's dominance is CPU-only: it is one 15361×n GEMM (~4ms on A100 fp64). The NNLS solve is sequential — it cannot be bought back with GPU parallelism — so its A100 share of the rect eval will be much larger than 20% (plausibly the dominant block of the 52.7ms/32.9ms rect rows; RAL re-profile will confirm).

PDIP probe on the real matrices (probe_pdip.py)

cond(Q_raw) ≈ 4.1e10 both meshes (Jacobi rescale does not reduce it: 1.7e11). Per iteration = one fresh n² Cholesky (factorize_kkt) + 2 triangular-solve pairs; iteration cost fully explains solve time (rect: 21 × 64.5ms; Delaunay: 19 × 59.8ms).

Variant	Rect iters	Rect speedup	Delaunay iters	Delaunay speedup	rel Δobjective
baseline (tol=n·eps·5e3 ≈ 1.7e-9)	21	1.00×	19	1.00×	—
tol 1e-8	20	0.95×*	18	1.04×	≤ 6e-15
tol 1e-6	18	1.06×	15	1.18×	≤ 3.8e-13
tol 1e-4	15	1.28×	13	1.32×	≤ 2.1e-10

* timing noise; iteration counts are the signal.

Because the log-det terms of the evidence are reconstruction-independent, Δlog_evidence from a looser solve equals the Δ of the NNLS quadratic objective — i.e. ≤ ~1e-8 at tol 1e-6 and ~1e-6 at tol 1e-4 for these systems. Both are far below sampler sensitivity; the Phase C parity gate will verify end-to-end.

vmap (8 lanes, real noise-covariance perturbations of q): confirmed max-lane effect — rect lanes spread 18–22 iterations and the whole batch pays 22. Tolerance loosening shifts the whole spread down (~uniform −3), so the saving survives batching. Per-lane early exit is not structurally available under vmap'd while_loop.

Note: the rect fit's saved NumPy-path reconstruction (fnnls_cholesky) differs from PDIP's by up to 0.27 (rel max) with matching objectives — active-set degeneracy of the near-empty rect outskirts, not a solver bug; Delaunay agrees to 1.7e-5.

Verdict

GO — scoped. Phase B = expose nnls_solver_tol and nnls_max_iter as config keys (defaults exactly reproduce current behavior: tol=n·eps·5e3, max_iter=50), with the measured recommendation that searches run at 1e-6 (rect −14%, Delaunay −21% of solve; ~5–15% of a whole A100 rect eval if NNLS's GPU share is as expected). jaxnnls does not expose these parameters, so the loop is vendored into autoarray/util/ (reusing jaxnnls.pdip.pdip_pc_step/initialize as imports, only the while_loop driver re-implemented — smallest possible fork surface) unless upstream takes a params PR quickly.

Parked (documented, pending A100 re-profile): cheaper PDIP iterations (PCG-on-fixed-factorization; the varying z/s diagonal degrades any fixed preconditioner — research-grade risk), alternative first-order solvers (cond 1e10 makes them iteration-hungry), fp32/mixed-precision PDIP (cond 1e10 vs fp32 eps).

The Delaunay question — is the mesh/scipy machinery still the target?

No. Post-#368 the qhull callback is ~9ms host-side (was ~53ms) and the Delaunay per-eval budget is now dominated by the same blocks as rect: F on CPU, NNLS (+ inversion setup) on GPU. Two Delaunay-specific residuals remain:

1. the callback still runs vmap_method="sequential", so under batched sampling it serializes per lane — this, not single-eval cost, is where Delaunay still pays vs rect (expected A100 vmap lane ~40–50ms vs rect 32.9ms; RAL re-profile will quantify). The fix (batch-parallel/threaded callback, plan item 3 of the 2026-07-09 review) stays parked until those numbers exist;
2. nothing else — mapping/interpolation moved to JAX in refactor: qhull-only Delaunay callback, exact JAX visibility-walk point location #368 and regularization H is ~2%.

So the right next Delaunay target = the same NNLS levers as rect (this task), then the vmap-callback question after the A100 re-profile.

[Jammy2211]

Interim update — A100 run queued, Phase B implemented + locally validated

A100 (RAL restored): all five HPC checkouts pulled to current main (PyAutoArray at the #368 merge), probe package + real extracted systems rsynced over, and SLURM job 330046 (submit_probe_nnls_a100) submitted: PDIP probe on both real systems + A100 re-runs of the two likelihood_breakdown scripts (also clears the outstanding post-#368 A100 re-profile from the Delaunay review). Currently pending — euclid-ral-gpu-[1-2] are still down/drained; it will start when the GPU nodes return.

Phase B (in the feature/nnls-solver-optimization worktree, unshipped):

• autoarray/util/jax_nnls.py — vendored while_loop driver with solver_tol / max_iter knobs; reuses all jaxnnls building blocks + the relaxed-KKT custom-vjp backward; knobs are static closure params behind an lru_cache so JIT tracing caches hit.
• inversion_util.reconstruction_positive_only_from reads general.yaml -> inversion -> nnls_solver_tol / nnls_max_iter (try/except fallbacks, same pattern as the existing nnls keys); packaged defaults reproduce upstream exactly.
• Validated on the real rect+Delaunay systems: defaults bit-identical to upstream jaxnnls (solutions and gradients), tol 1e-6 saves 3-4 iterations with rel drift ~1e-6, knob-cache stable. New NumPy-only unit tests pass (module-level JAX-import guard, packaged config defaults, numpy-path positive-only solve); full suite run in progress.

Ship gate: A100 numbers land in the PR benchmark table first.

[Jammy2211]

Laptop-GPU probe (RTX 2060 Max-Q 6GB, WSL2) — consumer feasibility datapoint

Full rect pixelization+MGE HST likelihood, real production configs, post-#368:

	Laptop CPU fp64	RTX 2060 fp64	RTX 2060 mp	GPU/CPU
Curvature matrix (F)	3.01s	0.397s	0.395s	7.6x
Inversion setup	1.58s	0.125s	0.077s	12.6x
NNLS reconstruction	1.21s	0.298s	0.294s	4.1x
whole eval	5.96s	0.881s	0.820s	6.8x

NNLS solver probe on the real systems (identical iteration counts to CPU, as expected): rect solve 336ms (3.4x CPU), Delaunay 329ms (2.7x); vmap-8 amortizes rect to 267ms/lane. Tolerance knob worth more on GPU for Delaunay (1e-6 -> 1.43x, 1e-4 -> 1.60x solve speedup).

Takeaways

• A consumer laptop GPU runs the full production likelihood ~7x faster than the same laptop's CPU — sub-second/eval — so laptop-GPU modeling is genuinely feasible for patient users.
• NNLS share grows 20% (CPU) -> 34% (2060) — confirms the solver as the right accelerator-side target; on A100 it should dominate outright (F's GEMM gets ~100x cheaper there, the sequential PDIP chain doesn't).
• Mixed precision is NOT the consumer rescue: F and the NNLS solve stay fp64 (cond(Q)~1e10), together 84% of the GPU eval. Log-evidence parity passed in both modes.
• UX trap for laptop users: jax silently picks CPU unless BOTH JAX_PLATFORM_NAME and JAX_PLATFORMS are unset/overridden — worth a docs note.
• mp run OOM'd at XLA_PYTHON_CLIENT_MEM_FRACTION=0.5 on 6GB (eager FFT mapping-matrix convolution, 398MB alloc); fine at 0.85.

A100 job 330046 still queued (euclid-ral-gpu-[1-2] down, "Not responding").

[Jammy2211]

Deeper-lever experiment: Gondzio multiple centrality corrections — tried, near-wash, parked

Anatomy first (RTX 2060, real rect system): the fresh Cholesky is the PDIP iteration (~14ms step; one extra KKT solve ~2ms) — so reusing the factorization for extra centrality-corrector solves (Gondzio MCC) had a real ceiling. Implemented in scratch/nnls_speedup/probe_gondzio.py (composite directions via KKT linearity; branch-free acceptance; correct: solutions match baseline to ~1e-8, all variants converge).

Result: the iteration reduction is real, the wall-time win is not.

(RTX 2060)	rect iters	rect single	rect vmap-8/lane	delaunay iters	delaunay single	delaunay vmap-8/lane
baseline	21	298ms	284ms	19	258ms	211ms
Gondzio K=2	18	1.09x	288ms (wash)	15	0.90x	222ms (worse)
tol 1e-6 alone	18	1.17x	247ms	15	1.29x	202ms
K=2 + tol 1e-6	16	1.23x	275ms	13	1.34x	196ms

The extra solves + line searches eat the saved factorizations, and under vmap (the production mode) every lane pays the corrector cost every iteration while acceptance diverges across lanes — net zero to negative. Stacking on the tolerance knob adds only ~3-6% over the knob alone.

Verdict: ship the tolerance/max-iter knobs (Phase B, done); park Gondzio with these numbers. Solver dead-end ledger for this problem class (cond(Q)1e10, n1600): warm starts (hurt), fp32/iterative refinement (cond x eps32 ~ 600, diverges), Gondzio MCC (iterations -15% but wall-time wash). One-command re-bench on A100 if its factorization:solve ratio turns out very different: python3 scratch/nnls_speedup/probe_gondzio.py --mesh rectangular.

[Jammy2211]

Outcome: findings shipped, implementation parked (user decision 2026-07-09)

The measured optimization (solver tolerance/iteration knobs, ~15-20% of the NNLS solve) is not being merged — the risk of a new solver code path was judged not worth the size of the win. The validated implementation is preserved on branch feature/nnls-solver-optimization (d8a1c84: vendored driver + config keys + tests, defaults bit-identical to upstream jaxnnls in solutions and gradients on the real production systems, full suite 876-passed) should the trade-off change — e.g. if the A100 numbers show a larger production share than expected, or the follow-up experiment (below) promotes a new solver anyway.

What this task established (the durable knowledge)

Per-eval breakdowns, production HST pixelization+MGE configs, post-#368:

step	laptop CPU fp64	RTX 2060 fp64	RTX 2060 mp
Curvature matrix (F)	3.01s (51%)	0.397s (45%)	0.395s (48%)
Inversion setup	1.58s (27%)	0.125s (14%)	0.077s (9%)
NNLS reconstruction	1.21s (20%)	0.298s (34%)	0.294s (36%)
whole eval	5.96s	0.881s	0.820s

• Consumer-laptop modeling is feasible: 6.8x over the same laptop's CPU, sub-second/eval, evidence parity passing. Mixed precision is not the consumer rescue (F + NNLS stay fp64 at cond(Q)~1e10 = 84% of the eval).
• NNLS share grows on accelerators (20% CPU -> 34% GPU) and the solve is latency-bound, not flops-bound (2060 does an iteration in ~14ms despite 1/100th the A100's fp64 flops).
• Delaunay verdict: post-refactor: qhull-only Delaunay callback, exact JAX visibility-walk point location #368 the mesh/scipy machinery is no longer the target — remaining Delaunay-specific cost is the callback's vmap serialization only (quantified by the queued A100 job).

The interior-point ledger is complete — measured on the real systems:

lever	result
tolerance 1e-6 / 1e-4	-3 to -6 of ~20 iters, Δlog_ev ~1e-8/1e-6 (parked branch exposes it)
warm start from unconstrained solve	HURTS (un-centers; cond ~1e10-1e16)
fp32 / iterative refinement	diverges (cond x eps32 ~ 600; KKT worsens near convergence)
Gondzio multiple centrality corrections	iters -15% but wall-time wash; worse under vmap
per-iteration anatomy	the fresh Cholesky IS the iteration (~14ms vs ~2ms/extra solve, RTX 2060)

Follow-up

Different-algorithm candidates (block principal pivoting; ADMM with fixed factorization) are filed as PyAutoMind/feature/autolens_profiling/nnls_bpp_admm_solver_experiment.md with hard promotion gates (>=2x at evidence parity, vmap-clean, PDIP fallback). Scratch harness + real extracted systems in autolens_profiling/scratch/nnls_speedup/ are the shared testbed.

Pending addendum

A100 job 330046 remains queued (euclid-ral-gpu-[1-2] down). When it runs it delivers: A100 per-step shares (the true production NNLS share), the post-#368 Delaunay re-profile, and on-A100 PDIP probe numbers. Results will be posted here; suggest closing this issue after that addendum.

[Jammy2211]

Ledger update from #370: BPP and ADMM both measured and negative (BPP machine-exact but 36-53 factorizations vs PDIP's 19-21; ADMM plateaus at rel dObj ~3e-6, four orders short of the evidence envelope). The dead-end ledger is now: warm starts, fp32/IR, Gondzio, BPP, ADMM. PDIP ~20 iters x Cholesky is near-optimal for this problem class — the remaining routes to a faster positive-only solve are problem-level (fewer source pixels, conditioning via regularization design, or relaxing positivity), not algorithmic.

[Jammy2211]

User decision update: the knobs are now shipping after all — reworked as per-fit Settings options (default off) per request, PR #371. The A100 addendum (job 330046) will be posted here when the GPU nodes return. Closing; the durable record is autolens_profiling/results/notes/nnls_solver_ledger.md.