feat: KNN-barycentric interpolator — drop-in JAX-native Delaunay alternative

[Jammy2211]

Overview

A pure-JAX wildcard for replacing scipy.spatial.Delaunay in PyAutoArray's source-plane interpolation. Pick the 3 nearest mesh vertices in source plane (existing brute-force kNN) and compute exact barycentric weights on that triangle — replacing the Wendland C4 kernel weights of the existing InterpolatorKNearestNeighbor. If validation passes (log-evidence rtol=1e-3 vs true Delaunay at the production HST fiducial), this is a drop-in replacement that eliminates the scipy callback bottleneck (~16.87 ms/element on A100, 24% of the production likelihood). Code is small (~50-100 lines); the hard part is the science validation.

Plan

• Implement a pure-JAX interpolator that picks the 3 nearest mesh vertices in source plane (k=3) and computes exact barycentric weights on that triangle.
• For queries outside their nearest-3 triangle (negative barycentric component): start with clip + renormalize; escalate to "take more neighbors, pick 3 with non-negative coords" only if validation fails.
• Register as a new KNNBarycentric(Delaunay) mesh class so users opt in by swapping mesh — no pipeline-code changes.
• Validate by comparing log_evidence against InterpolatorDelaunay at the HST imaging fiducial (EXPECTED_LOG_EVIDENCE_HST = 26288.321397232066). Targets: rtol=1e-4 → drop-in default; rtol=1e-3 → opt-in alternative; rtol=1e-2 fail → abandon and pursue split-callback approach instead.
• Measure speedup via the existing z_projects/profiling/scripts/delaunay_vmap_probe.py harness at PROBE_BATCH_SIZE=20. Expected: 69.5 → ~56 ms per element (~1.25× overall).

Detailed implementation plan

Affected Repositories

• PyAutoArray (primary, library)
• autolens_workspace_developer (regression follow-up)
• autolens_workspace_test (smoke follow-up)

Work Classification

Library (then workspace follow-up)

Branch Survey

Repository	Current Branch	Dirty?
./PyAutoArray	main	clean
./autolens_workspace_developer	main	dirty (pre-existing unrelated work; worktree gives clean origin/main copy)
./autolens_workspace_test	main	clean

Suggested branch: feature/knn-barycentric
Worktree root: ~/Code/PyAutoLabs-wt/knn-barycentric/ (created later by /start_library)

Coexistence note: Active task jit-regression-drift claims autolens_workspace_developer + autolens_profiling for drift triage of existing jit/ scripts. New regression scripts under delaunay_knn_barycentric.py are additive and won't collide with the constants being triaged, but the merge order should sequence the drift PR first if both land in the same window.

Implementation Steps

PyAutoArray (library — primary work)

1. autoarray/inversion/mesh/interpolator/knn.py

   • Add barycentric_weights_from_3_nearest(query_points, mesh_points, nearest_3_indices, xp):
     • Gathers the 3 vertices for each query (shape (Q, 3, 2))
     • Computes signed barycentric coordinates (reuse triangle_area_xp from delaunay.py)
     • Clips negatives to 0, renormalizes so the row sums to 1
     • Zeros weights on degenerate triangles (matches existing Delaunay code path)
     • Returns weights shape (Q, 3)
   • Add InterpolatorKNNBarycentric(InterpolatorKNearestNeighbor):
     • Overrides _mappings_sizes_weights to call get_interpolation_weights with k_neighbors=3 and pipe the indices through barycentric_weights_from_3_nearest instead of Wendland weights
     • _mappings_sizes_weights_split: same swap on the split-points path (regularization requires it)

2. autoarray/inversion/mesh/mesh/knn.py

   • Add KNNBarycentric(Delaunay) class mirroring KNearestNeighbor:
     • No k_neighbors / radius_scale knobs (hard-coded k=3)
     • Keeps areas_factor for split regularization
     • interpolator_cls returns InterpolatorKNNBarycentric

3. autoarray/inversion/mesh/__init__.py — export KNNBarycentric.

