gaussian kernel converted successfully

Author: Jammy2211

autoarray/inversion/inversion/abstract.py

@@ -663,6 +663,9 @@ def log_det_regularization_matrix_term(self) -> float:
         float
             The log determinant of the regularization matrix.
         """
+        if not self.has(cls=AbstractRegularization):
+            return 0.0
+
         try:
             return 2.0 * np.sum(
                 jnp.log(

autoarray/inversion/regularization/constant_zeroth.py

@@ -55,5 +55,4 @@ def regularization_matrix_from(self, linear_obj: LinearObj) -> np.ndarray:
             coefficient=self.coefficient_neighbor,
             coefficient_zeroth=self.coefficient_zeroth,
             neighbors=linear_obj.neighbors,
-            neighbors_sizes=linear_obj.neighbors.sizes,
         )

autoarray/inversion/regularization/gaussian_kernel.py

@@ -1,4 +1,5 @@
 from __future__ import annotations
+import jax.numpy as jnp
 import numpy as np
 from typing import TYPE_CHECKING
 
@@ -7,52 +8,45 @@
 
 from autoarray.inversion.regularization.abstract import AbstractRegularization
 
-from autoarray import numba_util
-
-
-@numba_util.jit()
 def gauss_cov_matrix_from(
     scale: float,
-    pixel_points: np.ndarray,
-) -> np.ndarray:
+    pixel_points: jnp.ndarray,  # shape (N, 2)
+) -> jnp.ndarray:
     """
-    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‐pixel Gaussian covariance matrix for regularization.
+
+    For N source‐pixels at coordinates (y_i, x_i), we define
 
-    the covariance matrix includes one non-linear parameters, the scale coefficient, which is used to
-    determine the typical scale of the regularization pattern.
+      C_ij = exp( -||p_i - p_j||^2 / (2 scale^2) )
+
+    plus a tiny diagonal “jitter” (1e-8) to ensure numerical stability.
 
     Parameters
     ----------
     scale
-        the typical scale of the regularization pattern .
+        The characteristic length scale of the Gaussian 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 pixel.
 
     Returns
     -------
-    np.ndarray
-        The source covariance matrix (2d array), shape [N_source_pixels, N_source_pixels].
+    cov : jnp.ndarray, shape (N, N)
+        The Gaussian covariance matrix.
     """
+    # Ensure array:
+    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)
 
-    pixels = len(pixel_points)
-    covariance_matrix = np.zeros(shape=(pixels, pixels))
-
-    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
+    # Gaussian kernel
+    cov  = jnp.exp(-d2 / (2.0 * scale**2))        # (N, N)
 
-            covariance_matrix[i, j] += np.exp(-1.0 * d_ij**2 / (2 * scale**2))
+    # Add tiny jitter on the diagonal
+    N    = pts.shape[0]
+    cov  = cov + jnp.eye(N, dtype=cov.dtype) * 1e-8
 
-    return covariance_matrix
+    return cov
 
 
 class GaussianKernel(AbstractRegularization):
@@ -117,7 +111,8 @@ def regularization_matrix_from(self, linear_obj: LinearObj) -> np.ndarray:
         The regularization matrix.
         """
         covariance_matrix = gauss_cov_matrix_from(
-            scale=self.scale, pixel_points=np.array(linear_obj.source_plane_mesh_grid)
+            scale=self.scale,
+            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/regularization_util.py

@@ -82,13 +82,10 @@ class in the module `autoarray.inversion.regularization`.
     )
 
 
-
-@numba_util.jit()
 def constant_zeroth_regularization_matrix_from(
     coefficient: float,
     coefficient_zeroth: float,
     neighbors: np.ndarray,
-    neighbors_sizes: np.ndarray,
 ) -> np.ndarray:
     """
     From the pixel-neighbors array, setup the regularization matrix using the instance regularization scheme.
@@ -113,24 +110,34 @@ class in the module ``autoarray.inversion.regularization``.
         The regularization matrix computed using Regularization where the effective regularization
         coefficient of every source pixel is the same.
     """
+    S, P = neighbors.shape
+    reg1 = coefficient**2
+    reg0 = coefficient_zeroth**2
+
+    # 1) Flatten (i,j) neighbor‐pairs
+    I = jnp.repeat(jnp.arange(S), P)       # (S*P,)
+    J = neighbors.reshape(-1)              # (S*P,)
 
