perf: reuse one Cholesky factor of F+H for the solve and the log det

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

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
+    )