diff --git a/draft/feature/autoarray/regularization_jax_gradient_gaps.md b/draft/feature/autoarray/regularization_jax_gradient_gaps.md index 754dcfa7..4566cbb2 100644 --- a/draft/feature/autoarray/regularization_jax_gradient_gaps.md +++ b/draft/feature/autoarray/regularization_jax_gradient_gaps.md @@ -5,7 +5,7 @@ Target: autoarray Repos: - PyAutoArray - autolens_workspace_test -Difficulty: medium +Difficulty: small Autonomy: supervised Priority: normal Status: draft @@ -53,6 +53,105 @@ Reformulation candidates (in the spirit of the opt-in slogdet, PyAutoArray#391): - Cheaper interim: scale the fixed `1e-8` diagonal jitter with the kernel's dynamic range / N, and expose it as a kwarg. +### Status of leg 2, re-derived against main `efaf3041` (2026-08-09) + +The two candidates above were written before #391 landed and are now partly +stale. Checked against source rather than assumed: + +- **`log det H` implicit — ALREADY SHIPPED.** `log_det_regularization_matrix_term_from` + exists on all four kernel schemes (`MaternKernel`, `MaternAdaptKernel`, + `GaussianKernel`, `ExponentialKernel`), gated behind + `Settings.log_det_method == "slogdet"` (PyAutoArray#391). Nothing owed. +- **jitter as a kwarg — ALREADY SHIPPED.** All four carry + `jitter: Optional[float] = None` + a `jitter_value` property. The *scaling* + was the remaining half — tracked as leg 2b and now shipped too (see its block + below). Note `gaussian_kernel.py` already trace-scales + (`1e-8 * abs(diag_mean)`), but on the formed `H` rather than on `C`; that + stabilisation is a separate guard on the explicit inverse and was left alone. +- **"Keep `H` implicit" in general — NOT ACHIEVABLE, do not re-open.** + `curvature_reg_matrix` (`inversion/abstract.py:366`) is a dense + `xp.add(curvature_matrix, regularization_matrix)` feeding the dense solve for + the reconstruction. `H` must still be formed there. Only the *evidence* terms + can avoid the explicit inverse. Any future attempt needs an iterative-solver + design first, which is a different task. + +**Leg 2a — `s^T H s` implicit — SHIPPED 2026-08-09** on branch +`claude/automind-task-planning-gm4flt` (PyAutoArray). The one genuinely open +piece: `regularization_term` (`inversion/abstract.py:698`) contracted the formed +`H`. For the kernel schemes the term is `coefficient * s^T C^-1 s`, from one +Cholesky solve against `s` instead of forming `C^-1` and contracting it. + +Built to the exact shape of the shipped log-det shortcut: a +`regularization_term_from` hook on `AbstractRegularization` returning `None` by +default, overridden by the four kernel schemes, consulted by +`AbstractInversion.regularization_term` behind a new opt-in +`Settings.regularization_term_method` (`"matmul"` default, `"cho_solve"` +opt-in). Deliberately a *separate* setting from `log_det_method` so the two +evidence terms can move onto their exact factorizations independently — that is +what makes an evidence shift attributable to one term rather than both. Default +evidence values unchanged, so archived comparability holds (the constraint the +#391 adversarial probe verdict imposed). + +Measured gain on a clustered fixture at `cond(C) = 3.2e9`, graded against an +exactly-known reference (`s = C v` makes `C^-1 s = v`, so the true form is +`s^T v`): relative error **5.99e-08 explicit vs 2.93e-16 implicit** — the +implicit path is at machine precision, confirming this prompt's +~1e-6..4e-5 noise-floor measurement was the explicit inverse. + +Two traps found and pinned by tests: `MaternAdaptKernel` passes +`coefficient=0.0` to `MaternKernel.__init__` (its weights live inside `C_w`), so +inheriting the parent term would silently zero it — it needs its own override; +and `GaussianKernel`'s formed matrix carries a symmetrisation + trace-scaled +jitter the shortcut excludes, matching how its log-det shortcut already behaves. +Schemes with no factorization return `None`, and one `None` falls the whole +computation back to the formed matrix, so mixed inversions stay correct. + +9 new tests in `test_autoarray/inversion/regularizations/test_kernel_regularization_term.py`. +Full inversion suite green (244 passed, 9 skipped). + +**Leg 2b — jitter scaling — SHIPPED 2026-08-09** on the same branch. + +Investigated before designing, and the prompt's framing ("scale with the +kernel's dynamic range / N") turned out to point at a real and more serious +problem than "cheaper interim" suggests. The jitter is a fixed absolute +`1e-8 * I`, which is only meaningful when `diag(C) ~ 1`. That holds for the +three unweighted kernels (`K(0) == 1`, measured) but NOT for +`MaternAdaptKernel`, whose `C_ii = w_i²` spans the adaptive-weight dynamic +range. Measured distortion of the faintest pixel on a 40-pixel fixture: + +| inner/outer | faintest `C_ii` | distortion | +|---|---|---| +| 1.0 / 1.0 | 1.0e+00 | 1.0e-08 | +| 0.5 / 4.0 | 3.9e-03 | 2.6e-06 | +| 1.0 / 20.0 | 6.3e-06 | 1.6e-03 | +| 0.1 / 100.0 | 1.0e-08 | **1.0e+00** | + +`inner_coefficient`/`outer_coefficient` are **free model parameters**, so a +sampler can walk into the bottom row mid-fit, at which point the jitter is 100% +of the faintest pixels' variance and their kernel structure is gone. Silent — +no exception, no NaN, just wrong smoothing. + +Fix: `jitter_relative=True` applies `jitter * diag(diag(C))` (`C_ii *= 1 + jitter`). +With `C = D^½ R D^½` for correlation matrix `R`, that is exactly +`D^½ (R + jitter·I) D^½` — the jitter lands on the *correlation* matrix, so every +pixel gets the same relative protection whatever its scale. + +One rejected design worth recording so it is not retried: `jitter = N·eps·max(diag)` +("as small as possible above the round-off floor") fixes the distortion but lets +`cond(C)` reach **3.2e15** on smooth clustered vertices — the edge of float64, +reintroducing exactly the noise leg 2 exists to remove. The correlation-relative +rule instead leaves conditioning *unchanged* (3.16e9 both ways, measured). + +Default `False` everywhere, byte-identical to previous behaviour. Threaded +through all three covariance call sites per scheme, not just the constructor +(pinned by a test). 7 new tests in `test_kernel_jitter_relative.py`. + +**This prompt now has no open work.** Both remaining legs are shipped; the file +is ready to advance to `complete/` once the branch merges. Note the JAX leg of +the gate (`jax_grad/regularization.py` re-passing) is still owed — it lives in +`autolens_workspace_test` and needs `tfp-nightly`, neither available in the +session that did this work. + Gate any change on the `regularization.py` jax_grad script re-passing and on FoM parity on the numpy path. @@ -70,7 +169,7 @@ implemented — the capability is absent by design. The `rectangular_adapt_constant_split_guard.md` merge question this section raises is therefore moot; that prompt was recorded as complete on 2026-08-09 ([[rectangular-adapt-constant-split-guard]]). Verified by the draft/ sweep against -main `efaf3041`. **Leg 2 is now the only open work in this prompt.*** +main `efaf3041`. **All legs are now shipped — see the status block under leg 2.*** `ConstantSplit`/`AdaptSplit`/`AdaptSplitZeroth` on a rectangular mesh fail with a raw broadcasting `TypeError` ((784,784) vs (3808,3808)): the split