Activity

1. Jammy2211 commented on Aug 31, 2026

   

   Jammy2211

   on Aug 31, 2026 – with Claude

   CollaboratorAuthor

   More actions

   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
2. Jammy2211 commented on Sep 1, 2026

   

   Jammy2211

   on Sep 1, 2026 – with Claude

   CollaboratorAuthor

   More actions

   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
3. Jammy2211

   closed this as completedon Sep 1, 2026