Feature/matern adaptive

autoarray/inversion/regularization/__init__.py

@@ -10,3 +10,4 @@
 from .gaussian_kernel import GaussianKernel
 from .exponential_kernel import ExponentialKernel
 from .matern_kernel import MaternKernel
+from .matern_adaptive_brightness_kernel import MaternAdaptiveBrightnessKernel

autoarray/inversion/regularization/matern_adaptive_brightness_kernel.py

@@ -0,0 +1,153 @@
+from __future__ import annotations
+import numpy as np
+from typing import TYPE_CHECKING
+
+from autoarray.inversion.regularization.matern_kernel import MaternKernel
+
+if TYPE_CHECKING:
+    from autoarray.inversion.linear_obj.linear_obj import LinearObj
+
+from autoarray.inversion.regularization.matern_kernel import matern_kernel
+
+
+def matern_cov_matrix_from(
+    scale: float,
+    nu: float,
+    pixel_points,
+    weights=None,
+    xp=np,
+):
+    """
+    Construct the regularization covariance matrix (N x N) using a Matérn kernel,
+    optionally modulated by per-pixel weights.
+
+    If `weights` is provided (shape [N]), the covariance is:
+        C_ij = K(d_ij; scale, nu) * w_i * w_j
+    with a small diagonal jitter added for numerical stability.
+
+    Parameters
+    ----------
+    scale
+        Typical correlation length of the Matérn kernel.
+    nu
+        Smoothness parameter of the Matérn kernel.
+    pixel_points
+        Array-like of shape [N, 2] with (y, x) coordinates (or any 2D coords; only distances matter).
+    weights
+        Optional array-like of shape [N]. If None, treated as all ones.
+    xp
+        Backend (numpy or jax.numpy).
+
+    Returns
+    -------
+    covariance_matrix
+        Array of shape [N, N].
+    """
+
+    # --------------------------------
+    # Pairwise distances (broadcasted)
+    # --------------------------------
+    diff = pixel_points[:, None, :] - pixel_points[None, :, :]  # (N, N, 2)
+    d_ij = xp.sqrt(diff[..., 0] ** 2 + diff[..., 1] ** 2)  # (N, N)
+
+    # --------------------------------
+    # Base Matérn covariance
+    # --------------------------------
+    covariance_matrix = matern_kernel(d_ij, l=scale, v=nu, xp=xp)  # (N, N)
+
+    # --------------------------------
+    # Apply weights: C_ij *= w_i * w_j
+    # (broadcasted outer product, JAX-safe)
+    # --------------------------------
+    if weights is not None:
+        w = xp.asarray(weights)
+        # Ensure shape (N,) -> outer product (N,1)*(1,N) -> (N,N)
+        covariance_matrix = covariance_matrix * (w[:, None] * w[None, :])
+
+    # --------------------------------
+    # Add diagonal jitter (JAX-safe)
+    # --------------------------------
+    pixels = pixel_points.shape[0]
+    covariance_matrix = covariance_matrix + 1e-8 * xp.eye(pixels)
+
+    return covariance_matrix
+
+
+class MaternAdaptiveBrightnessKernel(MaternKernel):
+    def __init__(
+        self,
+        coefficient: float = 1.0,
+        scale: float = 1.0,
+        nu: float = 0.5,
+        rho: float = 1.0,
+    ):
+        """
+        Regularization which uses a Matern smoothing kernel to regularize the solution with regularization weights
+        that adapt to the brightness of the source being reconstructed.
+
+        For this regularization scheme, every pixel is regularized with every other pixel. This contrasts many other
+        schemes, where regularization is based on neighboring (e.g. do the pixels share a Delaunay edge?) or computing
+        derivates around the center of the pixel (where nearby pixels are regularization locally in similar ways).
+
+        This makes the regularization matrix fully dense and therefore maybe change the run times of the solution.
+        It also leads to more overall smoothing which can lead to more stable linear inversions.
+
+        For the weighted regularization scheme, each pixel is given an 'effective regularization weight', which is
+        applied when each set of pixel neighbors are regularized with one another. The motivation of this is that
+        different regions of a pixelization's mesh require different levels of regularization (e.g., high smoothing where the
+        no signal is present and less smoothing where it is, see (Nightingale, Dye and Massey 2018)).
+
+        This scheme is not used by Vernardos et al. (2022): https://arxiv.org/abs/2202.09378, but it follows
+        a similar approach.
+
+        A full description of regularization and this matrix can be found in the parent `AbstractRegularization` class.
+
+        Parameters
+        ----------
+        coefficient
+            The regularization coefficient which controls the degree of smooth of the inversion reconstruction.
+        scale
+            The typical scale of the exponential regularization pattern.
+        nu
+            Controls the derivative of the regularization pattern (`nu=0.5` is a Gaussian).
+        """
+        super().__init__(coefficient=coefficient, scale=scale, nu=nu)
+        self.rho = rho
+
+    def covariance_kernel_weights_from(
+        self, linear_obj: LinearObj, xp=np
+    ) -> np.ndarray:
+        """
+        Returns per-pixel kernel weights that adapt to the reconstructed pixel brightness.
+        """
+        # Assumes linear_obj.pixel_signals_from is xp-aware elsewhere in the codebase.
+        pixel_signals = linear_obj.pixel_signals_from(signal_scale=1.0, xp=xp)
+
+        max_signal = xp.max(pixel_signals)
+        max_signal = xp.maximum(max_signal, 1e-8)  # avoid divide-by-zero (JAX-safe)
+
+        return xp.exp(-self.rho * (1.0 - pixel_signals / max_signal))
+
+    def regularization_matrix_from(self, linear_obj: LinearObj, xp=np) -> np.ndarray:
+        kernel_weights = self.covariance_kernel_weights_from(
+            linear_obj=linear_obj, xp=xp
+        )
+
+        # Follow the xp pattern used in the Matérn kernel module (often `.array` for grids).
+        pixel_points = linear_obj.source_plane_mesh_grid.array
+
+        covariance_matrix = matern_cov_matrix_from(
+            scale=self.scale,
+            pixel_points=pixel_points,
+            nu=self.nu,
+            weights=kernel_weights,
+            xp=xp,
+        )
+
+        return self.coefficient * xp.linalg.inv(covariance_matrix)
+
+    def regularization_weights_from(self, linear_obj: LinearObj, xp=np) -> np.ndarray:
+        """
+        Returns the regularization weights of this regularization scheme.
+        """
+        return 1.0 / self.covariance_kernel_weights_from(linear_obj=linear_obj, xp=xp)