4. Unit tests — test_autoarray/inversion/mesh/interpolator/test_knn_barycentric.py:

   • Triangle interior: weights match pixel_weights_delaunay_from to fp64 precision when the 3 nearest are the containing Delaunay triangle (constructed by hand)
   • Boundary (query outside triangle of 3 nearest): clipped weights ≥ 0, sum to 1
   • Degenerate (collinear 3 nearest): no NaN, weights finite
   • Stays numpy-only — no JAX in library unit tests

autolens_workspace_developer (regression — follow-up)

5. jax_profiling/jit/imaging/delaunay_knn_barycentric.py — parallel of delaunay.py:
   • Same dataset, same fiducial, same model setup
   • Mesh swapped to KNNBarycentric
   • Assertion: log_evidence within rtol=1e-3 of EXPECTED_LOG_EVIDENCE_HST = 26288.321397232066
   • If rtol=1e-4 also passes, note it in a comment — that's the drop-in-default condition

autolens_workspace_test (smoke — follow-up)

6. scripts/jax_assertions/knn_barycentric.py — mirror nnls.py pattern:
   • Well-conditioned case: finite output, weights sum to 1 per row
   • Ill-conditioned (caustic-crossing lens model from existing test suite): assert tolerance vs Delaunay's source reconstruction

Speed measurement (sanity check, no PR)

7. Re-run z_projects/profiling/scripts/delaunay_vmap_probe.py at PROBE_BATCH_SIZE=20 with mesh swapped from Delaunay to KNNBarycentric. Compare per-call timings to existing output/imaging/delaunay/hpc_a100_fp64_vmap_probe.json. Target: ~13 ms saving per element, 69.5 → ~56 ms.

Decision gate (validation outcome)

Outcome	Ship as
rtol=1e-4 passes	Default Delaunay-equivalent mesh; can deprecate scipy path
rtol=1e-3 passes, rtol=1e-4 fails	Opt-in alternative via config flag, default off
rtol=1e-2 fails	Abandon; fall back to split-callback approach in delaunay_research.md

Key Files

New code (PyAutoArray):

• autoarray/inversion/mesh/interpolator/knn.py — add barycentric_weights_from_3_nearest + InterpolatorKNNBarycentric
• autoarray/inversion/mesh/mesh/knn.py — add KNNBarycentric mesh class
• autoarray/inversion/mesh/__init__.py — export
• test_autoarray/inversion/mesh/interpolator/test_knn_barycentric.py — unit tests

New code (autolens_workspace_developer):

• jax_profiling/jit/imaging/delaunay_knn_barycentric.py — regression script at rtol=1e-3

New code (autolens_workspace_test):

• scripts/jax_assertions/knn_barycentric.py — smoke / cross-conditioning

Read-only reference:

• autoarray/inversion/mesh/interpolator/delaunay.py — pixel_weights_delaunay_from, triangle_area_xp (reuse), InterpolatorDelaunay
• autoarray/inversion/mesh/interpolator/knn.py — current get_interpolation_weights, InterpolatorKNearestNeighbor
• autoarray/inversion/mesh/mesh/knn.py — existing KNearestNeighbor(Delaunay) pattern
• autoarray/inversion/mappers/abstract.py:255 — Mapper.mapping_matrix call site
• z_projects/profiling/scripts/delaunay_vmap_probe.py — speed harness
• PyAutoPrompt/autoarray/delaunay_research.md — split-callback fallback strategy
• z_projects/profiling/FINDINGS_nnls_v2.md — full investigation background

Out of scope

• Pure-JAX Delaunay (separate research project — delaunay_research.md, ~3-6 months)
• Multiprocessing the scipy callback (separate optimization in delaunay_research.md)
• Changes to NNLS, log_ev, PSF FFT (not bottlenecks at production batch=20)
• Sampler-level Delaunay caching (PyAutoFit concern)

Original Prompt

Click to expand starting prompt

InterpolatorKNearestNeighbor variant: barycentric weights on top-3 nearest

A pure-JAX wildcard for replacing scipy.spatial.Delaunay in PyAutoArray's
source-plane interpolation. The idea is small and testable, the
potential payoff is large — if it holds the log-evidence to rtol=1e-4
against true Delaunay, it's a drop-in replacement that eliminates the
scipy callback bottleneck entirely.

Background from the nnls-vmap-speedup investigation

