Skip to content

perf: numba Sibson kernel so DelaunayNN matches Delaunay speed on CPU (#623 phase 1) - #624

Merged
Jammy2211 merged 1 commit into
mainfrom
feature/delaunay-nn-numba
Oct 15, 2026
Merged

Jammy2211 merged 1 commit into
mainfrom
feature/delaunay-nn-numba

Conversation

@Jammy2211

Copy link
Copy Markdown
Collaborator

Summary

Phase 1 of #623. DelaunayNN was ~40x slower than Delaunay on the NumPy (use_jax=False) path because sibson_mappings_weights_from_tables ran _sibson_single_from_tables once per query in a Python list comprehension (~1.5 ms × ~11k queries per likelihood). This adds a parallel numba kernel (prange over queries, module-level in interpolator/_sibson_numba.py, lazily imported so numba stays optional and the on-disk cache hits across processes) that ports the per-query code with the same arithmetic order: cavity walk + tolerance, cavity overflow, Watson contributions, outside-mesh / on-vertex / on-edge rows, degenerate and neighbour-overflow → NaN, and the six outputs with identical dtypes and padding. The NumPy branch dispatches to it when numba is importable and general.yaml numba.use_numba is true (float64 inputs); otherwise the unchanged reference loop runs. JAX path untouched. Phase 2 (routing audit of the numba imaging/interferometer likelihood paths) follows on the same issue.

Aris SMA interferometer (5.5M visibilities, threads=1, steady median per likelihood; same-run comparison):

mesh vertices before after Sibson step
Delaunay 1405 0.43 s 0.41 s –
DelaunayNN 1405 18.1 s 0.57 s (1.39× Delaunay) 8.46 s → 12.9 ms (11,116 queries)
Delaunay 5526 12.7 s 8.2 s –
DelaunayNN 5526 56 s 12.1 s 63 ms (27,600 queries)

logL identical to the reference loop (14092443.357190056). All four cases route to InversionInterferometerSparseNumba (DelaunayNN nnz/source column 17.8 @1405, 1.0 @5526, gate 60) — measured only; routing is phase 2. Numba cache: first process 18.8 s compile, second process cache hit, first call 0.53 s.

API Changes

None — internal changes only. (New sibson.sibson_numba_kernel() builder and private helpers; no public signature or behaviour change beyond speed. Weights agree with the reference to ≤1e-14, not bit-for-bit, because the reference groups vertex ids with an unstable argsort.)
See full details below.

Test Plan

  • New test_sibson_numba.py (12): kernel vs reference on random meshes (outside-hull, on-vertex, on-edge-midpoint queries), cocircular grid, genuine degenerate triangle, five small cavity/neighbour cap combos (overflow asserted), dispatch both ways — mappings/sizes/diagnostics identical, weights 1e-12
  • Full test_autoarray/: 2026 passed, 4 xfailed (local, worktree)
  • Emulated unittest-nojax leg (path finder blocks jax/jaxlib/…): interpolator dir 95 passed, 8 skipped, 3 xfailed
  • autolens_workspace_test: misc/jax_assertions/delaunay_nn.py, delaunay_nn_caps.py, delaunay_walk.py, imaging/mesh_interpolator_source_recovery.py — all exit 0 (DelaunayNN recovery pearson 0.99935)
  • Independent review (fresh-context Opus): no correctness bugs; two findings fixed (cross-process numba cache miss; availability probe now catches ImportError)
Full API Changes (for automation & release notes)

Added

  • autoarray.inversion.mesh.interpolator.sibson.sibson_numba_kernel() — returns the compiled parallel Sibson kernel (imports _sibson_numba on first use)
  • autoarray/inversion/mesh/interpolator/_sibson_numba.py — private module: _sibson_from_tables_parallel + pairwise-sum helpers

Changed Behaviour

  • sibson_mappings_weights_from_tables(..., xp=np) — dispatches to the numba kernel when available/enabled; outputs unchanged up to ≤1e-14 weight rounding

Heart RED override (development only)

Authorised by the live human in this session on 2026-10-15 for #623: "I override, tests will be handled once the epic refactor elsewhere is sorted". Permits commit, push and this PR only — merge needs a separate human /prm with required checks green; this PR does not claim to fix Heart.

  • Current pyauto-heart readiness --json: verdict stale (ts 2026-10-14T14:57:28Z) — "release validation incomplete: no rehearsal for current source".
  • Current Heart feed RED reasons (verbatim): "Heart: RED — Release RED; monitoring RED · 10/100 · 415 unresolved (updated 4h ago)"; "[red] release validation FAILED (stage integrate)"; "[red] PyAutoHands/workflows/Python Version Matrix: failure"; "[red] PyAutoHeart/workflows/Release Integrate: failure"; "[red] PyAutoMind/workflows/Lifecycle Drift: failure"; "[red] PyAutoArray/Tests: red"; "[red] PyAutoArray/Tests: red"; "[red] autofit_workspace_test/Smoke Tests: red"; "[red] autogalaxy_workspace_test/Smoke Tests: red"; "[red] autolens_workspace_test/Smoke Tests: red"; "[red] CI wall-clock: 26 gates · slowest PyAutoBrain Nightly Release 63m · 5 hang events"; "[red] Release readiness: release validation FAILED (stage integrate)"; "[red] autofit_test: red"; "[red] Release validation: NOT release_ready — v2026.10.15.1.dev81701 profile=release (2026-10-15T07:04:34+00:00)"; "[red] validation_report/stages/integrate: fail"; "[red] PyAutoArray: red"; "[red] autolens_inference: red"; "[red] autolens_profiling: red"; "[red] Worktree drift: 5 orphan / 0 missing / 4 dirty"; "[red] euclid_strong_lens_modeling_pipeline: red". (Note: PyAutoArray main Tests at ccddfba is green on all three legs.)
  • Branch gates passed: tests, emulated nojax, smoke scripts, independent review (above).

Generated by the PyAutoLabs agent workflow.

🤖 Generated with Claude Code

Replace the per-query Python loop in the NumPy branch of
sibson_mappings_weights_from_tables with a parallel numba kernel
(module-level in _sibson_numba.py, lazily imported, cross-process cache)
that ports _sibson_single_from_tables exactly: cavity walk and overflow,
Watson contributions, outside-mesh/on-vertex/on-edge rows, degenerate and
neighbour-overflow NaN rows and all six outputs. The reference loop stays
as the fallback when numba is unavailable or disabled. JAX path untouched.

Aris SMA interferometer, 1405 vertices: DelaunayNN 18.1 s -> 0.57 s per
likelihood (Delaunay 0.41 s); Sibson step 8.5 s -> 13 ms.

Phase 1 of #623.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@Jammy2211 Jammy2211 added the pending-release PR queued for the next release build label Oct 15, 2026
@Jammy2211
Jammy2211 merged commit 53d94b9 into main Oct 15, 2026
3 checks passed
@Jammy2211
Jammy2211 deleted the feature/delaunay-nn-numba branch October 15, 2026 16:13
@Jammy2211 Jammy2211 removed the pending-release PR queued for the next release build label Oct 16, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant