Merge pull request #437 from PyAutoLabs/claude/automind-task-planning-gm4flt

Kernel regularization: opt-in Cholesky-solve evidence term + relative jitter

Author: Jammy2211

autoarray/config/general.yaml

@@ -10,6 +10,7 @@ inversion:
   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). Under "slogdet" the kernel regularization schemes (Matern/Gaussian/Exponential) also compute the regularization log-det analytically from a Cholesky of their covariance instead of factorizing the formed inverse. See PyAutoArray#391.
+  regularization_term_method: matmul  # How the Bayesian-evidence regularization term s^T H s is computed. "matmul" (default) is the historical s @ (H @ s) against the explicitly formed regularization matrix; "cho_solve" evaluates coefficient * s^T C^-1 s for the kernel schemes (Matern/Gaussian/Exponential/MaternAdapt) via one Cholesky solve of their covariance C, avoiding the explicit inverse whose round-off is amplified by cond(C) (~1e9 on clustered traced mesh vertices). Opt-in, non-default; does not change the default evidence. Schemes with no such factorization fall back to the formed matrix.
 numba:
   use_numba: true
   cache: true

autoarray/inversion/inversion/abstract.py

@@ -691,10 +691,50 @@ def regularization_term(self) -> float:
 
         The above works include the regularization_matrix coefficient (lambda) in this calculation. In PyAutoLens,
         this is already in the regularization matrix and thus implicitly included in the matrix multiplication.
+
+        Under ``regularization_term_method == "cho_solve"`` (opt-in, default off), regularization schemes
+        which know a factorization of their own matrix may instead supply their contribution directly via
+        :meth:`AbstractRegularization.regularization_term_from` — the kernel schemes (``MaternKernel`` etc.)
+        return ``coefficient * s^T C^-1 s`` from a single Cholesky solve of their covariance ``C``, avoiding
+        the round-off of contracting the explicitly formed inverse (whose error is amplified by ``cond(C)``,
+        ~1e9 on clustered traced mesh vertices). Because ``regularization_matrix_reduced`` is the block
+        diagonal of the per-object matrices when every linear object is regularized, the term is the sum of
+        the per-object terms; if any scheme has no shortcut (returns ``None``) the whole computation falls
+        back to the formed matrix. The default ``"matmul"`` path never consults the shortcut, so default
+        evidence values are unchanged.
+
+        Returns
+        -------
+        float
+            The regularization term of the inversion.
         """
         if not self.has(cls=AbstractRegularization):
             return 0.0
 
+        if (
+            self.settings.regularization_term_method == "cho_solve"
+            and self.all_linear_obj_have_regularization
+        ):
+            # `reconstruction_reduced` is the full reconstruction here (the guard above is exactly the
+            # no-reduction case), so the per-object slices index it directly.
+            reconstruction = self.reconstruction_reduced
+
+            term_list = [
+                regularization.regularization_term_from(
+                    linear_obj=linear_obj,
+                    reconstruction=reconstruction[param_range[0] : param_range[1]],
+                    xp=self._xp,
+                )
+                for linear_obj, regularization, param_range in zip(
+                    self.linear_obj_list,
+                    self.regularization_list,
+                    self.param_range_list_from(cls=LinearObj),
+                )
+            ]
+
+            if all(term is not None for term in term_list):
+                return sum(term_list)
+
         return self._xp.matmul(
             self.reconstruction_reduced.T,
             self._xp.matmul(

autoarray/inversion/regularization/abstract.py

@@ -213,3 +213,47 @@ def log_det_regularization_matrix_term_from(
         has no factorization-aware shortcut.
         """
         return None
+
+    def regularization_term_from(
+        self, linear_obj: LinearObj, reconstruction: np.ndarray, xp=np
+    ) -> Optional[float]:
+        """
+        Returns this scheme's contribution to the regularization term ``s^T H s``
+        computed from a factorization the scheme itself knows about, or ``None`` when
+        no such shortcut exists (the default).
+
+        This is the ``s^T H s`` counterpart of
+        :meth:`log_det_regularization_matrix_term_from`, and exists for the same
+        reason. The kernel regularization schemes build ``H = coefficient * C^-1``
+        from a dense covariance ``C``, so their term is
+        ``coefficient * s^T C^-1 s`` — obtainable from a single Cholesky *solve*
+        against ``s`` rather than by forming ``C^-1`` and contracting it. Forming the
+        explicit inverse carries round-off amplified by ``cond(C)`` (~1e9 on the
+        clustered traced vertices of the kNN mesh families), which then enters the
+        evidence through this term.
+
+        Note this cannot remove the explicit inverse from the inversion altogether:
+        ``curvature_reg_matrix`` is a dense ``F + H`` feeding the dense solve for the
+        reconstruction, so ``H`` is still formed there. This shortcut removes the
+        formed inverse from the *evidence* terms only.
+
+        The inversion consumes this ONLY when
+        ``Settings.regularization_term_method == "cho_solve"`` — the default
+        ``"matmul"`` path never calls it, so default likelihood values are unchanged.
+        See ``AbstractInversion.regularization_term``.
+
+        Parameters
+        ----------
+        linear_obj
+            The linear object (e.g. a ``Mapper``) whose regularization matrix the
+            term is of.
+        reconstruction
+            The reconstructed values ``s`` of this linear object's parameters (the
+            slice of the inversion's reconstruction belonging to ``linear_obj``).
+
+        Returns
+        -------
+        The scalar ``s^T H s`` for this linear object, or ``None`` when this scheme
+        has no factorization-aware shortcut.
+        """
+        return None

