From 8e68c5e65f75aa7045df667ad7a73246cab4b45c Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Fri, 17 Jul 2026 17:48:42 +0100 Subject: [PATCH] feat: opt-in gradient-safe log-det via Settings.log_det_method (default unchanged) The two Bayesian-evidence log-det terms (log_det_curvature_reg_matrix_term and log_det_regularization_matrix_term) use 2*sum(log(diag(cholesky(M)))), which returns NaN on the JAX backend (raises on numpy) when M is not numerically positive-definite -- stalling gradient-based searches on ill-conditioned regularization matrices (autolens_workspace_developer#104). Add Settings.log_det_method (Optional -> conf general.yaml fallback, mirroring the use_positive_only_solver pattern): - "cholesky" (default) -- the historical computation, byte-for-byte unchanged including the test-mode LinAlgError guard. Both terms route through a shared _log_det_symmetric_from helper; the default branch is the identical expression. - "slogdet" -- logabsdet of xp.linalg.slogdet(M): identical where M is positive-definite, finite and differentiable where the Cholesky would NaN. Opt-in, non-default, general (serves the adaptive schemes too); intended for gradient-based work and for comparison against the Cholesky evidence. Scope decided by an independent adversarial review of the reg-logdet investigation: the current evidence's coefficient-dependence is correct to machine precision, so the DEFAULT is deliberately not changed. This adds an opt-in alternative only. Full test_autoarray suite 926 passed (every existing log_evidence test unchanged = byte-identical default). New tests: default==cholesky==numpy log det; slogdet==cholesky on a PD matrix; slogdet finite where cholesky fails on a non-PD matrix. No workspace config mirroring needed -- packaged-only keys layer through production config (verified, like nnls_target_kappa). Closes #391. Co-Authored-By: Claude Opus 4.8 --- autoarray/config/general.yaml | 1 + autoarray/inversion/inversion/abstract.py | 69 ++++++++++--------- autoarray/settings.py | 22 ++++++ .../inversion/inversion/test_abstract.py | 63 +++++++++++++++++ 4 files changed, 123 insertions(+), 32 deletions(-) diff --git a/autoarray/config/general.yaml b/autoarray/config/general.yaml index 53f42b9e..50bcff31 100644 --- a/autoarray/config/general.yaml +++ b/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 diff --git a/autoarray/inversion/inversion/abstract.py b/autoarray/inversion/inversion/abstract.py index d6bef799..7cacd4c6 100644 --- a/autoarray/inversion/inversion/abstract.py +++ b/autoarray/inversion/inversion/abstract.py @@ -702,34 +702,52 @@ 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: """ @@ -737,7 +755,8 @@ def log_det_regularization_matrix_term(self) -> float: 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: diff --git a/autoarray/settings.py b/autoarray/settings.py index 76b8f20e..663097eb 100644 --- a/autoarray/settings.py +++ b/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 diff --git a/test_autoarray/inversion/inversion/test_abstract.py b/test_autoarray/inversion/inversion/test_abstract.py index b5ae9c57..992555ce 100644 --- a/test_autoarray/inversion/inversion/test_abstract.py +++ b/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]])