Repository navigation
refactor: qhull-only Delaunay callback, JAX point location #367
Description
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 inautolens_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.
scipy_delaunay_tri_only(points_np)— new minimal host callback: qhull
Delaunay, then buildsimplices_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 onenp.uniqueover packed
directed edges. Pad widths are module constants (start K_INCIDENT=24,
K_NEIGHBOR=20 — measured maxima ~11/~13);assertdegree ≤ pad so overflow
fails loudly (no silent guards).- 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 (replacesnp.add.atinbarycentric_dual_area_fromon this path).
- Rewire
jax_delaunay: pure_callback →scipy_delaunay_tri_onlyonly;
then in JAX: locate data grid (seed = nearest vertex), dual areas, split
points via existingsplit_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_fromdownstream, unchanged. jax_delaunay_matern: same tri-only callback + data-grid locate; no
splits/areas (mirrorsscipy_delaunay_matern's reduced outputs).- Numpy path (
scipy_delaunay,scipy_delaunay_matern, thexp is np
branch ofInterpolatorDelaunay.delaunay): byte-identical, untouched. - Unit tests (numpy-only per repo rules) in
test_autoarray/inversion/mesh/: grouped-table construction; 2-ring locate
vsscipy.find_simplexon 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)
- Parity script(s) in the JAX/parity area (exact placement per repo layout at
start_workspace): Delaunay+Hilbert imaging config — assert
log_evidenceparity between eagerxp=npFitImaging(still the old
scipy path → the numerical reference),jax.jitround-trip, and
fitness._vmapbatch 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. - 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 changePyAutoArray/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_methodchanges, 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).
Activity
Jammy2211 commentedon Jul 9, 2026
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:
tri.neighbors, which enables the same visibility-walk algorithmfind_simplexitself 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 maskedlax.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:
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).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 commentedon Jul 9, 2026
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):autoarray/inversion/mesh/interpolator/delaunay.py+234/−37 (CRLF preserved to keep the diff minimal) + newtest_delaunay_walk.py.jax_delaunay/jax_delaunay_materncontracts unchanged; no downstream-repo suites required).The fork needing a human. Two things the contract won't let the run decide:
pending-releasePR.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 oldOptions.
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 commentedon Jul 9, 2026
Library PR created — #368
Ship executed on the in-session sign-off (commit
041d5f60, 873 tests passed,pending-releaselabel verified).Review faculty verdict: CLEAN (surface: 2 files, +365/−37,
python-sourceflag). 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:
autolens_workspace_testparity suite (eager vs jit vs vmap on the Delaunay configs, incl. near-caustic) proceeds on the already-claimedfeature/delaunay-qhull-callbackbranch behind the library-first merge gate.autolens_profilingwhen RAL returns (expected: Delaunay vmap lane 106.8 → ~40-50 ms).Jammy2211 commentedon Jul 9, 2026
Phase 2 complete — workspace_test parity suite validated; ship sign-off requested
All validation green under the new library path (
feature/delaunay-qhull-callbackworktree):delaunay.pyvmap pindelaunay.pyeager vs JIT (tightened 1e-4 → 1e-8)delaunay_near_caustic.pyeager vs JIT @ 1e-8Diff on the claimed branch:
delaunay.py+4/−1 (tightened assert + comment),delaunay_near_caustic.pyadded.delaunay_mge.pydeliberately 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 commentedon Jul 9, 2026
Phase 2 shipped — PyAutoLabs/autolens_workspace_test#155
Commit
efc5d13,pending-releaseverified, 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 commentedon Jul 9, 2026
Shipped and merged
89a5fbdb);feature/delaunay-qhull-callbackdeleted on the remote.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).