autoarray/inversion/regularization/exponential_kernel.py

@@ -12,6 +12,7 @@ def exp_cov_matrix_from(
     scale: float,
     pixel_points: np.ndarray,  # shape (N, 2)
     jitter: float = 1e-8,
+    jitter_relative: bool = False,
     xp=np,
 ) -> np.ndarray:  # shape (N, N)
     """
@@ -54,7 +55,9 @@ def exp_cov_matrix_from(
 
     # add a small jitter on the diagonal
     N = pts.shape[0]
-    cov = cov + xp.eye(N, dtype=cov.dtype) * jitter
+    from autoarray.inversion.regularization.matern_kernel import apply_jitter
+
+    cov = apply_jitter(cov, jitter=jitter, jitter_relative=jitter_relative, xp=xp)
 
     return cov
 
@@ -65,6 +68,7 @@ def __init__(
         coefficient: float = 1.0,
         scale: float = 1.0,
         jitter: Optional[float] = None,
+        jitter_relative: bool = False,
     ):
         """
         Regularization which uses an Exponential smoothing kernel to regularize the solution.
@@ -99,10 +103,17 @@ def __init__(
             ``None`` (default) uses the historical value 1e-8 — behaviour is identical
             to not having this parameter (it is a fixed setting, not a free model
             parameter, hence the ``None`` default).
+        jitter_relative
+            If ``True`` the jitter is applied *relative* to each pixel's own variance
+            (``C_ii *= 1 + jitter``) rather than as a fixed absolute ``jitter * I``.
+            ``False`` (default) preserves the historical behaviour exactly. The absolute
+            convention assumes ``C_ii ~ 1``, which holds for this unweighted kernel but not
+            for the adaptive one; see :func:`apply_jitter` for why and when to switch.
         """
         self.coefficient = coefficient
         self.scale = scale
         self.jitter = jitter
+        self.jitter_relative = jitter_relative
 
         super().__init__()
 
@@ -150,6 +161,7 @@ def regularization_matrix_from(self, linear_obj: LinearObj, xp=np) -> np.ndarray
             scale=self.scale,
             pixel_points=linear_obj.source_plane_mesh_grid.array,
             jitter=self.jitter_value,
+            jitter_relative=self.jitter_relative,
             xp=xp,
         )
 
@@ -173,6 +185,7 @@ def log_det_regularization_matrix_term_from(
             scale=self.scale,
             pixel_points=linear_obj.source_plane_mesh_grid.array,
             jitter=self.jitter_value,
+            jitter_relative=self.jitter_relative,
             xp=xp,
         )
 
@@ -181,3 +194,33 @@ def log_det_regularization_matrix_term_from(
         )
 
         return linear_obj.params * np.log(self.coefficient) - log_det_covariance
+
+    def regularization_term_from(
+        self, linear_obj: LinearObj, reconstruction: np.ndarray, xp=np
+    ) -> float:
+        """
+        The regularization term ``s^T H s`` from a single Cholesky solve of the kernel
+        covariance: ``H = coefficient * C^-1``, so
+        ``s^T H s = coefficient * s^T C^-1 s``, with the quadratic form evaluated by
+        solving ``C x = s`` rather than by forming ``C^-1``.
+
+        Consumed by the inversion only when
+        ``Settings.regularization_term_method == "cho_solve"`` (see
+        :meth:`AbstractRegularization.regularization_term_from`); the default
+        ``"matmul"`` path contracts the formed ``H`` and is unchanged.
+        """
+        from autoarray.inversion.regularization.matern_kernel import (
+            quadratic_form_via_cholesky,
+        )
+
+        covariance_matrix = exp_cov_matrix_from(
+            scale=self.scale,
+            pixel_points=linear_obj.source_plane_mesh_grid.array,
+            jitter=self.jitter_value,
+            jitter_relative=self.jitter_relative,
+            xp=xp,
+        )
+
+        return self.coefficient * quadratic_form_via_cholesky(
+            covariance_matrix, reconstruction, xp=xp
+        )

autoarray/inversion/regularization/gaussian_kernel.py

@@ -12,6 +12,7 @@ def gauss_cov_matrix_from(
     scale: float,
     pixel_points: np.ndarray,  # shape (N, 2)
     jitter: float = 1e-8,
+    jitter_relative: bool = False,
     xp=np,
 ) -> np.ndarray:
     """
@@ -46,7 +47,9 @@ def gauss_cov_matrix_from(
 
     # Add tiny jitter on the diagonal
     N = pts.shape[0]
-    cov = cov + xp.eye(N, dtype=cov.dtype) * jitter
+    from autoarray.inversion.regularization.matern_kernel import apply_jitter
+
+    cov = apply_jitter(cov, jitter=jitter, jitter_relative=jitter_relative, xp=xp)
 
     return cov
 
@@ -57,6 +60,7 @@ def __init__(
         coefficient: float = 1.0,
         scale: float = 1.0,
         jitter: Optional[float] = None,
+        jitter_relative: bool = False,
     ):
         """
         Regularization which uses a Gaussian smoothing kernel to regularize the solution.
@@ -89,10 +93,17 @@ def __init__(
             ``None`` (default) uses the historical value 1e-8 — behaviour is identical
             to not having this parameter (it is a fixed setting, not a free model
             parameter, hence the ``None`` default).
+        jitter_relative
+            If ``True`` the jitter is applied *relative* to each pixel's own variance
+            (``C_ii *= 1 + jitter``) rather than as a fixed absolute ``jitter * I``.
+            ``False`` (default) preserves the historical behaviour exactly. The absolute
+            convention assumes ``C_ii ~ 1``, which holds for this unweighted kernel but not
+            for the adaptive one; see :func:`apply_jitter` for why and when to switch.
         """
         self.coefficient = coefficient
         self.scale = scale
         self.jitter = jitter
+        self.jitter_relative = jitter_relative
         super().__init__()
 
     @property
@@ -139,6 +150,7 @@ def regularization_matrix_from(self, linear_obj: LinearObj, xp=np) -> np.ndarray
             scale=self.scale,
             pixel_points=linear_obj.source_plane_mesh_grid.array,
             jitter=self.jitter_value,
+            jitter_relative=self.jitter_relative,
             xp=xp,
         )
 
@@ -157,9 +169,10 @@ def regularization_matrix_from(self, linear_obj: LinearObj, xp=np) -> np.ndarray
         N = regularization_matrix.shape[0]
         diag_mean = xp.mean(xp.diag(regularization_matrix))
         h_jitter = 1e-8 * xp.abs(diag_mean)
-        regularization_matrix = regularization_matrix + xp.eye(
-            N, dtype=regularization_matrix.dtype
-        ) * h_jitter
+        regularization_matrix = (
+            regularization_matrix
+            + xp.eye(N, dtype=regularization_matrix.dtype) * h_jitter
+        )
 
         return regularization_matrix
 
@@ -185,6 +198,7 @@ def log_det_regularization_matrix_term_from(
             scale=self.scale,
             pixel_points=linear_obj.source_plane_mesh_grid.array,
             jitter=self.jitter_value,
+            jitter_relative=self.jitter_relative,
             xp=xp,
         )
 
@@ -193,3 +207,39 @@ def log_det_regularization_matrix_term_from(
         )
 
         return linear_obj.params * np.log(self.coefficient) - log_det_covariance
+
+    def regularization_term_from(
+        self, linear_obj: LinearObj, reconstruction: np.ndarray, xp=np
+    ) -> float:
+        """
+        The regularization term ``s^T H s`` from a single Cholesky solve of the kernel
+        covariance: ``H = coefficient * C^-1``, so
+        ``s^T H s = coefficient * s^T C^-1 s``, with the quadratic form evaluated by
+        solving ``C x = s`` rather than by forming ``C^-1``.
+
+        As with :meth:`log_det_regularization_matrix_term_from`, this is the term of
+        the analytic ``coefficient * C^-1``: it excludes both the symmetrisation and
+        the trace-scaled stabilisation jitter that :meth:`regularization_matrix_from`
+        applies to the formed matrix, since both exist only to guard the
+        factorization of the explicit inverse that this shortcut avoids entirely.
+
+        Consumed by the inversion only when
+        ``Settings.regularization_term_method == "cho_solve"`` (see
+        :meth:`AbstractRegularization.regularization_term_from`); the default
+        ``"matmul"`` path contracts the formed ``H`` and is unchanged.
+        """
+        from autoarray.inversion.regularization.matern_kernel import (
+            quadratic_form_via_cholesky,
+        )
+
+        covariance_matrix = gauss_cov_matrix_from(
+            scale=self.scale,
+            pixel_points=linear_obj.source_plane_mesh_grid.array,
+            jitter=self.jitter_value,
+            jitter_relative=self.jitter_relative,
+            xp=xp,
+        )
+
+        return self.coefficient * quadratic_form_via_cholesky(
+            covariance_matrix, reconstruction, xp=xp
+        )