black

Author: Jammy2211

autoarray/inversion/pixelization/mappers/delaunay.py

@@ -9,11 +9,11 @@
 
 
 def pix_indexes_for_sub_slim_index_delaunay_from(
-    source_plane_data_grid,          # (N_sub, 2)
+    source_plane_data_grid,  # (N_sub, 2)
     simplex_index_for_sub_slim_index,  # (N_sub,)
-    pix_indexes_for_simplex_index,   # (M, 3)
-    delaunay_points,                 # (N_points, 2)
-    xp=np,                           # <--- choose backend: np or jnp
+    pix_indexes_for_simplex_index,  # (M, 3)
+    delaunay_points,  # (N_points, 2)
+    xp=np,  # <--- choose backend: np or jnp
 ):
     """
     XP-compatible version of pix_indexes_for_sub_slim_index_delaunay_from.
@@ -34,7 +34,7 @@ def ones_like(x, dtype):
     N_sub = source_plane_data_grid.shape[0]
 
     # Boolean mask for points that fall inside simplices
-    inside_mask = simplex_index_for_sub_slim_index >= 0     # shape (N_sub,)
+    inside_mask = simplex_index_for_sub_slim_index >= 0  # shape (N_sub,)
     outside_mask = ~inside_mask
 
     # ----------------------------
@@ -56,7 +56,7 @@ def ones_like(x, dtype):
         axis=-1,
     )
 
-    nearest = xp.argmin(d2, axis=1).astype(np.int32)   # (N_sub,)
+    nearest = xp.argmin(d2, axis=1).astype(np.int32)  # (N_sub,)
 
     # (N_sub, 3) → [nearest, -1, -1]
     nn_triplets = xp.stack(
@@ -93,23 +93,24 @@ def ones_like(x, dtype):
 
 import numpy as np
 
+
 def triangle_area_xp(c0, c1, c2, xp):
     """
     Twice triangle area using vector cross product magnitude.
     Calling via xp ensures NumPy or JAX backend operation.
     """
-    v0 = c1 - c0   # (..., 2)
+    v0 = c1 - c0  # (..., 2)
     v1 = c2 - c0
     cross = v0[..., 0] * v1[..., 1] - v0[..., 1] * v1[..., 0]
     return xp.abs(cross)
 
 
 def pixel_weights_delaunay_from(
-    source_plane_data_grid,         # (N_sub, 2)
-    source_plane_mesh_grid,         # (N_pix, 2)
+    source_plane_data_grid,  # (N_sub, 2)
+    source_plane_mesh_grid,  # (N_pix, 2)
     slim_index_for_sub_slim_index,  # (N_sub,)  UNUSED? kept for signature compatibility
-    pix_indexes_for_sub_slim_index, # (N_sub, 3), padded with -1
-    xp=np,                          # backend: np (default) or jnp
+    pix_indexes_for_sub_slim_index,  # (N_sub, 3), padded with -1
+    xp=np,  # backend: np (default) or jnp
 ):
     """
     XP-compatible (NumPy/JAX) version of pixel_weights_delaunay_from.
