From 2c9a22bb99311e1635a345db9bd6b6fa03b7af39 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 20 Aug 2026 21:23:38 +0000 Subject: [PATCH 1/2] perf: in-place Cholesky buffer + copy-free numba solves for fnnls_cholesky MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The positive-only inversion solve (fnnls_cholesky) dominates the numba CPU Delaunay likelihood (~74% of a euclid eval, autolens_profiling#151): ~150 Bro-Jong active-set iterations, each rebuilding the ~1300^2 factor via np.insert/np.delete (two O(n^2) copies per change) and re-scanning/copying it inside scipy's solve_triangular/cho_solve. Two changes, same mathematics: - The factor now lives in the top-left k x k corner of one preallocated buffer, grown/shrunk in place (cholinsertlast_inplace / choldeleteindexes_inplace; the numba Givens kernels are unchanged and proven bitwise-stable on strided views). - The per-iteration triangular solves run through new copy-free numba kernels (_solve_upper_transposed_buffer, _cho_solve_buffer) that read the buffer's upper triangle directly — scipy on a strided view re-copies and finite-scans O(k^2) per call, which had cancelled the buffer's win. On the real euclid 1310-param system (Delaunay Hilbert-1250 fiducial, autolens_profiling PR #152): solve 3.29 s -> 1.40 s (2.35x), identical active set, solutions agree to ~1e-9 absolute. Results are not bitwise identical to the old path (LAPACK picks layout-dependent but equally valid kernels; the old implementation already mixed layouts iteration-to-iteration) — equivalence is tested at 1e-13 in the new test_autoarray/util/test_cholesky_inplace.py, plus exact factor-property checks and scipy.optimize.nnls cross-checks. The out-of-place cholinsertlast/choldeleteindexes remain for reference and tests. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01Vcc7MUBMnNU6n8qqS9ioVZ --- autoarray/util/cholesky_funcs.py | 103 +++++++++++ autoarray/util/fnnls.py | 55 ++++-- test_autoarray/util/test_cholesky_inplace.py | 169 +++++++++++++++++++ 3 files changed, 310 insertions(+), 17 deletions(-) create mode 100644 test_autoarray/util/test_cholesky_inplace.py diff --git a/autoarray/util/cholesky_funcs.py b/autoarray/util/cholesky_funcs.py index 60ff68016..fd94b1e1b 100644 --- a/autoarray/util/cholesky_funcs.py +++ b/autoarray/util/cholesky_funcs.py @@ -152,3 +152,106 @@ def choldeleteindexes(U, indexes): U = L return U + + +@numba_util.jit() +def _solve_upper_transposed_buffer(Ubuf, k, b): + """ + Solve ``U^T y = b`` for the active k x k upper factor ``Ubuf[:k, :k]``, + overwriting ``b`` with ``y`` (row-oriented forward substitution, touching + only contiguous row slices of the buffer's upper triangle). + + This replaces `scipy.linalg.solve_triangular(..., trans=1)` on the buffer + view: scipy copies a non-contiguous view into a fresh array and scans it + for non-finite values on every call, which at n ~ 1000 and ~150 calls per + fnnls solve re-creates the very memory traffic the in-place buffer + removes. + """ + for i in range(k): + yi = b[i] / Ubuf[i, i] + b[i] = yi + for j in range(i + 1, k): + b[j] -= Ubuf[i, j] * yi + return b + + +@numba_util.jit() +def _cho_solve_buffer(Ubuf, k, b): + """ + Solve ``(U^T U) s = b`` for the active k x k upper factor + ``Ubuf[:k, :k]``, overwriting ``b`` with ``s`` — LAPACK ``cho_solve`` for + an upper factor, reading only the buffer's upper triangle and making no + copies (see `_solve_upper_transposed_buffer` for why that matters). + """ + b = _solve_upper_transposed_buffer(Ubuf, k, b) + for i in range(k - 1, -1, -1): + b[i] = (b[i] - np.dot(Ubuf[i, i + 1 : k], b[i + 1 : k])) / Ubuf[i, i] + return b + + +def cholinsertlast_inplace(Ubuf, k, x): + """ + In-place variant of `cholinsertlast` for a factor held in a preallocated + buffer: the active k x k factor is `Ubuf[:k, :k]` and the new row/column is + written directly into row/column `k` of the buffer. + + Same arithmetic as `cholinsertlast` (a transposed-triangular solve for + the new column, the same `_pivot_from_schur` pivot, on the same values), + but with a copy-free numba substitution in place of scipy's + `solve_triangular` and without the two full O(k^2) `np.insert` + reallocations per call — which dominate the + positive-only inversion solve at n ~ 1000 (one such insertion per fnnls + active-set iteration, ~150 iterations per likelihood evaluation). + + Only the upper triangle of the active region is maintained; entries below + the diagonal are never written or read (every consumer — the solve and + update kernels above — references the upper triangle only). ``x[:k]`` is + overwritten by the solve. + + Returns the new active size ``k + 1``; the factor is ``Ubuf[:k+1, :k+1]``. + """ + S12 = _solve_upper_transposed_buffer(Ubuf, k, x[:k]) + + Ubuf[:k, k] = S12 + + Ubuf[k, k] = _pivot_from_schur( + schur=x[k] - S12.dot(S12), diagonal=x[k], index=k + ) + + return k + 1 + + +def choldeleteindexes_inplace(Ubuf, k, indexes): + """ + In-place variant of `choldeleteindexes` for a factor held in a + preallocated buffer: remove the given positions from the active k x k + factor `Ubuf[:k, :k]` by shifting the surviving rows/columns within the + buffer, then re-triangularize the trailing block with the same numba + Givens kernel (`_cholupdate`) on the same values — no `np.delete` + reallocations (two full-factor copies per deleted index). + + NumPy guarantees correct results for the overlapping same-array slice + assignments used for the shifts (it buffers the source when views + overlap). Only the upper triangle is maintained, as in + `cholinsertlast_inplace`. + + Returns the new active size; the factor is ``Ubuf[:k', :k']``. + """ + for index in sorted(indexes, reverse=True): + # The deleted row's tail is the rank-1 update vector for the trailing + # block — copied out before the shifts overwrite it. + x = Ubuf[index, index + 1 : k].copy() + + # Deleting row+column `index`: the top-right block shifts one column + # left; the trailing block shifts one step up-left along the diagonal. + Ubuf[:index, index : k - 1] = Ubuf[:index, index + 1 : k] + Ubuf[index : k - 1, index : k - 1] = Ubuf[index + 1 : k, index + 1 : k] + + k -= 1 + + # If the deleted index was at the end, the factor needs no update. + + if index < k: + _cholupdate(Ubuf[index:k, index:k], x) + + return k diff --git a/autoarray/util/fnnls.py b/autoarray/util/fnnls.py index a7ef9e65d..c4985d852 100644 --- a/autoarray/util/fnnls.py +++ b/autoarray/util/fnnls.py @@ -1,6 +1,10 @@ import numpy as np -from autoarray.util.cholesky_funcs import cholinsertlast, choldeleteindexes +from autoarray.util.cholesky_funcs import ( + _cho_solve_buffer, + cholinsertlast_inplace, + choldeleteindexes_inplace, +) from autoarray import exc @@ -50,6 +54,18 @@ def fnnls_cholesky( w = ZTx - (ZTZ) @ d s_chol = np.zeros(n) + # The Cholesky factor of ZTZ[passive][:, passive] lives in the top-left + # k_active x k_active corner of a single preallocated buffer, updated in + # place by cholinsertlast_inplace / choldeleteindexes_inplace as the + # active set changes. The buffer is allocated once (first factorisation) + # instead of the factor being rebuilt with np.insert/np.delete every + # iteration — the dominant cost of this solver at n ~ 1000. Zeroed, not + # empty: the update/solve kernels only ever read the upper triangle, but + # keeping the rest exactly zero costs one memset and keeps every k x k + # view a valid dense factor for inspection and tests. + U_buffer = np.zeros((n, n)) + k_active = 0 + if P_initial.shape[0] != 0: P_number = np.arange(len(P), dtype="int") P_inorder = P_number[P_initial] @@ -76,22 +92,27 @@ def fnnls_cholesky( if loop_count == 0: # We need to initialize the Cholesky factorisation, U, for the first loop. U = slg.cholesky(ZTZ[P_inorder][:, P_inorder]) + k_active = U.shape[0] + U_buffer[:k_active, :k_active] = U else: - U = cholinsertlast(U, ZTZ[idmax][P_inorder]) + k_active = cholinsertlast_inplace( + U_buffer, k_active, ZTZ[idmax][P_inorder] + ) - # solve the lstsq problem by cho_solve + # solve the lstsq problem via the copy-free buffer cho_solve - s_chol[P_inorder] = slg.cho_solve((U, False), ZTx[P_inorder]) + s_chol[P_inorder] = _cho_solve_buffer(U_buffer, k_active, ZTx[P_inorder]) P[idmax] = True while np.any(P) and np.min(s_chol[P]) <= tolerance: - s_chol, d, P, P_inorder, U = fix_constraint_cholesky( + s_chol, d, P, P_inorder, k_active = fix_constraint_cholesky( ZTx=ZTx, s_chol=s_chol, d=d, P=P, P_inorder=P_inorder, - U=U, + U_buffer=U_buffer, + k_active=k_active, tolerance=tolerance, ) @@ -132,18 +153,18 @@ def fnnls_cholesky( return d -def fix_constraint_cholesky(ZTx, s_chol, d, P, P_inorder, U, tolerance): +def fix_constraint_cholesky(ZTx, s_chol, d, P, P_inorder, U_buffer, k_active, tolerance): """ Similar to fix_constraint, but solve the lstsq by Cholesky factorisation. If this function is called, it means some solutions in the current passive sets needed to be taken out and put into the active set. So, this function involves 3 procedure: 1. Identifying what solutions should be taken out of the current passive set. - 2. Updating the P, P_inorder and the Cholesky factorisation U. - 3. Solving the lstsq by using the new Cholesky factorisation U. + 2. Updating the P, P_inorder and the Cholesky factorisation (the active + k_active x k_active corner of U_buffer, updated in place). + 3. Solving the lstsq by using the new Cholesky factorisation. As some solutions are taken out from the passive set, the Cholesky factorisation needs to be - updated by choldeleteindexes. To realize that, we call the `choldeleteindexes` from - cholesky_funcs. + updated in place by `choldeleteindexes_inplace` from cholesky_funcs. """ q = P * (s_chol <= tolerance) alpha = np.min(d[q] / (d[q] - s_chol[q])) @@ -153,20 +174,20 @@ def fix_constraint_cholesky(ZTx, s_chol, d, P, P_inorder, U, tolerance): id_delete = np.where(d[P_inorder] <= tolerance)[0] - U = choldeleteindexes(U, id_delete) # update the Cholesky factorisation + # update the Cholesky factorisation + + k_active = choldeleteindexes_inplace(U_buffer, k_active, id_delete) P_inorder = np.delete(P_inorder, id_delete) # update the P_inorder P[d <= tolerance] = False # update the P - # solve the lstsq problem by cho_solve + # solve the lstsq problem via the copy-free buffer cho_solve if len(P_inorder): - from scipy import linalg as slg - # there could be a case where P_inorder is empty. - s_chol[P_inorder] = slg.cho_solve((U, False), ZTx[P_inorder]) + s_chol[P_inorder] = _cho_solve_buffer(U_buffer, k_active, ZTx[P_inorder]) s_chol[~P] = 0.0 # set solutions taken out of the passive set to be 0 - return s_chol, d, P, P_inorder, U + return s_chol, d, P, P_inorder, k_active diff --git a/test_autoarray/util/test_cholesky_inplace.py b/test_autoarray/util/test_cholesky_inplace.py new file mode 100644 index 000000000..91879391a --- /dev/null +++ b/test_autoarray/util/test_cholesky_inplace.py @@ -0,0 +1,169 @@ +""" +The in-place Cholesky update path (`cholinsertlast_inplace` / +`choldeleteindexes_inplace` + the preallocated buffer in `fnnls_cholesky`) +must reproduce the out-of-place `cholinsertlast` / `choldeleteindexes` +results to the last few ulp, and the maintained factor must stay an exact +Cholesky factor of the active submatrix. + +Exact bitwise agreement between the two implementations is NOT required (and +does not hold): LAPACK picks different, equally valid dtrtrs/dpotrs +invocations depending on the input's memory layout (F-contiguous vs +C-contiguous vs strided view), producing last-ulp differences — the +out-of-place implementation already mixes layouts between its own iterations +(`scipy.linalg.cholesky` returns F-order, `np.insert`/`np.delete` return +C-order). The production tolerance for this solver's output is the profiling +pins' rtol=1e-6; the cross-implementation tolerance here is far tighter. +""" + +import numpy as np +import pytest +from scipy import linalg as slg +from scipy.optimize import nnls + +from autoarray.util.cholesky_funcs import ( + cholinsertlast, + cholinsertlast_inplace, + choldeleteindexes, + choldeleteindexes_inplace, +) +from autoarray.util.fnnls import fnnls_cholesky + + +def _random_spd(n, seed): + rng = np.random.default_rng(seed) + Z = rng.normal(size=(2 * n, n)) + return Z.T @ Z + n * np.eye(n) + + +def _buffer_from(U, n_max): + buffer = np.zeros((n_max, n_max)) + k = U.shape[0] + buffer[:k, :k] = U + return buffer, k + + +def _assert_factors_match(U_reference, U_inplace): + """The two factors agree to within a few ulp (see module docstring).""" + np.testing.assert_allclose( + np.triu(U_inplace), np.triu(U_reference), rtol=1e-13, atol=1e-13 + ) + + +def _assert_is_cholesky_of(U_view, A_sub): + """The maintained upper triangle is an exact Cholesky factor of the + active submatrix (the property every downstream cho_solve relies on).""" + R = np.triu(U_view) + np.testing.assert_allclose(R.T @ R, A_sub, rtol=1e-12, atol=1e-12) + + +@pytest.mark.parametrize("seed", [0, 1, 2]) +def test__cholinsertlast_inplace__matches_out_of_place(seed): + n = 12 + A = _random_spd(n, seed) + + k = 7 + U = slg.cholesky(A[:k, :k]) + x = A[k, : k + 1].copy() + + S = cholinsertlast(U.copy(), x.copy()) + + buffer, k_active = _buffer_from(U, n) + k_active = cholinsertlast_inplace(buffer, k_active, x.copy()) + + assert k_active == k + 1 + _assert_factors_match(S, buffer[:k_active, :k_active]) + _assert_is_cholesky_of(buffer[:k_active, :k_active], A[: k + 1, : k + 1]) + + +@pytest.mark.parametrize( + "indexes", + [[0], [7], [3, 5], [0, 1, 6], [2, 3, 4, 5], [0, 1, 2, 3, 4, 5, 6, 7]], +) +def test__choldeleteindexes_inplace__matches_out_of_place(indexes): + n = 12 + A = _random_spd(n, seed=3) + + k = 8 + U = slg.cholesky(A[:k, :k]) + + L = choldeleteindexes(U.copy(), list(indexes)) + + buffer, k_active = _buffer_from(U, n) + k_active = choldeleteindexes_inplace(buffer, k_active, list(indexes)) + + assert k_active == k - len(indexes) + _assert_factors_match(L, buffer[:k_active, :k_active]) + + keep = [i for i in range(k) if i not in indexes] + _assert_is_cholesky_of( + buffer[:k_active, :k_active], A[np.ix_(keep, keep)] + ) + + +def test__interleaved_inserts_and_deletes__match(): + n = 20 + A = _random_spd(n, seed=4) + + k = 4 + U = slg.cholesky(A[:k, :k]) + buffer, k_active = _buffer_from(U, n) + + # Mimic fnnls's usage: grow to the next leading size, shed some indexes, + # grow again — comparing the two implementations after every operation. + for op, arg in [ + ("insert", None), + ("insert", None), + ("delete", [1, 3]), + ("insert", None), + ("delete", [0]), + ("insert", None), + ("insert", None), + ]: + if op == "insert": + k_old = U.shape[0] + x = A[k_old, : k_old + 1].copy() + U = cholinsertlast(U, x.copy()) + k_active = cholinsertlast_inplace(buffer, k_active, x.copy()) + else: + U = choldeleteindexes(U, arg) + k_active = choldeleteindexes_inplace(buffer, k_active, arg) + + assert k_active == U.shape[0] + _assert_factors_match(U, buffer[:k_active, :k_active]) + + +@pytest.mark.parametrize("seed", [0, 1, 2, 3, 4]) +def test__fnnls_cholesky__matches_scipy_nnls(seed): + rng = np.random.default_rng(seed) + n = 30 + Z = rng.normal(size=(50, n)) + # A mixed-sign target so a substantial subset of the solution is clamped + # at zero and the delete path is exercised. + x = Z @ rng.normal(size=n) + rng.normal(size=50) + + ZTZ = Z.T @ Z + ZTx = Z.T @ x + + d = fnnls_cholesky(ZTZ, ZTx) + d_ref, _ = nnls(Z, x) + + assert np.all(d >= 0.0) + assert d == pytest.approx(d_ref, rel=1e-6, abs=1e-8) + + +@pytest.mark.parametrize("seed", [0, 1, 2]) +def test__fnnls_cholesky__warm_start_matches_cold_start(seed): + rng = np.random.default_rng(seed) + n = 30 + Z = rng.normal(size=(50, n)) + x = Z @ rng.normal(size=n) + rng.normal(size=50) + + ZTZ = Z.T @ Z + ZTx = Z.T @ x + + d_cold = fnnls_cholesky(ZTZ, ZTx) + + P_initial = np.where(slg.solve(ZTZ.copy(), ZTx.copy(), assume_a="pos") > 0)[0] + d_warm = fnnls_cholesky(ZTZ, ZTx, P_initial=P_initial) + + assert d_warm == pytest.approx(d_cold, rel=1e-8, abs=1e-10) From 89be276ce2c631415e3f07f79420f55ebf6dd63a Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 20 Aug 2026 21:36:05 +0000 Subject: [PATCH 2/2] perf: numba shift kernel for in-place Cholesky deletes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Profiling the new solver showed choldeleteindexes_inplace as the largest remaining cost (0.64 s of 1.45 s at euclid): numpy's overlapping slice assignments buffer the source region, re-introducing per-delete temp allocations. _choldelete_shift_buffer closes the deleted row/column gap with plain numba loops (destinations strictly behind sources, upper triangle only) — pure value movement, bitwise identical. Euclid 1310-param system: solve 1.40 s -> 1.16 s (2.83x total vs the np.insert/np.delete baseline's 3.29 s); euclid eval 2.34 s/call with the profiling pin passing. test_autoarray: same 1026 passed / 3 pre-existing pynufft failures. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01Vcc7MUBMnNU6n8qqS9ioVZ --- autoarray/util/cholesky_funcs.py | 35 ++++++++++++++++++++++---------- 1 file changed, 24 insertions(+), 11 deletions(-) diff --git a/autoarray/util/cholesky_funcs.py b/autoarray/util/cholesky_funcs.py index fd94b1e1b..e4e59c902 100644 --- a/autoarray/util/cholesky_funcs.py +++ b/autoarray/util/cholesky_funcs.py @@ -221,18 +221,34 @@ def cholinsertlast_inplace(Ubuf, k, x): return k + 1 +@numba_util.jit() +def _choldelete_shift_buffer(Ubuf, k, index): + """ + Shift the active k x k factor's rows/columns to close the gap left by + deleting row+column ``index``: the top-right block moves one column left, + the trailing block moves one step up-left along the diagonal. Pure value + movement (bitwise), touching only the upper triangle — the numba loops + write destinations strictly behind their sources, so no temporary is + needed (numpy's overlapping slice assignment buffers the source instead, + which at ~100 deletes per fnnls solve is real allocation traffic). + """ + for i in range(index): + for j in range(index, k - 1): + Ubuf[i, j] = Ubuf[i, j + 1] + for i in range(index, k - 1): + for j in range(i, k - 1): + Ubuf[i, j] = Ubuf[i + 1, j + 1] + + def choldeleteindexes_inplace(Ubuf, k, indexes): """ In-place variant of `choldeleteindexes` for a factor held in a preallocated buffer: remove the given positions from the active k x k factor `Ubuf[:k, :k]` by shifting the surviving rows/columns within the - buffer, then re-triangularize the trailing block with the same numba - Givens kernel (`_cholupdate`) on the same values — no `np.delete` - reallocations (two full-factor copies per deleted index). - - NumPy guarantees correct results for the overlapping same-array slice - assignments used for the shifts (it buffers the source when views - overlap). Only the upper triangle is maintained, as in + buffer (`_choldelete_shift_buffer`), then re-triangularize the trailing + block with the same numba Givens kernel (`_cholupdate`) on the same + values — no `np.delete` reallocations (two full-factor copies per + deleted index). Only the upper triangle is maintained, as in `cholinsertlast_inplace`. Returns the new active size; the factor is ``Ubuf[:k', :k']``. @@ -242,10 +258,7 @@ def choldeleteindexes_inplace(Ubuf, k, indexes): # block — copied out before the shifts overwrite it. x = Ubuf[index, index + 1 : k].copy() - # Deleting row+column `index`: the top-right block shifts one column - # left; the trailing block shifts one step up-left along the diagonal. - Ubuf[:index, index : k - 1] = Ubuf[:index, index + 1 : k] - Ubuf[index : k - 1, index : k - 1] = Ubuf[index + 1 : k, index + 1 : k] + _choldelete_shift_buffer(Ubuf, k, index) k -= 1