perf: in-place Cholesky buffer + copy-free numba solves for fnnls_cholesky

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 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Vcc7MUBMnNU6n8qqS9ioVZ

Author: claude

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

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

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)