From 25f5f5e63a46c486a8f013c09bb9d984ec2380d2 Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 26 Jul 2026 13:37:31 +0000 Subject: [PATCH 1/7] docs: final assessment of rectangular-mesh pixelized-source gradients Records the 2026-07-26 re-certification and linear-algebra precision probe in the gradient-audit README: all four jax_grad pixelization variants pass strict FD post-consolidation (production shape <= 6.2e-8 rel err); slogdet vs cholesky log-det gradients agree to 2e-15; the relaxed-KKT NNLS backward pass contributes <= 3e-10 relative gradient error at the default target_kappa; the os_pix=1 einstein_radius FD poisoning persists under the unconstrained solver, re-confirming PyAutoArray#377's solver exoneration. Verdict: no linear-algebra reformulation warranted for gradient precision. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_013FSqnkgZv97PU9JdkCcthy --- jax_profiling/gradient/README.md | 59 ++++++++++++++++++++++++++++++++ 1 file changed, 59 insertions(+) diff --git a/jax_profiling/gradient/README.md b/jax_profiling/gradient/README.md index 856c06d..9116dbe 100644 --- a/jax_profiling/gradient/README.md +++ b/jax_profiling/gradient/README.md @@ -163,8 +163,67 @@ Three structurally distinct obstacles, in increasing depth: `_jax_delaunay_tables` in `jax.custom_jvp` with a zero rule, then FD-validate the frozen-triangulation gradients the same way as the rectangular mesh. +## Final assessment: rectangular-mesh pixelized-source gradients (2026-07-26) + +Post-consolidation (PyAutoArray#403 — the kernel-CDF meshes now ARE +`RectangularAdaptDensity` / `RectangularAdaptImage`; the FD tests are variants +A–D of `autolens_workspace_test/scripts/imaging/jax_grad/pixelization.py`), a +final certification re-run plus a linear-algebra precision probe (fresh +environment, jax 0.10.2 CPU, float64) settles the two standing questions. + +**Do the gradients work?** Yes — re-certified. All four variants pass strict +FD on all 14 parameters: `RectangularUniform` rel err ≤ 2e-7, +`RectangularAdaptDensity` os_pix=4 ≤ 1.7e-7, and the full production shape +(`RectangularAdaptImage` + `reg.Adapt` + `AdaptImages` + border relocator, +os_pix=4) ≤ 6.2e-8, with typical parameters at 1e-9–1e-11. The one standing +exclusion (os_pix=1 `einstein_radius`) reproduces exactly as documented: all +three FD steps land on branch flips (FD values 5.4e8 / −5.4e7 / −5.3e6 against +a self-consistent AD of +1.5e5) while every other parameter matches at ≤ 4e-7. + +**Does the linear algebra need reformulating (à la the slogdet option)?** +No — measured, each candidate in turn on the os_pix=4 production-adjacent +config: + +- **`log_det_method="slogdet"` vs `"cholesky"`** (PyAutoArray#391/#392): LL + differs by 6.4e-10 and the 14-parameter AD gradients agree to **2.1e-15** + max relative — in the positive-definite regime the two formulations are + gradient-equivalent to machine precision. `slogdet` buys *robustness* + (finite + differentiable where the Cholesky NaNs, i.e. extreme + regularization coefficients during exploration), not precision; keeping it + opt-in for gradient-based searches is the right call. +- **Relaxed-KKT NNLS backward pass** (jaxnnls implicit differentiation): + sweeping `target_kappa` 1e-9 → 1e-13 around the 1e-11 default moves the + gradient by ≤ **3.1e-10** max relative — the backward-pass relaxation error + is 2–3 orders of magnitude below the FD certification floor. Tightening + `nnls_solver_tol` to 1e-12 moves gradients by 7.5e-15 (the forward solve is + already converged). +- **NNLS custom-VJP vs exact unconstrained solve**: with + `use_positive_only_solver=False` (plain `linalg.solve`, exact built-in + implicit autodiff) the AD-vs-FD floor is the same (max 1.1e-7 vs 1.7e-7) — + the positive-only solver's custom VJP loses nothing measurable against the + analytically exact reference. +- **os_pix=1 branch-flip cross-check**: the `einstein_radius` FD poisoning + persists unchanged under the unconstrained solver — re-confirming #377's + solver exoneration on the imaging config (the flips live in the fused JIT + graph, not the NNLS solve). AD remains the trustworthy value there. + +The AD-vs-FD residual (~1e-7 worst-case) is FD noise (step quantization + +micro-flips), not autodiff error: FD converges toward AD as steps shrink at +every clean step. **Verdict: gradients for rectangular-mesh pixelized sources +are correct and already at the measurable precision floor; no linear-algebra +reformulation is warranted.** What remains open is orthogonal to gradient +precision: the JIT-only branch-flip localization (#377 follow-up — an XLA +fusion ulp threshold, measure-thin, documented LL accuracy floor under jit), +and the bandwidth-default quality question +(`PyAutoMind/draft/research/autoarray/rectangular_kernel_bandwidth_defaults.md`), +which affects reconstruction quality, not differentiability. + ## Findings log +- **2026-07-26** (final assessment, this section above): certification + re-passes post-consolidation; slogdet/cholesky gradient-equivalent to 2e-15; + relaxed-KKT backward error ≤ 3e-10; solver re-exonerated at os_pix=1 on the + imaging config. No linear-algebra reformulation warranted. - **2026-07-10** (kernel-CDF certification, PyAutoArray#373/#374): the kernel-density CDF meshes pass strict FD on ALL parameters in every configuration, including the two previously-dead corners (imaging os_pix=1, From 5f6071e26001159af966549bde0de9c881b9a4b0 Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 26 Jul 2026 14:31:13 +0000 Subject: [PATCH 2/7] docs: KNN-mesh gradient certification rows + findings-log entry Records the 2026-07-26 jax_grad/knn.py certification in the gradient-audit README: KNearestNeighbor (Wendland) and KNNBarycentric FD-certified strict on all parameters under split-family regularization; neighbour-based regularization pinned as the scipy TracerArrayConversionError boundary; the KNNBarycentric science-gate caveat (PyAutoArray#317) carried alongside its gradient status. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_013FSqnkgZv97PU9JdkCcthy --- jax_profiling/gradient/README.md | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/jax_profiling/gradient/README.md b/jax_profiling/gradient/README.md index 9116dbe..f312a69 100644 --- a/jax_profiling/gradient/README.md +++ b/jax_profiling/gradient/README.md @@ -30,6 +30,8 @@ finite and non-zero. | Imaging, `RectangularKernelAdaptImage` + `reg.Adapt` + adapt images + border relocator (os_pix=4, bandwidth=0.1) | **works** | **yes** (2026-07-10: strict FD-step-sweep, all 14 params; variant G) | full production shape; FoM parity floor 6.3e-4 relative (intrinsic: the kernel smooths the adapt-image weights over its bandwidth; swept bandwidth×n_knots 2026-07-10) | | Interferometer, `RectangularKernelAdaptDensity` + `reg.Adapt` (sparse operator) | **works — the interferometer staircase now has its escape hatch** | **yes** (2026-07-10: strict FD-step-sweep, all 7 mass/shear params live; variant D of `jax_grad/interferometer.py`) | the production-shape adaptive mesh with usable gradients on the no-over-sampling path; FoM parity 3.9e-5 relative at default bandwidth | | Imaging, Delaunay pixelization | **hard error** (re-confirmed 2026-07-09: 3 PASS / 8 ERROR) | n/a | `jax.pure_callback` → `scipy.spatial.Delaunay` has no JVP rule (`PyAutoArray .../interpolator/delaunay.py:126`); see "Why Delaunay gradients are infeasible today" below (phase 2) | +| Imaging, `KNearestNeighbor` (Wendland kNN) + `reg.ConstantSplit` / `reg.AdaptSplit` | **works** | **yes** (2026-07-26: strict FD-step-sweep, all 14 params, rel err ≤ 3.3e-8; `jax_grad/knn.py` variants A/B) | the JAX-native Delaunay-family mesh: Hilbert image mesh + edge zeroing, no scipy callback anywhere in the graph — gradients flow through traced query points AND traced mesh vertices. **Split-family regularization only**: `reg.Constant`/`ConstantZeroth`/`Adapt` need `MeshGeometryDelaunay.neighbors` (a direct scipy call on the traced mesh grid) and raise `TracerArrayConversionError` under `jax.grad` — pinned as a negative test in the script. Science caveat: Wendland kNN historically underperforms Delaunay (kernel knobs, caustic smearing — see PyAutoArray#317 background) | +| Imaging, `KNNBarycentric` + `reg.ConstantSplit` | **works** (gradients only) | **yes** (2026-07-26: strict FD-step-sweep, all 14 params, rel err ≤ 4.1e-7; `jax_grad/knn.py` variant C) | 3-nearest barycentric weights; slightly noisier FD than Wendland (3-NN-set swaps move weights discontinuously — measure-zero jump sites). **Mesh failed its science gate as a Delaunay replacement** (PyAutoArray#317: 2.2% log-evidence drift, ~5% of vertices unreachable) — certified for gradient correctness, not for production science | | Point source, source-plane χ² (`FitPositionsSource`) | **works** (probe 4/4 PASS; forward `jax.jit` still blocked by the `Grid2DIrregular` xp gap) | **yes** (2026-07-09, rel err ≤ 5e-6; `jax_grad/point_source.py`) | includes magnification-via-Hessian term (3rd derivatives of the potential); flux/H0 legitimately zero in positions-only fits | | Point source, image-plane (`FitPositionsImagePairAll`) | prior probe: **not differentiable** | n/a | `PointSolver` triangle-tiling forward solve uses `jnp.where` masking + integer neighbour lookups | | Weak lensing (`FitWeak`, `xp=jnp`) | **works** | **yes** (2026-07-09, rel err ≤ 3e-9, plain + redshift-scaled; `jax_grad/weak.py`) | gradients through the deflection-Hessian shear derivation are correct; no step-by-step probe needed — full pipeline validated first try | @@ -220,6 +222,18 @@ which affects reconstruction quality, not differentiability. ## Findings log +- **2026-07-26** (KNN mesh certification, `jax_grad/knn.py` — new): both + k-nearest-neighbour meshes FD-certified strict on all 14 params + (`KNearestNeighbor` + ConstantSplit/AdaptSplit ≤ 3.3e-8 rel err; + `KNNBarycentric` ≤ 4.1e-7). Pure JAX end to end — no scipy callback, + gradients flow through traced mesh vertices as well as queries. Boundary + pinned: neighbour-based regularizations (`Constant`/`ConstantZeroth`/ + `Adapt`) raise `TracerArrayConversionError` (scipy `Delaunay` on the + traced mesh grid in `MeshGeometryDelaunay.neighbors`) — split-family and + kernel schemes are the JAX pairings. `AdaptSplit()` at default + coefficients (inner == outer == 1.0) is numerically identical to + `ConstantSplit(1.0)` — use asymmetric coefficients to exercise the + adaptive path. - **2026-07-26** (final assessment, this section above): certification re-passes post-consolidation; slogdet/cholesky gradient-equivalent to 2e-15; relaxed-KKT backward error ≤ 3e-10; solver re-exonerated at os_pix=1 on the From 98d25ed4bcb7553a43a107fe9fb9fc80d183af39 Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 26 Jul 2026 16:54:04 +0000 Subject: [PATCH 3/7] docs: regularization x mesh gradient matrix + Matern/tfp verdict Adds the 2026-07-26 sweep results to the gradient-audit README: the full al.reg x {rectangular, KNN} compatibility matrix (FD-certified / AD-live- but-FD-limited / hard-error / incompatible, with mechanisms), the Matern verdict (tfp bessel_kve gradients work; nu=2.5 strict-certified on the rectangular mesh), the kernel-scheme conditioning note (explicit C^-1 at cond ~1e9 on clustered traced vertices is the one linear-algebra reformulation candidate the sweep surfaced), and a findings-log entry. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_013FSqnkgZv97PU9JdkCcthy --- jax_profiling/gradient/README.md | 45 ++++++++++++++++++++++++++++++++ 1 file changed, 45 insertions(+) diff --git a/jax_profiling/gradient/README.md b/jax_profiling/gradient/README.md index f312a69..b987fdb 100644 --- a/jax_profiling/gradient/README.md +++ b/jax_profiling/gradient/README.md @@ -220,8 +220,53 @@ and the bandwidth-default quality question (`PyAutoMind/draft/research/autoarray/rectangular_kernel_bandwidth_defaults.md`), which affects reconstruction quality, not differentiability. +## Regularization × mesh gradient matrix (2026-07-26 sweep) + +Every ``al.reg`` scheme swept against the gradient-capable meshes +(``RectangularAdaptDensity`` os_pix=4; ``KNearestNeighbor`` / +``KNNBarycentric`` with Hilbert image mesh + edge zeroing, os_pix=1), on the +jax_test 14-parameter fiducial. Positive results are pinned by +``autolens_workspace_test/scripts/imaging/jax_grad/regularization.py``; +mesh-family negatives by ``jax_grad/knn.py``. + +| Regularization | Rectangular (kernel-CDF) | KNN meshes | Notes | +|---|---|---|---| +| `Constant` | **FD-certified** (jax_grad/pixelization.py) | **hard error** | rectangular neighbors are analytic/static; Delaunay-family neighbors call scipy on the traced mesh grid (`MeshGeometryDelaunay.neighbors`) → `TracerArrayConversionError` | +| `Adapt` | **FD-certified** (production config, pixelization.py variant D) | **hard error** | same neighbors split as `Constant` | +| `ConstantSplit` / `AdaptSplit` | **incompatible** (shape error: split machinery expects 4-cross-per-pixel, rectangular's shared 4-corner mappings are per-query) | **FD-certified** (jax_grad/knn.py, ≤ 3.3e-8) | the split family is the KNN/Delaunay-family production pairing; note `AdaptSplit()` at default inner==outer==1.0 ≡ `ConstantSplit(1.0)` | +| `Zeroth` | **FD-certified** (1.4e-7) | **FD-certified** (7.5e-8 / 2.9e-7) | neighbour-free, pure xp | +| `MaternKernel(nu=2.5)` | **FD-certified strict** (2.2e-4) | **works — FD-limited** (~2e-3, at the FD noise floor) | **the tfp question: YES, gradients flow** — `tfp.substrates.jax.math.bessel_kve` (tfp-nightly) ships a registered gradient w.r.t. its argument (`nu` is static), and the dense-covariance Cholesky inverse differentiates. See conditioning note below | +| `MaternKernel(nu=0.5)` / `MaternAdaptKernel` / `GaussianKernel` | **works — FD-limited** (1.4e-2 / 9.4e-2, at each variant's FD noise floor) | **works — FD-limited** (~2e-3 / 6.7e-3) | gradients finite + live everywhere; the likelihood itself carries a 1e-6..4e-5 absolute numerical noise floor (dense kernel inverse / bessel lowering), which central FD divides by the step — the "mismatch" equals that floor in every case, i.e. no evidence of wrong AD. `MaternAdaptKernel()` at default coefficients ≡ `MaternKernel` (uniform weights) | +| `BrightnessZeroth` | **hard error** | **hard error** | not xp-ported: numpy boolean ops on the traced pixel-signals array (`TracerArrayConversionError`) | +| `ExponentialKernel` | **hard error** | **hard error** | not xp-ported: numpy (N,N,2) pairwise-diff build (`TracerArrayConversionError`) — unlike `GaussianKernel`/`MaternKernel`, which are xp-threaded | +| `ConstantZeroth` | **broken** (numpy too) | **broken** | known dead code — missing `neighbors_sizes` argument (filed: `draft/bug/autoarray/constant_zeroth_broken_dead_code.md`) | +| `CurvatureMask` / `FourthOrderMask` | **incompatible** | **incompatible** | dpsi (potential-correction) schemes sized to the data grid (952), not source-mesh schemes (784/330) — shape error by construction | +| `AdaptSplitZeroth` | **incompatible** (split shape) | **hard error** (neighbors via its zeroth/adapt leg) | | + +**Kernel-scheme conditioning note.** The kernel regularizations form +``coefficient * C^-1`` explicitly (``inv_via_cholesky``). On the rectangular +mesh's well-spaced vertices cond(C) ≈ 3e5 (nu=2.5) and the scheme is +strict-FD-certifiable; on the KNN meshes' TRACED (clustered) vertices +cond(C) ≈ 1.4e9 (min pairwise separation 7e-3 vs median 9e-2), which puts a +~1e-6 absolute noise floor on the likelihood — the one genuine +linear-algebra reformulation candidate this sweep surfaced (avoid the +explicit inverse, e.g. keep H implicit through the Cholesky of C for the +``s^T H s`` and ``log det H`` terms, and/or scale the 1e-8 diagonal jitter +with the kernel's dynamic range). + ## Findings log +- **2026-07-26** (regularization × mesh sweep, table above + + `jax_grad/regularization.py` — new): every `al.reg` scheme swept on both + gradient-capable mesh families. New FD certifications: `Zeroth` (both + families) and `MaternKernel(nu=2.5)` on rectangular — settling the + tfp/Bessel question: **Matérn gradients work** (tfp-nightly `bessel_kve` + has a registered gradient). Kernel schemes elsewhere are AD-live but + FD-limited by the likelihood's own noise floor (dense `C^-1` at + cond ~1e9 on clustered traced vertices) — reformulation candidate noted. + xp-port gaps found: `BrightnessZeroth`, `ExponentialKernel` (hard error + under trace); split-family structurally incompatible with rectangular; + `CurvatureMask`/`FourthOrderMask` are dpsi-only. - **2026-07-26** (KNN mesh certification, `jax_grad/knn.py` — new): both k-nearest-neighbour meshes FD-certified strict on all 14 params (`KNearestNeighbor` + ConstantSplit/AdaptSplit ≤ 3.3e-8 rel err; From 2bb551c5ffd1298fb51913c5545e6440e413ebfb Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 26 Jul 2026 17:18:05 +0000 Subject: [PATCH 4/7] =?UTF-8?q?docs:=20ExponentialKernel=20matrix=20row=20?= =?UTF-8?q?=E2=80=94=20xp-ported,=20gradients=20verified?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Updates the regularization x mesh matrix: ExponentialKernel moves from hard-error to works/FD-limited on both mesh families after the PyAutoArray xp port (same kernel-family noise floor as Matern nu=0.5, which it equals mathematically). Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_013FSqnkgZv97PU9JdkCcthy --- jax_profiling/gradient/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/jax_profiling/gradient/README.md b/jax_profiling/gradient/README.md index b987fdb..3deae08 100644 --- a/jax_profiling/gradient/README.md +++ b/jax_profiling/gradient/README.md @@ -238,7 +238,7 @@ mesh-family negatives by ``jax_grad/knn.py``. | `MaternKernel(nu=2.5)` | **FD-certified strict** (2.2e-4) | **works — FD-limited** (~2e-3, at the FD noise floor) | **the tfp question: YES, gradients flow** — `tfp.substrates.jax.math.bessel_kve` (tfp-nightly) ships a registered gradient w.r.t. its argument (`nu` is static), and the dense-covariance Cholesky inverse differentiates. See conditioning note below | | `MaternKernel(nu=0.5)` / `MaternAdaptKernel` / `GaussianKernel` | **works — FD-limited** (1.4e-2 / 9.4e-2, at each variant's FD noise floor) | **works — FD-limited** (~2e-3 / 6.7e-3) | gradients finite + live everywhere; the likelihood itself carries a 1e-6..4e-5 absolute numerical noise floor (dense kernel inverse / bessel lowering), which central FD divides by the step — the "mismatch" equals that floor in every case, i.e. no evidence of wrong AD. `MaternAdaptKernel()` at default coefficients ≡ `MaternKernel` (uniform weights) | | `BrightnessZeroth` | **hard error** | **hard error** | not xp-ported: numpy boolean ops on the traced pixel-signals array (`TracerArrayConversionError`) | -| `ExponentialKernel` | **hard error** | **hard error** | not xp-ported: numpy (N,N,2) pairwise-diff build (`TracerArrayConversionError`) — unlike `GaussianKernel`/`MaternKernel`, which are xp-threaded | +| `ExponentialKernel` | **works — FD-limited** (kernel-family noise floor, as `Matérn ν=0.5` — the exponential kernel *is* Matérn ν=0.5) | **works — FD-limited** | xp-ported 2026-07-26 (was: numpy (N,N,2) pairwise-diff build + un-threaded `xp`; also switched to the NaN-safe `sqrt(d²+1e-20)` distance form — `linalg.norm`'s derivative is NaN at the zero diagonal) | | `ConstantZeroth` | **broken** (numpy too) | **broken** | known dead code — missing `neighbors_sizes` argument (filed: `draft/bug/autoarray/constant_zeroth_broken_dead_code.md`) | | `CurvatureMask` / `FourthOrderMask` | **incompatible** | **incompatible** | dpsi (potential-correction) schemes sized to the data grid (952), not source-mesh schemes (784/330) — shape error by construction | | `AdaptSplitZeroth` | **incompatible** (split shape) | **hard error** (neighbors via its zeroth/adapt leg) | | From c81566297a75b2bd6089f6e9936557c3f84c1ff7 Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 26 Jul 2026 18:27:55 +0000 Subject: [PATCH 5/7] =?UTF-8?q?docs:=20Delaunay=20frozen-tables=20gradient?= =?UTF-8?q?=20probe=20=E2=80=94=20the=20hard-error=20verdict=20is=20unbloc?= =?UTF-8?q?kable?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The 2026-07-26 probe shows stop_gradient on the tables pure_callback input unlocks jax.grad through the full Delaunay likelihood: with the visibility walk in-graph, only the int32 tables (piecewise-constant in the vertices) are frozen, so this is the exact almost-everywhere derivative. 14/14 params live on the production shape, FD median 9.6e-6 / max 2.1e-3, eager==jit. Updates the status row and findings log; ship task filed to PyAutoMind. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_013FSqnkgZv97PU9JdkCcthy --- jax_profiling/gradient/README.md | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/jax_profiling/gradient/README.md b/jax_profiling/gradient/README.md index 3deae08..f50e845 100644 --- a/jax_profiling/gradient/README.md +++ b/jax_profiling/gradient/README.md @@ -29,7 +29,7 @@ finite and non-zero. | Imaging, `RectangularKernelAdaptDensity` (os_pix=4) | **works** | **yes** (2026-07-10: strict FD-step-sweep, all 14 params; variant F) | FoM parity with linear `AdaptDensity` = 2.7e-5 relative at default bandwidth | | Imaging, `RectangularKernelAdaptImage` + `reg.Adapt` + adapt images + border relocator (os_pix=4, bandwidth=0.1) | **works** | **yes** (2026-07-10: strict FD-step-sweep, all 14 params; variant G) | full production shape; FoM parity floor 6.3e-4 relative (intrinsic: the kernel smooths the adapt-image weights over its bandwidth; swept bandwidth×n_knots 2026-07-10) | | Interferometer, `RectangularKernelAdaptDensity` + `reg.Adapt` (sparse operator) | **works — the interferometer staircase now has its escape hatch** | **yes** (2026-07-10: strict FD-step-sweep, all 7 mass/shear params live; variant D of `jax_grad/interferometer.py`) | the production-shape adaptive mesh with usable gradients on the no-over-sampling path; FoM parity 3.9e-5 relative at default bandwidth | -| Imaging, Delaunay pixelization | **hard error** (re-confirmed 2026-07-09: 3 PASS / 8 ERROR) | n/a | `jax.pure_callback` → `scipy.spatial.Delaunay` has no JVP rule (`PyAutoArray .../interpolator/delaunay.py:126`); see "Why Delaunay gradients are infeasible today" below (phase 2) | +| Imaging, Delaunay pixelization | **hard error — but UNBLOCKABLE by one line** (probed 2026-07-26) | probe: median 9.6e-6, max 2.1e-3 over 14 params (mass/shear live) | `jax.pure_callback` has no JVP rule — but the current architecture runs EVERYTHING differentiable in-graph (visibility-walk point location, barycentric weights, dual areas, split points); the callback returns only int32 tables, which are piecewise-constant in the vertices, so `stop_gradient` on the callback input yields the EXACT almost-everywhere derivative (nothing dropped). Probe passed on the full production shape (Hilbert + edge zeroing + AdaptSplit); residual FD scatter on mass = FD steps crossing re-wiring events. Ship task filed: `PyAutoMind/draft/feature/autoarray/delaunay_frozen_tables_gradient.md`. The 2026-07-09 "infeasible" section below predates the in-graph walk and is superseded on points 1–2 (point 3's a.e. analysis is what the probe confirms) | | Imaging, `KNearestNeighbor` (Wendland kNN) + `reg.ConstantSplit` / `reg.AdaptSplit` | **works** | **yes** (2026-07-26: strict FD-step-sweep, all 14 params, rel err ≤ 3.3e-8; `jax_grad/knn.py` variants A/B) | the JAX-native Delaunay-family mesh: Hilbert image mesh + edge zeroing, no scipy callback anywhere in the graph — gradients flow through traced query points AND traced mesh vertices. **Split-family regularization only**: `reg.Constant`/`ConstantZeroth`/`Adapt` need `MeshGeometryDelaunay.neighbors` (a direct scipy call on the traced mesh grid) and raise `TracerArrayConversionError` under `jax.grad` — pinned as a negative test in the script. Science caveat: Wendland kNN historically underperforms Delaunay (kernel knobs, caustic smearing — see PyAutoArray#317 background) | | Imaging, `KNNBarycentric` + `reg.ConstantSplit` | **works** (gradients only) | **yes** (2026-07-26: strict FD-step-sweep, all 14 params, rel err ≤ 4.1e-7; `jax_grad/knn.py` variant C) | 3-nearest barycentric weights; slightly noisier FD than Wendland (3-NN-set swaps move weights discontinuously — measure-zero jump sites). **Mesh failed its science gate as a Delaunay replacement** (PyAutoArray#317: 2.2% log-evidence drift, ~5% of vertices unreachable) — certified for gradient correctness, not for production science | | Point source, source-plane χ² (`FitPositionsSource`) | **works** (probe 4/4 PASS; forward `jax.jit` still blocked by the `Grid2DIrregular` xp gap) | **yes** (2026-07-09, rel err ≤ 5e-6; `jax_grad/point_source.py`) | includes magnification-via-Hessian term (3rd derivatives of the potential); flux/H0 legitimately zero in positions-only fits | @@ -256,6 +256,15 @@ with the kernel's dynamic range). ## Findings log +- **2026-07-26** (Delaunay frozen-tables gradient probe): `stop_gradient` on + the tables `pure_callback` input unlocks `jax.grad` through the full + Delaunay likelihood — the in-graph visibility walk means only the int32 + tables (true derivative: zero between re-wirings) are frozen, so this is + the exact a.e. gradient, not an approximation. Probe on the production + shape: 14/14 params live, FD median 9.6e-6 / max 2.1e-3, eager==jit. Ship + task filed to PyAutoMind. Remaining Delaunay-vs-KNN trade: the callback is + `vmap_method="sequential"` — one host qhull call per vmap lane — so the + KNN meshes stay the batched-throughput option. - **2026-07-26** (regularization × mesh sweep, table above + `jax_grad/regularization.py` — new): every `al.reg` scheme swept on both gradient-capable mesh families. New FD certifications: `Zeroth` (both From e545538c8e1d0d783ddd43218e9e0e4f8be7420a Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 26 Jul 2026 19:30:40 +0000 Subject: [PATCH 6/7] =?UTF-8?q?docs:=20Delaunay=20gradients=20shipped=20?= =?UTF-8?q?=E2=80=94=20frozen-tables=20section=20+=20status=20row=20update?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_013FSqnkgZv97PU9JdkCcthy --- jax_profiling/gradient/README.md | 26 ++++++++++++++++++++++++-- 1 file changed, 24 insertions(+), 2 deletions(-) diff --git a/jax_profiling/gradient/README.md b/jax_profiling/gradient/README.md index f50e845..cb34a8b 100644 --- a/jax_profiling/gradient/README.md +++ b/jax_profiling/gradient/README.md @@ -29,7 +29,7 @@ finite and non-zero. | Imaging, `RectangularKernelAdaptDensity` (os_pix=4) | **works** | **yes** (2026-07-10: strict FD-step-sweep, all 14 params; variant F) | FoM parity with linear `AdaptDensity` = 2.7e-5 relative at default bandwidth | | Imaging, `RectangularKernelAdaptImage` + `reg.Adapt` + adapt images + border relocator (os_pix=4, bandwidth=0.1) | **works** | **yes** (2026-07-10: strict FD-step-sweep, all 14 params; variant G) | full production shape; FoM parity floor 6.3e-4 relative (intrinsic: the kernel smooths the adapt-image weights over its bandwidth; swept bandwidth×n_knots 2026-07-10) | | Interferometer, `RectangularKernelAdaptDensity` + `reg.Adapt` (sparse operator) | **works — the interferometer staircase now has its escape hatch** | **yes** (2026-07-10: strict FD-step-sweep, all 7 mass/shear params live; variant D of `jax_grad/interferometer.py`) | the production-shape adaptive mesh with usable gradients on the no-over-sampling path; FoM parity 3.9e-5 relative at default bandwidth | -| Imaging, Delaunay pixelization | **hard error — but UNBLOCKABLE by one line** (probed 2026-07-26) | probe: median 9.6e-6, max 2.1e-3 over 14 params (mass/shear live) | `jax.pure_callback` has no JVP rule — but the current architecture runs EVERYTHING differentiable in-graph (visibility-walk point location, barycentric weights, dual areas, split points); the callback returns only int32 tables, which are piecewise-constant in the vertices, so `stop_gradient` on the callback input yields the EXACT almost-everywhere derivative (nothing dropped). Probe passed on the full production shape (Hilbert + edge zeroing + AdaptSplit); residual FD scatter on mass = FD steps crossing re-wiring events. Ship task filed: `PyAutoMind/draft/feature/autoarray/delaunay_frozen_tables_gradient.md`. The 2026-07-09 "infeasible" section below predates the in-graph walk and is superseded on points 1–2 (point 3's a.e. analysis is what the probe confirms) | +| Imaging, Delaunay pixelization | **works — frozen-tables a.e.-exact gradient** (SHIPPED 2026-07-26: `stop_gradient` on the tables callback input) | **yes** (2026-07-26: median 9.6e-6, max 2.1e-3 over 14 params at documented rtol=1e-2; `jax_grad/delaunay.py`) | everything differentiable runs in-graph (visibility-walk point location, barycentric weights, dual areas, split points); the callback returns only int32 tables, piecewise-constant in the vertices, so freezing them under differentiation drops nothing — the exact a.e. derivative. Residual FD scatter on mass/shear = FD steps crossing triangle-flip events (measure-zero likelihood discontinuities). Batched caveat: callback is `vmap_method="sequential"` — KNN meshes remain the vmap-throughput option. See "Delaunay gradients: SHIPPED" section below | | Imaging, `KNearestNeighbor` (Wendland kNN) + `reg.ConstantSplit` / `reg.AdaptSplit` | **works** | **yes** (2026-07-26: strict FD-step-sweep, all 14 params, rel err ≤ 3.3e-8; `jax_grad/knn.py` variants A/B) | the JAX-native Delaunay-family mesh: Hilbert image mesh + edge zeroing, no scipy callback anywhere in the graph — gradients flow through traced query points AND traced mesh vertices. **Split-family regularization only**: `reg.Constant`/`ConstantZeroth`/`Adapt` need `MeshGeometryDelaunay.neighbors` (a direct scipy call on the traced mesh grid) and raise `TracerArrayConversionError` under `jax.grad` — pinned as a negative test in the script. Science caveat: Wendland kNN historically underperforms Delaunay (kernel knobs, caustic smearing — see PyAutoArray#317 background) | | Imaging, `KNNBarycentric` + `reg.ConstantSplit` | **works** (gradients only) | **yes** (2026-07-26: strict FD-step-sweep, all 14 params, rel err ≤ 4.1e-7; `jax_grad/knn.py` variant C) | 3-nearest barycentric weights; slightly noisier FD than Wendland (3-NN-set swaps move weights discontinuously — measure-zero jump sites). **Mesh failed its science gate as a Delaunay replacement** (PyAutoArray#317: 2.2% log-evidence drift, ~5% of vertices unreachable) — certified for gradient correctness, not for production science | | Point source, source-plane χ² (`FitPositionsSource`) | **works** (probe 4/4 PASS; forward `jax.jit` still blocked by the `Grid2DIrregular` xp gap) | **yes** (2026-07-09, rel err ≤ 5e-6; `jax_grad/point_source.py`) | includes magnification-via-Hessian term (3rd derivatives of the potential); flux/H0 legitimately zero in positions-only fits | @@ -125,7 +125,29 @@ to autodiff — clean steps converge to AD at 1e-6..1e-9 relative, so a wrong AD fails every step. The flips deserve their own investigation (likely positive-only-solver / PDIP tie-breaks — NNLS-ledger territory). -## Why Delaunay gradients are infeasible today (phase 2a, 2026-07-09) +## Delaunay gradients: SHIPPED via frozen tables (2026-07-26) + +The section below this one records the 2026-07-09 verdict and is kept as +history; its premises are superseded. The in-graph visibility walk moved +every differentiable quantity (point location, barycentric weights, dual +areas, split points) inside the JIT program, 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 re-wiring events — so wrapping the callback input in +``stop_gradient`` (shipped in ``_jax_delaunay_tables``) yields the EXACT +almost-everywhere derivative, not an approximation (the audit's point 3 +below, now realised; points 1–2 described the pre-walk architecture). + +Certification: `autolens_workspace_test/scripts/imaging/jax_grad/delaunay.py` +(production shape: Hilbert + edge zeroing + AdaptSplit) — all 14 params live, +lens light at 1e-8..1e-10, mass/shear at 1e-5..2e-3 (FD steps straddling +triangle-flip events — measure-zero likelihood discontinuities where no +method has a gradient; AD differentiates the branch the point is on), at a +documented rtol=1e-2. Batched-sampler caveat: the callback remains +``vmap_method="sequential"`` (one host qhull call per vmap lane) — the KNN +meshes stay the batched-throughput option. + +## Why Delaunay gradients were infeasible before the walk (phase 2a, 2026-07-09 — superseded, see above) Re-confirmed on current mains via `imaging/delaunay.py`: **3 PASS / 8 ERROR**. The pre-inversion stages (ray-trace, blurred lens light, profile subtraction) are fully From 3b9c8cebcca282b9c1cf913e15a8fc319b530529 Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 26 Jul 2026 19:51:20 +0000 Subject: [PATCH 7/7] docs: BrightnessZeroth/AdaptSplitZeroth matrix rows post xp-port Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_013FSqnkgZv97PU9JdkCcthy --- jax_profiling/gradient/README.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/jax_profiling/gradient/README.md b/jax_profiling/gradient/README.md index cb34a8b..a062b4e 100644 --- a/jax_profiling/gradient/README.md +++ b/jax_profiling/gradient/README.md @@ -259,11 +259,11 @@ mesh-family negatives by ``jax_grad/knn.py``. | `Zeroth` | **FD-certified** (1.4e-7) | **FD-certified** (7.5e-8 / 2.9e-7) | neighbour-free, pure xp | | `MaternKernel(nu=2.5)` | **FD-certified strict** (2.2e-4) | **works — FD-limited** (~2e-3, at the FD noise floor) | **the tfp question: YES, gradients flow** — `tfp.substrates.jax.math.bessel_kve` (tfp-nightly) ships a registered gradient w.r.t. its argument (`nu` is static), and the dense-covariance Cholesky inverse differentiates. See conditioning note below | | `MaternKernel(nu=0.5)` / `MaternAdaptKernel` / `GaussianKernel` | **works — FD-limited** (1.4e-2 / 9.4e-2, at each variant's FD noise floor) | **works — FD-limited** (~2e-3 / 6.7e-3) | gradients finite + live everywhere; the likelihood itself carries a 1e-6..4e-5 absolute numerical noise floor (dense kernel inverse / bessel lowering), which central FD divides by the step — the "mismatch" equals that floor in every case, i.e. no evidence of wrong AD. `MaternAdaptKernel()` at default coefficients ≡ `MaternKernel` (uniform weights) | -| `BrightnessZeroth` | **hard error** | **hard error** | not xp-ported: numpy boolean ops on the traced pixel-signals array (`TracerArrayConversionError`) | +| `BrightnessZeroth` | **xp-ported 2026-07-26** — differentiates; standalone use resamples (supplement by design: weights → 0 on high-signal pixels → singular system) | same | was a dropped-`xp` call site (`pixel_signals_from` called without the backend — same bug class as `ExponentialKernel`); its gradient path is FD-certified through `AdaptSplitZeroth` | | `ExponentialKernel` | **works — FD-limited** (kernel-family noise floor, as `Matérn ν=0.5` — the exponential kernel *is* Matérn ν=0.5) | **works — FD-limited** | xp-ported 2026-07-26 (was: numpy (N,N,2) pairwise-diff build + un-threaded `xp`; also switched to the NaN-safe `sqrt(d²+1e-20)` distance form — `linalg.norm`'s derivative is NaN at the zero diagonal) | | `ConstantZeroth` | **broken** (numpy too) | **broken** | known dead code — missing `neighbors_sizes` argument (filed: `draft/bug/autoarray/constant_zeroth_broken_dead_code.md`) | | `CurvatureMask` / `FourthOrderMask` | **incompatible** | **incompatible** | dpsi (potential-correction) schemes sized to the data grid (952), not source-mesh schemes (784/330) — shape error by construction | -| `AdaptSplitZeroth` | **incompatible** (split shape) | **hard error** (neighbors via its zeroth/adapt leg) | | +| `AdaptSplitZeroth` | **incompatible** (split shape) | **FD-certified** (8.2e-8, post BrightnessZeroth xp-port 2026-07-26) | its KNN hard error was the `BrightnessZeroth` dropped-`xp` call in the zeroth leg | **Kernel-scheme conditioning note.** The kernel regularizations form ``coefficient * C^-1`` explicitly (``inv_via_cholesky``). On the rectangular