Merge pull request #469 from PyAutoLabs/claude/numerical-inversion-failures-7xsp1k

fix: form the reconstruction covariance via Cholesky, not an elementwise sqrt

Author: Jammy2211

autoarray/inversion/inversion/abstract.py

@@ -1,4 +1,5 @@
 import copy
+import warnings
 
 import numpy as np
 from typing import Dict, List, Optional, Type, Union
@@ -836,27 +837,99 @@ def log_det_regularization_matrix_term(self) -> float:
         return self._log_det_symmetric_from(self.regularization_matrix_reduced)
 
     @property
-    def reconstruction_noise_map_with_covariance(self) -> np.ndarray:
+    def reconstruction_covariance_matrix(self) -> np.ndarray:
         """
-        Returns the noise-map of the reconstruction as a two dimension matrix which accounts for the covariance
-        of the noise between pixels.
+        Returns the covariance matrix of the reconstruction, ``C = [F + reg_coeff*H]^-1``.
+
+        This is the inverse of the curvature matrix with regularization -- the same matrix used to solve for the
+        reconstruction via the linear inversion. Its diagonal holds the variance of each reconstructed pixel and
+        its off-diagonal entries the covariances between pixels; the off-diagonals are routinely negative, which is
+        a property of the covariance and not an error.
+
+        For the RMS standard deviation of each pixel (the quantity used for scientific analysis) use
+        `reconstruction_noise_map`, which takes the square root of this matrix's diagonal.
+
+        The inverse is formed from a Cholesky factorization rather than `np.linalg.inv`, for two reasons:
+
+        - `cho_factor` raises `LinAlgError` when the matrix is not positive-definite. `np.linalg.inv` raises only
+          on an exactly singular matrix, so an indefinite `curvature_reg_matrix` -- which the inversion does
+          encounter, hence `Settings.no_regularization_add_to_curvature_diag_value` -- previously returned a
+          plausible-looking covariance with no error and no warning. A covariance is only defined for a
+          positive-definite matrix, so failing here is correct and is handled by the callers in
+          `inversion/plot/inversion_plots.py`.
+        - `np.linalg.inv` is LU-based and exploits neither the symmetry nor the positive-definiteness this matrix
+          has. Its output drifts out of symmetry as conditioning worsens (measured at ~5e-7 absolute at
+          `cond ~ 1e12`, against ~3e-16 for the Cholesky solve), while a covariance matrix is symmetric by
+          definition.
+
+        Every failure mode raises `LinAlgError`, including a non-finite `curvature_reg_matrix`. That case is
+        checked explicitly because scipy would otherwise raise `ValueError`, which the callers -- here and
+        downstream -- do not catch; the previous implementation returned a silently all-`NaN` matrix instead.
+
+        The matrix is symmetrized on input. `cho_factor` reads only the upper triangle, so an asymmetric input
+        would otherwise be inverted as though its lower triangle matched, silently and with no diagnostic.
+        `curvature_reg_matrix` is `F + H` and symmetric by construction, so this is defensive only.
+
+        This property is NumPy-only: the input is coerced with `np.asarray`, so a JAX `curvature_reg_matrix`
+        forces a device-to-host transfer. It is a post-fit diagnostic, not part of the likelihood, so it is not on
+        the JIT path.
+
+        Returns
+        -------
+        The covariance matrix of the reconstruction, of shape [total_params, total_params].
+        """
+        from scipy.linalg import cho_factor, cho_solve
 
-        The diagonal of this matrix is the noise-map of the reconstruction, which can be used for analysing the
-        reconstruction with noise properties that are representative of the fit and therefore should be used
-        for any scientific analysis (e.g. source reconstructions of strong lenses).
+        matrix = np.asarray(self.curvature_reg_matrix)
 
-        This noise-map is defined as the RMS standard deviation of the noise in every pixel of the reconstruction.
-        This definition is identical to the `noise_map` attributes of dataset objects.
+        if not np.isfinite(matrix).all():
+            raise np.linalg.LinAlgError(
+                "The curvature_reg_matrix contains non-finite entries (NaN or inf), so the reconstruction "
+                "covariance is undefined. Raised as LinAlgError so the plotting and CSV callers, which guard "
+                "on LinAlgError, degrade gracefully rather than aborting the model-fit."
+            )
+
+        # cho_factor reads only the upper triangle; symmetrize so an asymmetric input cannot be silently
+        # inverted as though its lower triangle matched its upper.
+        matrix = 0.5 * (matrix + matrix.T)
+
+        covariance = cho_solve(
+            cho_factor(matrix, check_finite=False),
+            np.eye(matrix.shape[0], dtype=matrix.dtype),
+            check_finite=False,
+        )
 
