Merge pull request #392 from PyAutoLabs/feature/gradient-safe-logdet

feat: opt-in gradient-safe log-det via Settings (default unchanged)

Author: Jammy2211

autoarray/config/general.yaml

@@ -9,6 +9,7 @@ inversion:
   nnls_jacobi_preconditioning: true   # If True (default), the curvature matrix passed to jaxnnls.solve_nnls_primal is Jacobi-preconditioned (D Q D y = D q, x = D y). Fixes NaN backward-pass gradients on ill-conditioned Q and roughly halves forward solve time. Set False to restore the raw unpreconditioned solve.
   nnls_target_kappa: 1.0e-11          # Central-path relaxation parameter passed to jaxnnls.solve_nnls_primal. Larger values smooth the relaxed-KKT backward pass and prevent NaN gradients on ill-conditioned Q; smaller values tighten the primal solve. Verified finite gradients across all MGE/rectangular/delaunay pipelines (imaging + interferometer) with scale invariance over 5 orders of magnitude in noise. jaxnnls's own default (1e-3) is too aggressive for the backward pass.
   reconstruction_vmax_factor: 0.5     # Plots of an Inversion's reconstruction use the reconstructed data's bright value multiplied by this factor.
+  log_det_method: cholesky            # How the Bayesian-evidence log-determinant terms are computed. "cholesky" (default) is the historical 2*sum(log(diag(cholesky(M)))); "slogdet" uses logabsdet of slogdet(M), which is identical where M is positive-definite but finite (not NaN) where the Cholesky fails, for gradient-based searches (opt-in, non-default; does not change the default evidence). See PyAutoArray#391.
 numba:
   use_numba: true
   cache: true

autoarray/inversion/inversion/abstract.py

@@ -702,42 +702,61 @@ def regularization_term(self) -> float:
             ),
         )
 
-    @property
-    def log_det_curvature_reg_matrix_term(self) -> float:
+    def _log_det_symmetric_from(self, matrix) -> float:
         """
-        The log determinant of [F + reg_coeff*H] is used to determine the Bayesian evidence of the solution.
+        Log determinant of a symmetric positive-(semi)definite ``matrix``, used by the two evidence log-det terms.
+
+        The method is selected by ``self.settings.log_det_method`` (see :class:`Settings`):
 
-        This uses the Cholesky decomposition which is already computed before solving the reconstruction.
+        - ``"cholesky"`` (default) — ``2 * sum(log(diag(cholesky(matrix))))``. This is the historical
+          computation and is byte-for-byte unchanged, including the test-mode ``LinAlgError`` guard. On the
+          NumPy backend a non-positive-definite matrix raises; on the JAX backend the Cholesky returns NaN,
+          which propagates as a non-finite figure of merit and can stall gradient-based searches
+          (autolens_workspace_developer#104).
+        - ``"slogdet"`` — the ``logabsdet`` from ``xp.linalg.slogdet(matrix)``. Wherever ``matrix`` is
+          positive-definite this equals the Cholesky value exactly (the sign is +1), but where the Cholesky
+          would NaN it returns a finite, differentiable value instead, so it never stalls a gradient search.
+          It is **opt-in and non-default** — intended for gradient-based work and for comparison against the
+          Cholesky evidence, not as a replacement for it. See PyAutoArray#391.
         """
-        if not self.has(cls=AbstractRegularization):
-            return 0.0
+        if self.settings.log_det_method == "slogdet":
+            return self._xp.linalg.slogdet(matrix)[1]
 
         try:
             return 2.0 * self._xp.sum(
-                self._xp.log(
-                    self._xp.diag(
-                        self._xp.linalg.cholesky(self.curvature_reg_matrix_reduced)
-                    )
-                )
+                self._xp.log(self._xp.diag(self._xp.linalg.cholesky(matrix)))
             )
         except np.linalg.LinAlgError:
             if is_test_mode():
