From 13241001454ab7420630659d0f20477ad7cc8fbf Mon Sep 17 00:00:00 2001
From: Jammy2211 <JNightingale2211@gmail.com>
Date: Mon, 27 Jul 2026 15:40:56 +0100
Subject: [PATCH] test: pin Delaunay NaN lane isolation

---
 scripts/imaging/jax_likelihood/delaunay.py    | 82 +++++++++++++++++++
 .../interferometer/jax_likelihood/delaunay.py | 20 +++++
 2 files changed, 102 insertions(+)

diff --git a/scripts/imaging/jax_likelihood/delaunay.py b/scripts/imaging/jax_likelihood/delaunay.py
index 33ab5294..453c6ef2 100644
--- a/scripts/imaging/jax_likelihood/delaunay.py
+++ b/scripts/imaging/jax_likelihood/delaunay.py
@@ -302,6 +302,19 @@
     err_msg="delaunay: JAX vmap likelihood mismatch",
 )
 
+# A non-finite trajectory must not abort the sequential Delaunay callback or
+# contaminate its sibling lanes. Fitness deliberately converts the raw NaN
+# figure of merit to its configured resample value after the forward pass.
+poisoned_parameters = parameters.at[1, :].set(jnp.nan)
+poisoned_result = np.asarray(fitness._vmap(poisoned_parameters))
+finite_lanes = np.array([0, 2])
+
+np.testing.assert_array_equal(
+    poisoned_result[finite_lanes], np.asarray(result)[finite_lanes]
+)
+assert poisoned_result[1] == -1.0e99
+print("PASS: poisoned Delaunay vmap lane is isolated and resampled.")
+
 
 """
 __Path A: jit-wrap ``analysis.fit_from``__
@@ -339,3 +352,72 @@
     float(fit.log_likelihood), float(fit_np.log_likelihood), rtol=1e-8
 )
 print("PASS: jit(fit_from) round-trip matches NumPy scalar.")
+
+nan_instance = model.instance_from_vector(
+    vector=np.full(model.total_free_parameters, np.nan)
+)
+nan_fit = fit_jit_fn(nan_instance)
+assert np.isnan(float(nan_fit.log_likelihood))
+print("PASS: invalid Delaunay mesh reaches the raw imaging likelihood as NaN.")
+
+
+"""
+__Callback lane isolation + gradient parity__
+
+Exercise partial poisoning directly at the mesh boundary. The last-vertex and
+non-last-vertex cases pin the negative-index/IEEE ordering asymmetry that can
+otherwise turn identical sentinel tables into finite weights.
+"""
+from autoarray.inversion.mesh.interpolator.delaunay import (
+    jax_delaunay,
+    pixel_weights_delaunay_from,
+)
+
+rng = np.random.default_rng(7)
+mesh_points_np = rng.uniform(-1.0, 1.0, size=(40, 2))
+query_points = jnp.asarray(rng.uniform(-0.8, 0.8, size=(24, 2)))
+pixel_values = jnp.linspace(0.0, 1.0, mesh_points_np.shape[0])
+
+
+def interpolated_sum(mesh_points):
+    _, _, mappings, split_points, split_mappings = jax_delaunay(
+        mesh_points, query_points
+    )
+    weights = pixel_weights_delaunay_from(query_points, mesh_points, mappings, xp=jnp)
+    split_weights = pixel_weights_delaunay_from(
+        split_points, mesh_points, split_mappings, xp=jnp
+    )
+    mapped_values = pixel_values[mappings.clip(min=0)]
+    split_mapped_values = pixel_values[split_mappings.clip(min=0)]
+    return jnp.sum(weights * mapped_values) + jnp.sum(
+        split_weights * split_mapped_values
+    )
+
+
+mesh_value_and_grad = jax.jit(jax.vmap(jax.value_and_grad(interpolated_sum)))
+solo_mesh_value, solo_mesh_grad = jax.jit(jax.value_and_grad(interpolated_sum))(
+    jnp.asarray(mesh_points_np)
+)
+assert np.any(np.asarray(solo_mesh_grad) != 0.0)
+
+for label, poison_index in (("non-last", 7), ("last", 39), ("all", None)):
+    mesh_batch = np.repeat(mesh_points_np[None, :, :], 4, axis=0)
+    if poison_index is None:
+        mesh_batch[1, :, :] = np.nan
+    else:
+        mesh_batch[1, poison_index, 0] = np.nan
+
+    mesh_values, mesh_grads = mesh_value_and_grad(jnp.asarray(mesh_batch))
+    mesh_values = np.asarray(mesh_values)
+    mesh_grads = np.asarray(mesh_grads)
+
+    assert np.isnan(mesh_values[1]), label
+    np.testing.assert_array_equal(
+        mesh_values[[0, 2, 3]], np.repeat(np.asarray(solo_mesh_value)[None], 3)
+    )
+    np.testing.assert_array_equal(
+        mesh_grads[[0, 2, 3]],
+        np.repeat(np.asarray(solo_mesh_grad)[None, :, :], 3, axis=0),
+    )
+
+print("PASS: partial/all mesh poisoning preserves finite-lane values and gradients.")
diff --git a/scripts/interferometer/jax_likelihood/delaunay.py b/scripts/interferometer/jax_likelihood/delaunay.py
index 12715aa1..7b875ab3 100644
--- a/scripts/interferometer/jax_likelihood/delaunay.py
+++ b/scripts/interferometer/jax_likelihood/delaunay.py
@@ -243,6 +243,19 @@ class in **PyAutoFit**, which pairs the model with likelihood.
     err_msg="interferometer/delaunay: JAX vmap likelihood mismatch",
 )
 
+# The sparse inversion path must obey the same lane-isolation contract as
+# imaging. Fitness sees the raw NaN after the forward pass and converts only
+# that lane to its configured resample value.
+poisoned_parameters = parameters.at[1, :].set(jnp.nan)
+poisoned_result = np.asarray(fitness._vmap(poisoned_parameters))
+finite_lanes = np.array([0, 2])
+
+np.testing.assert_array_equal(
+    poisoned_result[finite_lanes], np.asarray(result)[finite_lanes]
+)
+assert poisoned_result[1] == -1.0e99
+print("PASS: poisoned sparse Delaunay vmap lane is isolated and resampled.")
+
 
 """
 __Path A: jit-wrap ``analysis.fit_from``__
@@ -278,6 +291,13 @@ class in **PyAutoFit**, which pairs the model with likelihood.
 )
 print("PASS: jit(fit_from) round-trip matches NumPy scalar.")
 
+nan_instance = model.instance_from_vector(
+    vector=np.full(model.total_free_parameters, np.nan)
+)
+nan_fit = fit_jit_fn(nan_instance)
+assert np.isnan(float(nan_fit.log_likelihood))
+print("PASS: invalid Delaunay mesh reaches the raw interferometer likelihood as NaN.")
+
 
 """
 __Path B: TransformerNUFFT cross-check__