Add KNNBarycentric mesh: pure-JAX barycentric weights on top-3 nearest

A JAX-native approximation to InterpolatorDelaunay that avoids the
scipy.spatial.Delaunay callback. Picks the 3 nearest mesh vertices in
source plane and computes locally-exact barycentric weights on the
triangle they form. When the 3 nearest are the containing Delaunay
triangle's vertices the weights are bit-identical to Delaunay;
otherwise weights are clipped non-negative and renormalized.
Degenerate (collinear) triangles fall back to nearest-neighbor.

Additive only: new `KNNBarycentric(KNearestNeighbor)` mesh class
selectable via `aa.mesh.KNNBarycentric(pixels=...)`, no existing
behaviour changed.

Author: Jammy2211

autoarray/inversion/mesh/interpolator/knn.py

@@ -241,3 +241,172 @@ def _mappings_sizes_weights_split(self):
     #         k=self.mesh.k_neighbors,
     #         radius_scale=self.mesh.radius_scale,
     #     )
+
+
+def barycentric_weights_from_3_nearest(
+    query_points,
+    mesh_points,
+    nearest_3_indices,
+    xp,
+):
+    """
+    Compute barycentric weights for each query point on the triangle formed by its
+    3 nearest mesh vertices.
+
+    Signed barycentric coordinates are computed, then clipped to be non-negative
+    and renormalized so each row sums to 1. Queries inside the triangle return
+    the exact Delaunay weights; queries outside return a clipped approximation
+    (a convex combination of the 3 nearest, biased toward whichever vertices are
+    on the same side of the triangle as the query).
+
+    Degenerate triangles (collinear vertices) get zero weights to avoid NaN.
+
+    Parameters
+    ----------
+    query_points : (Q, 2)
+        Query point (x, y) coordinates.
+    mesh_points : (N, 2)
+        Mesh vertex (x, y) coordinates.
+    nearest_3_indices : (Q, 3)
+        Indices into mesh_points of the 3 nearest vertices for each query.
+    xp : module
+        numpy or jax.numpy.
+
+    Returns
+    -------
+    weights : (Q, 3)
+        Barycentric weights, clipped non-negative and row-normalized.
+    """
+    vertices = mesh_points[nearest_3_indices]  # (Q, 3, 2)
+    p0 = vertices[:, 0]
+    p1 = vertices[:, 1]
+    p2 = vertices[:, 2]
+    q = query_points
+
+    def signed_cross(a, b, c):
+        return (b[..., 0] - a[..., 0]) * (c[..., 1] - a[..., 1]) - (
+            b[..., 1] - a[..., 1]
+        ) * (c[..., 0] - a[..., 0])
+
+    total = signed_cross(p0, p1, p2)
+    w0 = signed_cross(q, p1, p2)
+    w1 = signed_cross(p0, q, p2)
+    w2 = signed_cross(p0, p1, q)
+
+    eps = xp.asarray(1e-12, dtype=total.dtype)
+    safe_total = xp.where(xp.abs(total) > eps, total, 1.0)
+
+    bary = xp.stack([w0, w1, w2], axis=1) / safe_total[:, None]
+
+    clipped = xp.maximum(bary, 0.0)
+    row_sum = xp.sum(clipped, axis=1, keepdims=True)
+    safe_sum = xp.where(row_sum > eps, row_sum, 1.0)
+    weights = clipped / safe_sum
+
+    # Degenerate triangles fall back to nearest-neighbor (weight 1 on column 0,
+    # which `get_interpolation_weights` orders as the closest mesh vertex).
+    # Same fallback policy as `pix_indexes_for_sub_slim_index_delaunay_from`
+    # for outside-simplex points.
+    nearest_only = xp.asarray([1.0, 0.0, 0.0], dtype=weights.dtype)
+
+    degenerate = xp.abs(total) <= eps
+    weights = xp.where(degenerate[:, None], nearest_only[None, :], weights)
+
+    return weights
+
+
+class InterpolatorKNNBarycentric(InterpolatorKNearestNeighbor):
+    """
+    Interpolator that picks the 3 nearest mesh vertices in the source plane and
+    computes locally-exact barycentric weights on the triangle they form.
+
+    Approximates :class:`InterpolatorDelaunay` without the scipy.spatial.Delaunay
+    callback: when the 3 nearest are the containing Delaunay triangle's vertices,
+    the weights are bit-identical to Delaunay; otherwise they are clipped-and-
+    renormalized barycentric weights on whichever triangle the 3 nearest form.
+
+    The kNN connectivity knobs (``k_neighbors``, ``radius_scale``,
+    ``split_neighbor_division``) on the parent :class:`KNearestNeighbor` mesh are
+    inherited and still control the regularization-spacing computation via
+    ``distance_to_self``. Interpolation always uses k=3, irrespective of
+    ``mesh.k_neighbors``.
+    """
+
+    @cached_property
+    def _mappings_sizes_weights(self):
+
+        try:
+            query_points = self.data_grid.over_sampled.array
+        except AttributeError:
+            try:
+                query_points = self.data_grid.array
+            except AttributeError:
+                query_points = self.data_grid
+
+        mappings, _, _ = get_interpolation_weights(
+            points=self.mesh_grid_xy,
+            query_points=query_points,
+            k_neighbors=3,
+            radius_scale=1.0,
+        )
+
+        weights = barycentric_weights_from_3_nearest(
+            query_points=query_points,
+            mesh_points=self.mesh_grid_xy,
+            nearest_3_indices=mappings,
+            xp=self._xp,
+        )
+
+        mappings = self._xp.asarray(mappings)
+        weights = self._xp.asarray(weights)
+
+        sizes = self._xp.full(
+            (mappings.shape[0],),
+            mappings.shape[1],
+        )
+
+        return mappings, sizes, weights
+
+    @cached_property
+    def _mappings_sizes_weights_split(self):
+        """
+        Same spacing scheme as :class:`InterpolatorKNearestNeighbor` but the
+        split-point interpolator is :class:`InterpolatorKNNBarycentric` so the
+        split-regularization weights are also barycentric rather than Wendland.
+        """
+        from autoarray.inversion.regularization.regularization_util import (
+            split_points_from,
+        )
+
+        neighbor_index = int(self.mesh.k_neighbors) // self.mesh.split_neighbor_division
+
+        distance_to_self = self.distance_to_self
+        others = distance_to_self[:, 1:]
+        idx = int(neighbor_index) - 1
+        idx = max(0, min(idx, others.shape[1] - 1))
+        r_k = others[:, idx]
+
+        split_step = self.mesh.areas_factor * r_k
+
+        split_points = split_points_from(
+            points=self.mesh_grid.array,
+            area_weights=split_step,
+            xp=self._xp,
+        )
+
+        interpolator = InterpolatorKNNBarycentric(
+            mesh=self.mesh,
+            mesh_grid=self.mesh_grid,
+            data_grid=split_points,
+            xp=self._xp,
+        )
+
+        mappings = interpolator.mappings
+        weights = interpolator.weights
+
+        sizes = self._xp.full(
+            (mappings.shape[0],),
+            mappings.shape[1],
+        )
+
+        return mappings, sizes, weights

