refactor: qhull-only Delaunay callback, JAX point location

[Jammy2211]

Overview

Shrink the JAX Delaunay pure_callback to qhull-only and move point location into the JIT program. The callback currently does triangulation + find_simplex (data grid + split points) + numpy assembly on the host — 40–320 ms of host-serial work per likelihood eval that runs sequentially per vmap lane (vmap_method="sequential"), the measured reason the A100 Delaunay cell sits at 106.8 ms/lane vs rectangular's 32.9 ms. A validated PoC (autolens_profiling/scratch/delaunay_speedup/demo_1_callback_shrink.py) shows 4.4× less host-serial work at sub_size 1 and ~30× at sub_size 4, with 99.99% identical mappings and the residual tail falling back to the existing outside-hull semantics. Hard requirement: the likelihood must be numerically unchanged — proven by parity tests in autolens_workspace_test.

Plan

• Replace the JAX-path callback so scipy does only the qhull triangulation, returning the padded simplices plus two small fixed-shape adjacency tables (vertex→incident simplices, vertex→neighbor vertices).
• Reimplement "which triangle contains this point" inside JAX: nearest-vertex seed + full 2-ring signed-barycentric containment test; points contained in none fall back to the nearest-vertex mapping, exactly the existing outside-hull behaviour.
• Move dual areas (scatter-add), split-point generation, and split-point location into JAX; seed split points at their own nearest vertex; same treatment for the Matérn variant.
• The numpy/scipy path is untouched and remains the numerical reference.
• Numpy-only unit tests in PyAutoArray for the new helpers vs scipy.find_simplex; full likelihood-parity suite in autolens_workspace_test (eager vs jit vs vmap, incl. a near-caustic config).
• Ship as two PRs behind the library-first gate: PyAutoArray first, workspace_test parity suite second. (The Feature Agent's split-into-phases recommendation maps onto this built-in two-step; no extra prompts filed.)

Detailed implementation plan

Affected Repositories

• PyAutoArray (primary)
• autolens_workspace_test (parity tests, phase 2, library-first merge gate)

Branch Survey

Repository	Current Branch	Dirty?
./PyAutoArray	main (b0f368c)	clean
./autolens_workspace_test	fix/visualization-jit-test-mode-output-path (unrelated, open PR)	clean

Suggested branch: feature/delaunay-qhull-callback
Worktree root: ~/Code/PyAutoLabs-wt/delaunay-qhull-callback/

Implementation Steps — Phase 1 (PyAutoArray)

All in autoarray/inversion/mesh/interpolator/delaunay.py unless noted. The
return contract of jax_delaunay / jax_delaunay_matern is unchanged
(points, simplices_padded, mappings, split_points, splitted_mappings — same
shapes/dtypes), so DelaunayInterface, _mappings_sizes_weights[_split] and
everything downstream need no edits; only the internals move from host to JAX.

1. scipy_delaunay_tri_only(points_np) — new minimal host callback: qhull
   Delaunay, then build simplices_padded (2N,3 int32, −1 pad; as today),
   incident (N, K_INCIDENT) via one stable-argsort grouping pass, and
   neighbors (N, K_NEIGHBOR, self included) via one np.unique over packed
   directed edges. Pad widths are module constants (start K_INCIDENT=24,
   K_NEIGHBOR=20 — measured maxima ~11/~13); assert degree ≤ pad so overflow
   fails loudly (no silent guards).
2. JAX-side helpers (xp-generic where practical so they unit-test under numpy):
   • nearest_vertex_from(query, points, chunk=1024) — lax.map-chunked
     squared-distance argmin (chunking bounds the (chunk, N) intermediate on
     GPU under vmap).
   • pix_indexes_delaunay_2ring_from(query, seed, points, simplices_padded, incident, neighbors) — candidates = incident[neighbors[seed]]
     (the full 2-ring, ~K_N×K_I padded); signed barycentric weights via cross
     products with −1-pad masking (−inf min-weight); winner = max min-weight;
     inside iff min-weight ≥ −1e-12; else fallback row [seed, −1, −1].
     Do not use a greedy walk — the PoC measured walk stalls (98.6%
     coverage vs 99.99% for the full 2-ring).
   • dual_areas_from(points, simplices_padded) — triangle areas → .at[].add
     scatter (replaces np.add.at in barycentric_dual_area_from on this path).
3. Rewire jax_delaunay: pure_callback → scipy_delaunay_tri_only only;
   then in JAX: locate data grid (seed = nearest vertex), dual areas, split
   points via existing split_points_from(..., xp=jnp) with offsets
   areas_factor*sqrt(areas), locate split points seeded at their own
   nearest vertex (PoC lesson: parent seeding breaks fallback parity with
   the KDTree outside-hull rule). Weights for both grids continue to come from
   pixel_weights_delaunay_from downstream, unchanged.
4. jax_delaunay_matern: same tri-only callback + data-grid locate; no
   splits/areas (mirrors scipy_delaunay_matern's reduced outputs).
5. Numpy path (scipy_delaunay, scipy_delaunay_matern, the xp is np
   branch of InterpolatorDelaunay.delaunay): byte-identical, untouched.
6. Unit tests (numpy-only per repo rules) in
   test_autoarray/inversion/mesh/: grouped-table construction; 2-ring locate
   vs scipy.find_simplex on random + blob/ring adaptive meshes including
   outside-hull points (agreement modulo the documented ≤0.04% fallback tail,
   which must map to the nearest vertex); pad-overflow assert fires.

Implementation Steps — Phase 2 (autolens_workspace_test, after library PR)

7. Parity script(s) in the JAX/parity area (exact placement per repo layout at
   start_workspace): Delaunay+Hilbert imaging config — assert
   log_evidence parity between eager xp=np FitImaging (still the old
   scipy path → the numerical reference), jax.jit round-trip, and
   fitness._vmap batch eval; rtol 1e-8 target (fall back to the pinned-value
   1e-4 only if fp-drift through Cholesky/log-det justifies it, documented).
   Include a near-caustic configuration (high ellipticity + shear) so folded /
   skinny source-plane triangles are exercised. Interferometer Delaunay config
   if the repo covers one.
8. Local validation: run
   autolens_profiling/likelihood_runtime/imaging/delaunay.py (CPU) — pinned
   log-evidence 29110.92085793 must hold; record before/after timings.

Key Files

• PyAutoArray/autoarray/inversion/mesh/interpolator/delaunay.py — the change
• PyAutoArray/autoarray/inversion/regularization/regularization_util.py — split_points_from (already xp-generic, reused not modified)
• autolens_profiling/scratch/delaunay_speedup/demo_1_callback_shrink.py — validated PoC to port (incl. saved outputs)
• autolens_profiling/results/notes/sparse_vs_dense_inversion_path.md — A100 baselines to beat

Risks / trade-offs

• ≤0.04% of points map via the nearest-vertex fallback where scipy found a
  distant triangle → bounded by the workspace_test parity rtol; if exceeded,
  widen the candidate set to the 3-ring before any other change.
• Pathological meshes could overflow the pad constants → loud assert by design.
• Explicitly out of scope: the fnnls/NNLS solver (separate follow-up),
  the numpy path, vmap_method changes, sparse w-tilde.

Autonomy

--auto launch 2026-07-09; effective level supervised (header supervised,
refactor cap safe). Plan written here per the contract; ship sign-off will park
as a batched question on this issue. Merge stays human.

Original Prompt

Click to expand starting prompt

Speed up the JAX Delaunay likelihood by shrinking the scipy pure_callback to qhull-only. Currently InterpolatorDelaunay (PyAutoArray autoarray/inversion/mesh/interpolator/delaunay.py) runs scipy_delaunay through jax.pure_callback with vmap_method=sequential: triangulation + find_simplex on the ~18k oversampled data points + split-point find_simplex + numpy assembly, ~40-320ms of host-serial work per likelihood eval (scales with over-sampling), serialized per vmap lane — the measured reason the A100 Delaunay cell is 106.8ms/lane vs rectangular's 32.9ms. Proven PoC in autolens_profiling/scratch/delaunay_speedup/demo_1_callback_shrink.py: callback returns only simplices_padded + two fixed-shape adjacency tables (vertex->incident simplices, vertex->neighbor vertices, one vectorized numpy pass); point location moves into JAX as nearest-vertex seed + full 2-ring signed-barycentric containment (max-min-weight winner, misses fall back to nearest-vertex mapping = existing outside-hull semantics); dual areas via scatter-add, split points and their location also JAX-side (seed split points at their own nearest vertex, NOT the parent, to match KDTree fallback). PoC parity: 99.99% identical mappings (the 0.01-0.04% tail is the fallback), split grid 100%, matched rows 1e-15; host-serial 4.4x less at sub_size 1, 30x at sub_size 4. Implement fully in PyAutoArray: JAX branch of InterpolatorDelaunay.delaunay only (numpy/scipy path unchanged), including the jax_delaunay_matern variant, fixed global pad widths for the adjacency tables (assert in callback, fail loudly), chunked nearest-vertex distance search (VRAM under vmap). Hard requirement: the likelihood must be numerically unchanged — add parity testing in autolens_workspace_test: jax.jit round-trip + fitness._vmap batch eval comparing new vs old path log-evidence on the Delaunay imaging configs (and interferometer if covered), rtol at the pinned-value level (e.g. 1e-8 on log_evidence, plus exact mapping-matrix comparison modulo the documented fallback tail), exercised near-caustic so folded/skinny triangles are covered. Also keep the eager (xp=np) FitImaging reference assert. Library unit tests stay numpy-only per repo rules: unit-test the new point-location helpers with xp=np against scipy find_simplex on random + adaptive meshes. Do NOT change the fnnls/NNLS solver (separate follow-up task).

[Jammy2211]

Phase 1 implementation update — locator redesigned after a failed parity gate

The plan's 2-ring point locator is out; an exact visibility walk is in. What happened:

1. Implemented the qhull-only callback + 2-ring locator per the plan; new numpy-only unit tests passed and the full suite stayed green (875 passed).
2. The in-session parity gate then failed the hard requirement: on the production HST config the new JAX path gave log-evidence 29110.151 vs eager 29110.921 (Δ = −0.77). Probing the real border-relocated source-plane grids showed why the PoC estimate didn't transfer: only ~31% of query points are inside the traced mesh hull, and 5 inside points had their containing simplex beyond the nearest vertex's 2-ring — caustic-compressed geometry defeats ring-based coverage arguments.
3. Fix: qhull already returns tri.neighbors, which enables the same visibility-walk algorithm find_simplex itself uses — start at a simplex incident to the nearest vertex, step through the neighbor opposite the most-negative barycentric weight, exit at a hull edge → nearest-vertex fallback (identical outside-hull semantics). Fixed-shape in JAX via a masked lax.fori_loop (cap 128 steps; production resolves in ≤64), chunked over query points for VRAM under vmap. Measured on the production grids: 100.0000% identical mappings, 0 differences.

Parity gate now passes: new-path JIT − eager = 5.4e-9, old-path JIT − eager = 5.2e-9 — fp-noise-identical, pinned value 29110.92085793 holds. This is exactness by construction (same algorithm), not statistical coverage.

Notes for review:

• The callback now returns simplices_padded + tri.neighbors (padded) + one incident simplex per vertex — the 2-ring adjacency tables and pad-width constants from the plan are gone (simpler than planned).
• Known trade-off: single-thread CPU JAX-path per-eval rose ~20% (14.4 s vs 11.8 s locally) since CPU-JAX runs the walk slower than scipy's C. GPU/vmap is the target of this refactor; the eager numpy path is byte-identical.
• Unit tests rewritten for the walk (test_delaunay_walk.py): exact find_simplex agreement modulo fp edge ties that must still contain the query, KDTree-fallback equality outside the hull, vertex/edge-midpoint degeneracies.

Next: full-suite rerun, then ship_library (PR-open per the autonomy contract), then phase 2 (autolens_workspace_test parity suite) behind the library-first gate.

[Jammy2211]

Ship sign-off requested (supervised --auto checkpoint)

Where things stand. Phase 1 implementation is complete and validated in the worktree (feature/delaunay-qhull-callback, uncommitted pending this sign-off):

• Diff: autoarray/inversion/mesh/interpolator/delaunay.py +234/−37 (CRLF preserved to keep the diff minimal) + new test_delaunay_walk.py.
• Tests: full suite 873 passed in the worktree (was 875 with the interim 2-ring tests; −6 removed, +4 walk tests).
• Parity (the hard requirement): production HST Delaunay config — new-path JIT vs eager Δ = 5.4e-9 (old path: 5.2e-9); mapping identity on production grids 100.0000%; pinned log-evidence 29110.92085793 holds. Eager numpy path byte-identical.
• API: none — internal changes only (jax_delaunay/jax_delaunay_matern contracts unchanged; no downstream-repo suites required).
• Review faculty: runs post-commit (no reviewable diff until then) — will run before PR-open.

The fork needing a human. Two things the contract won't let the run decide:

1. Ship sign-off — commit, push, and open the pending-release PR.
2. Heart YELLOW acknowledgement — current reasons (all chronic organism-state, none branch-specific):
   workspace validation not passing (5 failed, 2026-07-08) · 58 stale parked scripts · autolens_assistant pinned BEHIND · install verification not run · no release validation for current source · PyAutoMind open PR 9d old

Options.

• (a) Approve ship + acknowledge the YELLOW reason set above → run review faculty, commit, push, open PR (stops at PR-open; merge stays human), then post the workspace-impact recommendation and queue phase 2 (workspace_test parity suite).
• (b) Request changes → say what; the branch stays parked.
• (c) Hold → task stays awaiting-input.

Recommendation: (a). The YELLOW reasons predate this task and are tracked elsewhere; the branch's own evidence (tests + exact-parity gate) is complete.

[Jammy2211]

Library PR created — #368

Ship executed on the in-session sign-off (commit 041d5f60, 873 tests passed, pending-release label verified).

Review faculty verdict: CLEAN (surface: 2 files, +365/−37, python-source flag). Re-reviewed the walk-loop carry dtypes, chunk-padding semantics, fallback parity with the KDTree convention, the untouched numpy path, and the chunked-argmin VRAM shape under vmap batch 62 — nothing must change. Known trade-offs (CPU-JAX ~20% slower per-eval; 128-step walk cap with nearest-vertex fallback) are disclosed in the PR body.

Workspace impact: none — additive follow-up only. API changes are internal-only; grep across all workspaces finds no code use of the changed symbols (one docstring mention in autolens_workspace_developer). The curated smoke lists have no active Delaunay entry (the one candidate, jax_likelihood_functions/imaging/delaunay_mge.py, is disabled for unrelated jax-0.7 reasons), so the smoke leg is n/a-stated — the production-scale parity gate run in-session (JIT−eager 5.4e-9, 100.0000% mapping identity) is the stronger check.

Next:

1. Merge is yours (PR refactor: qhull-only Delaunay callback, exact JAX visibility-walk point location #368; run stops at PR-open per the contract).
2. Phase 2 — autolens_workspace_test parity suite (eager vs jit vs vmap on the Delaunay configs, incl. near-caustic) proceeds on the already-claimed feature/delaunay-qhull-callback branch behind the library-first merge gate.
3. A100 re-profile via autolens_profiling when RAL returns (expected: Delaunay vmap lane 106.8 → ~40-50 ms).

[Jammy2211]

Phase 2 complete — workspace_test parity suite validated; ship sign-off requested

All validation green under the new library path (feature/delaunay-qhull-callback worktree):

Check	Config	Result
Existing delaunay.py vmap pin	AdaptSplit + 30 zeroed edge px + PowerLaw	−22205.87818084 exact match, 3 lanes
delaunay.py eager vs JIT (tightened 1e-4 → 1e-8)	same	Δ relative ~7e-15
NEW delaunay_near_caustic.py eager vs JIT @ 1e-8	axis-ratio ~0.5 PowerLaw + γ=0.11 shear, ConstantSplit, folds/slivers	Δ relative 3.3e-10
near-caustic vmap lanes vs eager figure_of_merit @ 1e-8	same	Δ relative 2.5e-10

Diff on the claimed branch: delaunay.py +4/−1 (tightened assert + comment), delaunay_near_caustic.py added. delaunay_mge.py deliberately stays at 1e-4 (mixed-precision config — fp32-scale deltas by design).

Sign-off options: (a) ship — commit, push, open the workspace PR (pending-release, merge gated library-first behind PR #368); (b) request changes; (c) hold. Recommendation: (a). Heart reason set unchanged from the acknowledged launch list.

[Jammy2211]

Phase 2 shipped — PyAutoLabs/autolens_workspace_test#155

Commit efc5d13, pending-release verified, cross-referenced on PR #368. Merge order: PyAutoArray #368 first, then workspace_test #155 (library-first gate; the 1e-8 asserts were validated against the new library).

Both PRs are now at PR-open with merge in human hands — the full task deliverable is complete. Remaining post-merge follow-ups tracked here: A100 re-profile via autolens_profiling when RAL returns (expected Delaunay vmap lane 106.8 → ~40-50 ms), and the separate fnnls/NNLS deep-dive task.

[Jammy2211]

Shipped and merged

• PyAutoArray refactor: qhull-only Delaunay callback, exact JAX visibility-walk point location #368 merged (89a5fbdb); feature/delaunay-qhull-callback deleted on the remote.
• autolens_workspace_test fix issue causing ndarray to not be a type when accessed via jax wrapper #155 auto-merge armed, lands when its smoke checks finish (library-first order preserved).

Task complete. The JAX Delaunay callback is now qhull-only with exact in-JIT visibility-walk point location; likelihood numerically unchanged (5.4e-9 vs eager, guarded permanently at rtol 1e-8 incl. the near-caustic config).

Post-merge follow-ups (tracked, not this issue): A100 re-profile via autolens_profiling when RAL returns (expected Delaunay vmap lane 106.8 → ~40-50 ms); fnnls/NNLS solver deep-dive (separate PyAutoMind task to be filed at pickup).