autoarray/inversion/regularization/matern_kernel.py

@@ -44,6 +44,20 @@ def kv_xp(v, z, xp=np):
             )
 
 
+def gamma_xp(x, xp=np):
+    """
+    XP-compatible Gamma(x).
+    """
+    if xp is np:
+        import scipy.special as sc
+
+        return sc.gamma(x)
+    else:
+        import jax.scipy.special as jsp
+
+        return jsp.gamma(x)
+
+
 def matern_kernel(r, l: float = 1.0, v: float = 0.5, xp=np):
     """
     XP-compatible Matérn kernel.
@@ -55,7 +69,7 @@ def matern_kernel(r, l: float = 1.0, v: float = 0.5, xp=np):
 
     z = xp.sqrt(2.0 * v) * r / l
 
-    part1 = 2.0 ** (1.0 - v) / math.gamma(v)  # scalar constant
+    part1 = 2.0 ** (1.0 - v) / gamma_xp(v, xp)  # scalar constant
     part2 = z**v
     part3 = kv_xp(v, z, xp)
 
@@ -141,8 +155,8 @@ def __init__(self, coefficient: float = 1.0, scale: float = 1.0, nu: float = 0.5
         """
 
         self.coefficient = coefficient
-        self.scale = float(scale)
-        self.nu = float(nu)
+        self.scale = scale
+        self.nu = nu
         super().__init__()
 
     def regularization_weights_from(self, linear_obj: LinearObj, xp=np) -> np.ndarray:

autoarray/preloads.py

@@ -26,6 +26,7 @@ def __init__(
         linear_light_profile_blurred_mapping_matrix=None,
         use_voronoi_areas: bool = True,
         areas_factor: float = 0.5,
+        skip_areas: bool = False,
     ):
         """
         Stores preloaded arrays and matrices used during pixelized linear inversions, improving both performance
@@ -123,3 +124,4 @@ def __init__(
 
         self.use_voronoi_areas = use_voronoi_areas
         self.areas_factor = areas_factor
+        self.skip_areas = skip_areas

autoarray/structures/arrays/kernel_2d.py

@@ -612,7 +612,7 @@ def convolved_image_from(
         image,
         blurring_image,
         jax_method="direct",
-        use_mixed_precision : bool = False,
+        use_mixed_precision: bool = False,
         xp=np,
     ):
         """

autoarray/structures/mesh/delaunay_2d.py

