move utils to their specific modules

Author: Jammy2211

autoarray/inversion/regularization/__init__.py

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

autoarray/inversion/regularization/adaptive_brightness.py

@@ -1,14 +1,121 @@
 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.inversion.regularization import regularization_util
 
+def adaptive_regularization_weights_from(
+    inner_coefficient: float, outer_coefficient: float, pixel_signals: jnp.ndarray
+) -> jnp.ndarray:
+    """
+    Returns the regularization weights for the adaptive regularization scheme (e.g. ``AdaptiveBrightness``).
+
+    The weights define the effective regularization coefficient of every mesh parameter (typically pixels
+    of a ``Mapper``).
+
+    They are computed using an estimate of the expected signal in each pixel.
+
+    Two regularization coefficients are used, corresponding to the:
+
+    1) pixel_signals: pixels with a high pixel-signal (i.e. where the signal is located in the pixelization).
+    2) 1.0 - pixel_signals: pixels with a low pixel-signal (i.e. where the signal is not located in the pixelization).
+
+    Parameters
+    ----------
+    inner_coefficient
+        The inner regularization coefficients which controls the degree of smoothing of the inversion reconstruction
+        in the inner regions of a mesh's reconstruction.
+    outer_coefficient
+        The outer regularization coefficients which controls the degree of smoothing of the inversion reconstruction
+        in the outer regions of a mesh's reconstruction.
+    pixel_signals
+        The estimated signal in every pixelization pixel, used to change the regularization weighting of high signal
+        and low signal pixelizations.
+
+    Returns
+    -------
+    jnp.ndarray
+        The adaptive regularization weights which act as the effective regularization coefficients of
+        every source pixel.
+    """
+    return (
+        inner_coefficient * pixel_signals + outer_coefficient * (1.0 - pixel_signals)
+    ) ** 2.0
+
+
+def weighted_regularization_matrix_from(
+    regularization_weights: jnp.ndarray,
+    neighbors: jnp.ndarray,
+) -> jnp.ndarray:
+    """
+    Returns the regularization matrix of the adaptive regularization scheme (e.g. ``AdaptiveBrightness``).
+
+    This matrix is computed using the regularization weights of every mesh pixel, which are computed using the
+    function ``adaptive_regularization_weights_from``. These act as the effective regularization coefficients of
+    every mesh pixel.
+
+    The regularization matrix is computed using the pixel-neighbors array, which is setup using the appropriate
+    neighbor calculation of the corresponding ``Mapper`` class.
+
+    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).
+    neighbors_sizes
+        An array of length (total_pixels) which gives the number of neighbors of every pixel in the
+        Voronoi grid.
+
+    Returns
+    -------
+    jnp.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 = 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,)
+
+    # 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)
+
+    # 5) Scatter into the diagonal:
+    #    - the tiny 1e-8 floor on each i < S
+    #    - 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)
+    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)
+
+    # 6) Scatter the off‐diagonal subtractions:
+    mat = mat.at[I, J].add(-w_ij)
+    mat = mat.at[J, I].add(-w_ij)
+
+    # 7) Drop the extra row/column S and return the S×S result
+    return mat[:S, :S]
 
 class AdaptiveBrightness(AbstractRegularization):
     def __init__(
@@ -70,7 +177,7 @@ def __init__(
         self.outer_coefficient = outer_coefficient
         self.signal_scale = signal_scale
 
-    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.
 
@@ -91,13 +198,13 @@ def regularization_weights_from(self, linear_obj: LinearObj) -> np.ndarray:
         """
         pixel_signals = linear_obj.pixel_signals_from(signal_scale=self.signal_scale)
 
-        return regularization_util.adaptive_regularization_weights_from(
+        return adaptive_regularization_weights_from(
             inner_coefficient=self.inner_coefficient,
             outer_coefficient=self.outer_coefficient,
             pixel_signals=pixel_signals,
         )
 
-    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].
 