-        It is computed as the square root of the inverse of the curvature matrix with regularization, which is the
-        same matrix used to solve for the reconstruction via the linear inversion.
+        # cho_solve is accurate but not bitwise symmetric; a covariance matrix is symmetric by definition.
+        return 0.5 * (covariance + covariance.T)
+
+    @property
+    def reconstruction_noise_map_with_covariance(self) -> np.ndarray:
+        """
+        Deprecated alias of `reconstruction_covariance_matrix`.
+
+        This property previously returned ``np.sqrt(np.linalg.inv(curvature_reg_matrix))`` -- an elementwise square
+        root of the whole covariance matrix. Because the off-diagonal entries of a covariance matrix are routinely
+        negative, every such entry was `NaN` by construction, for any input matrix however well-conditioned, and
+        each call emitted `RuntimeWarning: invalid value encountered in sqrt`.
+
+        It now returns the covariance matrix itself, so the values differ: the diagonal holds variances rather
+        than standard deviations, and the off-diagonals hold covariances rather than `NaN`.
 
         Returns
         -------
-        The noise-map of the reconstruction as a two dimension matrix which accounts for the covariance of the noise
-        between pixels.
+        The covariance matrix of the reconstruction (see `reconstruction_covariance_matrix`).
         """
-        return np.sqrt(np.linalg.inv(self.curvature_reg_matrix))
+        warnings.warn(
+            "`reconstruction_noise_map_with_covariance` is deprecated; use "
+            "`reconstruction_covariance_matrix` instead. Note the values have changed: it now returns the "
+            "covariance matrix, so its diagonal holds variances rather than standard deviations (the previous "
+            "elementwise square root made every off-diagonal NaN). For the RMS noise of each pixel use "
+            "`reconstruction_noise_map`.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+
+        return self.reconstruction_covariance_matrix
 
     @property
     def reconstruction_noise_map(self):
@@ -870,15 +943,21 @@ def reconstruction_noise_map(self):
         The noise-map of the reconstruction is the RMS standard deviation of the noise in every pixel of the
         reconstruction. This definition is identical to the `noise_map` attributes of dataset objects.
 
-        It is computed as the square root of the diagonal of the `reconstruction_noise_map_with_covariance` matrix,
-        which is the same matrix used to solve for the reconstruction via the linear inversion.
+        It is computed as the square root of the diagonal of `reconstruction_covariance_matrix`, which is the
+        inverse of the same matrix used to solve for the reconstruction via the linear inversion.
+
+        This previously took the diagonal of an elementwise-square-rooted matrix. The two are algebraically
+        identical -- `np.sqrt` is elementwise, so it commutes with taking the diagonal -- but only numerically
+        equivalent, since the covariance is now formed by Cholesky rather than LU. The difference is
+        conditioning-limited roundoff, measured at ~7e-15 relative at `cond ~ 1e3` rising to ~4e-5 at
+        `cond ~ 1e13`; neither result is the more correct one.
 
         Returns
         -------
         The noise-map of the reconstruction as a one dimensional ndarray, which does not account for the covariance
         of the noise between pixels.
         """
-        return np.diagonal(self.reconstruction_noise_map_with_covariance)
+        return np.sqrt(np.diag(self.reconstruction_covariance_matrix))
 
     @property
     def reconstruction_noise_map_dict(self) -> Dict[LinearObj, np.ndarray]:

test_autoarray/inversion/inversion/test_abstract.py

@@ -1,3 +1,5 @@
+import warnings
+
 import numpy as np
 import pytest
 
@@ -676,19 +678,143 @@ def test__log_det_method__slogdet_is_finite_where_cholesky_fails_on_non_positive
     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():
+def test__reconstruction_noise_map__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]])
 
     inversion = aa.m.MockInversion(curvature_reg_matrix=curvature_reg_matrix)
 
-    assert inversion.reconstruction_noise_map_with_covariance[0, 0] == pytest.approx(
-        np.sqrt(2.5), 1.0e-2
-    )
+    assert inversion.reconstruction_covariance_matrix[0, 0] == pytest.approx(2.5, 1.0e-2)
     assert inversion.reconstruction_noise_map == pytest.approx(
         np.sqrt(np.array([2.5, 1.0, 0.5])), 1.0e-3
     )
 
 
+def test__reconstruction_covariance_matrix__off_diagonals_are_finite_and_negative():
+    """
+    The off-diagonal entries of a covariance matrix are covariances and are routinely negative.
+
+    `reconstruction_covariance_matrix` previously applied `np.sqrt` elementwise to the whole inverse, so every
+    negative off-diagonal became NaN by construction -- for any matrix, however well-conditioned -- while
+    emitting `RuntimeWarning: invalid value encountered in sqrt`. Only the [0, 0] diagonal element was asserted,
+    so nothing caught it.
+    """
+    curvature_reg_matrix = np.array([[1.0, 1.0, 1.0], [1.0, 2.0, 1.0], [1.0, 1.0, 3.0]])
+
+    inversion = aa.m.MockInversion(curvature_reg_matrix=curvature_reg_matrix)
+
+    with warnings.catch_warnings():
+        warnings.simplefilter("error", RuntimeWarning)
+        covariance = inversion.reconstruction_covariance_matrix
+
+    assert np.all(np.isfinite(covariance))
+
+    # this matrix has anti-correlated pixels, so the off-diagonals are genuinely negative
+    assert covariance[0, 1] < 0.0
+    assert covariance == pytest.approx(np.linalg.inv(curvature_reg_matrix), 1.0e-8)
+
+
+def test__reconstruction_covariance_matrix__is_accurate_and_symmetric_when_ill_conditioned():
+    """
+    Ground truth is exact by construction: for `A = Q diag(w) Q.T` the inverse is `Q diag(1/w) Q.T`.
+
+    The symmetry half only guards the symmetrization line -- `0.5 * (C + C.T)` is bitwise symmetric for any C --
+    so the accuracy assertion against the constructed truth is what tests the factorization itself.
+    """
+    rng = np.random.default_rng(1234)
+    q, _ = np.linalg.qr(rng.standard_normal((25, 25)))
+    eigenvalues = np.logspace(0, 9, 25)
+
+    curvature_reg_matrix = (q * eigenvalues) @ q.T
+    curvature_reg_matrix = 0.5 * (curvature_reg_matrix + curvature_reg_matrix.T)
+
+    covariance_true = (q * (1.0 / eigenvalues)) @ q.T
+
+    inversion = aa.m.MockInversion(curvature_reg_matrix=curvature_reg_matrix)
+
+    covariance = inversion.reconstruction_covariance_matrix
+
+    # cond ~ 1e9, so the achievable accuracy is eps * cond ~ 2e-7; the measured error is ~3e-9. This is not a
+    # claim that Cholesky beats LU here -- it does not, `np.linalg.inv` measures ~7e-10 on this matrix.
+    assert covariance == pytest.approx(covariance_true, abs=1.0e-7)
+    assert covariance == pytest.approx(covariance.T, abs=1.0e-15)
+
+
+def test__reconstruction_covariance_matrix__asymmetric_input_is_symmetrized_not_silently_upper_triangle():
+    """
+    `cho_factor` reads only the upper triangle, so an asymmetric input would be inverted as though its lower
+    triangle matched its upper -- silently, and differing from the true inverse.
+    """
+    curvature_reg_matrix = np.array([[2.0, 0.5], [0.1, 2.0]])
+    symmetrized = 0.5 * (curvature_reg_matrix + curvature_reg_matrix.T)
+
+    inversion = aa.m.MockInversion(curvature_reg_matrix=curvature_reg_matrix)
+
+    assert inversion.reconstruction_covariance_matrix == pytest.approx(
+        np.linalg.inv(symmetrized), 1.0e-8
+    )
+
+
+def test__reconstruction_covariance_matrix__non_finite_matrix_raises_lin_alg_error():
+    """
+    scipy raises `ValueError` on a non-finite matrix, which the plotting and CSV callers do not catch -- they
+    guard on `LinAlgError`. The CSV writer explicitly promises not to abort the enclosing model-fit, so the
+    non-finite case is converted rather than allowed to escape.
+    """
+    curvature_reg_matrix = np.array([[1.0, np.nan], [np.nan, 2.0]])
+
+    inversion = aa.m.MockInversion(curvature_reg_matrix=curvature_reg_matrix)
+
+    with pytest.raises(np.linalg.LinAlgError, match="non-finite"):
+        inversion.reconstruction_covariance_matrix
+
+
+def test__reconstruction_noise_map__is_sqrt_of_covariance_diagonal():
+    """
+    The invariant, asserted directly rather than via hand-computed values.
+
+    `reconstruction_noise_map` used to be `np.diagonal(...)` of an already-square-rooted matrix, which was
+    correct only incidentally -- because `np.sqrt` is elementwise. It now takes the square root of the
+    covariance diagonal itself, so the relationship is stated rather than emergent.
+    """
+    curvature_reg_matrix = np.array([[1.0, 1.0, 1.0], [1.0, 2.0, 1.0], [1.0, 1.0, 3.0]])
+
+    inversion = aa.m.MockInversion(curvature_reg_matrix=curvature_reg_matrix)
+
+    assert inversion.reconstruction_noise_map == pytest.approx(
+        np.sqrt(np.diag(inversion.reconstruction_covariance_matrix)), 1.0e-12
+    )
+
+
+def test__reconstruction_covariance_matrix__raises_on_a_non_positive_definite_matrix():
+    """
+    A covariance is only defined for a positive-definite matrix.
+
+    `np.linalg.inv` raises only on an exactly singular matrix, so an indefinite `curvature_reg_matrix` returned
+    a plausible-looking covariance with no error and no warning. The Cholesky factorization rejects it, and the
+    plotting and CSV callers already catch `LinAlgError`.
+    """
+    # symmetric, non-singular, but indefinite (eigenvalues +1 and -1)
+    curvature_reg_matrix = np.array([[0.0, 1.0], [1.0, 0.0]])
+
+    inversion = aa.m.MockInversion(curvature_reg_matrix=curvature_reg_matrix)
+
+    assert np.isfinite(np.linalg.inv(curvature_reg_matrix)).all()  # inv is silent here
+
+    with pytest.raises(np.linalg.LinAlgError):
+        inversion.reconstruction_covariance_matrix
+
+
+def test__reconstruction_noise_map_with_covariance__is_deprecated_alias():
+    curvature_reg_matrix = np.array([[1.0, 1.0, 1.0], [1.0, 2.0, 1.0], [1.0, 1.0, 3.0]])
+
+    inversion = aa.m.MockInversion(curvature_reg_matrix=curvature_reg_matrix)
+
+    with pytest.warns(DeprecationWarning, match="reconstruction_covariance_matrix"):
+        covariance = inversion.reconstruction_noise_map_with_covariance
+
+    assert covariance == pytest.approx(inversion.reconstruction_covariance_matrix, 1.0e-12)
+
+
 def test__max_pixel_list_from_and_centre__returns_top_pixels_and_brightest_centre():
 
     source_plane_mesh_grid = aa.Grid2DIrregular(

test_autoarray/inversion/plot/test_inversion_plotters.py

@@ -79,8 +79,8 @@ def test__inversion_subplot_of_mapper__singular_curvature_reg_matrix(
 
     monkeypatch.setattr(
         type(inversion),
-        "reconstruction_noise_map_with_covariance",
-        property(lambda self: np.sqrt(np.linalg.inv(np.zeros((params, params))))),
+        "reconstruction_covariance_matrix",
+        property(lambda self: np.linalg.inv(np.zeros((params, params)))),
     )
 
     with pytest.raises(np.linalg.LinAlgError):
@@ -107,8 +107,8 @@ def test__save_reconstruction_csv__singular_curvature_reg_matrix(
 
     monkeypatch.setattr(
         type(inversion),
-        "reconstruction_noise_map_with_covariance",
-        property(lambda self: np.sqrt(np.linalg.inv(np.zeros((params, params))))),
+        "reconstruction_covariance_matrix",
+        property(lambda self: np.linalg.inv(np.zeros((params, params)))),
     )
 
     with pytest.raises(np.linalg.LinAlgError):