@@ -124,7 +125,7 @@ def pixel_weights_delaunay_from(
     # -----------------------------
     # If pix_indexes_for_sub_slim_index[sub][1] == -1 → NOT in simplex
     has_simplex = pix_indexes_for_sub_slim_index[:, 1] != -1  # (N_sub,)
-    no_simplex  = xp.logical_not(has_simplex)
+    no_simplex = xp.logical_not(has_simplex)
 
     # -----------------------------
     # GATHER TRIANGLE VERTICES
@@ -135,12 +136,12 @@ def pixel_weights_delaunay_from(
     # (N_sub, 3, 2)
     vertices = source_plane_mesh_grid[safe_indices]
 
-    p0 = vertices[:, 0]   # (N_sub, 2)
+    p0 = vertices[:, 0]  # (N_sub, 2)
     p1 = vertices[:, 1]
     p2 = vertices[:, 2]
 
     # Query points
-    q = source_plane_data_grid   # (N_sub, 2)
+    q = source_plane_data_grid  # (N_sub, 2)
 
     # -----------------------------
     # TRIANGLE AREAS (barycentric numerators)
@@ -164,23 +165,17 @@ def pixel_weights_delaunay_from(
             xp.zeros(N_sub),
             xp.zeros(N_sub),
         ],
-        axis=1
+        axis=1,
     )
 
     # -----------------------------
     # SELECT BETWEEN CASES
     # -----------------------------
-    pixel_weights = xp.where(
-        has_simplex[:, None],
-        weights_bary,
-        weights_nn
-    )
+    pixel_weights = xp.where(has_simplex[:, None], weights_bary, weights_nn)
 
     return pixel_weights
 
 
-
-
 class MapperDelaunay(AbstractMapper):
     """
     To understand a `Mapper` one must be familiar `Mesh` objects and the `mesh` and `pixelization` packages, where
@@ -285,14 +280,12 @@ def pix_sub_weights(self) -> PixSubWeights:
         )
         pix_indexes_for_simplex_index = delaunay.simplices
 
-        mappings, sizes = (
-            pix_indexes_for_sub_slim_index_delaunay_from(
-                source_plane_data_grid=self.source_plane_data_grid.over_sampled,
-                simplex_index_for_sub_slim_index=simplex_index_for_sub_slim_index,
-                pix_indexes_for_simplex_index=pix_indexes_for_simplex_index,
-                delaunay_points=delaunay.points,
-                xp=self._xp,
-            )
+        mappings, sizes = pix_indexes_for_sub_slim_index_delaunay_from(
+            source_plane_data_grid=self.source_plane_data_grid.over_sampled,
+            simplex_index_for_sub_slim_index=simplex_index_for_sub_slim_index,
+            pix_indexes_for_simplex_index=pix_indexes_for_simplex_index,
+            delaunay_points=delaunay.points,
+            xp=self._xp,
         )
 
         mappings = mappings.astype("int")
@@ -351,7 +344,9 @@ def pix_sub_weights_split_cross(self) -> PixSubWeights:
         append_line_float = np.zeros((len(splitted_weights), 1), dtype="float")
 
         return PixSubWeights(
-            mappings=self._xp.hstack((splitted_mappings.astype(self._xp.int32), append_line_int)),
+            mappings=self._xp.hstack(
+                (splitted_mappings.astype(self._xp.int32), append_line_int)
+            ),
             sizes=splitted_sizes.astype(self._xp.int32),
             weights=self._xp.hstack((splitted_weights, append_line_float)),
         )

autoarray/inversion/pixelization/mesh/delaunay.py

@@ -54,6 +54,4 @@ def mesh_grid_from(
             Settings controlling the pixelization for example if a border is used to relocate its exterior coordinates.
         """
 
-        return Mesh2DDelaunay(
-            values=source_plane_mesh_grid, _xp=xp
-        )
+        return Mesh2DDelaunay(values=source_plane_mesh_grid, _xp=xp)

autoarray/inversion/pixelization/mesh/triangulation.py

@@ -64,21 +64,21 @@ def mapper_grids_from(
         relocated_grid = self.relocated_grid_from(
             border_relocator=border_relocator,
             source_plane_data_grid=source_plane_data_grid,
-            xp=xp
+            xp=xp,
         )
 
         relocated_mesh_grid = self.relocated_mesh_grid_from(
             border_relocator=border_relocator,
             source_plane_data_grid=relocated_grid.over_sampled,
             source_plane_mesh_grid=source_plane_mesh_grid,
-            xp=xp
+            xp=xp,
         )
 
         try:
             source_plane_mesh_grid = self.mesh_grid_from(
                 source_plane_data_grid=relocated_grid.over_sampled,
                 source_plane_mesh_grid=relocated_mesh_grid,
-            xp=xp
+                xp=xp,
             )
         except ValueError as e:
             raise e

autoarray/inversion/regularization/adaptive_brightness_split.py

@@ -104,13 +104,13 @@ def regularization_matrix_from(self, linear_obj: LinearObj, xp=np) -> np.ndarray
             splitted_mappings=pix_sub_weights_split_cross.mappings,
             splitted_sizes=pix_sub_weights_split_cross.sizes,
             splitted_weights=pix_sub_weights_split_cross.weights,
-            xp=xp
+            xp=xp,
         )
 
         return regularization_util.pixel_splitted_regularization_matrix_from(
             regularization_weights=regularization_weights,
             splitted_mappings=splitted_mappings,
             splitted_sizes=splitted_sizes,
             splitted_weights=splitted_weights,
-            xp=xp
+            xp=xp,
         )

autoarray/inversion/regularization/adaptive_brightness_split_zeroth.py

@@ -106,7 +106,7 @@ def regularization_matrix_from(self, linear_obj: LinearObj, xp=np) -> np.ndarray
             splitted_mappings=pix_sub_weights_split_cross.mappings,
             splitted_sizes=pix_sub_weights_split_cross.sizes,
             splitted_weights=pix_sub_weights_split_cross.weights,
-            xp=xp
+            xp=xp,
         )
 
         regularization_matrix = (
@@ -115,7 +115,7 @@ def regularization_matrix_from(self, linear_obj: LinearObj, xp=np) -> np.ndarray
                 splitted_mappings=splitted_mappings,
                 splitted_sizes=splitted_sizes,
                 splitted_weights=splitted_weights,
-                xp=xp
+                xp=xp,
             )
         )
 