(Full findings: z_projects/profiling/FINDINGS_nnls_v2.md. Companion
research doc: PyAutoPrompt/autoarray/delaunay_research.md. Closed
issue: #307.)

At production batch=20 on A100, the Delaunay imaging likelihood costs
69.5 ms per element. Decomposition:

scipy.spatial.Delaunay via pure_callback = 16.87 ms (24%)  <-- this prompt
other JAX-traced inversion setup         = ~25 ms (36%)
PSF FFT convolution                      = ~9 ms (13%)
log_ev (slogdet + matmul)                = 12 ms (17%)
NNLS reconstruction                      = 6.2 ms (9%)
misc                                     = ~0.5 ms (1%)

pure_callback with vmap_method="sequential" invokes scipy serially
per batch element → 16.87 × 20 = 337 ms wall per batched likelihood
call is held on a single CPU running scipy/Qhull. The barycentric
weight computation (0.01 ms) and the sparse mapping matrix scatter
(0.15 ms) are essentially free — the 16.87 ms IS the scipy work.

Of that 16.87 ms, ~52% is find_simplex (point location for 15361
query points), called twice, and ~14% is the actual Delaunay
triangulation on 1231 mesh points. The bottleneck is point location
on CPU, not the triangulation per se.

The science context the user supplied

PyAutoArray's premier interpolator is InterpolatorDelaunay. It uses
scipy.spatial.Delaunay to triangulate the source-plane mesh vertices
and find which triangle each over-sampled image-plane data pixel falls
in. The barycentric coordinates of the data pixel inside its triangle
give the three interpolation weights to the three triangle vertices.

InterpolatorKNearestNeighbor was a previous JAX-friendly attempt at
the same problem (autoarray/inversion/mesh/interpolator/knn.py). It
uses brute-force k-nearest-neighbor search in JAX (no scipy callback)
and weighs the neighbors with a Wendland C4 kernel (smoothed
inverse-distance, compact support).

Per the user, kNN performed scientifically much worse than Delaunay
because:

• The Wendland kernel has knobs (k_neighbors, radius_scale) that
  are hard to set to work for all lenses.
• Delaunay adapts cleanly to local mesh density; kNN smears across
  density gradients.
• At caustic crossings (folds in the source-plane mapping), Delaunay's
  triangulation correctly spans the fold; kNN smears across it.

So Delaunay is scientifically the right method but the least JAX-friendly.

The wildcard idea

Use kNN to pick the top-3 nearest neighbors in source plane, then
compute exact barycentric coordinates on the triangle those 3 form.

This replaces Wendland kernel weights with locally-exact barycentric
weights — the same weights Delaunay uses, just on a triangle chosen
by Euclidean nearest-neighbor instead of by Delaunay triangulation.

When the 3 nearest mesh vertices happen to be the 3 vertices of the
containing Delaunay triangle (the common case for "interior" query
points), the result is bit-identical to Delaunay. When they aren't
(boundary points, density-gradient regions, caustic-crossing points),
the result is an approximation — but one that still respects local
mesh topology rather than smearing with a global kernel.

The potentially-killer subtlety: when the query point is OUTSIDE the
triangle formed by its 3 nearest neighbors, the barycentric coords
have a negative component. Two options:

1. Clip and renormalize: max(bary, 0), then normalize so they sum
   to 1. Gives a valid convex combination but slightly distorts.
2. Fall back to kNN-Wendland for those query points. Hybrid.
3. Take more neighbors (k=4, 5, 6) and pick the 3 with non-negative
   barycentric coords. Most expressive but more complex.

Start with option 1. If log-evidence rtol fails, escalate to option 3.

Why this might work where pure kNN-Wendland didn't

• Barycentric weights on 3 nearest are locally exact when the 3 are
  the correct triangle vertices. The Wendland kernel never is.
• The "knobs" go away: no radius_scale, no kernel shape. Just k=3.
• At caustic crossings, the 3 nearest still bracket the local source-plane
  structure better than a smoothly-decaying kernel.

Where it might still fail:

• Boundary points (mesh edge) where the 3 nearest don't form a
  containing triangle, AND clip-renormalize doesn't approximate well.
• Highly anisotropic mesh density where the 3 nearest are nearly
  collinear (degenerate triangle → numerical instability in barycentric).

Both failure modes are testable.

Performance expectations

InterpolatorKNearestNeighbor is already pure JAX with lax.fori_loop
over blocks of mesh points (brute-force kNN, no scipy callback). Adding
barycentric on top-3 is a small post-processing step on the kNN output.

Estimated runtime under vmap=20 on A100:

• kNN search (current code, k=3 instead of k=10): ~1-3 ms per element
  (brute-force, 1231 mesh × 15361 queries × 3 = 5.7e7 ops, easy GPU work)
• Barycentric weight computation on the 3 picked vertices: ~0.1 ms
  per element
• Total: ~2-4 ms per element

Replacing InterpolatorDelaunay's 16.87 ms per element callback with
~3 ms pure JAX gives ~13 ms saving per element, full pipeline drops
from 69.5 → 56 ms per element, batch_time 1.4 → 1.12 sec.
~1.25× speedup overall, just from this swap.

If it works, it's also vmap-parallel (no sequential callback), so
larger batch sizes will continue to amortize well — unlike the scipy
callback which is sequential-per-element forever.

Decision criteria up front

This is a science-validated speedup wildcard. Code is small (~50-100
lines new). The hard part is the validation: does kNN-barycentric give
the same scientific answer as Delaunay across the production lens
modeling matrix?

If yes — best possible outcome: drop scipy.spatial.Delaunay entirely,
get ~1.25× speedup on Delaunay production, zero external dep on scipy
in the inversion path.

If no — small sunk cost, fall back to the more conservative
split-the-callback approach in delaunay_research.md.

Either way the data tells us something useful.

[Jammy2211]

Library PR Created — 2026-05-16

PR

• Add KNNBarycentric mesh: JAX-native Delaunay-class interpolator #318

Status

Workspace changes needed. Next: /start_workspace

API Changes

Purely additive — new aa.mesh.KNNBarycentric(pixels=...) mesh class (inherits from KNearestNeighbor), new InterpolatorKNNBarycentric interpolator class, new barycentric_weights_from_3_nearest() helper in autoarray.inversion.mesh.interpolator.knn. No removals, no signature changes, no behaviour changes to existing classes. Users opt in by swapping aa.mesh.Delaunay(...) for aa.mesh.KNNBarycentric(...).

Library validation

• 7/7 new unit tests pass (interior-triangle exact match to Delaunay, edge split, outside-triangle clipping, degenerate fallback, 200-sample direct comparison to pixel_weights_delaunay_from).
• Full test_autoarray/ suite: 787 passed, 0 failed.
• Direct comparison test confirms barycentric weights are bit-identical to pixel_weights_delaunay_from when the 3 nearest are the containing Delaunay triangle's vertices.

What workspace work covers

1. autolens_workspace_developer/jax_profiling/jit/imaging/delaunay_knn_barycentric.py — regression at HST imaging fiducial, assertion rtol=1e-3 vs Delaunay's EXPECTED_LOG_EVIDENCE_HST = 26288.321397232066. This is the science gate that determines whether this ships as drop-in default (rtol=1e-4), opt-in alternative (rtol=1e-3), or gets abandoned in favor of the split-callback approach.
2. autolens_workspace_test/scripts/jax_assertions/knn_barycentric.py — well-conditioned + ill-conditioned (caustic-crossing) smoke vs Delaunay.

Implementation notes

• KNNBarycentric inherits all kNN connectivity knobs (k_neighbors=10, radius_scale=1.5, etc.) — these still control the regularization-spacing computation via distance_to_self. Interpolation always uses k=3, irrespective of mesh.k_neighbors.
• Degenerate fallback policy: when the 3 nearest are collinear (zero-area triangle), weight falls back to [1, 0, 0] on the closest vertex. Mirrors pix_indexes_for_sub_slim_index_delaunay_from's outside-simplex behavior. A real case of this was caught in the smoke check.

[Jammy2211]

Shipped — NEGATIVE Result — 2026-05-16

PRs

• Library: Add KNNBarycentric mesh: JAX-native Delaunay-class interpolator #318 (follow-up commit 58b314ce added writability + split-padding bugfixes surfaced during workspace integration)
• Developer regression: Add jit/imaging/delaunay_knn_barycentric.py — NEGATIVE result documenting wildcard failure autolens_workspace_developer#70
• Smoke: Add jax_assertions/knn_barycentric.py — smoke for new KNNBarycentric mesh autolens_workspace_test#99

Science gate verdict: WILDCARD FAILED

At the HST imaging fiducial (39×39 overlay → 1231 mesh vertices, MGE-60 lens):

log_evidence (aa.mesh.KNNBarycentric)	26872.33
log_evidence (aa.mesh.Delaunay reference)	26288.32
Absolute difference	+584 nats (kNN-bary higher)
Relative error	2.22%
Gate outcome	FAILS rtol=1e-2 (the abandon threshold)

Per the decision tree in PyAutoPrompt/issued/knn_barycentric.md: rtol=1e-2 fail → abandon; fall back to split-callback approach in PyAutoPrompt/issued/delaunay_research.md.

Root cause (more interesting than the gate result itself)

~5% of mesh vertices (60/1291) are never used by the kNN-barycentric interpolator. They're never in any over-sampled query's nearest-3, nor in any split-point's nearest-3. The regularization matrix H is therefore structurally rank-deficient by 60, and the production FitImaging.log_evidence quietly absorbs the zero modes via its log_det_curvature_reg_matrix_term.

This is a fundamental property of using kNN connectivity instead of a triangulation. scipy's Delaunay guarantees every mesh vertex belongs to at least one simplex → every vertex contributes to the reconstruction. kNN makes no such guarantee — sparse/edge vertices get orphaned. Beyond the log-evidence drift, this means the user is paying for 5% of mesh vertices and getting nothing back.

What stays in the repo (additive infrastructure)

The library code works correctly and is covered by 8 unit tests:

• aa.mesh.KNNBarycentric mesh class
• InterpolatorKNNBarycentric interpolator (k=3 + clip-and-renormalize barycentric)
• barycentric_weights_from_3_nearest() helper
• Regression-catch on observed log_evidence (26872.33 at rtol=1e-4) in delaunay_knn_barycentric.py

These components may be useful for future research (the building blocks — signed barycentric coords, structurally-rank-aware fallback — could be reused), and the regression script + smoke serve as documentary evidence so future investigators reproduce the verdict before re-attempting.

Recommended next step

Move to the split-callback approach in PyAutoPrompt/issued/delaunay_research.md (option A): keep scipy.spatial.Delaunay for the cheap CPU triangulation (~5ms for 1231 pts), replace the two expensive find_simplex calls (~21ms × 2, the actual bottleneck) with a JAX-native brute-force point-in-triangle implementation that vmaps across the batch. Expected ~1.4× speedup on the full pipeline at batch=20, with the science kept correct by construction (still uses Delaunay's triangulation).

Implementation notes worth keeping

• _mappings_sizes_weights must use np.array() not np.asarray() on the numpy path — the regularization code does in-place assignment and a read-only view of a jax.Array fails with "assignment destination is read-only".
• _mappings_sizes_weights_split must pad mappings to k+1 columns — reg_split_np_from writes splitted_mappings[i][j+1] = pixel_index for the central-pixel insertion.
• Both gotchas now have regression tests in test_knn_barycentric.py.

[Jammy2211]

Closing as completed (NEGATIVE result). All three PRs merged:

• Library: Add KNNBarycentric mesh: JAX-native Delaunay-class interpolator #318
• Developer regression: Add jit/imaging/delaunay_knn_barycentric.py — NEGATIVE result documenting wildcard failure autolens_workspace_developer#70
• Smoke: Add jax_assertions/knn_barycentric.py — smoke for new KNNBarycentric mesh autolens_workspace_test#99

Wildcard failed the science gate (2.22% drift vs Delaunay, fails rtol=1e-2). Library infrastructure stays in the repo as additive code that other research may reuse; regression + smoke serve as documentary evidence so the negative result survives.

Next step recommended in PyAutoPrompt/autoarray/delaunay_research.md: resume option A (split-callback) from the already-implemented branch feature/delaunay-jax-find-simplex (commit eda747c, shelved 2026-05-11 at measured 1.19x-1.23x). With the wildcard ruled out, A is the only realistic JAX-native lever and worth the half-day to ship.