autoarray/inversion/mesh/mesh/__init__.py

@@ -6,3 +6,4 @@
 from .rectangular_uniform import RectangularUniform
 from .delaunay import Delaunay
 from .knn import KNearestNeighbor
+from .knn import KNNBarycentric

autoarray/inversion/mesh/mesh/knn.py

@@ -77,3 +77,28 @@ def interpolator_cls(self):
         )
 
         return InterpolatorKNearestNeighbor
+
+
+class KNNBarycentric(KNearestNeighbor):
+    """
+    A mesh that inherits k-nearest-neighbour connectivity from
+    :class:`KNearestNeighbor` but uses :class:`InterpolatorKNNBarycentric` to
+    compute interpolation weights as locally-exact barycentric coordinates on
+    the triangle formed by the 3 nearest source-plane mesh vertices, rather
+    than a Wendland kernel.
+
+    This is a JAX-native approximation to Delaunay barycentric interpolation
+    that avoids the scipy.spatial.Delaunay callback. The kNN connectivity knobs
+    (``k_neighbors``, ``radius_scale``, ``split_neighbor_division``) are
+    inherited and still control the regularization-spacing computation, but the
+    *interpolation* weights always use k=3 + barycentric and ignore them.
+    """
+
+    @property
+    def interpolator_cls(self):
+
+        from autoarray.inversion.mesh.interpolator.knn import (
+            InterpolatorKNNBarycentric,
+        )
+
+        return InterpolatorKNNBarycentric

test_autoarray/inversion/pixelization/interpolator/test_knn_barycentric.py