autoarray/inversion/regularization/constant_split.py

@@ -66,7 +66,7 @@ def regularization_matrix_from(self, linear_obj: LinearObj, xp=np) -> np.ndarray
             splitted_mappings=pix_sub_weights_split_cross.mappings,
             splitted_sizes=pix_sub_weights_split_cross.sizes,
             splitted_weights=pix_sub_weights_split_cross.weights,
-            xp=xp
+            xp=xp,
         )
 
         pixels = int(len(splitted_mappings) / 4)
@@ -77,5 +77,5 @@ def regularization_matrix_from(self, linear_obj: LinearObj, xp=np) -> np.ndarray
             splitted_mappings=splitted_mappings,
             splitted_sizes=splitted_sizes,
             splitted_weights=splitted_weights,
-            xp=xp
+            xp=xp,
         )

autoarray/inversion/regularization/regularization_util.py

@@ -82,6 +82,7 @@ def reg_split_np_from(
 
     return splitted_mappings, splitted_sizes, splitted_weights
 
+
 def reg_split_from(
     splitted_mappings: np.ndarray,
     splitted_sizes: np.ndarray,
@@ -143,27 +144,27 @@ def reg_split_from(
     # 2. Pixel index for each row: i // 4
     # -------------------------------------------------------------
     pixel_index = (jnp.arange(N) // 4).astype(mappings.dtype)  # (N,)
-    pix_b = pixel_index[:, None]                                # (N,1)
+    pix_b = pixel_index[:, None]  # (N,1)
 
     # -------------------------------------------------------------
     # 3. Mask of valid columns j < size[i]
     # -------------------------------------------------------------
-    cols = jnp.arange(K)[None, :]         # (1,4)
-    valid_mask = cols < sizes[:, None]    # (N,4)
+    cols = jnp.arange(K)[None, :]  # (1,4)
+    valid_mask = cols < sizes[:, None]  # (N,4)
 
     # -------------------------------------------------------------
     # 4. Self match: mapping[i,j] == pixel_index AND j is valid
     # -------------------------------------------------------------
-    self_mask = (mappings == pix_b) & valid_mask               # (N,4)
-    row_has_self = jnp.any(self_mask, axis=1)                  # (N,)
+    self_mask = (mappings == pix_b) & valid_mask  # (N,4)
+    row_has_self = jnp.any(self_mask, axis=1)  # (N,)
 
     # Position of self per row
-    self_pos = jnp.argmax(self_mask, axis=1)                   # (N,)
+    self_pos = jnp.argmax(self_mask, axis=1)  # (N,)
 
     # -------------------------------------------------------------
     # 5. Add +1 weight at self_pos where row_has_self == True
     # -------------------------------------------------------------
-    one_hot = jnn.one_hot(self_pos, K, dtype=weights.dtype)    # (N,4)
+    one_hot = jnn.one_hot(self_pos, K, dtype=weights.dtype)  # (N,4)
     weights = weights + one_hot * row_has_self[:, None]
 
     # -------------------------------------------------------------
@@ -172,21 +173,19 @@ def reg_split_from(
     no_self = ~row_has_self
 
     # Insert position = sizes[i]
-    insert_pos = sizes                                          # (N,)
+    insert_pos = sizes  # (N,)
     insert_mask = no_self[:, None] & (cols == sizes[:, None])
 
     # New mappings and weights
     mappings = jnp.where(insert_mask, pix_b, mappings)
-    weights  = jnp.where(insert_mask, jnp.array(1.0, weights.dtype), weights)
+    weights = jnp.where(insert_mask, jnp.array(1.0, weights.dtype), weights)
 
     # Updated sizes: +1 if no self detected
     sizes_new = sizes + no_self.astype(sizes.dtype)
 
     return mappings, sizes_new, weights
 
 
-
-
 def pixel_splitted_regularization_matrix_np_from(
     regularization_weights: np.ndarray,
     splitted_mappings: np.ndarray,
@@ -228,12 +227,11 @@ def pixel_splitted_regularization_matrix_np_from(
     return regularization_matrix
 
 
-
 def pixel_splitted_regularization_matrix_from(
-    regularization_weights: np.ndarray,   # (P,)
-    splitted_mappings: np.ndarray,        # (4P, 4)
-    splitted_sizes: np.ndarray,           # (4P,)
-    splitted_weights: np.ndarray,          # (4P, 4)
+    regularization_weights: np.ndarray,  # (P,)
+    splitted_mappings: np.ndarray,  # (4P, 4)
+    splitted_sizes: np.ndarray,  # (4P,)
+    splitted_weights: np.ndarray,  # (4P, 4)
     xp=np,
 ):
     """
@@ -279,33 +277,33 @@ def pixel_splitted_regularization_matrix_from(
     P = splitted_mappings.shape[0] // 4
 
     # Square, positive regularization weights
-    reg_w = regularization_weights**2.0              # (P,)
+    reg_w = regularization_weights**2.0  # (P,)
 
     # Add diagonal jitter (2e-8)
-    reg_mat = jnp.eye(P) * 2e-8                      # (P, P)
+    reg_mat = jnp.eye(P) * 2e-8  # (P, P)
 
     # ----- Build all 4P contributions at once -----
 
     # Mask away padded entries (where mapping = -1)
-    valid = splitted_mappings != -1                  # (4P, 4)
+    valid = splitted_mappings != -1  # (4P, 4)
 
     # Extract valid mapping rows and weights
     # BUT keep fixed shape (4) and just zero out invalid ones
     map_fixed = jnp.where(valid, splitted_mappings, 0)  # (4P, 4)
-    w_fixed   = jnp.where(valid, splitted_weights,  0.) # (4P, 4)
+    w_fixed = jnp.where(valid, splitted_weights, 0.0)  # (4P, 4)
 
     # Compute all outer products of weights
     # w_fixed[:, :, None] * w_fixed[:, None, :]  → (4P, 4, 4)
-    outer = w_fixed[:, :, None] * w_fixed[:, None, :]   # (4P, 4, 4)
+    outer = w_fixed[:, :, None] * w_fixed[:, None, :]  # (4P, 4, 4)
 
     # Build corresponding row and col index grids
-    rows = map_fixed[:, :, None]   # (4P, 4, 1)
-    cols = map_fixed[:, None, :]   # (4P, 1, 4)
+    rows = map_fixed[:, :, None]  # (4P, 4, 1)
+    cols = map_fixed[:, None, :]  # (4P, 1, 4)
 
     # Multiply each 4x4 block by its pixel’s regularization weight
     # Rows 0–3 belong to pixel 0, rows 4–7 to pixel 1, etc.
-    pixel_index = jnp.arange(4 * P) // 4            # (4P,)
-    block_scale = reg_w[pixel_index]                # (4P,)
+    pixel_index = jnp.arange(4 * P) // 4  # (4P,)
+    block_scale = reg_w[pixel_index]  # (4P,)
     outer_scaled = outer * block_scale[:, None, None]
 
     # Now scatter-add all entries into the (P,P) matrix