@@ -339,6 +339,65 @@ def pix_indexes_for_sub_slim_index_delaunay_from(
     return out
 
 
+def scipy_delaunay_matern(points_np, query_points_np):
+    """
+    Minimal SciPy Delaunay callback for Matérn regularization.
+
+    Returns only what’s needed for mapping:
+      - points (tri.points)
+      - simplices_padded
+      - mappings (pix indexes for each query point)
+    """
+
+    max_simplices = 2 * points_np.shape[0]
+
+    # --- Delaunay mesh ---
+    tri = Delaunay(points_np)
+
+    points = tri.points.astype(points_np.dtype)
+    simplices = tri.simplices.astype(np.int32)
+
+    # --- Pad simplices to fixed shape for JAX ---
+    simplices_padded = -np.ones((max_simplices, 3), dtype=np.int32)
+    simplices_padded[: simplices.shape[0]] = simplices
+
+    # --- find_simplex for query points ---
+    simplex_idx = tri.find_simplex(query_points_np).astype(np.int32)  # (Q,)
+
+    mappings = pix_indexes_for_sub_slim_index_delaunay_from(
+        source_plane_data_grid=query_points_np,
+        simplex_index_for_sub_slim_index=simplex_idx,
+        pix_indexes_for_simplex_index=simplices,
+        delaunay_points=points_np,
+    )
+
+    return points, simplices_padded, mappings
+
+
+def jax_delaunay_matern(points, query_points):
+    """
+    JAX wrapper using pure_callback to run SciPy Delaunay on CPU,
+    returning only the minimal outputs needed for Matérn usage.
+    """
+    import jax
+    import jax.numpy as jnp
+
+    N = points.shape[0]
+    Q = query_points.shape[0]
+    max_simplices = 2 * N
+
+    points_shape = jax.ShapeDtypeStruct((N, 2), points.dtype)
+    simplices_padded_shape = jax.ShapeDtypeStruct((max_simplices, 3), jnp.int32)
+    mappings_shape = jax.ShapeDtypeStruct((Q, 3), jnp.int32)
+
+    return jax.pure_callback(
+        lambda pts, qpts: scipy_delaunay_matern(np.asarray(pts), np.asarray(qpts)),
+        (points_shape, simplices_padded_shape, mappings_shape),
+        points,
+        query_points,
+    )
+
+
 class DelaunayInterface:
 
     def __init__(
@@ -466,33 +525,60 @@ def delaunay(self) -> "scipy.spatial.Delaunay":
 
             use_voronoi_areas = self.preloads.use_voronoi_areas
             areas_factor = self.preloads.areas_factor
+            skip_areas = self.preloads.skip_areas
 
         else:
 
             use_voronoi_areas = True
             areas_factor = 0.5
+            skip_areas = False
 
-        if self._xp.__name__.startswith("jax"):
+        if not skip_areas:
 
-            import jax.numpy as jnp
+            if self._xp.__name__.startswith("jax"):
 
-            points, simplices, mappings, split_points, splitted_mappings = jax_delaunay(
-                points=self.mesh_grid_xy,
-                query_points=self._source_plane_data_grid_over_sampled,
-                use_voronoi_areas=use_voronoi_areas,
-                areas_factor=areas_factor,
-            )
+                import jax.numpy as jnp
+
+                points, simplices, mappings, split_points, splitted_mappings = (
+                    jax_delaunay(
+                        points=self.mesh_grid_xy,
+                        query_points=self._source_plane_data_grid_over_sampled,
+                        use_voronoi_areas=use_voronoi_areas,
+                        areas_factor=areas_factor,
+                    )
+                )
+
+            else:
+
+                points, simplices, mappings, split_points, splitted_mappings = (
+                    scipy_delaunay(
+                        points_np=self.mesh_grid_xy,
+                        query_points_np=self._source_plane_data_grid_over_sampled,
+                        use_voronoi_areas=use_voronoi_areas,
+                        areas_factor=areas_factor,
+                    )
+                )
 
         else:
 
-            points, simplices, mappings, split_points, splitted_mappings = (
-                scipy_delaunay(
+            if self._xp.__name__.startswith("jax"):
+
+                import jax.numpy as jnp
+
+                points, simplices, mappings = jax_delaunay_matern(
+                    points=self.mesh_grid_xy,
+                    query_points=self._source_plane_data_grid_over_sampled,
+                )
+
+            else:
+
+                points, simplices, mappings = scipy_delaunay_matern(
                     points_np=self.mesh_grid_xy,
                     query_points_np=self._source_plane_data_grid_over_sampled,
-                    use_voronoi_areas=use_voronoi_areas,
-                    areas_factor=areas_factor,
                 )
-            )
+
+            split_points = None
+            splitted_mappings = None
 
         return DelaunayInterface(
             points=points,