-                # Singular curvature-reg matrix from a fabricated test-mode model;
-                # this evidence term is discarded. Return a benign 0.0 so unguarded
-                # test-mode paths do not crash (matches the reconstruction guard in
-                # inversion_util). Normal runs re-raise unchanged. numpy-only: the
-                # JAX cholesky returns NaN rather than raising.
+                # Singular matrix from a fabricated test-mode model; this evidence term is discarded.
+                # Return a benign 0.0 so unguarded test-mode paths do not crash (matches the reconstruction
+                # guard in inversion_util). Normal runs re-raise unchanged. numpy-only: the JAX cholesky
+                # returns NaN rather than raising (which is exactly what "slogdet" exists to avoid).
                 return 0.0
             raise
 
+    @property
+    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`).
+        """
+        if not self.has(cls=AbstractRegularization):
+            return 0.0
+
+        return self._log_det_symmetric_from(self.curvature_reg_matrix_reduced)
+
     @property
     def log_det_regularization_matrix_term(self) -> float:
         """
         The Bayesian evidence of an inversion which quantifies its overall goodness-of-fit uses the log determinant
         of regularization matrix, Log[Det[Lambda*H]].
 
         Unlike the determinant of the curvature reg matrix, which uses an existing preloading Cholesky decomposition
-        used for the source reconstruction, this uses scipy sparse linear algebra to solve the determinant efficiently.
+        used for the source reconstruction, this uses scipy sparse linear algebra to solve the determinant efficiently
+        (or ``slogdet`` when ``Settings.log_det_method == "slogdet"`` — see :meth:`_log_det_symmetric_from`).
 
         Returns
         -------
@@ -747,21 +766,7 @@ def log_det_regularization_matrix_term(self) -> float:
         if not self.has(cls=AbstractRegularization):
             return 0.0
 
-        try:
-            return 2.0 * self._xp.sum(
-                self._xp.log(
-                    self._xp.diag(
-                        self._xp.linalg.cholesky(self.regularization_matrix_reduced)
-                    )
-                )
-            )
-        except np.linalg.LinAlgError:
-            if is_test_mode():
-                # Singular regularization matrix from a fabricated test-mode model;
-                # this evidence term is discarded. Benign 0.0 in test mode, unchanged
-                # (raise) in normal operation. numpy-only (JAX cholesky returns NaN).
-                return 0.0
-            raise
+        return self._log_det_symmetric_from(self.regularization_matrix_reduced)
 
     @property
     def reconstruction_noise_map_with_covariance(self) -> np.ndarray:

autoarray/settings.py

@@ -17,6 +17,7 @@ def __init__(
         no_regularization_add_to_curvature_diag_value: float = None,
         nnls_solver_tol: Optional[float] = None,
         nnls_max_iter: Optional[int] = None,
+        log_det_method: Optional[str] = None,
     ):
         """
         The settings of an Inversion, customizing how a linear set of equations are solved for.
@@ -95,6 +96,19 @@ def __init__(
             jaxnnls's own hard-coded cap of 50 (production HST pixelization+MGE systems converge in ~19-21
             iterations). Under `vmap` the solve runs until the slowest lane in the batch converges, so this
             also caps the worst-case batched cost. Only the JAX (`xp=jnp`) path honors this.
+        log_det_method
+            Which computation is used for the two Bayesian-evidence log-determinant terms
+            (``log_det_curvature_reg_matrix_term`` and ``log_det_regularization_matrix_term``). `None`
+            (default) reads the packaged value ``"cholesky"``.
+
+            - ``"cholesky"`` (default) — ``2 * sum(log(diag(cholesky(M))))``, the historical computation.
+              On a non-positive-definite matrix the NumPy backend raises and the JAX backend returns NaN.
+            - ``"slogdet"`` — the ``logabsdet`` of ``xp.linalg.slogdet(M)``. Where ``M`` is positive-definite
+              this equals the Cholesky value exactly; where the Cholesky would NaN it returns a finite,
+              differentiable value instead, so it never stalls a gradient-based search
+              (autolens_workspace_developer#104). This is an **opt-in, non-default** alternative intended for
+              gradient-based work and for comparison against the Cholesky evidence — it is not a replacement
+              for the default and the default evidence path is unchanged. See PyAutoArray#391.
         """
         self.use_mixed_precision = use_mixed_precision
         self.nnls_solver_tol = nnls_solver_tol
@@ -105,6 +119,7 @@ def __init__(
         self._no_regularization_add_to_curvature_diag_value = (
             no_regularization_add_to_curvature_diag_value
         )
+        self._log_det_method = log_det_method
 
     @property
     def use_positive_only_solver(self):
@@ -135,3 +150,10 @@ def no_regularization_add_to_curvature_diag_value(self):
             ]
 
         return self._no_regularization_add_to_curvature_diag_value
+
+    @property
+    def log_det_method(self):
+        if self._log_det_method is None:
+            return conf.instance["general"]["inversion"]["log_det_method"]
+
+        return self._log_det_method

test_autoarray/inversion/inversion/test_abstract.py

@@ -613,6 +613,69 @@ def test__determinant_of_positive_definite_matrix_via_cholesky__tridiagonal_matr
     )
 
 
+def test__log_det_method__default_is_cholesky_and_matches_numpy_log_det():
+    # The default log_det_method must reproduce the historical Cholesky value exactly.
+    assert aa.Settings().log_det_method == "cholesky"
+
+    matrix = np.array([[2.0, -1.0, 0.0], [-1.0, 2.0, -1.0], [0.0, -1.0, 2.0]])
+
+    inversion = aa.m.MockInversion(
+        linear_obj_list=[aa.m.MockLinearObj(regularization=aa.m.MockRegularization())],
+        curvature_reg_matrix=matrix,
+    )
+
+    assert inversion.log_det_curvature_reg_matrix_term == pytest.approx(
+        np.log(np.linalg.det(matrix)), 1.0e-4
+    )
+
+
+def test__log_det_method__slogdet_matches_cholesky_on_positive_definite_matrix():
+    # Where the matrix is positive-definite, "slogdet" is identical to "cholesky".
+    matrix = np.array([[2.0, -1.0, 0.0], [-1.0, 2.0, -1.0], [0.0, -1.0, 2.0]])
+
+    cholesky_inversion = aa.m.MockInversion(
+        linear_obj_list=[aa.m.MockLinearObj(regularization=aa.m.MockRegularization())],
+        curvature_reg_matrix=matrix,
+        regularization_matrix=matrix,
+    )
+    slogdet_inversion = aa.m.MockInversion(
+        linear_obj_list=[aa.m.MockLinearObj(regularization=aa.m.MockRegularization())],
+        curvature_reg_matrix=matrix,
+        regularization_matrix=matrix,
+        settings=aa.Settings(log_det_method="slogdet"),
+    )
+
+    expected = np.log(np.linalg.det(matrix))
+
+    assert cholesky_inversion.log_det_curvature_reg_matrix_term == pytest.approx(
+        expected, 1.0e-8
+    )
+    assert slogdet_inversion.log_det_curvature_reg_matrix_term == pytest.approx(
+        cholesky_inversion.log_det_curvature_reg_matrix_term, 1.0e-8
+    )
+    assert slogdet_inversion.log_det_regularization_matrix_term == pytest.approx(
+        cholesky_inversion.log_det_regularization_matrix_term, 1.0e-8
+    )
+
+
+def test__log_det_method__slogdet_is_finite_where_cholesky_fails_on_non_positive_definite():
+    # A symmetric matrix with a negative eigenvalue: Cholesky cannot factor it (raises /
+    # returns NaN), but slogdet returns the finite log|det|. This is the property that
+    # keeps a gradient-based search from stalling (autolens_workspace_developer#104).
+    matrix = np.array([[1.0, 2.0, 0.0], [2.0, 1.0, 0.0], [0.0, 0.0, 3.0]])
+    assert np.linalg.eigvalsh(matrix).min() < 0.0  # confirm it is not positive-definite
+
+    inversion = aa.m.MockInversion(
+        linear_obj_list=[aa.m.MockLinearObj(regularization=aa.m.MockRegularization())],
+        curvature_reg_matrix=matrix,
+        settings=aa.Settings(log_det_method="slogdet"),
+    )
+
+    result = inversion.log_det_curvature_reg_matrix_term
+    assert np.isfinite(result)
+    assert result == pytest.approx(np.linalg.slogdet(matrix)[1], 1.0e-8)
+
+
 def test__reconstruction_noise_map__asymmetric_curvature_reg_matrix__correct_diagonal_noise_values():
     curvature_reg_matrix = np.array([[1.0, 1.0, 1.0], [1.0, 2.0, 1.0], [1.0, 1.0, 3.0]])