end to end matern tests and skip areas added

Author: Jammy2211

autoarray/inversion/regularization/matern_adaptive_brightness_kernel.py

@@ -9,6 +9,7 @@
 
 from autoarray.inversion.regularization.matern_kernel import matern_kernel
 
+
 def matern_cov_matrix_from(
     scale: float,
     nu: float,
@@ -47,7 +48,7 @@ def matern_cov_matrix_from(
     # 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)
+    d_ij = xp.sqrt(diff[..., 0] ** 2 + diff[..., 1] ** 2)  # (N, N)
 
     # --------------------------------
     # Base Matérn covariance
@@ -113,7 +114,9 @@ def __init__(
         super().__init__(coefficient=coefficient, scale=scale, nu=nu)
         self.rho = rho
 
-    def covariance_kernel_weights_from(self, linear_obj: LinearObj, xp=np) -> np.ndarray:
+    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.
         """
@@ -126,10 +129,12 @@ def covariance_kernel_weights_from(self, linear_obj: LinearObj, xp=np) -> np.nda
         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)
+        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
+        pixel_points = linear_obj.source_plane_mesh_grid.array
 
         covariance_matrix = matern_cov_matrix_from(
             scale=self.scale,
@@ -145,4 +150,4 @@ def regularization_weights_from(self, linear_obj: LinearObj, xp=np) -> np.ndarra
         """
         Returns the regularization weights of this regularization scheme.
         """
-        return 1.0 / self.covariance_kernel_weights_from(linear_obj=linear_obj, xp=xp)
+        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/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,

test_autoarray/inversion/regularizations/test_matern_adaptive_brightness_kernel.py

@@ -8,7 +8,9 @@
 
 def test__regularization_matrix():
 
-    reg = aa.reg.MaternAdaptiveBrightnessKernel(coefficient=1.0, scale=2.0, nu=2.0, rho=1.0)
+    reg = aa.reg.MaternAdaptiveBrightnessKernel(
+        coefficient=1.0, scale=2.0, nu=2.0, rho=1.0
+    )
 
     neighbors = np.array(
         [
@@ -41,7 +43,9 @@ def test__regularization_matrix():
     assert regularization_matrix[0, 0] == pytest.approx(18.7439565009, 1.0e-4)
     assert regularization_matrix[0, 1] == pytest.approx(-8.786547368, 1.0e-4)
 
-    reg = aa.reg.MaternAdaptiveBrightnessKernel(coefficient=1.5, scale=2.5, nu=2.5, rho=1.5)
+    reg = aa.reg.MaternAdaptiveBrightnessKernel(
+        coefficient=1.5, scale=2.5, nu=2.5, rho=1.5
+    )
 
     pixel_signals = np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0])
 
@@ -52,4 +56,4 @@ def test__regularization_matrix():
     regularization_matrix = reg.regularization_matrix_from(linear_obj=mapper)
 
     assert regularization_matrix[0, 0] == pytest.approx(121.0190770, 1.0e-4)
-    assert regularization_matrix[0, 1] == pytest.approx(-66.9580331, 1.0e-4)
+    assert regularization_matrix[0, 1] == pytest.approx(-66.9580331, 1.0e-4)