regulsirztion refactor complete and rectangular works

Author: Jammy2211

autoarray/inversion/pixelization/mesh/mesh_util.py

@@ -110,6 +110,7 @@ def rectangular_corner_neighbors(
 
     return neighbors, neighbors_sizes
 
+
 def rectangular_top_edge_neighbors(
     neighbors: np.ndarray, neighbors_sizes: np.ndarray, shape_native: Tuple[int, int]
 ) -> Tuple[np.ndarray, np.ndarray]:
@@ -145,6 +146,7 @@ def rectangular_top_edge_neighbors(
 
     return neighbors, neighbors_sizes
 
+
 def rectangular_left_edge_neighbors(
     neighbors: np.ndarray, neighbors_sizes: np.ndarray, shape_native: Tuple[int, int]
 ) -> Tuple[np.ndarray, np.ndarray]:
@@ -180,6 +182,7 @@ def rectangular_left_edge_neighbors(
 
     return neighbors, neighbors_sizes
 
+
 def rectangular_right_edge_neighbors(
     neighbors: np.ndarray, neighbors_sizes: np.ndarray, shape_native: Tuple[int, int]
 ) -> Tuple[np.ndarray, np.ndarray]:
@@ -215,6 +218,7 @@ def rectangular_right_edge_neighbors(
 
     return neighbors, neighbors_sizes
 
+
 def rectangular_bottom_edge_neighbors(
     neighbors: np.ndarray, neighbors_sizes: np.ndarray, shape_native: Tuple[int, int]
 ) -> Tuple[np.ndarray, np.ndarray]:
@@ -288,14 +292,15 @@ def rectangular_central_neighbors(
 
     # Compute neighbor indices
     neighbors[pixel_indices, 0] = pixel_indices - n_cols  # Up
-    neighbors[pixel_indices, 1] = pixel_indices - 1       # Left
-    neighbors[pixel_indices, 2] = pixel_indices + 1       # Right
+    neighbors[pixel_indices, 1] = pixel_indices - 1  # Left
+    neighbors[pixel_indices, 2] = pixel_indices + 1  # Right
     neighbors[pixel_indices, 3] = pixel_indices + n_cols  # Down
 
     neighbors_sizes[pixel_indices] = 4
 
     return neighbors, neighbors_sizes
 
+
 def rectangular_edge_pixel_list_from(shape_native: Tuple[int, int]) -> List[int]:
     """
     Returns a list of the 1D indices of all pixels on the edge of a rectangular pixelization,

autoarray/inversion/regularization/__init__.py

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

autoarray/inversion/regularization/adaptive_brightness.py

@@ -80,19 +80,19 @@ def weighted_regularization_matrix_from(
         coefficient of every source pixel is different.
     """
     S, P = neighbors.shape
-    reg_w = regularization_weights ** 2
+    reg_w = regularization_weights**2
 
     # 1) Flatten the (i→j) neighbor pairs
-    I = jnp.repeat(jnp.arange(S), P)            # (S*P,)
-    J = neighbors.reshape(-1)                   # (S*P,)
+    I = jnp.repeat(jnp.arange(S), P)  # (S*P,)
+    J = neighbors.reshape(-1)  # (S*P,)
 
     # 2) Remap “no neighbor” entries to an extra slot S, whose weight=0
     OUT = S
     J = jnp.where(J < 0, OUT, J)
 
     # 3) Build an extended weight vector with a zero at index S
     reg_w_ext = jnp.concatenate([reg_w, jnp.zeros((1,))], axis=0)
-    w_ij = reg_w_ext[J]                         # (S*P,)
+    w_ij = reg_w_ext[J]  # (S*P,)
 
     # 4) Start with zeros on an (S+1)x(S+1) canvas so we can scatter into row S safely
     mat = jnp.zeros((S + 1, S + 1), dtype=regularization_weights.dtype)
@@ -102,10 +102,9 @@ def weighted_regularization_matrix_from(
     #    - sum_j reg_w[j] into diag[i]
     #    - sum contributions reg_w[j] into diag[j]
     #    (diagonal at OUT=S picks up zeros only)
-    diag_updates_i = jnp.concatenate([
-        jnp.full((S,), 1e-8),
-        jnp.zeros((1,))  # out‐of‐bounds slot stays zero
-    ], axis=0)
+    diag_updates_i = jnp.concatenate(
+        [jnp.full((S,), 1e-8), jnp.zeros((1,))], axis=0  # out‐of‐bounds slot stays zero
+    )
     mat = mat.at[jnp.diag_indices(S + 1)].add(diag_updates_i)
     mat = mat.at[I, I].add(w_ij)
     mat = mat.at[J, J].add(w_ij)
@@ -117,6 +116,7 @@ def weighted_regularization_matrix_from(
     # 7) Drop the extra row/column S and return the S×S result
     return mat[:S, :S]
 
+
 class AdaptiveBrightness(AbstractRegularization):
     def __init__(
         self,

autoarray/inversion/regularization/brightness_zeroth.py

@@ -42,6 +42,7 @@ def brightness_zeroth_regularization_weights_from(
     """
     return coefficient * (1.0 - pixel_signals)
 
+
 def brightness_zeroth_regularization_matrix_from(
     regularization_weights: jnp.ndarray,
 ) -> jnp.ndarray:
@@ -63,7 +64,6 @@ def brightness_zeroth_regularization_matrix_from(
     return jnp.diag(regularization_weight_squared)
 
 
-
 class BrightnessZeroth(AbstractRegularization):
     def __init__(
         self,

autoarray/inversion/regularization/constant.py

@@ -52,7 +52,9 @@ class in the module `autoarray.inversion.regularization`.
     # This ensures that JAX can efficiently drop these entries during matrix updates.
     neighbors = jnp.where(neighbors == -1, OUT_OF_BOUND_IDX, neighbors)
     return (
-        jnp.diag(1e-8 + regularization_coefficient * neighbors_sizes).at[I_IDX, neighbors]
+        jnp.diag(1e-8 + regularization_coefficient * neighbors_sizes).at[
+            I_IDX, neighbors
+        ]
         # unique indices should be guranteed by neighbors-spec
         .add(-regularization_coefficient, mode="drop", unique_indices=True)
     )

autoarray/inversion/regularization/constant_zeroth.py

@@ -50,12 +50,14 @@ class in the module ``autoarray.inversion.regularization``.
     # This ensures that JAX can efficiently drop these entries during matrix updates.
     neighbors = jnp.where(neighbors == -1, OUT_OF_BOUND_IDX, neighbors)
     const = (
-        jnp.diag(1e-8 + regularization_coefficient * neighbors_sizes).at[I_IDX, neighbors]
+        jnp.diag(1e-8 + regularization_coefficient * neighbors_sizes).at[
+            I_IDX, neighbors
+        ]
         # unique indices should be guranteed by neighbors-spec
         .add(-regularization_coefficient, mode="drop", unique_indices=True)
     )
 
-    reg_coeff = coefficient_zeroth ** 2.0
+    reg_coeff = coefficient_zeroth**2.0
     # Identity matrix scaled by reg_coeff does exactly ∑_i reg_coeff * e_i e_i^T
     zeroth = jnp.eye(P) * reg_coeff
 

autoarray/inversion/regularization/exponential_kernel.py

@@ -1,58 +1,50 @@
 from __future__ import annotations
-import numpy as np
+import jax.numpy as jnp
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
     from autoarray.inversion.linear_obj.linear_obj import LinearObj
 
 from autoarray.inversion.regularization.abstract import AbstractRegularization
 
-from autoarray import numba_util
 
-
-@numba_util.jit()
 def exp_cov_matrix_from(
     scale: float,
-    pixel_points: np.ndarray,
-) -> np.ndarray:
+    pixel_points: jnp.ndarray,  # shape (N, 2)
+) -> jnp.ndarray:  # shape (N, N)
     """
-    Consutruct the source brightness covariance matrix, which is used to determined the regularization
-    pattern (i.e, how the different  source pixels are smoothed).
+    Construct the source brightness covariance matrix using an exponential kernel:
+
+        cov[i,j] = exp(- d_{ij} / scale)
 
-    The covariance matrix includes one non-linear parameters, the scale coefficient, which is used to determine
-    the typical scale of the regularization pattern.
+    with a tiny jitter 1e-8 added on the diagonal for numerical stability.
 
     Parameters
     ----------
     scale
-        The typical scale of the regularization pattern .
+        The length‐scale of the exponential kernel.
     pixel_points
-        An 2d array with shape [N_source_pixels, 2], which save the source pixelization coordinates (on source plane).
-        Something like [[y1,x1], [y2,x2], ...]
+        Array of shape (N, 2) giving the (y,x) coordinates of each source‐plane pixel.
 
     Returns
     -------
-    np.ndarray
-        The source covariance matrix (2d array), shape [N_source_pixels, N_source_pixels].
+    jnp.ndarray, shape (N, N)
+        The exponential covariance matrix.
     """
+    # pairwise differences: shape (N, N, 2)
+    diff = pixel_points[:, None, :] - pixel_points[None, :, :]
 
-    pixels = len(pixel_points)
-    covariance_matrix = np.zeros(shape=(pixels, pixels))
+    # Euclidean distances: shape (N, N)
+    d = jnp.linalg.norm(diff, axis=-1)
 
-    for i in range(pixels):
-        covariance_matrix[i, i] += 1e-8
-        for j in range(pixels):
-            xi = pixel_points[i, 1]
-            yi = pixel_points[i, 0]
-            xj = pixel_points[j, 1]
-            yj = pixel_points[j, 0]
-            d_ij = np.sqrt(
-                (xi - xj) ** 2 + (yi - yj) ** 2
-            )  # distance between the pixel i and j
+    # exponential kernel
+    cov = jnp.exp(-d / scale)
 
-            covariance_matrix[i, j] += np.exp(-1.0 * d_ij / scale)
+    # add a small jitter on the diagonal
+    N = pixel_points.shape[0]
+    cov = cov + jnp.eye(N) * 1e-8
 
-    return covariance_matrix
+    return cov
 
 
 class ExponentialKernel(AbstractRegularization):
@@ -83,7 +75,7 @@ def __init__(self, coefficient: float = 1.0, scale: float = 1.0):
 
         super().__init__()
 
-    def regularization_weights_from(self, linear_obj: LinearObj) -> np.ndarray:
+    def regularization_weights_from(self, linear_obj: LinearObj) -> jnp.ndarray:
         """
         Returns the regularization weights of this regularization scheme.
 
@@ -102,9 +94,9 @@ def regularization_weights_from(self, linear_obj: LinearObj) -> np.ndarray:
         -------
         The regularization weights.
         """
-        return self.coefficient * np.ones(linear_obj.params)
+        return self.coefficient * jnp.ones(linear_obj.params)
 
-    def regularization_matrix_from(self, linear_obj: LinearObj) -> np.ndarray:
+    def regularization_matrix_from(self, linear_obj: LinearObj) -> jnp.ndarray:
         """
         Returns the regularization matrix with shape [pixels, pixels].
 
@@ -119,7 +111,7 @@ def regularization_matrix_from(self, linear_obj: LinearObj) -> np.ndarray:
         """
         covariance_matrix = exp_cov_matrix_from(
             scale=self.scale,
-            pixel_points=np.array(linear_obj.source_plane_mesh_grid),
+            pixel_points=linear_obj.source_plane_mesh_grid.array,
         )
 
-        return self.coefficient * np.linalg.inv(covariance_matrix)
+        return self.coefficient * jnp.linalg.inv(covariance_matrix)

autoarray/inversion/regularization/gaussian_kernel.py

@@ -8,6 +8,7 @@
 
 from autoarray.inversion.regularization.abstract import AbstractRegularization
 
+
 def gauss_cov_matrix_from(
     scale: float,
     pixel_points: jnp.ndarray,  # shape (N, 2)
@@ -34,17 +35,17 @@ def gauss_cov_matrix_from(
         The Gaussian covariance matrix.
     """
     # Ensure array:
-    pts = jnp.asarray(pixel_points)               # (N, 2)
+    pts = jnp.asarray(pixel_points)  # (N, 2)
     # Compute squared distances: ||p_i - p_j||^2
-    diffs = pts[:, None, :] - pts[None, :, :]     # (N, N, 2)
-    d2   = jnp.sum(diffs**2, axis=-1)             # (N, N)
+    diffs = pts[:, None, :] - pts[None, :, :]  # (N, N, 2)
+    d2 = jnp.sum(diffs**2, axis=-1)  # (N, N)
 
     # Gaussian kernel
-    cov  = jnp.exp(-d2 / (2.0 * scale**2))        # (N, N)
+    cov = jnp.exp(-d2 / (2.0 * scale**2))  # (N, N)
 
     # Add tiny jitter on the diagonal
-    N    = pts.shape[0]
-    cov  = cov + jnp.eye(N, dtype=cov.dtype) * 1e-8
+    N = pts.shape[0]
+    cov = cov + jnp.eye(N, dtype=cov.dtype) * 1e-8
 
     return cov
 
@@ -111,8 +112,7 @@ def regularization_matrix_from(self, linear_obj: LinearObj) -> np.ndarray:
         The regularization matrix.
         """
         covariance_matrix = gauss_cov_matrix_from(
-            scale=self.scale,
-            pixel_points=linear_obj.source_plane_mesh_grid.array
+            scale=self.scale, pixel_points=linear_obj.source_plane_mesh_grid.array
         )
 
         return self.coefficient * jnp.linalg.inv(covariance_matrix)

autoarray/inversion/regularization/regularization_util.py

@@ -3,10 +3,24 @@
 
 from autoarray import exc
 
-from autoarray.inversion.regularization.adaptive_brightness import adaptive_regularization_weights_from
-from autoarray.inversion.regularization.brightness_zeroth import brightness_zeroth_regularization_matrix_from
-from autoarray.inversion.regularization.brightness_zeroth import brightness_zeroth_regularization_weights_from
-from autoarray.inversion.regularization.constant import constant_regularization_matrix_from
+from autoarray.inversion.regularization.adaptive_brightness import (
+    adaptive_regularization_weights_from,
+)
+from autoarray.inversion.regularization.adaptive_brightness import (
+    weighted_regularization_matrix_from,
+)
+from autoarray.inversion.regularization.brightness_zeroth import (
+    brightness_zeroth_regularization_matrix_from,
+)
+from autoarray.inversion.regularization.brightness_zeroth import (
+    brightness_zeroth_regularization_weights_from,
+)
+from autoarray.inversion.regularization.constant import (
+    constant_regularization_matrix_from,
+)
+from autoarray.inversion.regularization.constant_zeroth import (
+    constant_zeroth_regularization_matrix_from,
+)
 from autoarray.inversion.regularization.exponential_kernel import exp_cov_matrix_from
 from autoarray.inversion.regularization.gaussian_kernel import gauss_cov_matrix_from
 from autoarray.inversion.regularization.matern_kernel import matern_kernel

autoarray/inversion/regularization/zeroth.py

@@ -30,7 +30,7 @@ class in the module `autoarray.inversion.regularization`.
         coefficient of every source pixel is the same.
     """
 
-    reg_coeff = coefficient ** 2.0
+    reg_coeff = coefficient**2.0
 
     # Identity matrix scaled by reg_coeff does exactly ∑_i reg_coeff * e_i e_i^T