fix: isolate invalid Delaunay callback lanes

Author: Jammy2211

autoarray/inversion/mesh/interpolator/delaunay.py

@@ -96,20 +96,27 @@ def scipy_delaunay_tri_only(points_np):
         coplanar/duplicate (the walk then starts from simplex 0 instead,
         which is equally valid — the walk converges from any start).
     """
+    N = points_np.shape[0]
+    simplices_padded = -np.ones((2 * N, 3), dtype=np.int32)
+    simplex_neighbors = -np.ones((2 * N, 3), dtype=np.int32)
+    vertex_simplex = -np.ones(N, dtype=np.int32)
+
+    # This function is the body of a sequentially-vmapped host callback.  A
+    # qhull exception in one lane aborts the entire batch, so keep non-finite
+    # input lane-local and encode it with the callback's existing padding
+    # convention.  The JAX-side walk and weights preserve the lane as NaN.
+    if not np.isfinite(points_np).all():
+        return simplices_padded, simplex_neighbors, vertex_simplex
+
     from scipy.spatial import Delaunay
 
-    N = points_np.shape[0]
     tri = Delaunay(points_np)
     simplices = tri.simplices.astype(np.int32)  # (T, 3)
     T = simplices.shape[0]
 
-    simplices_padded = -np.ones((2 * N, 3), dtype=np.int32)
     simplices_padded[:T] = simplices
-
-    simplex_neighbors = -np.ones((2 * N, 3), dtype=np.int32)
     simplex_neighbors[:T] = tri.neighbors.astype(np.int32)
 
-    vertex_simplex = -np.ones(N, dtype=np.int32)
     simplex_ids = np.arange(T, dtype=np.int32)
     for k in range(3):
         vertex_simplex[simplices[:, k]] = simplex_ids
@@ -185,9 +192,11 @@ def cross(u, v):
 
     def weights_of(cur, q_chunk):
         verts = simplices_padded[cur]  # (chunk, 3)
-        a = points[verts[:, 0]]
-        b = points[verts[:, 1]]
-        c = points[verts[:, 2]]
+        valid_simplex = (verts >= 0).all(axis=1)
+        safe_verts = verts.clip(min=0)
+        a = points[safe_verts[:, 0]]
+        b = points[safe_verts[:, 1]]
+        c = points[safe_verts[:, 2]]
         den = cross(b - a, c - a)
         den = xp.where(den != 0.0, den, 1.0)
         w = (
@@ -201,13 +210,14 @@ def weights_of(cur, q_chunk):
             )
             / den[:, None]
         )
-        return verts, w
+        return verts, w, valid_simplex
 
     def walk_step(carry, q_chunk):
         cur, done, outside = carry
-        _, w = weights_of(cur, q_chunk)
+        _, w, valid_simplex = weights_of(cur, q_chunk)
         minw = w.min(axis=1)
         opposite = w.argmin(axis=1)
+        outside = outside | (~done & ~valid_simplex)
         done = done | (~outside & (minw >= -1.0e-12))
         nxt = simplex_neighbors[cur, opposite]
         outside = outside | (~done & (nxt < 0))
@@ -551,6 +561,10 @@ def pixel_weights_delaunay_from(
     # -----------------------------
     pixel_weights = xp.where(has_simplex[:, None], weights_bary, weights_nn)
 
+    if xp is not np:
+        mesh_is_finite = xp.isfinite(mesh_grid).all()
+        pixel_weights = xp.where(mesh_is_finite, pixel_weights, xp.nan)
+
     return pixel_weights
 
 

test_autoarray/inversion/pixelization/interpolator/test_delaunay_walk.py

@@ -52,6 +52,68 @@ def test__tri_only_callback__tables_match_triangulation():
         assert v in tri.simplices[vertex_simplex[v]]
 
 
+def test__tri_only_callback__non_finite_points_return_sentinel_tables(monkeypatch):
+    def qhull_must_not_run(_):
+        raise AssertionError("Delaunay called for non-finite points")
+
+    monkeypatch.setattr("scipy.spatial.Delaunay", qhull_must_not_run)
+
+    rng = np.random.default_rng(5)
+    finite_points = rng.uniform(size=(40, 2))
+    poisoned_points = []
+
+    for poison_index in (7, 39):
+        points = finite_points.copy()
+        points[poison_index, 0] = np.nan
+        poisoned_points.append(points)
+
+    poisoned_points.append(np.full_like(finite_points, np.nan))
+
+    for points in poisoned_points:
+        simplices, neighbors, vertex_simplex = scipy_delaunay_tri_only(points)
+
+        assert simplices.shape == (80, 3)
+        assert neighbors.shape == (80, 3)
+        assert vertex_simplex.shape == (40,)
+        assert simplices.dtype == np.int32
+        assert neighbors.dtype == np.int32
+        assert vertex_simplex.dtype == np.int32
+        assert (simplices == -1).all()
+        assert (neighbors == -1).all()
+        assert (vertex_simplex == -1).all()
+
+
+def test__walk_locator__sentinel_tables_keep_nearest_vertex_fallback_live():
+    rng = np.random.default_rng(6)
+    finite_points = rng.uniform(size=(40, 2))
+    query = rng.uniform(size=(8, 2))
+
+    poisoned_points = []
+    for poison_index in (7, 39):
+        points = finite_points.copy()
+        points[poison_index, 0] = np.nan
+        poisoned_points.append(points)
+    poisoned_points.append(np.full_like(finite_points, np.nan))
+
+    simplices = -np.ones((80, 3), dtype=np.int32)
+    neighbors = -np.ones((80, 3), dtype=np.int32)
+    vertex_simplex = -np.ones(40, dtype=np.int32)
+
+    for points in poisoned_points:
+        mappings = pix_indexes_delaunay_walk_from(
+            query_points=query,
+            points=points,
+            simplices_padded=simplices,
+            simplex_neighbors=neighbors,
+            vertex_simplex=vertex_simplex,
+            xp=np,
+        )
+
+        assert (mappings[:, 0] >= 0).all()
+        assert (mappings[:, 0] < points.shape[0]).all()
+        assert (mappings[:, 1:] == -1).all()
+
+
 def _assert_matches_find_simplex(points, query):
     """The walk must reproduce find_simplex + KDTree-fallback semantics: rows
     identical, except at most fp edge ties where the returned simplex still