Merge pull request #453 from PyAutoLabs/feature/fnnls-inplace-cholesky-buffer

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

Author: Jammy2211

autoarray/util/cholesky_funcs.py

@@ -152,3 +152,119 @@ 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
+
+
+@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 (`_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']``.
+    """
+    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()
+
+        _choldelete_shift_buffer(Ubuf, k, index)
+
+        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