diff --git a/autoarray/inversion/inversion/abstract.py b/autoarray/inversion/inversion/abstract.py index dbe50f50d..5d29741e6 100644 --- a/autoarray/inversion/inversion/abstract.py +++ b/autoarray/inversion/inversion/abstract.py @@ -88,6 +88,15 @@ def __init__( self._preloads = preloads + # The Cholesky factor of the passive submatrix that the positive-only solve builds, and + # the global indices it was built on (`None` when the solve took the full system). + # Published by `reconstruction` and read by `log_det_curvature_reg_matrix_term`, which + # reads `log det(F + lambda*H)` off it instead of factorising that matrix a second time. + # Per-instance and never a module global: one inversion's factor must never be able to + # answer another inversion's log determinant. + self._nnls_factor = None + self._nnls_factor_ids = None + @property def _xp(self): if self.use_jax: @@ -355,15 +364,23 @@ def regularization_matrix_reduced(self) -> Optional[np.ndarray]: # Zero rows and columns in the matrix we want to ignore return self.regularization_matrix[ids_to_keep][:, ids_to_keep] - @property + @cached_property def curvature_reg_matrix(self) -> np.ndarray: """ The linear system of equations solves for F + regularization_coefficient*H, which is computed below. - For a single mapper, this function overwrites the cached `curvature_matrix`, because for large matrices this - avoids overheads in memory allocation. The `curvature_matrix` is removed as a cached property as a result, - to ensure if we access it after computing the `curvature_reg_matrix` it is correctly recalculated in a new - array of memory. + The addition is **out-of-place**: `curvature_matrix` is left intact and a new `[params, params]` array is + returned. An earlier version of this property added `H` into `F` in place and then deleted the cached + `curvature_matrix` to force its recomputation, and this docstring described that; the in-place form was + removed in `0766edd4` (2025-06-26), so nothing here invalidates anything (and on the imaging inversions, + where the matrix is large, `curvature_matrix` is itself a `cached_property`). + + Cached because the likelihood reaches it more than once per evaluation (`reconstruction` and + `curvature_reg_matrix_reduced`, and previously `log_det_curvature_reg_matrix_term` through the latter), + and each access re-ran the add -- ~2 ms per access of an (n, n) sum at n = 1500. It was a + `cached_property` until the blanket `e819fa12` (2025-11-05) sweep turned it into a plain `property` along + with a dozen others. The cost of the cache is one extra (n, n) array (18 MB at n = 1500, on top of `F` + and `H`) held for the inversion's lifetime. """ if not self.has(cls=AbstractRegularization): return self.curvature_matrix @@ -606,6 +623,13 @@ def reconstruction(self) -> np.ndarray: # so it is safe to consult here and nowhere higher up. ids_to_keep = self.solve_ids_to_keep + # The solve's Cholesky factor of its passive submatrix, published through this + # out-dict and stored on the instance below so that + # `log_det_curvature_reg_matrix_term` can read `log det(F + lambda*H)` off it. + # `reconstruction_positive_only_from` clears it on entry, so it is empty whenever + # no solve completed. + factor = {} + if ids_to_keep is not None: # Use advanced indexing to select rows/columns @@ -624,9 +648,13 @@ def reconstruction(self) -> np.ndarray: fingerprint=self._nnls_warm_start_fingerprint( ids_to_keep=ids_to_keep ), + factor=factor, ) ) + self._nnls_factor = factor + self._nnls_factor_ids = ids_to_keep + # Allocate full solution array reconstruction = self._xp.zeros(self.data_vector.shape[0]) @@ -642,14 +670,20 @@ def reconstruction(self) -> np.ndarray: else: - return inversion_util.reconstruction_positive_only_from( + reconstruction = inversion_util.reconstruction_positive_only_from( data_vector=self.data_vector, curvature_reg_matrix=self.curvature_reg_matrix, settings=self.settings, xp=self._xp, fingerprint=self._nnls_warm_start_fingerprint(), + factor=factor, ) + self._nnls_factor = factor + self._nnls_factor_ids = None + + return reconstruction + return inversion_util.reconstruction_positive_negative_from( data_vector=self.data_vector, curvature_reg_matrix=self.curvature_reg_matrix, @@ -898,14 +932,102 @@ def log_det_curvature_reg_matrix_term(self) -> float: """ The log determinant of [F + reg_coeff*H] is used to determine the Bayesian evidence of the solution. - This uses the Cholesky decomposition which is already computed before solving the reconstruction - (or ``slogdet`` when ``Settings.log_det_method == "slogdet"`` — see :meth:`_log_det_symmetric_from`). + On the NumPy path with the positive-only solver and the default ``"cholesky"`` method this term is read + off the Cholesky factor the reconstruction's non-negative least squares solve has already built, rather + than factorizing ``F + reg_coeff*H`` a second time. + + ``fnnls_cholesky`` finishes holding a factor ``U`` of ``M[P][:, P]``, where ``M`` is + ``curvature_reg_matrix_reduced`` and ``P`` is the solve's final passive set -- 1485 of the 1500 columns + on the production HST Delaunay ``AdaptSplit`` system. The block-determinant identity then gives the + full log determinant from that factor plus a Schur complement over the ~15 remaining columns (see + :func:`inversion_util.log_det_from_passive_cholesky_from`): ~5 ms in place of a ~40 ms dense + factorization of the same matrix, agreeing with it to <= 2e-12 nats. + + Because the factor comes from the solve, **this property evaluates** :attr:`reconstruction` (a + ``cached_property``) before reading it. The likelihood is unaffected -- it already evaluates + ``regularization_term``, and through it the reconstruction, before this term -- but a diagnostic that + asks for this term alone now pays for the reconstruction solve, where before it paid only for one + Cholesky. + + The fast path is taken only when every one of these holds; anything else, and any + ``np.linalg.LinAlgError`` raised while reading the factor, takes the dense route below unchanged (which + carries the test-mode guard in :meth:`_log_det_symmetric_from`): + + - the NumPy backend (the factor is a NumPy buffer built by a numba solver); + - ``Settings.log_det_method == "cholesky"`` (``"slogdet"`` is a different value by construction, and is + reached through the dense route); + - ``Settings.use_positive_only_solver`` (nothing else builds a factor, and this also keeps the + positive-negative route from evaluating the reconstruction it would not otherwise need); + - every linear object is regularized, so ``curvature_reg_matrix_reduced`` is the whole + ``curvature_reg_matrix`` the solve was handed rather than a ``mapper_indices`` subset of it; + - a factor is present with a non-empty passive block, and its matrix shape matches; + - the solve covered the whole reduced system. Note this is "no subset **or** the identity subset", not + "no subset": with ``use_edge_zeroed_pixels`` on but no edge pixels to zero, ``solve_ids_to_keep`` + returns ``arange(n)`` rather than ``None``, which is the production configuration of the system + above. """ if not self.has(cls=AbstractRegularization): return 0.0 + if ( + self._xp is np + and self.settings.log_det_method == "cholesky" + and self.settings.use_positive_only_solver + and self.all_linear_obj_have_regularization + ): + matrix = self.curvature_reg_matrix_reduced + + # Evaluates the solve (a `cached_property`), which is what publishes the factor + # read below. Bound to `_` because the value is not what is wanted here -- the + # side effect of the solve having run is. + _ = self.reconstruction + + factor = self._nnls_factor + ids = self._nnls_factor_ids + + if ( + factor + and factor.get("U_buffer") is not None + and factor.get("k_active", 0) > 0 + and factor.get("matrix_shape") == tuple(matrix.shape) + and self._nnls_factor_ids_cover(ids=ids, total=matrix.shape[0]) + ): + try: + return inversion_util.log_det_from_passive_cholesky_from( + matrix=matrix, + U_buffer=factor["U_buffer"], + k_active=factor["k_active"], + passive_set=factor["passive_set"], + ) + except np.linalg.LinAlgError: + # The Schur complement of the passive block was not positive-definite. That + # means this matrix is at the edge of what fp64 resolves as positive-definite, + # and the dense factorization below is the authority on it: it either returns + # its own value or raises the same error, where the test-mode guard and the + # likelihood's `FitException` resampling handle it as they always have. It must + # not be answered with a fabricated 0.0 here. + pass + return self._log_det_symmetric_from(self.curvature_reg_matrix_reduced) + @staticmethod + def _nnls_factor_ids_cover(ids: Optional[np.ndarray], total: int) -> bool: + """ + Did the solve whose factor was published cover every parameter of the reduced system? + + `None` means the solve took the full system. An explicit index array covers it only if it is the + identity `arange(total)`, which `solve_ids_to_keep` does return: `use_edge_zeroed_pixels` with no edge + pixels to zero yields `arange(n)`, not `None`, and that is the production configuration of the HST + Delaunay system this fast path was measured on. Any other index array is a proper subset, whose factor + says nothing about the full matrix's determinant. + """ + if ids is None: + return True + + ids = np.asarray(ids) + + return ids.shape == (total,) and np.array_equal(ids, np.arange(total)) + @property def log_det_regularization_matrix_term(self) -> float: """ diff --git a/autoarray/inversion/inversion/imaging/mapping.py b/autoarray/inversion/inversion/imaging/mapping.py index 017fd1931..65e70f26d 100644 --- a/autoarray/inversion/inversion/imaging/mapping.py +++ b/autoarray/inversion/inversion/imaging/mapping.py @@ -83,10 +83,11 @@ def curvature_matrix(self): concatenated ensuring their `curvature_matrix` values are solved for simultaneously. This includes all diagonal and off-diagonal terms describing the covariances between linear objects. - The `curvature_matrix` computed here is overwritten in memory when the regularization matrix is added to it, - because for large matrices this avoids overhead. For this reason, `curvature_matrix` is not a cached property - to ensure if we access it after computing the `curvature_reg_matrix` it is correctly recalculated in a new - array of memory. + This `curvature_matrix` is a `cached_property` and is NOT overwritten in memory: + `curvature_reg_matrix` adds the regularization matrix to it out-of-place and leaves it intact. An earlier + version of that property added `H` into `F` in place and deleted this cached value to force its + recomputation, which is what this docstring described; the in-place form was removed in `0766edd4` + (2025-06-26). """ return inversion_util.curvature_matrix_via_mapping_matrix_from( diff --git a/autoarray/inversion/inversion/imaging/sparse.py b/autoarray/inversion/inversion/imaging/sparse.py index 52f04fde0..78d57d433 100644 --- a/autoarray/inversion/inversion/imaging/sparse.py +++ b/autoarray/inversion/inversion/imaging/sparse.py @@ -241,10 +241,11 @@ def curvature_matrix(self) -> np.ndarray: for simultaneously. In the w-tilde formalism this requires us to consider the mappings between data and every linear object, meaning that the linear alegbra has both on and off diagonal terms. - The `curvature_matrix` computed here is overwritten in memory when the regularization matrix is added to it, - because for large matrices this avoids overhead. For this reason, `curvature_matrix` is not a cached property - to ensure if we access it after computing the `curvature_reg_matrix` it is correctly recalculated in a new - array of memory. + This `curvature_matrix` is a `cached_property` and is NOT overwritten in memory: + `curvature_reg_matrix` adds the regularization matrix to it out-of-place and leaves it intact. An earlier + version of that property added `H` into `F` in place and deleted this cached value to force its + recomputation, which is what this docstring described; the in-place form was removed in `0766edd4` + (2025-06-26). """ if self.has(cls=AbstractLinearObjFuncList): curvature_matrix = self._curvature_matrix_func_list_and_mapper diff --git a/autoarray/inversion/inversion/inversion_util.py b/autoarray/inversion/inversion/inversion_util.py index 16de7e72a..3dd67fc7b 100644 --- a/autoarray/inversion/inversion/inversion_util.py +++ b/autoarray/inversion/inversion/inversion_util.py @@ -259,6 +259,7 @@ def reconstruction_positive_only_from( settings: Settings = None, xp=np, fingerprint=None, + factor: Optional[dict] = None, ): """ Solve the linear system Eq.(2) (in terms of minimizing the quadratic value) of @@ -300,6 +301,13 @@ def reconstruction_positive_only_from( fingerprint Identifies the index space this solve's passive set lives in, enabling the cross-evaluation warm-start memo (`Settings.nnls_warm_start_memo`) on the NumPy path. `None` disables the memo for this call. + factor + If a dict is passed it is cleared on entry and filled with the Cholesky factor of the solve's passive + submatrix that `fnnls_cholesky` built anyway -- see that function's `factor` parameter for the keys, and + `log_det_from_passive_cholesky_from` for the caller that reads them. It is handed to *both* + `fnnls_cholesky` calls below, so a memo-seeded attempt that raises cannot leave the factor of a solve + whose result was discarded behind for the retry's caller to read. Purely observational: the returned + reconstruction is byte-identical whether or not it is passed. Notes ----- @@ -315,6 +323,13 @@ def reconstruction_positive_only_from( ------- Non-negative S that minimizes the Eq.(2) of https://arxiv.org/pdf/astro-ph/0302587.pdf. """ + if factor is not None: + # Cleared before anything else, so that every route which does not reach a successful + # `fnnls_cholesky` -- the JAX branch, a memo-seeded attempt that raises, the test-mode + # dummy -- leaves the caller reading "no factor" rather than a factor belonging to some + # other matrix. `fnnls_cholesky` publishes only on a successful return. + factor.clear() + if xp.__name__.startswith("jax"): from autonerves import conf @@ -416,6 +431,7 @@ def reconstruction_positive_only_from( (data_vector).T, P_initial=entry.passive_set, stats=stats, + factor=factor, ) stats["seed_source"] = "memo" except (RuntimeError, np.linalg.LinAlgError, ValueError): @@ -434,6 +450,7 @@ def reconstruction_positive_only_from( (data_vector).T, P_initial=np.linalg.solve(curvature_reg_matrix, data_vector) > 0, stats=stats, + factor=factor, ) stats["seed_source"] = "dense" @@ -592,6 +609,99 @@ def param_range_list_from(cls: Type, linear_obj_list) -> List[List[int]]: SPARSE_LOG_DET_MIN_PIXELS = 256 +def log_det_from_passive_cholesky_from( + matrix: np.ndarray, + U_buffer: np.ndarray, + k_active: int, + passive_set: np.ndarray, +) -> float: + """ + Returns `log det M` of a symmetric positive-definite `matrix` M from the Cholesky factor of + one of its principal submatrices, via the Schur complement of the remaining block. + + The positive-only reconstruction is a non-negative least squares solve, and `fnnls_cholesky` + finishes it holding a Cholesky factor `U` of `M[P][:, P]`, where `P` is the solve's final + *passive* set (the pixels whose reconstructed value is non-zero). On the production HST + Delaunay `pixels=1500` `AdaptSplit` system that is 1485 of the 1500 columns. The Bayesian + evidence then needs `log det M` of the full system, and until this function existed it + factorised the whole `[pixels, pixels]` matrix a second time from scratch -- ~40 ms of a + ~270 ms NumPy likelihood call, redoing 99% of the work `fnnls_cholesky` had just done. + + Partitioning M by the passive set `P` and the active set `A` (its complement), + + log det M = log det M_PP + log det(M_AA - M_AP M_PP^-1 M_PA) + + -- the block-determinant identity, where the second term is the Schur complement of the + already-factorised block. The first term is `2 * sum(log(diag(U)))`, free. The second needs + one triangular solve `U^T Y = M_PA` with `|A|` right-hand sides and one `|A| x |A|` + Cholesky, because `Y^T Y == M_AP M_PP^-1 M_PA`. With `|A|` of order ten that is ~5 ms + instead of ~40 ms, and it agrees with the dense factorisation to <= 2e-12 nats (measured + across six instances of the production system; the fiducial fit's log evidence is + bit-identical). + + Both blocks must be positive-definite for M to be, so a failure of either Cholesky raises + `np.linalg.LinAlgError` -- the same error the dense route raises on the same matrix, so the + caller's test-mode guard and `FitException` resampling apply unchanged. This function never + silently substitutes a value for a matrix it could not factorise. + + Parameters + ---------- + matrix + The `[pixels, pixels]` symmetric positive-definite matrix M whose log determinant is + computed, e.g. an inversion's `curvature_reg_matrix_reduced`. + U_buffer + The buffer `fnnls_cholesky` published, whose leading `k_active x k_active` upper + triangle is a Cholesky factor `U` with `matrix[P][:, P] == U.T @ U`. Only that corner is + read; the rest of the buffer is zero and meaningless. + k_active + The size of that valid leading block, i.e. `len(passive_set)`. + passive_set + The passive indices **in the order the factor's rows and columns are in**. They are + neither sorted nor contiguous (the active-set solver appends and deletes), so this order + is what pairs `U` with `matrix`, and it must be used for the off-diagonal block too. + + Returns + ------- + The log determinant of `matrix`. + """ + from scipy.linalg import cholesky, solve_triangular + + passive_set = np.asarray(passive_set) + + upper_passive = U_buffer[:k_active, :k_active] + + log_det = 2.0 * float(np.sum(np.log(np.diag(upper_passive)))) + + active = np.setdiff1d(np.arange(matrix.shape[0]), passive_set) + + if active.size == 0: + # Every column is passive: the factor is of the whole matrix already. + return log_det + + # `trans=1` solves `U^T Y = M_PA` against the upper factor in place of transposing it, + # which would hand LAPACK a non-contiguous array to copy. All `|A|` right-hand sides go + # in one call: `cholesky_funcs._solve_upper_transposed_buffer` is the copy-free kernel the + # solver itself uses, but it takes a single vector and overwrites it in place, so reaching + # for it here would mean `|A|` numba calls over `|A|` writable copies to save a LAPACK + # call on a 1485 x 15 block. It is not a clean fit at this size. + y = solve_triangular( + upper_passive, + matrix[np.ix_(passive_set, active)], + trans=1, + lower=False, + check_finite=False, + ) + + schur = matrix[np.ix_(active, active)] - y.T @ y + + # `cholesky` raises `np.linalg.LinAlgError` if this block is not positive-definite, which + # is propagated deliberately: it means `matrix` is not either, and the caller must treat + # that exactly as it treats a failure of the dense factorisation. + return log_det + 2.0 * float( + np.sum(np.log(np.diag(cholesky(schur, lower=False, check_finite=False)))) + ) + + def log_det_sparse_spd_from(matrix: np.ndarray) -> Optional[float]: """ Returns `log det M` of a sparse, symmetric positive-definite `matrix` via a sparse LU diff --git a/autoarray/util/fnnls.py b/autoarray/util/fnnls.py index 6d726aac6..0fe51dd28 100644 --- a/autoarray/util/fnnls.py +++ b/autoarray/util/fnnls.py @@ -29,6 +29,7 @@ def fnnls_cholesky( ZTx, P_initial=np.zeros(0, dtype=int), stats: Optional[dict] = None, + factor: Optional[dict] = None, ): """ Similar to fnnls, but use solving the lstsq problem by updating Cholesky factorisation. @@ -44,6 +45,26 @@ def fnnls_cholesky( indices, in the order they were added), `n_passive` and `warm_start_errors` (how many entries the warm start got wrong). Purely observational -- the returned solution is unaffected. + factor + If a dict is passed it is filled on return with the Cholesky factor this solve + already built, so a caller that needs `log det(ZTZ)` does not have to factorise + the same matrix a second time (see + `inversion_util.log_det_from_passive_cholesky_from`): + + - `U_buffer` -- the buffer itself (NOT a copy; it is local to this call and is + not referenced again after the return, so handing it over is free). Its + leading `k_active x k_active` upper triangle is the factor `U` with + `ZTZ[P][:, P] == U.T @ U`, where `P` is `passive_set`. Everything outside that + corner, and the whole lower triangle, is zero and meaningless. + - `k_active` -- the size of that valid leading block. + - `passive_set` -- the passive indices in the order the factor's rows and + columns are in (the same array as `stats["passive_set"]`, but published + independently of `stats`). + - `matrix_shape` -- `ZTZ.shape`, so a caller can check the factor it reads + belongs to the system it is asking about. + + Like `stats` this is purely observational: nothing in the solve reads `factor` + and the returned solution is byte-identical whether or not it is passed. """ from scipy import linalg as slg @@ -237,6 +258,17 @@ def fnnls_cholesky( stats["n_passive"] = int(P_inorder.size) stats["warm_start_errors"] = int(np.count_nonzero(P_initial_mask != P)) + if factor is not None: + # Hand the factor this solve already built to the caller. `U_buffer` is passed + # by reference deliberately: it is allocated inside this call and nothing here + # touches it after the return, so there is no copy to make (an (n, n) copy at + # n ~ 1500 would be 18 MB and ~5 ms, which is the whole saving the caller is + # after). Only `U_buffer[:k_active, :k_active]` is a valid factor. + factor["U_buffer"] = U_buffer + factor["k_active"] = int(k_active) + factor["passive_set"] = P_inorder.copy() + factor["matrix_shape"] = tuple(np.shape(ZTZ)) + return d diff --git a/test_autoarray/inversion/inversion/test_log_det_shared_factor.py b/test_autoarray/inversion/inversion/test_log_det_shared_factor.py new file mode 100644 index 000000000..6a4d09ec8 --- /dev/null +++ b/test_autoarray/inversion/inversion/test_log_det_shared_factor.py @@ -0,0 +1,569 @@ +""" +``log det(F + lambda*H)`` read off the NNLS Cholesky factor, on the NumPy path. + +One positive-only likelihood evaluation used to factorize ``F + lambda*H`` twice. +``fnnls_cholesky`` finishes the non-negative least squares solve holding a Cholesky factor of +``M[P][:, P]``, where ``M`` is ``curvature_reg_matrix_reduced`` and ``P`` is the solve's final +passive set -- 1485 of the 1500 columns on the production HST Delaunay ``pixels=1500`` +``AdaptSplit`` system -- and then threw it away; ``log_det_curvature_reg_matrix_term`` did a +fresh full Cholesky of the same matrix, ~40 ms of a ~270 ms NumPy likelihood call. + +``fnnls_cholesky`` now publishes that factor through an optional ``factor`` out-dict (the solve +itself is byte-identical -- nothing reads the dict), and the term reads the full log determinant +off it with the block-determinant identity + + log det M = log det M_PP + log det(M_AA - M_AP M_PP^-1 M_PA) + +(:func:`inversion_util.log_det_from_passive_cholesky_from`): a triangular solve with ``|A|`` +right-hand sides plus an ``|A| x |A|`` Cholesky, where ``A`` is the active set. + +These tests pin four things: the **identity** (the Schur route against a dense Cholesky of the +same matrix), the **factor contract** (what ``fnnls_cholesky`` publishes really is a factor of +the passive submatrix, and publishing it does not perturb the solution), the **end-to-end value +and the fact that the fast path is what produced it**, and every **fallback** -- because a +log determinant read off the wrong factor would be a silently wrong Bayesian evidence, not a +crash. + +On tolerances +------------- +Unlike the sparse regularization route (see ``test_log_det_sparse.py``, where two different +fp64 factorizations of an ill-conditioned matrix agree only to its conditioning), this is the +*same* factorization arithmetic rearranged, so it is pinned at ``rtol=1e-12`` throughout. The +production system measures <= 1.8e-12 nats from the dense ``potrf`` across six instances and a +bit-identical log evidence on the fiducial one. +""" + +import numpy as np +import pytest +import scipy.linalg + +import autoarray as aa + +from autoarray import fixtures +from autoarray.inversion.inversion import inversion_util +from autoarray.inversion.inversion.dataset_interface import DatasetInterface +from autoarray.util.fnnls import fnnls_cholesky + + +def _dense_log_det(matrix): + """The pre-lever computation: ``2 * sum(log(diag(cholesky(M))))``.""" + return 2.0 * np.sum(np.log(np.diag(np.linalg.cholesky(matrix)))) + + +def _random_spd(n, seed): + """A well-conditioned symmetric positive-definite matrix.""" + rng = np.random.default_rng(seed) + + matrix = rng.normal(size=(n, n)) + matrix = matrix @ matrix.T + n * np.eye(n) + + return 0.5 * (matrix + matrix.T) + + +def _factor_of_passive_set(matrix, passive_set): + """ + The buffer `fnnls_cholesky` would have published for this passive set: an ``(n, n)`` array + whose leading ``k x k`` upper triangle is `cholesky(matrix[P][:, P])` and whose every other + entry is zero, exactly as the solver's preallocated buffer is. + """ + n = matrix.shape[0] + + upper = scipy.linalg.cholesky(matrix[np.ix_(passive_set, passive_set)], lower=False) + + U_buffer = np.zeros((n, n)) + U_buffer[: upper.shape[0], : upper.shape[1]] = upper + + return U_buffer + + +# =================================================================== +# The identity: log_det_from_passive_cholesky_from vs a dense Cholesky +# =================================================================== + + +@pytest.mark.parametrize("n", [64, 256, 512]) +@pytest.mark.parametrize("n_active", [0, 1, 5, 40]) +def test__log_det_from_passive_cholesky_from__matches_dense_cholesky(n, n_active): + """ + The passive set is deliberately **neither sorted nor contiguous**: the active-set solver + appends and deletes columns, so the factor's row order is the order indices entered the + passive set and the off-diagonal block ``M_PA`` has to be taken in that same order. Sorting + ``P`` here would pass while hiding a route that only works for sorted sets. + """ + matrix = _random_spd(n=n, seed=n + n_active) + + rng = np.random.default_rng(n * 1000 + n_active) + + permutation = rng.permutation(n) + + passive_set = permutation[: n - n_active] + + log_det = inversion_util.log_det_from_passive_cholesky_from( + matrix=matrix, + U_buffer=_factor_of_passive_set(matrix=matrix, passive_set=passive_set), + k_active=passive_set.size, + passive_set=passive_set, + ) + + assert log_det == pytest.approx(_dense_log_det(matrix), rel=1.0e-12) + + +def test__log_det_from_passive_cholesky_from__empty_active_set__is_the_passive_factor_alone(): + """ + Every column passive: the published factor is already a factor of the whole matrix, so the + Schur term is skipped entirely (there is no ``|A| x |A|`` block to form). + """ + matrix = _random_spd(n=128, seed=11) + + passive_set = np.random.default_rng(11).permutation(128) + + U_buffer = _factor_of_passive_set(matrix=matrix, passive_set=passive_set) + + log_det = inversion_util.log_det_from_passive_cholesky_from( + matrix=matrix, + U_buffer=U_buffer, + k_active=128, + passive_set=passive_set, + ) + + assert log_det == pytest.approx( + 2.0 * np.sum(np.log(np.diag(U_buffer[:128, :128]))), rel=1.0e-12 + ) + assert log_det == pytest.approx(_dense_log_det(matrix), rel=1.0e-12) + + +def test__log_det_from_passive_cholesky_from__indefinite_schur_block__raises(): + """ + A matrix which is not positive-definite must raise, exactly as `np.linalg.cholesky` does on + it, so that the caller's test-mode guard and `FitException` resampling behave as they always + have. Returning 0.0 or a NaN here would put a fabricated number into the Bayesian evidence. + + The passive block is left positive-definite and the *active* block is broken, so it is the + Schur complement that fails rather than the published factor. + """ + matrix = _random_spd(n=64, seed=5) + + passive_set = np.arange(60) + + U_buffer = _factor_of_passive_set(matrix=matrix, passive_set=passive_set) + + matrix[60:, 60:] -= 10.0 * np.eye(4) * np.max(np.abs(matrix)) + + with pytest.raises(np.linalg.LinAlgError): + inversion_util.log_det_from_passive_cholesky_from( + matrix=matrix, + U_buffer=U_buffer, + k_active=passive_set.size, + passive_set=passive_set, + ) + + +# =================================================================== +# The factor contract: what fnnls_cholesky publishes +# =================================================================== + + +def _nnls_problem_with_negatives(n=40, n_data=90, seed=0): + """ + Normal equations whose unconstrained solution has negative components, so the solve's + passive set is a strict subset and the published factor is of a genuine submatrix. + """ + rng = np.random.default_rng(seed) + + Z = rng.normal(size=(n_data, n)) + x = Z @ rng.normal(size=n) + 3.0 * rng.normal(size=n_data) + + return Z.T @ Z, Z.T @ x + + +@pytest.mark.parametrize("seed", [0, 1, 2, 3]) +def test__fnnls_cholesky__published_factor_is_a_factor_of_the_passive_submatrix(seed): + """ + The real solver, not a reconstruction of it: `U_buffer[:k, :k]` transposed-squared must + reproduce `ZTZ[P][:, P]` in the published order, and `k_active` must be the passive count. + """ + ZTZ, ZTx = _nnls_problem_with_negatives(seed=seed) + + factor = {} + + reconstruction = fnnls_cholesky(ZTZ, ZTx, factor=factor) + + passive_set = factor["passive_set"] + k_active = factor["k_active"] + + # a strict subset, or the test is not exercising the thing it is about + assert 0 < k_active < ZTZ.shape[0] + assert passive_set.shape == (k_active,) + assert factor["matrix_shape"] == ZTZ.shape + + # the solution's support is exactly the passive set + assert set(np.flatnonzero(reconstruction).tolist()) == set(passive_set.tolist()) + + upper = factor["U_buffer"][:k_active, :k_active] + + assert upper.T @ upper == pytest.approx( + ZTZ[np.ix_(passive_set, passive_set)], abs=1.0e-10 + ) + + # the buffer outside the active corner is zero and meaningless, as documented + assert np.count_nonzero(np.tril(upper, k=-1)) == 0 + + +@pytest.mark.parametrize("seed", [0, 1, 2, 3]) +def test__fnnls_cholesky__publishing_the_factor_does_not_change_the_solution(seed): + """ + Bit-identical, not approximately equal: the `factor` out-dict adds a publish after the + solve and changes no arithmetic in it, and that is the property the whole lever rests on. + """ + ZTZ, ZTx = _nnls_problem_with_negatives(seed=seed) + + assert np.array_equal( + fnnls_cholesky(ZTZ, ZTx), + fnnls_cholesky(ZTZ, ZTx, factor={}), + ) + + +def test__reconstruction_positive_only_from__clears_a_stale_factor_on_entry(): + """ + A dict handed in carrying another matrix's factor must not survive the call, or a caller + could read a factor belonging to a system it is not asking about. + """ + ZTZ, ZTx = _nnls_problem_with_negatives(seed=0) + + factor = { + "U_buffer": np.eye(3), + "k_active": 3, + "passive_set": np.arange(3), + "matrix_shape": (3, 3), + "stale": True, + } + + inversion_util.reconstruction_positive_only_from( + data_vector=ZTx, + curvature_reg_matrix=ZTZ, + settings=aa.Settings(use_positive_only_solver=True), + fingerprint=None, + factor=factor, + ) + + assert "stale" not in factor + assert factor["matrix_shape"] == ZTZ.shape + + +# =================================================================== +# End to end on a real NumPy-path inversion +# =================================================================== + + +def _split_mapper(regularization, zeroed_pixels=0): + """ + The `delaunay_mapper_9_3x3` fixture's mesh, carrying a real split regularization scheme (the + production family) rather than `Constant`. + """ + source_plane_mesh_grid = aa.Grid2D.no_mask( + values=[ + [0.6, -0.3], + [0.5, -0.8], + [0.2, 0.1], + [0.0, 0.5], + [-0.3, -0.8], + [-0.6, -0.5], + [-0.4, -1.1], + [-1.2, 0.8], + [-1.5, 0.9], + ], + shape_native=(3, 3), + pixel_scales=1.0, + ) + + interpolator = aa.mesh.Delaunay( + pixels=9, zeroed_pixels=zeroed_pixels + ).interpolator_from( + source_plane_data_grid=fixtures.make_grid_2d_sub_2_7x7(), + source_plane_mesh_grid=source_plane_mesh_grid, + adapt_data=aa.Array2D.ones(shape_native=(3, 3), pixel_scales=0.1), + ) + + return aa.Mapper( + interpolator=interpolator, + regularization=regularization, + image_plane_mesh_grid=aa.Grid2D.uniform(shape_native=(3, 3), pixel_scales=0.1), + ) + + +def _dataset(seed): + """ + The `masked_imaging_7x7` fixture with seeded random data in place of its flat ones. + + The flat fixture's reconstruction is entirely positive, so its solve leaves an *empty* + active set and the Schur term is never formed. Signed data drives some source pixels to + zero, which is what puts columns in the active set -- the regime the fast path is for. + """ + dataset = fixtures.make_masked_imaging_7x7() + + data = aa.Array2D( + values=np.random.default_rng(seed).normal(0.0, 1.0, size=9), + mask=dataset.mask, + ) + + return DatasetInterface( + data=data, + noise_map=dataset.noise_map, + grids=dataset.grids, + psf=dataset.psf, + ) + + +def _counting_schur(monkeypatch): + """Count the fast path's calls without changing what it returns.""" + original = inversion_util.log_det_from_passive_cholesky_from + + calls = [] + + def _wrapped(**kwargs): + calls.append(kwargs) + return original(**kwargs) + + monkeypatch.setattr(inversion_util, "log_det_from_passive_cholesky_from", _wrapped) + + return calls + + +@pytest.mark.parametrize( + "regularization", + [ + aa.reg.ConstantSplit(coefficient=1.0), + aa.reg.AdaptSplit( + inner_coefficient=0.1, outer_coefficient=10.0, signal_scale=0.1 + ), + ], +) +@pytest.mark.parametrize("seed", [1, 2, 5]) +def test__log_det_curvature_reg_matrix_term__real_inversion__fast_path_matches_dense( + monkeypatch, regularization, seed +): + """ + The value, and the fact that the fast path is what produced it. Asserting only the value + would pass with the fast path silently never firing, which is exactly how this lever would + fail: it would be correct and do nothing. + """ + calls = _counting_schur(monkeypatch) + + inversion = aa.Inversion( + dataset=_dataset(seed), + linear_obj_list=[_split_mapper(regularization)], + settings=aa.Settings(use_positive_only_solver=True), + ) + + matrix = np.asarray(inversion.curvature_reg_matrix_reduced) + + term = inversion.log_det_curvature_reg_matrix_term + + assert len(calls) == 1 + assert term == pytest.approx(_dense_log_det(matrix), rel=1.0e-12) + + # the solve really did leave columns in the active set, so the Schur block was formed + k_active = inversion._nnls_factor["k_active"] + + assert 0 < k_active < matrix.shape[0] + assert np.count_nonzero(np.asarray(inversion.reconstruction)) == k_active + + +def test__log_det_curvature_reg_matrix_term__edge_zeroing_with_no_edge_pixels__fast_path_fires( + monkeypatch, +): + """ + The index-coverage trap. With `use_edge_zeroed_pixels` on and no edge pixels to zero, + `solve_ids_to_keep` returns `arange(n)` rather than `None` -- which is the configuration of + the production HST Delaunay system this lever was measured on. A coverage guard written as + "the solve took no subset" (`ids is None`) would decline every production evaluation and + save nothing, while still passing a value-only assertion. + """ + calls = _counting_schur(monkeypatch) + + inversion = aa.Inversion( + dataset=_dataset(2), + linear_obj_list=[_split_mapper(aa.reg.ConstantSplit(coefficient=1.0))], + settings=aa.Settings( + use_positive_only_solver=True, use_edge_zeroed_pixels=True + ), + ) + + assert np.array_equal(inversion.solve_ids_to_keep, np.arange(9)) + + matrix = np.asarray(inversion.curvature_reg_matrix_reduced) + + assert inversion.log_det_curvature_reg_matrix_term == pytest.approx( + _dense_log_det(matrix), rel=1.0e-12 + ) + assert len(calls) == 1 + + +def test__curvature_reg_matrix__is_cached__same_object_on_repeat_access(): + """ + Reached more than once per likelihood evaluation (the solve and the reduced matrix), and + each access re-ran an out-of-place `(n, n)` add -- ~2 ms per access at n = 1500. It was a + `cached_property` until the blanket `e819fa12` sweep. + """ + inversion = aa.Inversion( + dataset=_dataset(1), + linear_obj_list=[_split_mapper(aa.reg.ConstantSplit(coefficient=1.0))], + settings=aa.Settings(use_positive_only_solver=True), + ) + + matrix = inversion.curvature_reg_matrix + + assert inversion.curvature_reg_matrix is matrix + + # and it is the out-of-place sum, leaving the cached `curvature_matrix` intact + assert np.asarray(matrix) == pytest.approx( + np.asarray(inversion.curvature_matrix) + + np.asarray(inversion.regularization_matrix), + rel=1.0e-12, + ) + + +# =================================================================== +# Fallbacks: everything that must take the unchanged dense route +# =================================================================== + + +def test__log_det_curvature_reg_matrix_term__positive_negative_solver__dense_route( + monkeypatch, +): + """ + Route (c): no positive-only solve, so no factor exists. The term must not reach the fast + path -- and must not evaluate the reconstruction it would not otherwise need. + """ + calls = _counting_schur(monkeypatch) + + inversion = aa.Inversion( + dataset=_dataset(2), + linear_obj_list=[_split_mapper(aa.reg.ConstantSplit(coefficient=1.0))], + settings=aa.Settings(use_positive_only_solver=False), + ) + + matrix = np.asarray(inversion.curvature_reg_matrix_reduced) + + assert inversion.log_det_curvature_reg_matrix_term == pytest.approx( + _dense_log_det(matrix), rel=1.0e-12 + ) + assert calls == [] + assert inversion._nnls_factor is None + + +def test__log_det_curvature_reg_matrix_term__slogdet_method__unchanged(monkeypatch): + """The ``"slogdet"`` opt-in path is a different value by construction and is untouched.""" + calls = _counting_schur(monkeypatch) + + inversion = aa.Inversion( + dataset=_dataset(2), + linear_obj_list=[_split_mapper(aa.reg.ConstantSplit(coefficient=1.0))], + settings=aa.Settings(use_positive_only_solver=True, log_det_method="slogdet"), + ) + + matrix = np.asarray(inversion.curvature_reg_matrix_reduced) + + assert inversion.log_det_curvature_reg_matrix_term == pytest.approx( + np.linalg.slogdet(matrix)[1], rel=1.0e-12 + ) + assert calls == [] + + +def test__log_det_curvature_reg_matrix_term__fast_path_raises__falls_back_to_dense( + monkeypatch, +): + """ + A `LinAlgError` out of the Schur route means this matrix is at the edge of what fp64 + resolves as positive-definite. The dense factorization below is then the authority on it + (it either returns its own value or raises with the test-mode guard in place); the fast path + must never answer with a fabricated 0.0. + """ + + def raise_lin_alg_error(**kwargs): + raise np.linalg.LinAlgError("forced") + + monkeypatch.setattr( + inversion_util, "log_det_from_passive_cholesky_from", raise_lin_alg_error + ) + + inversion = aa.Inversion( + dataset=_dataset(2), + linear_obj_list=[_split_mapper(aa.reg.ConstantSplit(coefficient=1.0))], + settings=aa.Settings(use_positive_only_solver=True), + ) + + matrix = np.asarray(inversion.curvature_reg_matrix_reduced) + + assert inversion.log_det_curvature_reg_matrix_term == pytest.approx( + _dense_log_det(matrix), rel=1.0e-12 + ) + + +def test__factor_published_is_the_fallback_solve_s__when_a_memo_seed_raises(): + """ + A memo-seeded attempt that raises is retried from the dense-sign start + (`warm_start_fallback`), and it is the *retry's* factor that must be published: the failed + attempt's factor belongs to a solve whose result was discarded. + + The seeded attempt is forced to raise by patching the solver for the first call only, which + is how `reconstruction_positive_only_from`'s own retry is reached. + """ + from autoarray.inversion.inversion import nnls_memo + import autoarray.util.fnnls as fnnls_mod + + nnls_memo._nnls_passive_set_memo.clear() + + ZTZ, ZTx = _nnls_problem_with_negatives(seed=1) + + key = nnls_memo.memo_key(n=ZTx.shape[0], fingerprint="mesh") + + nnls_memo.passive_set_put( + key=key, passive_set=np.arange(ZTx.shape[0]), dense_error_fraction=0.5 + ) + + original = fnnls_mod.fnnls_cholesky + calls = [] + + def _raise_first( + ZTZ, ZTx, P_initial=np.zeros(0, dtype=int), stats=None, factor=None + ): + calls.append(P_initial) + + if len(calls) == 1: + raise np.linalg.LinAlgError("forced failure of the seeded attempt") + + return original(ZTZ, ZTx, P_initial, stats=stats, factor=factor) + + factor = {} + + try: + fnnls_mod.fnnls_cholesky = _raise_first + + reconstruction = inversion_util.reconstruction_positive_only_from( + data_vector=ZTx, + curvature_reg_matrix=ZTZ, + settings=aa.Settings( + use_positive_only_solver=True, nnls_warm_start_memo=True + ), + fingerprint="mesh", + factor=factor, + ) + finally: + fnnls_mod.fnnls_cholesky = original + nnls_memo._nnls_passive_set_memo.clear() + + # the seeded attempt raised and the dense-seed retry ran + assert len(calls) == 2 + + # the published factor is the retry's, and it describes the solution that was returned + assert factor["k_active"] == np.count_nonzero(reconstruction) + assert set(factor["passive_set"].tolist()) == set( + np.flatnonzero(reconstruction).tolist() + ) + + upper = factor["U_buffer"][: factor["k_active"], : factor["k_active"]] + + assert upper.T @ upper == pytest.approx( + ZTZ[np.ix_(factor["passive_set"], factor["passive_set"])], abs=1.0e-10 + ) diff --git a/test_autoarray/util/test_jax_nnls.py b/test_autoarray/util/test_jax_nnls.py index 8ee16ca74..9dac66580 100644 --- a/test_autoarray/util/test_jax_nnls.py +++ b/test_autoarray/util/test_jax_nnls.py @@ -226,9 +226,12 @@ def _solve_capturing_stats(monkeypatch, settings, fingerprint="mesh"): original = fnnls_mod.fnnls_cholesky captured = [] - def _wrapped(ZTZ, ZTx, P_initial=np.zeros(0, dtype=int), stats=None): + def _wrapped(ZTZ, ZTx, P_initial=np.zeros(0, dtype=int), stats=None, factor=None): + # `factor` is forwarded rather than swallowed: this double stands in for + # `fnnls_cholesky` and must publish everything the real one publishes, or the + # caller's log-determinant fast path silently loses its factor under the patch. captured.append(stats) - return original(ZTZ, ZTx, P_initial, stats=stats) + return original(ZTZ, ZTx, P_initial, stats=stats, factor=factor) monkeypatch.setattr(fnnls_mod, "fnnls_cholesky", _wrapped)