diff --git a/complete/2026/07/delaunay-frozen-tables-gradient.md b/complete/2026/07/delaunay-frozen-tables-gradient.md new file mode 100644 index 00000000..c80faeb6 --- /dev/null +++ b/complete/2026/07/delaunay-frozen-tables-gradient.md @@ -0,0 +1,133 @@ +## delaunay-frozen-tables-gradient +- completed: 2026-07-26 +- branch: `claude/rectangular-mesh-gradients-mh1j0z` (PyAutoArray + autolens_workspace_test + autolens_workspace_developer; PR/merge pending with the rest of the session's mesh-gradients work — no issue was opened, the task was conceived and executed in the same session) +- key commits: PyAutoArray `feat: unlock jax.grad through the Delaunay mesh via frozen integer tables` + `docs: state the sense in which the Delaunay mesh is autodifferentiable`; autolens_workspace_test `test: jax_grad FD certification for the Delaunay mesh` +- verdict: SHIPPED — the Delaunay mesh is autodifferentiable (a.e.-exact frozen-tables gradient) +- summary: | + The 2026-07-09 audit's "Delaunay is hard-undifferentiable" verdict was + stale: the visibility-walk refactor had moved every differentiable + quantity (point location, barycentric weights, dual areas, split points) + in-graph, leaving the host qhull `pure_callback` returning only int32 + connectivity tables. Those tables are piecewise-constant in the vertex + positions — their true derivative is zero between triangle-flip events — + so wrapping the callback input in `jax.lax.stop_gradient` (one line in + `_jax_delaunay_tables`) yields the EXACT almost-everywhere derivative, + not an approximation. Values are bit-identical (primal untouched): the + jax_likelihood delaunay regression literal passes unchanged. + + Certification (new `autolens_workspace_test/scripts/imaging/jax_grad/delaunay.py`, + production shape: Hilbert image mesh + circle edge zeroing + + reg.AdaptSplit(0.1, 10), 14-param truth-centred lens): all params live, + eager==jit, FD-step-sweep with lens light at 1e-8..1e-10 and mass/shear + at 1e-5..2e-3 (documented rtol=1e-2 — the scatter is FD steps straddling + flip events, the measure-zero jump seams where no method has a gradient; + AD differentiates the branch the point is on). + + THE SENSE IN WHICH IT IS AUTODIFFERENTIABLE (now in the Delaunay class + docstring): ReLU-network territory — piecewise-smooth, smooth within + each triangulation topology, measure-zero jump discontinuities at flips. + Contrast the kernel-CDF rectangular meshes: C-infinity by construction, + no seams — still the cleanest gradient-inference choice; Delaunay is the + scientifically exact piecewise-smooth alternative. + + REMAINING TRADE (documented, not fixed): the tables callback is + `vmap_method="sequential"` — one host qhull call per vmap lane — so the + KNN meshes keep the batched-throughput niche until the callback is + batched (delaunay_research.md option B territory). + + Lore: the audit README's "Why Delaunay gradients are infeasible" section + assessed the PRE-walk architecture (callback returned float split_points + /mappings, where freezing would have dropped real terms); its point 3 + (the frozen gradient is the correct a.e. derivative) is what this ships. + Probe pattern for the unlock: monkeypatch stop_gradient, value_and_grad, + FD sweep — 30 minutes from hypothesis to evidence. + +## Original prompt + +# Delaunay frozen-tables gradient — unlock jax.grad with stop_gradient on the tables callback + +Type: feature +Target: autoarray +Repos: +- PyAutoArray +- autolens_workspace_test +- autolens_workspace_developer +Difficulty: medium +Autonomy: supervised +Priority: high +Status: draft + +## The finding (probed 2026-07-26, wrap-up of the mesh-gradients assessment) + +The Delaunay mesh's "hard-undifferentiable" verdict (2026-07-09 audit) is +**stale**: since the visibility-walk landed, the JAX path +(`jax_delaunay` in `autoarray/inversion/mesh/interpolator/delaunay.py`) runs +everything differentiable **in-graph** — point location, barycentric weights, +dual areas, split points — and the `pure_callback` returns only int32 +triangulation tables (`simplices`, `neighbors`, `vertex_simplex`). Those +tables are piecewise-constant in the vertex positions (their true derivative +is zero between re-wiring events), so freezing them is NOT an approximation: +it is the exact almost-everywhere derivative of the likelihood. + +The unlock is one line — the callback input just needs to be excluded from +differentiation: + +```python +def _tables_stopgrad(points): + return _orig_tables(jax.lax.stop_gradient(points)) +``` + +Probe (monkeypatch, jax_test 14-param fiducial, full production shape: +Hilbert image mesh + circle edge zeroing + `reg.AdaptSplit(0.1, 10)`): + +- `value_and_grad` finite, non-zero, mass/shear live on all 14 params. +- eager == jit (rtol 1e-10). +- FD-step-sweep: **median rel err 9.6e-6, max 2.1e-3** (worst: + `mass.ell_comps_0`) — the residual scatter on mass params is FD stepping + across re-wiring events (measure-zero discontinuities, same class as the + KNN meshes' neighbour swaps and the certified branch flips of + PyAutoArray#377), while lens-light params match at 1e-8–1e-10. + +Probe script pattern: `jax_grad/knn.py` composition with +`mesh=al.mesh.Delaunay(pixels, zeroed_pixels)`; the monkeypatch above. + +## Task + +1. Ship the unlock in `_jax_delaunay_tables` (stop_gradient on `points` + before the callback — or an equivalent `custom_jvp` zero rule if + stop_gradient interacts badly with anything). Values are bit-identical + (the primal is untouched); only differentiation behaviour changes, from + "raises" to "exact a.e. gradient". +2. Certification: add a Delaunay variant to + `autolens_workspace_test/scripts/imaging/jax_grad/` (mirror `knn.py`; + FD-step-sweep, documented tolerance for the re-wiring FD scatter — + probe suggests rtol 3e-3 or per-param exclusions à la the os_pix=1 + einstein_radius precedent). Re-run the delaunay jax_likelihood scripts + (imaging, interferometer, datacube, multi) to confirm values unchanged. +3. Update `imaging/jax_grad/knn.py`'s docstring + the audit README's + Delaunay narrative ("Why Delaunay gradients are infeasible today" — + points 1–2 are superseded; point 3's a.e. analysis is what the probe + confirms) — partially done 2026-07-26 (row + findings-log updated; + the section prose still says infeasible). +4. Samplers caveat to document: gradient-based searches see measure-zero + value discontinuities at re-wiring events (the interpolant jumps when a + containing triangle's diagonal flips) — same practical class as the + XLA branch flips (#377), but geometric in origin and mesh-density + dependent. + +## Why this matters + +Delaunay is the flagship source reconstruction (exact spatial locality, +every vertex guaranteed simplex membership — the property whose absence +sank the KNN-barycentric wildcard, PyAutoArray#317). This makes it +gradient-capable at zero science cost. The remaining Delaunay-vs-KNN trade +is batched throughput only: the tables callback is +`vmap_method="sequential"` (one host qhull call per vmap lane), so KNN +meshes stay preferable for heavily-vmapped samplers until the callback is +batched (delaunay_research.md option B territory). + +## Constraints + +- The unlock must not change values (it cannot — primal untouched); the + jax_likelihood regression literals are the gate. +- Library unit tests numpy-only; JAX validation via workspace_test jax_grad. diff --git a/complete/index.md b/complete/index.md index 02b9620b..3a77dec6 100644 --- a/complete/index.md +++ b/complete/index.md @@ -6,7 +6,7 @@ Token-light navigation over the finished-work records (schema: only then grep a dated bucket. Curators: edit the band between the CURATED markers; everything below GENERATED is rebuilt. -739 records across 6 buckets. +741 records across 6 buckets. ## Highlights @@ -95,6 +95,7 @@ _(curate hard-won records here — survives regeneration.)_ - [cti-resurrection-phase5](2026/07/cti-resurrection-phase5.md) — closed - [database-latent-wheel-load](2026/07/database-latent-wheel-load.md) — database scripts failed on release wheels loading latent variables → AssertionError. Release-tail item D. Corr… - [defunct-grids-yaml-removal](2026/07/defunct-grids-yaml-removal.md) — Deleted the 10 orphan `config/grids.yaml` copies plus 3 references (2 library `config/README.md` bullets, 1 `a… +- [delaunay-frozen-tables-gradient](2026/07/delaunay-frozen-tables-gradient.md) - [delaunay-qhull-callback](2026/07/delaunay-qhull-callback.md) - [delete_pyautoheart_shim](2026/07/delete_pyautoheart_shim.md) - [desktop-acceptance-mcp](2026/07/desktop-acceptance-mcp.md) @@ -259,6 +260,7 @@ _(curate hard-won records here — survives regeneration.)_ - [prior-config-cache](2026/07/prior-config-cache.md) — The aggregator arc's "deeper follow-up" — cProfile showed 77% of per-result summary/model load was Model.__ini… - [prior-width-safety](2026/07/prior-width-safety.md) — Phase 2 of #1331 — RelativeWidthModifier abs(mean)+opt-in absolute_floor, default-path width guard (parameter-… - [priors-messages-fixes](2026/07/priors-messages-fixes.md) — Phase 1 of #1331 priors/messages batch — 7 fixes (01/02/04/10 + decisions 1/3/4), 12 regression tests, Fable m… +- [profiling-dataset-auto-simulate](2026/07/profiling-dataset-auto-simulate.md) — PR #89 - [profiling-mirror-taxonomy](2026/07/profiling-mirror-taxonomy.md) — profiling#85 + Brain#161 - [profiling-preopt-campaign](2026/07/profiling-preopt-campaign.md) — Phase-3 (likelihood_runtime) leg of the polish.md PreOptimizationTimes campaign, frozen "done enough" at user … - [psf-convolution-docstring](2026/07/psf-convolution-docstring.md) — Normalized PSF-convolution docs across 41 simulator.py (both workspaces). Tier A (3 reference files: imaging/s… diff --git a/draft/feature/autoarray/regularization_jax_gradient_gaps.md b/draft/feature/autoarray/regularization_jax_gradient_gaps.md new file mode 100644 index 00000000..6a609dad --- /dev/null +++ b/draft/feature/autoarray/regularization_jax_gradient_gaps.md @@ -0,0 +1,80 @@ +# Regularization JAX gradient gaps — xp-ports + kernel-scheme linear algebra + +Type: feature +Target: autoarray +Repos: +- PyAutoArray +- autolens_workspace_test +Difficulty: medium +Autonomy: supervised +Priority: normal +Status: draft + +## Context (2026-07-26 regularization × mesh gradient sweep) + +A full sweep of every `al.reg` scheme against the gradient-capable meshes +(`RectangularAdaptDensity` os_pix=4; `KNearestNeighbor`/`KNNBarycentric` +Hilbert + edge zeroing) mapped the JAX-gradient compatibility surface. The +matrix and measurements live in +`autolens_workspace_developer/jax_profiling/gradient/README.md` +("Regularization × mesh gradient matrix"); positive certifications are +pinned by `autolens_workspace_test/scripts/imaging/jax_grad/regularization.py` +and the mesh-family negatives by `jax_grad/knn.py`. Three actionable gaps +fell out — none blocks current production paths, so this is one prompt to +be split or trimmed at start-dev if any leg grows: + +## 1. xp-ports — DONE + +*Both legs shipped 2026-07-26 on `claude/rectangular-mesh-gradients-mh1j0z`: +`ExponentialKernel` (xp threaded through the covariance build, NaN-safe +`sqrt(d²+1e-20)` distances) and `BrightnessZeroth` (the `pixel_signals_from` +call site now threads `xp`). JAX gradients verified on both. Legs 2 and 3 +below remain the open work of this prompt.* + +## 2. Kernel-scheme linear algebra: avoid the explicit `C^-1` (the real one) + +`MaternKernel`/`MaternAdaptKernel`/`GaussianKernel` build +`H = coefficient * inv_via_cholesky(C)` — an explicit dense inverse. On +well-spaced vertices (rectangular mesh: cond(C) ≈ 3e5 at nu=2.5) this is +fine and `MaternKernel(nu=2.5)` is strict-FD-certified (2.2e-4). On TRACED +(clustered) mesh vertices (KNN meshes: min pairwise separation ~7e-3 vs +median ~9e-2 → cond(C) ≈ 1.4e9 at nu=2.5) the explicit inverse puts a +~1e-6..4e-5 absolute numerical noise floor on the likelihood itself +(measured as eager-vs-jit LL differences), which caps FD verifiability at +~1e-3..1e-2 relative and adds the same noise to sampler-visible likelihoods. + +Reformulation candidates (in the spirit of the opt-in slogdet, PyAutoArray#391): + +- Keep `H` implicit through the Cholesky of `C`: `s^T H s` via + `cho_solve(L, s)`, `log det H = -2 Σ log diag L` — no explicit inverse, + one factorization, strictly more accurate and faster. Requires the + inversion interface to accept an implicit/functional `H` (today it + consumes a dense matrix — check `curvature_reg_matrix` assembly). +- Cheaper interim: scale the fixed `1e-8` diagonal jitter with the kernel's + dynamic range / N, and expose it as a kwarg. + +Gate any change on the `regularization.py` jax_grad script re-passing and +on FoM parity on the numpy path. + +## 3. Split-family shape guard on rectangular meshes (papercut) + +`ConstantSplit`/`AdaptSplit`/`AdaptSplitZeroth` on a rectangular mesh fail +with a raw broadcasting `TypeError` ((784,784) vs (3808,3808)): the split +machinery assumes 4-cross-per-pixel splits, while the rectangular +interpolator reuses its per-query 4-corner mappings for +`_mappings_sizes_weights_split`. Either raise a clear +"split regularization requires a Delaunay-family mesh" exception at +composition time, or implement true pixel-centre crosses for the +rectangular geometry (relates to +`draft/feature/autoarray/rectangular_adapt_constant_split_guard.md` if that +covers the same surface — merge at intake if so). + +## Out of scope + +- `ConstantZeroth` dead code — already filed + (`draft/bug/autoarray/constant_zeroth_broken_dead_code.md`). +- `CurvatureMask`/`FourthOrderMask` — dpsi (potential-correction) schemes, + correctly incompatible with source meshes; nothing to fix. +- Making neighbour-based schemes (`Constant`/`Adapt`) JAX-traceable on the + Delaunay mesh family (kNN-derived neighbours) — bigger design question, + only worth filing if a production pipeline needs it.