@@ -0,0 +1,129 @@
+import numpy as np
+import pytest
+
+from autoarray.inversion.mesh.interpolator.knn import (
+    barycentric_weights_from_3_nearest,
+)
+
+
+def test__weights__interior_query_matches_delaunay_exactly():
+    mesh = np.array([[0.0, 0.0], [1.0, 0.0], [0.0, 1.0]])
+    query = np.array([[0.25, 0.25]])
+    indices = np.array([[0, 1, 2]])
+
+    weights = barycentric_weights_from_3_nearest(
+        query_points=query,
+        mesh_points=mesh,
+        nearest_3_indices=indices,
+        xp=np,
+    )
+
+    assert weights[0] == pytest.approx([0.5, 0.25, 0.25], rel=1e-12)
+    assert weights[0].sum() == pytest.approx(1.0, rel=1e-12)
+
+
+def test__weights__query_on_edge_splits_50_50():
+    mesh = np.array([[0.0, 0.0], [1.0, 0.0], [0.0, 1.0]])
+    query = np.array([[0.5, 0.5]])
+    indices = np.array([[0, 1, 2]])
+
+    weights = barycentric_weights_from_3_nearest(
+        query_points=query,
+        mesh_points=mesh,
+        nearest_3_indices=indices,
+        xp=np,
+    )
+
+    assert weights[0] == pytest.approx([0.0, 0.5, 0.5], abs=1e-12)
+
+
+def test__weights__outside_query_clipped_and_renormalized():
+    mesh = np.array([[0.0, 0.0], [1.0, 0.0], [0.0, 1.0]])
+    query = np.array([[-0.5, -0.5]])
+    indices = np.array([[0, 1, 2]])
+
+    weights = barycentric_weights_from_3_nearest(
+        query_points=query,
+        mesh_points=mesh,
+        nearest_3_indices=indices,
+        xp=np,
+    )
+
+    assert (weights[0] >= 0.0).all()
+    assert weights[0].sum() == pytest.approx(1.0, rel=1e-12)
+    # p0 is opposite both edges the query is "outside" of → keeps all weight
+    assert weights[0, 0] == pytest.approx(1.0, rel=1e-12)
+    assert weights[0, 1] == pytest.approx(0.0, abs=1e-12)
+    assert weights[0, 2] == pytest.approx(0.0, abs=1e-12)
+
+
+def test__weights__degenerate_triangle_falls_back_to_nearest_neighbor():
+    mesh = np.array([[0.0, 0.0], [1.0, 0.0], [2.0, 0.0]])
+    query = np.array([[1.0, 0.0]])
+    # column 0 = closest by construction (set by get_interpolation_weights)
+    indices = np.array([[0, 1, 2]])
+
+    weights = barycentric_weights_from_3_nearest(
+        query_points=query,
+        mesh_points=mesh,
+        nearest_3_indices=indices,
+        xp=np,
+    )
+
+    assert np.isfinite(weights).all()
+    # Degenerate → fall back to nearest-neighbor: [1, 0, 0]
+    assert weights[0] == pytest.approx([1.0, 0.0, 0.0], abs=1e-12)
+
+
+def test__weights__matches_pixel_weights_delaunay_from__for_inside_triangle_grid():
+    from autoarray.inversion.mesh.interpolator.delaunay import (
+        pixel_weights_delaunay_from,
+    )
+
+    mesh = np.array([[0.0, 0.0], [1.0, 0.0], [0.0, 1.0]])
+    rng = np.random.default_rng(seed=42)
+    # 200 random barycentric-coordinate samples → guaranteed inside-triangle queries
+    u = rng.uniform(0, 1, size=200)
+    v = rng.uniform(0, 1, size=200)
+    inside = u + v <= 1.0
+    u, v = u[inside], v[inside]
+    queries = np.stack([u, v], axis=1)  # (M, 2) — inside the unit-right triangle
+
+    indices = np.tile(np.array([[0, 1, 2]]), (queries.shape[0], 1))
+
+    bary = barycentric_weights_from_3_nearest(
+        query_points=queries,
+        mesh_points=mesh,
+        nearest_3_indices=indices,
+        xp=np,
+    )
+
+    delaunay = pixel_weights_delaunay_from(
+        data_grid=queries,
+        mesh_grid=mesh,
+        pix_indexes_for_sub_slim_index=indices,
+        xp=np,
+    )
+
+    assert bary == pytest.approx(delaunay, rel=1e-12, abs=1e-12)
+
+
+def test__mesh_class__interpolator_cls_is_knn_barycentric():
+    import autoarray as aa
+    from autoarray.inversion.mesh.interpolator.knn import (
+        InterpolatorKNNBarycentric,
+    )
+
+    mesh = aa.mesh.KNNBarycentric(pixels=10)
+    assert mesh.interpolator_cls is InterpolatorKNNBarycentric
+
+
+def test__mesh_class__inherits_knn_knobs():
+    import autoarray as aa
+
+    mesh = aa.mesh.KNNBarycentric(pixels=10)
+    # Knobs inherited from KNearestNeighbor still control regularization spacing
+    assert mesh.k_neighbors == 10
+    assert mesh.radius_scale == 1.5
+    assert mesh.areas_factor == 0.5
+    assert mesh.split_neighbor_division == 2