-    pixels = len(neighbors)
+    # 2) Remap “no neighbor” = -1 → OUT = S
+    OUT = S
+    J = jnp.where(J < 0, OUT, J)
 
-    regularization_matrix = np.zeros(shape=(pixels, pixels))
+    # 3) Start on an (S+1)x(S+1) zero canvas
+    M = jnp.zeros((S+1, S+1), dtype=jnp.float32)
 
-    regularization_coefficient = coefficient**2.0
-    regularization_coefficient_zeroth = coefficient_zeroth**2.0
+    # 4) Diagonal baseline: 1e-8 + reg0  for i in [0..S-1]
+    diag_base = jnp.concatenate([jnp.full((S,), 1e-8 + reg0), jnp.zeros((1,))])
+    M = M.at[jnp.diag_indices(S+1)].add(diag_base)
 
-    for i in range(pixels):
-        regularization_matrix[i, i] += 1e-8
-        regularization_matrix[i, i] += regularization_coefficient_zeroth
-        for j in range(neighbors_sizes[i]):
-            neighbor_index = neighbors[i, j]
-            regularization_matrix[i, i] += regularization_coefficient
-            regularization_matrix[i, neighbor_index] -= regularization_coefficient
+    # 5) Scatter the first-order reg1 into diag[i] for each neighbor (i→j):
+    #    M[i,i] += reg1
+    M = M.at[I, I].add(reg1)
 
-    return regularization_matrix
+    # 6) Scatter the off-diagonals: M[i,j] -= reg1
+    M = M.at[I, J].add(-reg1)
 
+    # 7) Return only the top‐left S×S block
+    return M[:S, :S]
 
 def adaptive_regularization_weights_from(
     inner_coefficient: float, outer_coefficient: float, pixel_signals: np.ndarray

autoarray/preloads.py

@@ -14,6 +14,7 @@ def __init__(
         self,
         mapper_indices: np.ndarray = None,
         source_pixel_zeroed_indices: np.ndarray = None,
+        linear_light_profile_blurred_mapping_matrix = None,
     ):
         """
         Stores preloaded arrays and matrices used during pixelized linear inversions, improving both performance
@@ -37,11 +38,17 @@ def __init__(
             Indices of source pixels that should be set to zero in the reconstruction. These typically correspond to
             outer-edge source-plane regions with no image-plane mapping (e.g. outside a circular mask), helping
             separate the lens light from the pixelized source model.
+        linear_light_profile_blurred_mapping_matrix
+            The evaluated images of the linear light profiles that make up the blurred mapping matrix component of the
+            inversion, with the other component being the pixelization's pixels. These are fixed when the lens light
+            is fixed to the maximum likelihood solution, allowing the blurred mapping matrix to be preloaded, but
+            the intensity values will still be solved for during the inversion.
         """
 
         self.mapper_indices = None
         self.source_pixel_zeroed_indices = None
         self.source_pixel_zeroed_indices_to_keep = None
+        self.linear_light_profile_blurred_mapping_matrix = None
 
         if mapper_indices is not None:
 
@@ -58,3 +65,9 @@ def __init__(
 
             # Get the indices where values_to_solve is True
             self.source_pixel_zeroed_indices_to_keep = jnp.where(values_to_solve)[0]
+
+        if linear_light_profile_blurred_mapping_matrix is not None:
+
+            self.linear_light_profile_blurred_mapping_matrix = jnp.array(
+                linear_light_profile_blurred_mapping_matrix
+            )

test_autoarray/inversion/regularizations/test_regularization_util.py

@@ -147,14 +147,11 @@ def test__constant_regularization_matrix_from():
 def test__constant_zeroth_regularization_matrix_from():
     neighbors = np.array([[1, 2, -1], [0, -1, -1], [0, -1, -1]])
 
-    neighbors_sizes = np.array([2, 1, 1])
-
     regularization_matrix = (
         aa.util.regularization.constant_zeroth_regularization_matrix_from(
             coefficient=2.0,
             coefficient_zeroth=0.5,
             neighbors=neighbors,
-            neighbors_sizes=neighbors_sizes,
         )
     )