Repository navigation
perf: in-place Cholesky buffer + copy-free numba solves for fnnls_cholesky #452
Description
Overview
The positive-only inversion solve (fnnls_cholesky) dominates the numba CPU Delaunay likelihood — ~74% of a euclid-resolution evaluation on the production campaign fiducial (Delaunay + Hilbert-1250, PyAutoLabs/autolens_profiling#151). Profiling showed the cost is implementation, not linear algebra: ~150 Bro–Jong active-set iterations, each rebuilding the ~1300² Cholesky factor via np.insert/np.delete (two O(n²) copies per change) and re-copying/finite-scanning it inside scipy's solve_triangular/cho_solve. This task makes the solver maintain its factor in one preallocated buffer with copy-free numba triangular solves — same mathematics, a fraction of the memory traffic.
Plan
- Hold the factor in the top-left k×k corner of a preallocated buffer, grown/shrunk in place as the active set changes (no
np.insert/np.delete). - Replace the per-iteration scipy
solve_triangular/cho_solvecalls with numba kernels that read the buffer's upper triangle directly (no copies, no per-call finite scans). - Keep the existing numba Givens up/down-date kernels (proven bitwise-stable on strided views) and the out-of-place functions for reference and tests.
- Validate against the pinned autolens_profiling delaunay_numba cells and the full test suite.
Detailed implementation plan
Affected Repositories
- PyAutoArray (primary)
Branch Survey
| Repository | Current Branch | Dirty? |
|---|---|---|
| ./PyAutoArray | main | clean |
Suggested branch: feature/fnnls-inplace-cholesky-buffer
Implementation Steps
autoarray/util/cholesky_funcs.py: addcholinsertlast_inplace/choldeleteindexes_inplace(buffer-view variants of the existing out-of-place functions) and the numba kernels_solve_upper_transposed_buffer(row-oriented forward substitution for U^T y = b) and_cho_solve_buffer(full (U^T U) s = b solve).autoarray/util/fnnls.py: allocateU_buffer = np.zeros((n, n))once per solve; thread(U_buffer, k_active)through the outer loop andfix_constraint_choleskyin place of the reallocatedU.test_autoarray/util/test_cholesky_inplace.py: cross-implementation agreement at 1e-13, exact factor-property checks (U'U = active submatrix),scipy.optimize.nnlscross-checks, warm-start = cold-start.
Key Files
autoarray/util/fnnls.py— the Bro–Jong active-set loopautoarray/util/cholesky_funcs.py— factor update + solve kernelsautoarray/inversion/inversion/inversion_util.py—reconstruction_positive_only_from(caller; unchanged)
Results (validated 2026-08-20, 4-core cloud container)
- Real euclid 1310-param system: solve 3.29 s → 1.40 s (2.35×), identical active set, solutions agree to ~1e-9 absolute.
- Pinned autolens_profiling cells (PR feat: numba CPU sparse-operator likelihood profiling (runtime, breakdown, multiprocessing scaling) autolens_profiling#152, rtol 1e-6) pass: euclid eval 4.92 → 2.68 s (1.83× end-to-end), solve step 3.61 → 1.57 s; hst 3.95 → 3.43 s.
test_autoarray/: 1026 passed; the only failures are 3 pre-existing pynufft transformer tests that fail identically on stockmainin this environment.- Results are not bitwise identical to the old path (LAPACK picks layout-dependent but equally valid kernels; the old implementation already mixed layouts between its own iterations) — equivalence is enforced at 1e-13 in the new tests.
Original Prompt
Click to expand starting prompt
PyAutoMind draft/feature/autoarray/numba_cpu_likelihood_positive_only_solver_speedup.md — "Numba CPU likelihood: positive-only solver speed-up (the Delaunay ~74% euclid lever)". Candidate directions in win/effort order: (1) in-place factor buffer [this issue], (2) block pivoting, (3) cross-evaluation warm starts, (4) BLAS-thread guidance. Full context and the instrumented probe table live in the prompt file and the PyAutoLabs/autolens_profiling#151 findings trail.
Activity
Jammy2211 commentedon Aug 20, 2026
Update: solve now 2.83× (3.29 → 1.16 s); phase-2 block strategies prototyped and both rejected — findings below.
Second commit on PR #453 (
89be276): profiling the new solver showed the delete path as the largest remaining cost (0.64 s of 1.45 s) — numpy's overlapping slice assignments buffer the source region, re-introducing per-delete temp allocations. A numba shift kernel (_choldelete_shift_buffer) removes that; bitwise-identical value movement. Euclid eval is now 2.34 s/call (was 4.92 s stock — 2.10× end-to-end), pin passing.Phase-2 prototypes (on the saved real euclid 1310-param system) — both naive block strategies fail on this system, and the failure modes are instructive:
Block principal pivoting (PJV/TNT-NN style) — cycles. The infeasibility count oscillates in the 10–20 range and revisits states (verified repeating F-sets); the p-counter/Murty anti-cycling fallback degenerates to one exchange per full 47 ms refactorization, i.e. worse than Bro–Jong. Root cause: the near-degenerate mesh columns (the same degeneracy behind the solver's documented ~1e-8 bistability) make whole-block exchanges overshoot persistently.
Bro–Jong with block insertion (top-m violators per outer iteration, reusing the new in-place machinery) — fast (0.74 s at m=16, 4 outer iterations) but wrong: the alpha-step constraint fix deletes most of each inserted block,
Preturns to its previous state, and theno_update >= max_repetitionssafeguard then terminates with KKT badly violated (max dual over the active set ≈ 3.7e3; objective worse by 5e-4 relative — far outside the 1e-6 pin tolerance). The safeguard masks non-convergence under block insertion.So a correct phase-2 needs real machinery (e.g. Benthem–Keenan fast combinatorial NNLS, or PJV with exact anti-cycling + degeneracy handling), not a loop tweak — filed as future work in the PyAutoMind prompt. The remaining 1.16 s is now roughly: Givens up/down-dates ~0.4 s, the 254 triangular-solve sweeps ~0.3 s (near the memory-bandwidth floor for 150 iterations at k≈1200), full-vector
wmatvecs ~0.25 s, insert-column solves ~0.15 s. The biggest remaining lever is cutting the iteration count (cross-eval warm starts, option 3 in the prompt).Generated by Claude Code