rectangular grid uses interpolation

Author: Jammy2211

autoarray/inversion/pixelization/mappers/mapper_util.py

@@ -144,6 +144,97 @@ def data_slim_to_pixelization_unique_from(
     return data_to_pix_unique, data_weights, pix_lengths
 
 
+def rectangular_mappings_weights_via_interpolation_from(
+    shape_native: Tuple[int, int],
+    source_plane_data_grid: np.ndarray,
+    source_plane_mesh_grid: np.ndarray,
+):
+    """
+    Compute bilinear interpolation weights and corresponding rectangular mesh indices for an irregular grid.
+
+    Given a flattened regular rectangular mesh grid and an irregular grid of data points, this function
+    determines for each irregular point:
+    - the indices of the 4 nearest rectangular mesh pixels (top-left, top-right, bottom-left, bottom-right), and
+    - the bilinear interpolation weights with respect to those pixels.
+
+    The function supports JAX and is compatible with JIT compilation.
+
+    Parameters
+    ----------
+    shape_native
+        The shape (Ny, Nx) of the original rectangular mesh grid before flattening.
+    source_plane_data_grid
+        The irregular grid of (y, x) points to interpolate.
+    source_plane_mesh_grid
+        The flattened regular rectangular mesh grid of (y, x) coordinates.
+
+    Returns
+    -------
+    mappings : np.ndarray of shape (N, 4)
+        Indices of the four nearest rectangular mesh pixels in the flattened mesh grid.
+        Order is: top-left, top-right, bottom-left, bottom-right.
+    weights : np.ndarray of shape (N, 4)
+        Bilinear interpolation weights corresponding to the four nearest mesh pixels.
+
+    Notes
+    -----
+    - Assumes the mesh grid is uniformly spaced.
+    - The weights sum to 1 for each irregular point.
+    - Uses bilinear interpolation in the (y, x) coordinate system.
+    """
+    source_plane_mesh_grid = source_plane_mesh_grid.reshape(*shape_native, 2)
+
+    # Assume mesh is shaped (Ny, Nx, 2)
+    Ny, Nx = source_plane_mesh_grid.shape[:2]
+
+    # Get mesh spacings and lower corner
+    y_coords = source_plane_mesh_grid[:, 0, 0]  # shape (Ny,)
+    x_coords = source_plane_mesh_grid[0, :, 1]  # shape (Nx,)
+
+    dy = y_coords[1] - y_coords[0]
+    dx = x_coords[1] - x_coords[0]
+
+    y_min = y_coords[0]
+    x_min = x_coords[0]
+
+    # shape (N_irregular, 2)
+    irregular = source_plane_data_grid
+
+    # Compute normalized mesh coordinates (floating indices)
+    fy = (irregular[:, 0] - y_min) / dy
+    fx = (irregular[:, 1] - x_min) / dx
+
+    # Integer indices of top-left corners
+    ix = np.floor(fx).astype(np.int32)
+    iy = np.floor(fy).astype(np.int32)
+
+    # Clip to stay within bounds
+    ix = np.clip(ix, 0, Nx - 2)
+    iy = np.clip(iy, 0, Ny - 2)
+
+    # Local coordinates inside the cell (0 <= tx, ty <= 1)
+    tx = fx - ix
+    ty = fy - iy
+
+    # Bilinear weights
+    w00 = (1 - tx) * (1 - ty)
+    w10 = tx * (1 - ty)
+    w01 = (1 - tx) * ty
+    w11 = tx * ty
+
+    weights = np.stack([w00, w10, w01, w11], axis=1)  # shape (N_irregular, 4)
+
+    # Compute indices of 4 surrounding pixels in the flattened mesh
+    i00 = iy * Nx + ix
+    i10 = iy * Nx + (ix + 1)
+    i01 = (iy + 1) * Nx + ix
+    i11 = (iy + 1) * Nx + (ix + 1)
+
+    mappings = np.stack([i00, i10, i01, i11], axis=1)  # shape (N_irregular, 4)
+
+    return mappings, weights
+
+
 @numba_util.jit()
 def pix_indexes_for_sub_slim_index_delaunay_from(
     source_plane_data_grid,

autoarray/inversion/pixelization/mappers/rectangular.py

@@ -9,6 +9,7 @@
 from autoarray.numba_util import profile_func
 from autoarray.geometry import geometry_util
 
+from autoarray.inversion.pixelization.mappers import mapper_util
 
 class MapperRectangular(AbstractMapper):
     """
@@ -99,19 +100,17 @@ def pix_sub_weights(self) -> PixSubWeights:
         dimension of the array `pix_indexes_for_sub_slim_index` 1 and all entries in `pix_weights_for_sub_slim_index`
         are equal to 1.0.
         """
-        mappings = geometry_util.grid_pixel_indexes_2d_slim_from(
-            grid_scaled_2d_slim=np.array(self.source_plane_data_grid.over_sampled),
-            shape_native=self.source_plane_mesh_grid.shape_native,
-            pixel_scales=self.source_plane_mesh_grid.pixel_scales,
-            origin=self.source_plane_mesh_grid.origin,
-        ).astype("int")
 
-        mappings = mappings.reshape((len(mappings), 1))
+        mappings, weights = (
+            mapper_util.rectangular_mappings_weights_via_interpolation_from(
+                shape_native=self.shape_native,
+                source_plane_mesh_grid=self.source_plane_mesh_grid.array,
+                source_plane_data_grid=self.source_plane_data_grid.over_sampled,
+            )
+        )
 
         return PixSubWeights(
             mappings=mappings,
-            sizes=np.ones(len(mappings), dtype="int"),
-            weights=np.ones(
-                (len(self.source_plane_data_grid.over_sampled), 1), dtype="int"
-            ),
+            sizes=4 * np.ones(len(mappings), dtype="int"),
+            weights=weights,
         )