remove mesh_tuil

Author: Jammy2211

autoarray/inversion/pixelization/mappers/mapper_util.py

@@ -3,359 +3,10 @@
 from typing import Tuple
 
 
-def forward_interp(xp, yp, x):
 
-    import jax
-    import jax.numpy as jnp
 
-    return jax.vmap(jnp.interp, in_axes=(1, 1, 1, None, None), out_axes=(1))(
-        x, xp, yp, 0, 1
-    )
 
 
-def reverse_interp(xp, yp, x):
-    import jax
-    import jax.numpy as jnp
-
-    return jax.vmap(jnp.interp, in_axes=(1, 1, 1), out_axes=(1))(x, xp, yp)
-
-
-def forward_interp_np(xp, yp, x):
-    """
-    xp: (N, M)
-    yp: (N, M)
-    x : (M,)  ← one x per column
-    """
-
-    if yp.ndim == 1 and xp.ndim == 2:
-        yp = np.broadcast_to(yp[:, None], xp.shape)
-
-    K, M = x.shape
-
-    out = np.empty((K, 2), dtype=xp.dtype)
-
-    for j in range(2):
-        out[:, j] = np.interp(x[:, j], xp[:, j], yp[:, j], left=0, right=1)
-
-    return out
-
-
-def reverse_interp_np(xp, yp, x):
-    """
-    xp : (N,) or (N, M)
-    yp : (N, M)
-    x  : (K, M)   query points per column
-    """
-
-    # Ensure xp is 2D: (N, M)
-    if xp.ndim == 1 and yp.ndim == 2:  # (N, 1)
-        xp = np.broadcast_to(xp[:, None], yp.shape)
-
-    # Shapes
-    K, M = x.shape
-
-    # Output
-    out = np.empty((K, 2), dtype=yp.dtype)
-
-    # Column-wise interpolation (cannot avoid this loop in pure NumPy)
-    for j in range(2):
-        out[:, j] = np.interp(x[:, j], xp[:, j], yp[:, j])
-
-    return out
-
-
-def create_transforms(traced_points, mesh_weight_map=None, xp=np):
-
-    N = traced_points.shape[0]  # // 2
-
-    if mesh_weight_map is None:
-        t = xp.arange(1, N + 1) / (N + 1)
-        t = xp.stack([t, t], axis=1)
-        sort_points = xp.sort(traced_points, axis=0)  # [::2]
-    else:
-        sdx = xp.argsort(traced_points, axis=0)
-        sort_points = xp.take_along_axis(traced_points, sdx, axis=0)
-        t = xp.stack([mesh_weight_map, mesh_weight_map], axis=1)
-        t = xp.take_along_axis(t, sdx, axis=0)
-        t = xp.cumsum(t, axis=0)
-
-    if xp.__name__.startswith("jax"):
-        transform = partial(forward_interp, sort_points, t)
-        inv_transform = partial(reverse_interp, t, sort_points)
-        return transform, inv_transform
-
-    transform = partial(forward_interp_np, sort_points, t)
-    inv_transform = partial(reverse_interp_np, t, sort_points)
-    return transform, inv_transform
-
-
-def adaptive_rectangular_transformed_grid_from(
-    source_plane_data_grid, grid, mesh_weight_map=None, xp=np
-):
-
-    mu = source_plane_data_grid.mean(axis=0)
-    scale = source_plane_data_grid.std(axis=0).min()
-    source_grid_scaled = (source_plane_data_grid - mu) / scale
-
-    transform, inv_transform = create_transforms(
-        source_grid_scaled, mesh_weight_map=mesh_weight_map, xp=xp
-    )
-
-    def inv_full(U):
-        return inv_transform(U) * scale + mu
-
-    return inv_full(grid)
-
-
-def adaptive_rectangular_areas_from(
-    source_grid_shape, source_plane_data_grid, mesh_weight_map=None, xp=np
-):
-
-    edges_y = xp.linspace(1, 0, source_grid_shape[0] + 1)
-    edges_x = xp.linspace(0, 1, source_grid_shape[1] + 1)
-
-    mu = source_plane_data_grid.mean(axis=0)
-    scale = source_plane_data_grid.std(axis=0).min()
-    source_grid_scaled = (source_plane_data_grid - mu) / scale
-
-    transform, inv_transform = create_transforms(
-        source_grid_scaled, mesh_weight_map=mesh_weight_map, xp=xp
-    )
-
-    def inv_full(U):
-        return inv_transform(U) * scale + mu
-
-    pixel_edges = inv_full(xp.stack([edges_y, edges_x]).T)
-    pixel_lengths = xp.diff(pixel_edges, axis=0).squeeze()  # shape (N_source, 2)
-
-    dy = pixel_lengths[:, 0]
-    dx = pixel_lengths[:, 1]
-
-    return xp.abs(xp.outer(dy, dx).flatten())
-
-
-def adaptive_rectangular_mappings_weights_via_interpolation_from(
-    source_grid_size: int,
-    source_plane_data_grid,
-    source_plane_data_grid_over_sampled,
-    mesh_weight_map=None,
-    xp=np,
-):
-    """
-    Compute bilinear interpolation indices and weights for mapping an oversampled
-    source-plane grid onto a regular rectangular pixelization.
-
-    This function takes a set of irregularly-sampled source-plane coordinates and
-    builds an adaptive mapping onto a `source_grid_size x source_grid_size` rectangular
-    pixelization using bilinear interpolation. The interpolation is expressed as:
-
-        f(x, y) ≈ w_bl * f(ix_down, iy_down) +
-                  w_br * f(ix_up,   iy_down) +
-                  w_tl * f(ix_down, iy_up) +
-                  w_tr * f(ix_up,   iy_up)
-
-    where `(ix_down, ix_up, iy_down, iy_up)` are the integer grid coordinates
-    surrounding the continuous position `(x, y)`.
-
-    Steps performed:
-      1. Normalize the source-plane grid by subtracting its mean and dividing by
-         the minimum axis standard deviation (to balance scaling).
-      2. Construct forward/inverse transforms which map the grid into the unit square [0,1]^2.
-      3. Transform the oversampled source-plane grid into [0,1]^2, then scale it
-         to index space `[0, source_grid_size)`.
-      4. Compute floor/ceil along x and y axes to find the enclosing rectangular cell.
-      5. Build the four corner indices: bottom-left (bl), bottom-right (br),
-         top-left (tl), and top-right (tr).
-      6. Flatten the 2D indices into 1D indices suitable for scatter operations,
-         with a flipped row-major convention: row = source_grid_size - i, col = j.
-      7. Compute bilinear interpolation weights (`w_bl, w_br, w_tl, w_tr`).
-      8. Return arrays of flattened indices and weights of shape `(N, 4)`, where
-         `N` is the number of oversampled coordinates.
-
-    Parameters
-    ----------
-    source_grid_size : int
-        The number of pixels along one dimension of the rectangular pixelization.
-        The grid is square: (source_grid_size x source_grid_size).
-    source_plane_data_grid : (M, 2) ndarray
-        The base source-plane coordinates, used to define normalization and transforms.
-    source_plane_data_grid_over_sampled : (N, 2) ndarray
-        Oversampled source-plane coordinates to be interpolated onto the rectangular grid.
-    mesh_weight_map
-        The weight map used to weight the creation of the rectangular mesh grid, which is used for the
-        `RectangularBrightness` mesh which adapts the size of its pixels to where the source is reconstructed.
-
-    Returns
-    -------
-    flat_indices : (N, 4) int ndarray
-        The flattened indices of the four neighboring pixel corners for each oversampled point.
-        Order: [bl, br, tl, tr].
-    weights : (N, 4) float ndarray
-        The bilinear interpolation weights for each of the four neighboring pixels.
-        Order: [w_bl, w_br, w_tl, w_tr].
-    """
-
-    # --- Step 1. Normalize grid ---
-    mu = source_plane_data_grid.mean(axis=0)
-    scale = source_plane_data_grid.std(axis=0).min()
-    source_grid_scaled = (source_plane_data_grid - mu) / scale
-
-    # --- Step 2. Build transforms ---
-    transform, inv_transform = create_transforms(
-        source_grid_scaled, mesh_weight_map=mesh_weight_map, xp=xp
-    )
-
-    # --- Step 3. Transform oversampled grid into index space ---
-    grid_over_sampled_scaled = (source_plane_data_grid_over_sampled - mu) / scale
-    grid_over_sampled_transformed = transform(grid_over_sampled_scaled)
-    grid_over_index = (source_grid_size - 3) * grid_over_sampled_transformed + 1
-
-    # --- Step 4. Floor/ceil indices ---
-    ix_down = xp.floor(grid_over_index[:, 0])
-    ix_up = xp.ceil(grid_over_index[:, 0])
-    iy_down = xp.floor(grid_over_index[:, 1])
-    iy_up = xp.ceil(grid_over_index[:, 1])
-
-    # --- Step 5. Four corners ---
-    idx_tl = xp.stack([ix_up, iy_down], axis=1)
-    idx_tr = xp.stack([ix_up, iy_up], axis=1)
-    idx_br = xp.stack([ix_down, iy_up], axis=1)
-    idx_bl = xp.stack([ix_down, iy_down], axis=1)
-
-    # --- Step 6. Flatten indices ---
-    def flatten(idx, n):
-        row = n - idx[:, 0]
-        col = idx[:, 1]
-        return row * n + col
-
-    flat_tl = flatten(idx_tl, source_grid_size)
-    flat_tr = flatten(idx_tr, source_grid_size)
-    flat_bl = flatten(idx_bl, source_grid_size)
-    flat_br = flatten(idx_br, source_grid_size)
-
-    flat_indices = xp.stack([flat_tl, flat_tr, flat_bl, flat_br], axis=1).astype(
-        "int64"
-    )
-
-    # --- Step 7. Bilinear interpolation weights ---
-    t_row = (grid_over_index[:, 0] - ix_down) / (ix_up - ix_down + 1e-12)
-    t_col = (grid_over_index[:, 1] - iy_down) / (iy_up - iy_down + 1e-12)
-
-    # Weights
-    w_tl = (1 - t_row) * (1 - t_col)
-    w_tr = (1 - t_row) * t_col
-    w_bl = t_row * (1 - t_col)
-    w_br = t_row * t_col
-    weights = xp.stack([w_tl, w_tr, w_bl, w_br], axis=1)
-
-    return flat_indices, weights
-
-
-def rectangular_mappings_weights_via_interpolation_from(
-    shape_native: Tuple[int, int],
-    source_plane_data_grid: np.ndarray,
-    source_plane_mesh_grid: np.ndarray,
-    xp=np,
-):
-    """
-    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 = xp.floor(fx).astype(xp.int32)
-    iy = xp.floor(fy).astype(xp.int32)
-
-    # Clip to stay within bounds
-    ix = xp.clip(ix, 0, Nx - 2)
-    iy = xp.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 = xp.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 = xp.stack([i00, i10, i01, i11], axis=1)  # shape (N_irregular, 4)
-
-    return mappings, weights
-
-
-def nearest_pixelization_index_for_slim_index_from_kdtree(grid, mesh_grid):
-    from scipy.spatial import cKDTree
-
-    kdtree = cKDTree(mesh_grid)
-
-    sparse_index_for_slim_index = []
-
-    for i in range(grid.shape[0]):
-        input_point = [grid[i, [0]], grid[i, 1]]
-        index = kdtree.query(input_point)[1]
-        sparse_index_for_slim_index.append(index)
-
-    return sparse_index_for_slim_index
 
 
 def adaptive_pixel_signals_from(