diff --git a/autoarray/util/cholesky_funcs.py b/autoarray/util/cholesky_funcs.py index 60ff6801..e4e59c90 100644 --- a/autoarray/util/cholesky_funcs.py +++ b/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 diff --git a/autoarray/util/fnnls.py b/autoarray/util/fnnls.py index a7ef9e65..c4985d85 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 00000000..91879391 --- /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)