Merge pull request #512 from PyAutoLabs/feature/adapt-linear-regularization

feat: AdaptPower regularization — power input + factor-2 scatter fix

Author: James Nightingale

autoarray/inversion/regularization/__init__.py

@@ -4,12 +4,16 @@
 from .constant_zeroth import ConstantZeroth
 from .constant_split import ConstantSplit
 from .adapt import Adapt
+from .adapt_power import AdaptPower
 from .adapt_split import AdaptSplit
+from .adapt_split_power import AdaptSplitPower
 from .brightness_zeroth import BrightnessZeroth
 from .adapt_split_zeroth import AdaptSplitZeroth
+from .adapt_split_zeroth_power import AdaptSplitZerothPower
 from .curvature_mask import CurvatureMask
 from .fourth_order_mask import FourthOrderMask
 from .gaussian_kernel import GaussianKernel
 from .exponential_kernel import ExponentialKernel
 from .matern_kernel import MaternKernel
 from .matern_adapt_kernel import MaternAdaptKernel
+from .matern_adapt_power_kernel import MaternAdaptPowerKernel

autoarray/inversion/regularization/adapt.py

@@ -10,7 +10,10 @@
 
 
 def adapt_regularization_weights_from(
-    inner_coefficient: float, outer_coefficient: float, pixel_signals: np.ndarray
+    inner_coefficient: float,
+    outer_coefficient: float,
+    pixel_signals: np.ndarray,
+    power: float = 2.0,
 ) -> np.ndarray:
     """
     Returns the regularization weights for the adaptive regularization scheme (e.g. ``Adapt``).
@@ -36,6 +39,14 @@ def adapt_regularization_weights_from(
     pixel_signals
         The estimated signal in every pixelization pixel, used to change the regularization weighting of high signal
         and low signal pixelizations.
+    power
+        The exponent the interpolated coefficient is raised to. The matrix builders square the returned weights
+        again, so the coefficient enters the regularization matrix at the power ``2 * power``.
+
+        The default ``2.0`` is the historical ``Adapt`` convention (a fourth-power coefficient dependence) and is
+        what the legacy ``Adapt``, ``AdaptSplit``, ``AdaptSplitZeroth`` and ``MaternAdaptKernel`` classes pass.
+        ``power=1.0`` gives the squared-once convention shared with ``Constant`` and is what the ``*Power``
+        classes (e.g. ``AdaptPower``) pass by default.
 
     Returns
     -------
@@ -45,7 +56,7 @@ def adapt_regularization_weights_from(
     """
     return (
         inner_coefficient * pixel_signals + outer_coefficient * (1.0 - pixel_signals)
-    ) ** 2.0
+    ) ** power
 
 
 def weighted_regularization_matrix_from(
@@ -129,6 +140,91 @@ def weighted_regularization_matrix_from(
     return mat[:S, :S]
 
 
+def weighted_regularization_matrix_single_scatter_from(
+    regularization_weights: np.ndarray,
+    neighbors: np.ndarray,
+    xp=np,
+) -> np.ndarray:
+    """
+    Returns the regularization matrix of the adaptive regularization scheme, scattering every mesh edge
+    **once** so that uniform weights reproduce ``Constant`` regularization exactly.
+
+    This is the corrected sibling of ``weighted_regularization_matrix_from``, used by the ``*Power``
+    classes (e.g. ``AdaptPower``). The two differ only in how often each mesh edge is scattered:
+
+    - ``weighted_regularization_matrix_from`` adds every ordered neighbor pair to **both** the
+      ``(i, j)`` and ``(j, i)`` entries. Because the neighbor list already holds each unordered edge
+      twice (once in row ``i``, once in row ``j``), every edge lands four times, which is exactly twice
+      what ``constant_regularization_matrix_from`` does. ``Adapt(inner=outer=c)`` is therefore
+      ``2 x`` ``Constant(c)``, not equal to it.
+    - This function adds every ordered neighbor pair once, to ``(i, i)`` and ``(i, j)`` only -- the same
+      bookkeeping ``Constant`` uses -- so ``AdaptPower(inner=outer=c)`` equals ``Constant(c)`` exactly.
+
+    The edge weight is the mean of the two endpoints' squared regularization weights,
+    ``0.5 * (w_i ** 2 + w_j ** 2)``. This is symmetric in ``(i, j)`` (so the matrix is symmetric even
+    for wildly varying adaptive weights), reduces to ``w ** 2`` when the weights are uniform, and makes
+    the result a weighted graph Laplacian plus a ``1e-8`` diagonal floor -- so its rows sum to the floor
+    and it is positive semi-definite by construction.
+
+    The legacy builder is left untouched: halving it in place would silently change the effective
+    regularization of every ``Adapt`` fit ever run.
+
+    Parameters
+    ----------
+    regularization_weights
+        The regularization weight of each pixel, adaptively governing the degree of gradient regularization
+        applied to each inversion parameter (e.g. mesh pixels of a ``Mapper``).
+    neighbors
+        An array of length (total_pixels) which provides the index of all neighbors of every pixel in
+        the mesh grid (entries of -1 correspond to no neighbor).
+
+    Returns
+    -------
+    np.ndarray
+        The regularization matrix computed using an adaptive regularization scheme where the effective
+        regularization coefficient of every source pixel is different.
+    """
+    S, P = neighbors.shape
+
+    reg_w = regularization_weights**2
+
+    # 1) Flatten the (i->j) neighbor pairs
+    I = xp.repeat(xp.arange(S), P)  # (S*P,)
+    J_raw = neighbors.reshape(-1)  # (S*P,)
+
+    # 2) Remap "no neighbor" entries to an extra slot S, whose weight = 0
+    OUT = S
+    valid = J_raw >= 0
+    J = xp.where(valid, J_raw, OUT)
+
+    # 3) Build an extended weight vector with a zero at index S
+    reg_w_ext = xp.concatenate([reg_w, xp.zeros((1,))], axis=0)
+
+    # 4) Symmetric per-edge weight: the mean of the two endpoints' squared weights. Padded entries are
+    #    masked to zero (their I endpoint is a real pixel, so the mean alone would not vanish).
+    w_ij = 0.5 * (reg_w_ext[I] + reg_w_ext[J]) * valid
+
+    # 5) Start with zeros on an (S+1)x(S+1) canvas so we can scatter into row S safely
+    mat = xp.zeros((S + 1, S + 1), dtype=regularization_weights.dtype)
+
+    diag_updates_i = xp.concatenate(
+        [xp.full((S,), 1e-8), xp.zeros((1,))], axis=0  # out-of-bounds slot stays zero
+    )
+
+    # 6) Scatter each ordered pair exactly once: onto the diagonal of i and the (i, j) off-diagonal
+    if xp.__name__.startswith("jax"):
+        mat = mat.at[xp.diag_indices(S + 1)].add(diag_updates_i)
+        mat = mat.at[I, I].add(w_ij)
+        mat = mat.at[I, J].add(-w_ij)
+    else:
+        np.add.at(mat, np.diag_indices(S + 1), diag_updates_i)
+        np.add.at(mat, (I, I), w_ij)
+        np.add.at(mat, (I, J), -w_ij)
+
+    # 7) Drop the extra row/column S and return the SxS result
+    return mat[:S, :S]
+
+
 class Adapt(AbstractRegularization):
     def __init__(
         self,
@@ -180,7 +276,28 @@ def __init__(
         neighbors come from a direct scipy call on the traced mesh grid (use
         ``AdaptSplit`` there). Note the defaults
         ``inner_coefficient == outer_coefficient == 1.0`` make the weighting
-        uniform — numerically identical to ``Constant(coefficient=1.0)``.
+        uniform — but *not* numerically identical to ``Constant(coefficient=1.0)``;
+        see the coefficient-convention note below.
+
+        **Coefficient convention (legacy, ``lambda^4``).** The coefficients are squared twice before they
+        reach the regularization matrix -- once by ``adapt_regularization_weights_from`` and once by the
+        matrix builder -- so the matrix scales as the *fourth* power of the coefficient, while
+        ``Constant`` scales as the second. Both carry the same ``LogUniform(1e-6, 1e6)`` prior, so this
+        scheme explores a far wider effective smoothing range and reaches a numerically non
+        positive-definite matrix from ``c ~ 1e4`` where ``Constant`` survives to ``c ~ 1e6``.
+
+        **It is also 2x ``Constant``, not equal to it.** The matrix builder scatters every mesh edge in
+        both directions, and the neighbor list already holds each unordered edge twice, so each edge
+        lands four times where ``Constant`` lands it twice.
+        ``Adapt(inner_coefficient=1.0, outer_coefficient=1.0)`` is therefore exactly ``2 x``
+        ``Constant(coefficient=1.0)``.
+
+        This behaviour is preserved deliberately: changing it would alter the coefficient scale of every
+        adaptive fit ever run. **New work should prefer ``AdaptPower``**, which takes a ``power`` argument
+        (default ``1.0``, giving the ``Constant``-matching ``lambda^2`` convention) and scatters each edge
+        once, so ``AdaptPower(inner=outer=c)`` equals ``Constant(c)`` exactly and is more robust to
+        gradient / NaN pathologies. The migration is ``c_new = c_old ** 2``, and
+        ``AdaptPower(power=2.0)`` reproduces this class's coefficient scaling exactly.
 
         Parameters
         ----------

autoarray/inversion/regularization/adapt_power.py

@@ -0,0 +1,137 @@
+from __future__ import annotations
+import numpy as np
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from autoarray.inversion.linear_obj.linear_obj import LinearObj
+
+from autoarray.inversion.regularization.adapt import Adapt
+from autoarray.inversion.regularization.adapt import adapt_regularization_weights_from
+from autoarray.inversion.regularization.adapt import (
+    weighted_regularization_matrix_single_scatter_from,
+)
+
+
+class AdaptPower(Adapt):
+    def __init__(
+        self,
+        inner_coefficient: float = 1.0,
+        outer_coefficient: float = 1.0,
+        signal_scale: float = 1.0,
+        power: float = 1.0,
+    ):
+        """
+        Regularization which uses the neighbors of the mesh (e.g. shared Delaunay vertexes) and values adapted to
+        the data being fitted to smooth an inversion's solution, with the coefficient convention of ``Constant``.
+
+        This is the corrected sibling of ``Adapt``. It reconstructs the same adaptive weighting -- high smoothing
+        where there is no signal, less where there is (Nightingale, Dye and Massey 2018) -- but fixes the two ways
+        in which ``Adapt`` diverges from ``Constant``:
+
+        1. **The coefficient enters at ``lambda^2``, not ``lambda^4``.** ``Adapt`` squares its coefficients twice
+           (once when interpolating them into per-pixel weights, once in the matrix builder), so its matrix scales
+           as the fourth power of the coefficient. This class raises the interpolated coefficient to ``power``
+           before the builder squares it, so the effective exponent is ``2 * power`` and the default ``power=1.0``
+           matches ``Constant``. Under the shared ``LogUniform(1e-6, 1e6)`` prior that means the prior now spans
+           ``lambda^2``, and the regularization matrix stays positive-definite to ``c ~ 1e6`` rather than
+           collapsing from ``c ~ 1e4`` -- the fragility that produced the likelihood-overflow floods seen in
+           free-coefficient adaptive fits.
+        2. **Every mesh edge is scattered once.** ``Adapt`` adds each ordered neighbor pair to both the ``(i, j)``
+           and ``(j, i)`` entries, and the neighbor list already holds each unordered edge twice, so its matrix is
+           exactly ``2 x`` ``Constant``'s. This class uses
+           ``weighted_regularization_matrix_single_scatter_from``, which scatters each ordered pair once with the
+           symmetric edge weight ``0.5 * (w_i ** 2 + w_j ** 2)`` -- a weighted graph Laplacian, so still symmetric
+           and positive semi-definite.
+
+        Together these make ``AdaptPower(inner_coefficient=c, outer_coefficient=c)`` **exactly equal** to
+        ``Constant(coefficient=c)`` for any ``c``, which is what ``Adapt``'s docstring always claimed and never
+        delivered.
+
+        A full description of regularization and this matrix can be found in the parent ``AbstractRegularization``
+        class; the ``B`` matrix construction is described on ``Adapt``.
+
+        **Migration from ``Adapt``.** The coefficient scale is squared: ``c_new = c_old ** 2``. To reproduce the
+        legacy class exactly, pass ``power=2.0`` -- but note the factor-2 scatter is fixed regardless, so
+        ``AdaptPower(power=2.0)`` is ``0.5 x`` ``Adapt`` with the same coefficients.
+
+        **JAX & gradient support**: as for ``Adapt`` -- JAX-differentiable and FD-certified on the rectangular
+        mesh family, but raising ``TracerArrayConversionError`` on the Delaunay mesh family, whose neighbors come
+        from a direct scipy call on the traced mesh grid (use ``AdaptSplitPower`` there).
+
+        Parameters
+        ----------
+        inner_coefficient
+            The inner regularization coefficient which controls the degree of smoothing of the inversion
+            reconstruction in the inner (high signal) regions of a mesh's reconstruction.
+        outer_coefficient
+            The outer regularization coefficient which controls the degree of smoothing of the inversion
+            reconstruction in the outer (low signal) regions of a mesh's reconstruction.
+        signal_scale
+            A factor which controls how rapidly the smoothness of regularization varies from high signal regions
+            to low signal regions.
+        power
+            The exponent the interpolated coefficient is raised to before the matrix builder squares it, so the
+            coefficient enters the regularization matrix at the power ``2 * power``. The default ``1.0`` is the
+            ``Constant`` convention; ``2.0`` is the legacy ``Adapt`` convention. This is a convention switch, not
+            a model parameter -- the shipped prior config fixes it as a ``Constant`` prior so a search never
+            samples it.
+        """
+        super().__init__(
+            inner_coefficient=inner_coefficient,
+            outer_coefficient=outer_coefficient,
+            signal_scale=signal_scale,
+        )
+
+        self.power = power
+
+    def regularization_weights_from(self, linear_obj: LinearObj, xp=np) -> np.ndarray:
+        """
+        Returns the regularization weights of this regularization scheme.
+
+        These are the interpolated inner / outer coefficients raised to ``self.power`` (default ``1.0``), as
+        opposed to ``Adapt``, which squares them.
+
+        Parameters
+        ----------
+        linear_obj
+            The linear object (e.g. a ``Mapper``) which uses these weights when performing regularization.
+
+        Returns
+        -------
+        The regularization weights.
+        """
+        pixel_signals = linear_obj.pixel_signals_from(
+            signal_scale=self.signal_scale, xp=xp
+        )
+
+        return adapt_regularization_weights_from(
+            inner_coefficient=self.inner_coefficient,
+            outer_coefficient=self.outer_coefficient,
+            pixel_signals=pixel_signals,
+            power=self.power,
+        )
+
+    def regularization_matrix_from(self, linear_obj: LinearObj, xp=np) -> np.ndarray:
+        """
+        Returns the regularization matrix with shape [pixels, pixels].
+
+        Every mesh edge is scattered once (unlike ``Adapt``), so uniform weights reproduce ``Constant`` exactly.
+
+        Parameters
+        ----------
+        linear_obj
+            The linear object (e.g. a ``Mapper``) which uses this matrix to perform regularization.
+
+        Returns
+        -------
+        The regularization matrix.
+        """
+        regularization_weights = self.regularization_weights_from(
+            linear_obj=linear_obj, xp=xp
+        )
+
+        return weighted_regularization_matrix_single_scatter_from(
+            regularization_weights=regularization_weights,
+            neighbors=linear_obj.mesh_geometry.neighbors,
+            xp=xp,
+        )

autoarray/inversion/regularization/adapt_split.py

@@ -67,7 +67,26 @@ def __init__(
         mesh family (e.g. the KNN meshes), structurally incompatible with the
         rectangular meshes. Note the defaults
         ``inner_coefficient == outer_coefficient == 1.0`` make the weighting
-        uniform — numerically identical to ``ConstantSplit(coefficient=1.0)``.
+        uniform — but *not* numerically identical to ``ConstantSplit(coefficient=1.0)``;
+        see the coefficient-convention note below.
+
+        **Coefficient convention (legacy, ``lambda^4``).** The coefficients are squared twice before they
+        reach the regularization matrix -- once by ``adapt_regularization_weights_from`` and once by the
+        matrix builder -- so the matrix scales as the *fourth* power of the coefficient, while
+        ``Constant`` scales as the second. Both carry the same ``LogUniform(1e-6, 1e6)`` prior, so this
+        scheme explores a far wider effective smoothing range and reaches a numerically non
+        positive-definite matrix from ``c ~ 1e4`` where ``Constant`` survives to ``c ~ 1e6``.
+
+        The split family does **not** carry the factor-2 scatter asymmetry of ``Adapt``:
+        it shares ``pixel_splitted_regularization_matrix_from`` with ``ConstantSplit``, so the
+        coefficient exponent is the only difference between the two.
+
+        This behaviour is preserved deliberately: changing it would alter the coefficient scale of every
+        adaptive fit ever run. **New work should prefer ``AdaptSplitPower``**, which takes a ``power`` argument
+        (default ``1.0``, giving the ``Constant``-matching ``lambda^2`` convention), so
+        ``AdaptSplitPower(inner=outer=c)`` equals ``ConstantSplit(c)`` exactly and is more robust to
+        gradient / NaN pathologies. The migration is ``c_new = c_old ** 2``, and
+        ``AdaptSplitPower(power=2.0)`` reproduces this class's coefficient scaling exactly.
 
         Parameters
         ----------