@@ -112,7 +219,7 @@ def regularization_matrix_from(self, linear_obj: LinearObj) -> np.ndarray:
         """
         regularization_weights = self.regularization_weights_from(linear_obj=linear_obj)
 
-        return regularization_util.weighted_regularization_matrix_from(
+        return weighted_regularization_matrix_from(
             regularization_weights=regularization_weights,
             neighbors=linear_obj.source_plane_mesh_grid.neighbors,
         )

autoarray/inversion/regularization/brightness_zeroth.py

@@ -1,5 +1,5 @@
 from __future__ import annotations
-import numpy as np
+import jax.numpy as jnp
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
@@ -10,6 +10,60 @@
 from autoarray.inversion.regularization import regularization_util
 
 
+def brightness_zeroth_regularization_weights_from(
+    coefficient: float, pixel_signals: jnp.ndarray
+) -> jnp.ndarray:
+    """
+    Returns the regularization weights for the brightness zeroth regularization scheme (e.g. ``BrightnessZeroth``).
+
+    The weights define the level of zeroth order regularization applied to every mesh parameter (typically pixels
+    of a ``Mapper``).
+
+    They are computed using an estimate of the expected signal in each pixel.
+
+    The zeroth order regularization coefficients is applied in combination with  1.0 - pixel_signals, which are
+    the pixels with a low pixel-signal (i.e. where the signal is not located near the source being reconstructed in
+    the pixelization).
+
+    Parameters
+    ----------
+    coefficient
+        The level of zeroth order regularization applied to every mesh parameter (typically pixels of a ``Mapper``),
+        with the degree applied varying based on the ``pixel_signals``.
+    pixel_signals
+        The estimated signal in every pixelization pixel, used to change the regularization weighting of high signal
+        and low signal pixelizations.
+
+    Returns
+    -------
+    jnp.ndarray
+        The zeroth order regularization weights which act as the effective level of zeroth order regularization
+        applied to every mesh parameter.
+    """
+    return coefficient * (1.0 - pixel_signals)
+
+def brightness_zeroth_regularization_matrix_from(
+    regularization_weights: jnp.ndarray,
+) -> jnp.ndarray:
+    """
+    Returns the regularization matrix for the zeroth-order brightness regularization scheme.
+
+    Parameters
+    ----------
+    regularization_weights
+        The regularization weights for each pixel, governing the strength of zeroth-order
+        regularization applied per inversion parameter.
+
+    Returns
+    -------
+    A diagonal regularization matrix where each diagonal element is the squared regularization weight
+    for that pixel.
+    """
+    regularization_weight_squared = regularization_weights**2.0
+    return jnp.diag(regularization_weight_squared)
+
+
+
 class BrightnessZeroth(AbstractRegularization):
     def __init__(
         self,
@@ -45,7 +99,7 @@ def __init__(
         self.coefficient = coefficient
         self.signal_scale = signal_scale
 
-    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 the ``BrightnessZeroth`` regularization scheme.
 
@@ -65,11 +119,11 @@ def regularization_weights_from(self, linear_obj: LinearObj) -> np.ndarray:
         """
         pixel_signals = linear_obj.pixel_signals_from(signal_scale=self.signal_scale)
 
-        return regularization_util.brightness_zeroth_regularization_weights_from(
+        return brightness_zeroth_regularization_weights_from(
             coefficient=self.coefficient, pixel_signals=pixel_signals
         )
 
-    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].
 
@@ -84,6 +138,6 @@ def regularization_matrix_from(self, linear_obj: LinearObj) -> np.ndarray:
         """
         regularization_weights = self.regularization_weights_from(linear_obj=linear_obj)
 
-        return regularization_util.brightness_zeroth_regularization_matrix_from(
+        return brightness_zeroth_regularization_matrix_from(
             regularization_weights=regularization_weights
         )

autoarray/inversion/regularization/constant.py

@@ -1,13 +1,61 @@
 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.inversion.regularization import regularization_util
+
+def constant_regularization_matrix_from(
+    coefficient: float,
+    neighbors: jnp.ndarray[[int, int], jnp.int64],
+    neighbors_sizes: jnp.ndarray[[int], jnp.int64],
+) -> jnp.ndarray[[int, int], jnp.float64]:
+    """
+    From the pixel-neighbors array, setup the regularization matrix using the instance regularization scheme.
+
+    A complete description of regularizatin and the `regularization_matrix` can be found in the `Regularization`
+    class in the module `autoarray.inversion.regularization`.
+
+    Memory requirement: 2SP + S^2
+    FLOPS: 1 + 2S + 2SP
+
+    Parameters
+    ----------
+    coefficient
+        The regularization coefficients which controls the degree of smoothing of the inversion reconstruction.
+    neighbors : ndarray, shape (S, P), dtype=int64
+        An array of length (total_pixels) which provides the index of all neighbors of every pixel in
+        the Voronoi grid (entries of -1 correspond to no neighbor).
+    neighbors_sizes : ndarray, shape (S,), dtype=int64
+        An array of length (total_pixels) which gives the number of neighbors of every pixel in the
+        Voronoi grid.
+
+    Returns
+    -------
+    regularization_matrix : ndarray, shape (S, S), dtype=float64
+        The regularization matrix computed using Regularization where the effective regularization
+        coefficient of every source pixel is the same.
+    """
+    S, P = neighbors.shape
+    # as the regularization matrix is S by S, S would be out of bound (any out of bound index would do)
+    OUT_OF_BOUND_IDX = S
+    regularization_coefficient = coefficient * coefficient
+
+    # flatten it for feeding into the matrix as j indices
+    neighbors = neighbors.flatten()
+    # now create the corresponding i indices
+    I_IDX = jnp.repeat(jnp.arange(S), P)
+    # Entries of `-1` in `neighbors` (indicating no neighbor) are replaced with an out-of-bounds index.
+    # 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]
+        # unique indices should be guranteed by neighbors-spec
+        .add(-regularization_coefficient, mode="drop", unique_indices=True)
+    )
 
 
 class Constant(AbstractRegularization):
@@ -38,7 +86,7 @@ def __init__(self, coefficient: 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.
 
@@ -57,9 +105,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].
 
@@ -73,7 +121,7 @@ def regularization_matrix_from(self, linear_obj: LinearObj) -> np.ndarray:
         The regularization matrix.
         """
 
-        return regularization_util.constant_regularization_matrix_from(
+        return constant_regularization_matrix_from(
             coefficient=self.coefficient,
             neighbors=linear_obj.neighbors,
             neighbors_sizes=linear_obj.neighbors.sizes,