From 02aa2e48dd0240bb597acedd7175ded268f412aa Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Thu, 15 Oct 2026 23:34:47 +0100 Subject: [PATCH] perf: reuse one Cholesky factor of F+H for the solve and the log det On the NumPy backend with use_positive_only_solver=False, a likelihood call factorized the same SPD matrix F + H twice: np.linalg.solve (LU) for the reconstruction and np.linalg.cholesky for log_det_curvature_reg_matrix_term. A new cached_property `_curvature_reg_matrix_cho_factor` computes one scipy.linalg.cho_factor (a copy; curvature_reg_matrix is not overwritten). `reconstruction` solves with cho_solve, and the log det reads 2 * sum(log(diag(L))) off the same factor when every linear object is regularized. The factor is None (unchanged routes) on JAX, with the positive-only solver, under log_det_method="slogdet", on LinAlgError, or when its diagonal is not finite. An indefinite-but-nonsingular F + H is therefore still solved by the LU and still makes the log det raise (or return 0.0 in test mode), and a singular one still makes the solve raise (or return its test-mode dummy). Tests pin all of these against the two-factorization values. Closes #633 Co-Authored-By: Claude Opus 5.5 --- autoarray/inversion/inversion/abstract.py | 83 +++++- .../test_positive_negative_shared_cholesky.py | 280 ++++++++++++++++++ 2 files changed, 362 insertions(+), 1 deletion(-) create mode 100644 test_autoarray/inversion/inversion/test_positive_negative_shared_cholesky.py diff --git a/autoarray/inversion/inversion/abstract.py b/autoarray/inversion/inversion/abstract.py index 771f999d0..94c40ca04 100644 --- a/autoarray/inversion/inversion/abstract.py +++ b/autoarray/inversion/inversion/abstract.py @@ -761,12 +761,80 @@ def reconstruction(self) -> np.ndarray: return reconstruction + # One Cholesky factor of `F + H`, shared with `log_det_curvature_reg_matrix_term`. `None` + # when the gate in `_curvature_reg_matrix_cho_factor` is closed or the factorization failed, + # in which case the unchanged LU solve below runs (with its test-mode guard). + cho_factor = self._curvature_reg_matrix_cho_factor + + if cho_factor is not None: + import scipy.linalg + + return scipy.linalg.cho_solve( + cho_factor, np.asarray(self.data_vector), check_finite=False + ) + return inversion_util.reconstruction_positive_negative_from( data_vector=self.data_vector, curvature_reg_matrix=self.curvature_reg_matrix, xp=self._xp, ) + @cached_property + def _curvature_reg_matrix_cho_factor(self): + """ + One Cholesky factor of ``F + reg_coeff*H`` for the positive-negative solve, shared by + :attr:`reconstruction` and :attr:`log_det_curvature_reg_matrix_term`. + + Without it, a NumPy positive-negative likelihood call factorizes the same symmetric + positive-definite matrix twice: ``np.linalg.solve`` (an LU) for the reconstruction and + ``np.linalg.cholesky`` for the log determinant. On the Aris SPT2319-55 interferometer setup + those two were 31% of a call at 1,405 source pixels and 72% at 2,947 (PyAutoArray#633, + campaign PyAutoPulse#62). One ``cho_factor`` serves both: ``cho_solve`` gives the + reconstruction and ``2 * sum(log(diag(L)))`` the log determinant. + + Returns the ``(c, lower)`` tuple of :func:`scipy.linalg.cho_factor`, or ``None`` -- meaning + "take the unchanged routes" -- when any of these holds: + + - the JAX backend (its solve and log determinant are traced and stay as they are); + - ``Settings.use_positive_only_solver`` (that solve builds and reuses its own factor); + - ``Settings.log_det_method != "cholesky"`` (the ``"slogdet"`` opt-in is untouched); + - the factorization raises ``LinAlgError`` (``scipy.linalg.LinAlgError`` is + ``numpy.linalg.LinAlgError``), i.e. the matrix is not positive-definite; + - the factor's diagonal is not finite (non-finite input). + + The ``None`` fallbacks keep failure behaviour identical to before: an indefinite but + nonsingular ``F + H`` is still solved by the LU and still makes the log determinant raise + (or return 0.0 in test mode) through :meth:`_log_det_symmetric_from`; a singular one still + makes the LU raise (or return its test-mode dummy). + + ``curvature_reg_matrix`` is a cached array that other properties read (the reduced matrix, + the covariance), so it is copied, never factorized in place. The factor is one further + ``(n, n)`` array held for the inversion's lifetime. + """ + if ( + self._xp is not np + or self.settings.use_positive_only_solver + or self.settings.log_det_method != "cholesky" + ): + return None + + import scipy.linalg + + try: + cho_factor = scipy.linalg.cho_factor( + np.asarray(self.curvature_reg_matrix), + lower=True, + overwrite_a=False, + check_finite=False, + ) + except np.linalg.LinAlgError: + return None + + if not np.all(np.isfinite(np.diag(cho_factor[0]))): + return None + + return cho_factor + @cached_property def reconstruction_reduced(self) -> np.ndarray: """ @@ -1019,7 +1087,9 @@ def log_det_curvature_reg_matrix_term(self) -> float: 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. + than factorizing ``F + reg_coeff*H`` a second time. On the NumPy path with the positive-negative solver it + is read off the Cholesky factor the reconstruction solved with (:attr:`_curvature_reg_matrix_cho_factor`), + for the same reason. ``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 @@ -1093,6 +1163,17 @@ def log_det_curvature_reg_matrix_term(self) -> float: # not be answered with a fabricated 0.0 here. pass + # Positive-negative NumPy path: read the log determinant off the Cholesky factor the + # reconstruction solves with (see `_curvature_reg_matrix_cho_factor`), rather than + # factorizing the same matrix again. Only when every linear object is regularized, so the + # reduced matrix is the full `curvature_reg_matrix` that was factorized. Reading the factor + # does not evaluate the reconstruction. + if self.all_linear_obj_have_regularization: + cho_factor = self._curvature_reg_matrix_cho_factor + + if cho_factor is not None: + return 2.0 * np.sum(np.log(np.diag(cho_factor[0]))) + return self._log_det_symmetric_from(self.curvature_reg_matrix_reduced) @staticmethod diff --git a/test_autoarray/inversion/inversion/test_positive_negative_shared_cholesky.py b/test_autoarray/inversion/inversion/test_positive_negative_shared_cholesky.py new file mode 100644 index 000000000..13228e3ee --- /dev/null +++ b/test_autoarray/inversion/inversion/test_positive_negative_shared_cholesky.py @@ -0,0 +1,280 @@ +""" +One Cholesky factor of ``F + lambda*H`` shared by the positive-negative solve and its log determinant. + +On the NumPy backend with ``use_positive_only_solver=False`` a likelihood call factorized the same +symmetric positive-definite matrix twice: ``np.linalg.solve`` (an LU) for ``reconstruction`` and +``np.linalg.cholesky`` for ``log_det_curvature_reg_matrix_term`` (PyAutoArray#633). Both now read +``AbstractInversion._curvature_reg_matrix_cho_factor``. + +These tests pin, against the two-factorization values computed here independently: + +- the reconstruction (rtol 1e-10) and the log determinant (rtol 1e-9), on imaging and + interferometer inversions, and that one factorization serves both; +- the failure semantics: an indefinite-but-nonsingular ``F + H`` is still solved by the LU and still + makes the log determinant raise (or return 0.0 in test mode); a singular one still makes the solve + raise (or return its test-mode dummy); +- the routes that must not use the factor: ``"slogdet"``, the positive-only solver, and a partially + regularized system's log determinant. +""" + +import numpy as np +import pytest +import scipy.linalg + +import autoarray as aa + +from autoarray import fixtures + + +def _two_factorization_values(inversion): + """The pre-#633 computation: an LU solve plus a separate dense Cholesky.""" + matrix = np.asarray(inversion.curvature_reg_matrix) + + reconstruction = np.linalg.solve(matrix, np.asarray(inversion.data_vector)) + + reduced = np.asarray(inversion.curvature_reg_matrix_reduced) + + log_det = 2.0 * np.sum(np.log(np.diag(np.linalg.cholesky(reduced)))) + + return reconstruction, log_det + + +def _count_cho_factor(monkeypatch): + original = scipy.linalg.cho_factor + + calls = [] + + def _wrapped(*args, **kwargs): + calls.append(1) + return original(*args, **kwargs) + + monkeypatch.setattr(scipy.linalg, "cho_factor", _wrapped) + + return calls + + +def _imaging_inversion(settings, linear_obj_list=None): + return aa.Inversion( + dataset=fixtures.make_masked_imaging_7x7(), + linear_obj_list=linear_obj_list or [fixtures.make_rectangular_mapper_7x7_3x3()], + settings=settings, + ) + + +def _interferometer_inversion(settings): + return aa.Inversion( + dataset=fixtures.make_interferometer_7(), + linear_obj_list=[fixtures.make_rectangular_mapper_7x7_3x3()], + settings=settings, + ) + + +def _with_matrix(inversion, matrix, data_vector): + """Seed the cached `curvature_reg_matrix` / `data_vector` with a chosen system.""" + inversion.__dict__["curvature_reg_matrix"] = matrix + inversion.__dict__["data_vector"] = data_vector + return inversion + + +@pytest.mark.parametrize( + "inversion_from", + [_imaging_inversion, _interferometer_inversion], + ids=["imaging", "interferometer"], +) +def test__shared_factor__matches_two_factorizations__one_factorization( + monkeypatch, inversion_from +): + calls = _count_cho_factor(monkeypatch) + + inversion = inversion_from(aa.Settings(use_positive_only_solver=False)) + + reconstruction_ref, log_det_ref = _two_factorization_values(inversion) + + # log det first: it must read the factor without evaluating the reconstruction. + log_det = inversion.log_det_curvature_reg_matrix_term + + assert "reconstruction" not in inversion.__dict__ + + reconstruction = np.asarray(inversion.reconstruction) + + assert len(calls) == 1 + assert inversion._curvature_reg_matrix_cho_factor is not None + + assert reconstruction == pytest.approx(reconstruction_ref, rel=1.0e-10, abs=0.0) + assert log_det == pytest.approx(log_det_ref, rel=1.0e-9) + + # the factor is a copy: `F + H` and `F` are left intact for the reduced matrix / covariance + assert np.asarray(inversion.curvature_reg_matrix) == pytest.approx( + np.asarray(inversion.curvature_matrix) + + np.asarray(inversion.regularization_matrix), + rel=1.0e-12, + ) + + +def test__shared_factor__random_spd_system__matches_solve_and_dense_cholesky(): + rng = np.random.default_rng(3) + + n = 9 # the rectangular 3x3 mapper's parameter count + + a = rng.normal(size=(n, n)) + matrix = a @ a.T + 1.0e-3 * np.eye(n) + data_vector = rng.normal(size=n) + + inversion = _with_matrix( + _imaging_inversion(aa.Settings(use_positive_only_solver=False)), + matrix=matrix, + data_vector=data_vector, + ) + + assert np.asarray(inversion.reconstruction) == pytest.approx( + np.linalg.solve(matrix, data_vector), rel=1.0e-10, abs=0.0 + ) + assert inversion.log_det_curvature_reg_matrix_term == pytest.approx( + 2.0 * np.sum(np.log(np.diag(np.linalg.cholesky(matrix)))), rel=1.0e-9 + ) + + +def _indefinite_nonsingular(n=9): + matrix = np.eye(n) + matrix[0, 0] = -1.0 + return matrix + + +def test__indefinite_matrix__lu_solve_kept__log_det_raises(monkeypatch): + monkeypatch.delenv("PYAUTO_TEST_MODE", raising=False) + + matrix = _indefinite_nonsingular() + data_vector = np.arange(1.0, 10.0) + + inversion = _with_matrix( + _imaging_inversion(aa.Settings(use_positive_only_solver=False)), + matrix=matrix, + data_vector=data_vector, + ) + + assert inversion._curvature_reg_matrix_cho_factor is None + + # as before: the LU solves an indefinite nonsingular system ... + assert np.asarray(inversion.reconstruction) == pytest.approx( + np.linalg.solve(matrix, data_vector), rel=1.0e-12 + ) + + # ... and the dense Cholesky log determinant rejects it + with pytest.raises(np.linalg.LinAlgError): + inversion.log_det_curvature_reg_matrix_term + + +def test__indefinite_matrix__test_mode__log_det_guard_returns_zero(monkeypatch): + monkeypatch.setenv("PYAUTO_TEST_MODE", "1") + + inversion = _with_matrix( + _imaging_inversion(aa.Settings(use_positive_only_solver=False)), + matrix=_indefinite_nonsingular(), + data_vector=np.ones(9), + ) + + assert inversion.log_det_curvature_reg_matrix_term == 0.0 + + +def test__singular_matrix__solve_raises__test_mode_dummy(monkeypatch): + matrix = np.ones((9, 9)) + + monkeypatch.delenv("PYAUTO_TEST_MODE", raising=False) + + inversion = _with_matrix( + _imaging_inversion(aa.Settings(use_positive_only_solver=False)), + matrix=matrix, + data_vector=np.arange(1.0, 10.0), + ) + + with pytest.raises(np.linalg.LinAlgError): + inversion.reconstruction + + monkeypatch.setenv("PYAUTO_TEST_MODE", "1") + + inversion = _with_matrix( + _imaging_inversion(aa.Settings(use_positive_only_solver=False)), + matrix=matrix, + data_vector=np.arange(1.0, 10.0), + ) + + assert (np.asarray(inversion.reconstruction) == 1.0).all() + + +def test__non_finite_matrix__no_factor__unchanged_routes(monkeypatch): + monkeypatch.delenv("PYAUTO_TEST_MODE", raising=False) + + matrix = np.eye(9) + matrix[4, 4] = np.nan + + inversion = _with_matrix( + _imaging_inversion(aa.Settings(use_positive_only_solver=False)), + matrix=matrix, + data_vector=np.ones(9), + ) + + assert inversion._curvature_reg_matrix_cho_factor is None + + +def test__slogdet_method__no_factor__unchanged(monkeypatch): + calls = _count_cho_factor(monkeypatch) + + inversion = _imaging_inversion( + aa.Settings(use_positive_only_solver=False, 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 np.asarray(inversion.reconstruction) == pytest.approx( + np.linalg.solve(matrix, np.asarray(inversion.data_vector)), rel=1.0e-12 + ) + assert inversion._curvature_reg_matrix_cho_factor is None + assert calls == [] + + +def test__positive_only_solver__no_factor(monkeypatch): + calls = _count_cho_factor(monkeypatch) + + inversion = _imaging_inversion(aa.Settings(use_positive_only_solver=True)) + + _ = inversion.reconstruction + _ = inversion.log_det_curvature_reg_matrix_term + + assert inversion._curvature_reg_matrix_cho_factor is None + assert calls == [] + + +def test__partially_regularized__solve_uses_factor__log_det_dense_reduced(): + dataset = fixtures.make_masked_imaging_7x7() + + grid = aa.Grid2D.from_mask(mask=dataset.mask) + + mapping_matrix = np.full(fill_value=0.5, shape=(9, 2)) + mapping_matrix[0, 0] = 0.8 + mapping_matrix[1, 1] = 0.4 + + linear_obj = aa.m.MockLinearObjFuncList( + parameters=2, grid=grid, mapping_matrix=mapping_matrix + ) + + inversion = _imaging_inversion( + aa.Settings(use_positive_only_solver=False), + linear_obj_list=[linear_obj, fixtures.make_rectangular_mapper_7x7_3x3()], + ) + + assert not inversion.all_linear_obj_have_regularization + + reconstruction_ref, log_det_ref = _two_factorization_values(inversion) + + assert np.asarray(inversion.reconstruction) == pytest.approx( + reconstruction_ref, rel=1.0e-10, abs=0.0 + ) + assert inversion._curvature_reg_matrix_cho_factor is not None + + # the factor is of the full matrix; the log det is of the mapper-only reduced matrix + assert inversion.log_det_curvature_reg_matrix_term == pytest.approx( + log_det_ref, rel=1.0e-9 + )