fix: AdaptSplit regularization log-det parity between numba and JAX

Overview

At identical parameters, the NumPy/numba and JAX likelihoods of a Delaunay + AdaptSplit interferometer inversion differ by ~20–270 in log likelihood on the 125×125 Aris mask and ~1 on the 512×512 mask. Chi-squared, curvature matrix, data vector and the curvature log-det are identical; the entire difference is log_det_regularization_matrix_term (2443.22 numba vs 2561.98 JAX). The two backends factorize H differently: NumPy offers it to the SuperLU sparse LU route (inversion_util.log_det_sparse_spd_from, diagonal pivoting only), JAX takes a dense Cholesky. On an extremely ill-conditioned AdaptSplit matrix (entries summing to ~1.4e-5) these can disagree — or the two backends may be building different H matrices. Log evidences are therefore not comparable across backends.

Plan

Reproduce on the 125×125 cache at one parameter vector; first compare the two backends' regularization matrices H directly, then the log-det of the same H under dense NumPy Cholesky, SuperLU, JAX Cholesky, slogdet and an mpmath high-precision reference.
Locate the discrepancy (construction vs factorization) and make both backends compute the term the same numerically robust way (diagonal scaling H = D^{1/2} Ĥ D^{1/2} is the leading candidate: log det H = Σ log d + log det Ĥ, well-conditioned, works for both backends and stays differentiable).
Check the non-positive-definite behaviour (numba raises LinAlgError → resample; JAX returns -inf) is consistent in effect and document it.
Add a regression test pinning the value on an ill-conditioned AdaptSplit matrix against the high-precision reference, on both backends.
Decide and record whether AdaptSplit's matrix construction should be better conditioned (scope note only unless small).

Tier: glance — merge mode: auto-merge on green if the Witness passes, post-merge summary

Detailed implementation plan

Affected Repositories

PyAutoArray (primary)

Branch Survey

Repository	Current Branch	Dirty?
array/PyAutoArray	main (`a84f724`)	clean

Suggested branch: feature/adaptsplit-logdet-backend-parity

Implementation Steps

Reproduce with the session scratchpad copy of aris_evidence/step7/bench.py --check (125×125 inputs in aris_evidence/aris_ds and aris_evidence/data) at a fixed parameter vector; dump regularization_matrix_reduced from both backends and np.max(np.abs(H_np - H_jax)).
If H differs: trace AdaptSplit.regularization_matrix_from → regularization_util numba vs JAX implementations (split-point weights, neighbour ordering, dtype) and fix the divergent one; the regression test then pins H parity.
If H agrees: compute log det of that one matrix five ways (np dense Cholesky, log_det_sparse_spd_from, jnp Cholesky, slogdet, mpmath at 50 digits) and identify which is wrong. Implement the fix in AbstractInversion._log_det_symmetric_from / log_det_regularization_matrix_term (abstract.py ~970–1150) and/or inversion_util.log_det_sparse_spd_from (inversion_util.py ~878): symmetric diagonal scaling before factorization on both backends is the preferred fix; if SuperLU without row pivoting is the culprit, scale or fall back to dense Cholesky above a pivot-ratio threshold.
Non-PD consistency: keep NumPy LinAlgError (resample) and JAX NaN/-inf semantics as documented in _log_det_symmetric_from, but verify the scaled factorization does not change which draws are non-PD; note in the PR body.
Tests: test_autoarray/inversion/inversion/test_log_det_sparse.py + a new ill-conditioned AdaptSplit test (small Delaunay mesh with extreme adapt weights) asserting NumPy and JAX log_det_regularization_matrix_term agree to <1e-3 and match the mpmath reference to <1e-6; keep the existing test_log_det_shared_factor.py and curvature log-det parity green.
Witness: re-run bench.py --check on the 125×125 setup; record numba vs JAX log_det_regularization_matrix_term before (2443.22 / 2561.98) and after (agree <1e-3) in the PR body, plus the 512×512 1,459-vertex pair (2576.59 / 2574.73).

Key Files

autoarray/inversion/inversion/abstract.py — _log_det_symmetric_from, log_det_regularization_matrix_term
autoarray/inversion/inversion/inversion_util.py — log_det_sparse_spd_from
autoarray/inversion/regularization/adapt_split.py, regularization_util.py — AdaptSplit matrix construction (numba + JAX)
test_autoarray/inversion/inversion/test_log_det_sparse.py — regression tests

Evidence

PROGRESS3/PROGRESS4 (step8.3 numerics): all other evidence terms identical to 1e-16 relative; JAX eager vs jit agree to 1e-8; ~1 in 4 random draws non-PD. Evidence copied to this session's scratchpad aris_evidence/ (temporary).

Heart RED development override

Heart was RED at launch (2026-10-15 evening): Heart: RED — Release RED; monitoring RED · 10/100 · 415 unresolved; [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] autofit_workspace_test/Smoke Tests: red; [red] autogalaxy_workspace_test/Smoke Tests: red; [red] autolens_workspace_test/Smoke Tests: red; [red] Release validation: NOT release_ready — v2026.10.15.1.dev81701; [red] Worktree drift: 5 orphan / 0 missing / 4 dirty; [red] autolens_inference: red; [red] autolens_profiling: red; [red] euclid_strong_lens_modeling_pipeline: red.

The live human authorized the development-only override in session 3e3dea38 (AskUserQuestion answer: "Yes, override for development (Recommended)") for the three PyAutoArray prompts filed from the Aris interferometer profiling. Permitted: commit, push, PR-open behind the autonomous-ship gate (tests, smoke, review must still pass). Merge stays a separate human act. Heart remains RED for release purposes. Launched --auto.

Original Prompt

Click to expand starting prompt

AdaptSplit regularization log-det differs between numba and JAX backends

Type: bug
Target: PyAutoArray
Repos:

PyAutoArray
Difficulty: large
Autonomy: supervised
Priority: high
Memory: reading-queue.md; wiki/lensing/sources/dark-matter-substructure.md; wiki/galaxies/sources/massive-ellipticals.md
Status: formalised
Consequence: glance
Witness: at identical parameters on the Aris interferometer setup, numba and JAX log_det_regularization_matrix_term agree to <1e-3 (currently 2443.22 vs 2561.98), with a unit test pinning the value on an ill-conditioned AdaptSplit matrix
Review-minutes: 3
Unattended: ready

At identical parameter vectors, the NumPy/numba and JAX likelihoods of a Delaunay + AdaptSplit
interferometer inversion differ by ~20-270 in log likelihood (old 125x125 mask, 1405-5526 vertices)
and ~1 on a 512x512 mask. Chi-squared, curvature matrix, data vector and curvature log-det are
identical; the whole difference is log_det_regularization_matrix_term (e.g. 2443.22 numba vs
2561.98 JAX). The AdaptSplit regularization matrix is extremely ill-conditioned (entries sum to
~1.4e-5), so its log-determinant is backend dependent. JAX eager vs jit agree to 1e-8.

Consequence: log evidences are not comparable across backends; a user mixing backends across
pipeline phases (or comparing models) gets inconsistent evidences. Also ~1 in 4 random draws is
non-positive-definite: numba raises LinAlgError (resample), JAX returns -inf — same behaviour
should be checked for consistency.

Do: locate where each backend computes the regularization log-det (Cholesky vs slogdet, jitter /
diagonal regularisation, dtype), decide the numerically correct value (e.g. compare to an
mpmath/high-precision reference on a small matrix), make both backends compute it the same robust
way, and add a regression test. Consider whether AdaptSplit's matrix construction itself should be
better conditioned.

Evidence: scratchpad /tmp/claude-1000/-home-jammy-Code-PyAutoLabs/50c8c6cd-6ac2-47c2-9290-7ee73f9c4315/scratchpad/aris/PROGRESS3.md, PROGRESS4.md, step7/bench.py (--check mode), step8/. PyAutoArray ccddfba.

Campaign membership: this task is logged and monitored under the PyAutoPulse interferometer campaign as task interferometer_large_mask_aris — PyAutoLabs/PyAutoPulse#62 (ledger rows queued in Mind behind the PyAutoPulse claim; the issue holds the campaign intent and baseline numbers until they land). Record the before/after witness table there at close-out.

Root cause found (read-only investigation, 2026-10-15) — plan amended

The log-det factorizations are not the cause. On the 125×125 Aris setup (1,405 vertices, one fixed vector) all five methods (NumPy dense Cholesky, SuperLU log_det_sparse_spd_from, JAX Cholesky x64, slogdet, diagonal scaling) agree with a 50-digit mpmath reference to ≤ 2e-8 on each backend's own H. The whole 118.76 gap is because the two backends build different H (max|H_np − H_jax| = 3.21 on max|H| 31.9; same sparsity pattern).

Chain: mapper_util.adaptive_pixel_signals_from (mapper_util.py ~54–74) divides each pixel's summed signal by pixel_counts, and that count includes mapping entries with zero weight. In this setup the mesh is a stride-2 subsample of the data grid, so 1,375 data points sit exactly on vertices with weights (0, 0, 1). For such a point any surrounding triangle is valid; scipy find_simplex and the JAX walk (pix_indexes_delaunay_walk_from, delaunay.py ~303) choose different ones in 1,114 of 5,496 rows — both geometrically correct (residual 2e-16). The differing zero-weight neighbours change the counts → pixel_signals (max diff 0.23) → AdaptSplit weights (0.53) → H (3.2). Zero weights add nothing to the mapping matrix, which is why chi², curvature, data vector and curvature log-det stay identical. Given identical weights, the split-stencil tables and the numba/JAX H builders already agree to 1e-14.

Simulated fix (count only non-zero-weight entries): signals agree to 4e-15, H to 4.8e-13, log-det 2468.446318 on both backends (Δ 4e-8). Over 20 prior draws: today every draw differs (1.4–822); with the fix 12/20 agree to ≤1e-6 and 2 more to <1e-3; the remaining 6 (max(coeff)⁴ from 5e6 to 3.4e9) still differ by 5e-3–2.2 because the absolute 1e-8 diagonal jitter makes the constant mode's eigenvalue round-off-dominated (cond 1e16–1e17) — exact-arithmetic log-dets of the two near-identical H differ, and PD-ness itself is decided by round-off. No factorization or scaling fixes that regime. No non-PD failures occurred at 125×125 (0/20 either backend); the "~1 in 4" figure was the 512×512 runs.

Amended plan

Fix: in adaptive_pixel_signals_from, require a non-zero weight in valid (& (pixel_weights != 0)), feeding both the masked weights and the counts. Works on both backends, differentiable, matches the docstring's "number of mapped data points". The NumPy evidence changes wherever zero-weight entries exist (Delaunay points on vertices/edges, possibly rectangular bilinear corners) — state this in the PR.
Leave the log-det code alone (log_det_sparse_spd_from, JAX dense Cholesky are accurate wherever H is well-posed).
Tests: (a) test_mapper_util.py: pixel signals from the scipy and JAX tables agree to atol 1e-12 on the synthetic 16×16 grid with a stride-2 mesh and 0.01*sin(7*grid) distortion (today 48/256 rows differ, signals differ by 0.10); plus a pure-NumPy case of one on-vertex point described by two triangles with weights (0,0,1). (b) test_log_det_sparse.py (JAX-marked): AdaptSplit on that geometry with moderate coefficients (inner 0.4, outer 1.7, signal_scale 0.3): H parity rtol 1e-12, NumPy and JAX log_det_regularization_matrix_term agree <1e-6 and match a high-precision reference <1e-6. Do not pin parity at high coefficients.
Scope note (not in this PR): the split builders' absolute jitter (regularization_util.py ~375/416/549/615) could become relative (1e-8 * max(regularization_weights**2)); that changes evidence values, so it is separate work. The 512×512 pair (2576.59 vs 2574.73) may be partly this effect; not run (RAM).

Witness unchanged: numba vs JAX agree <1e-3 on the Aris setup at moderate coefficients. Scripts: session 3e3dea38 scratchpad ld_investigation/ (temporary; the synthetic case will be committed as the test).

Patch drafted (not yet applied — PyAutoArray claim held by #626)

Built and verified in a throwaway copy of main a84f724 (no repo touched). Fix in autoarray/inversion/mappers/mapper_util.py adaptive_pixel_signals_from:

-    valid = (xp.arange(B)[None, :] < pix_size_for_sub_slim_index[:, None]).reshape(-1)
+    in_size = xp.arange(B)[None, :] < pix_size_for_sub_slim_index[:, None]
+    valid = in_size.reshape(-1) & (flat_weights != 0)

(masked on the flattened weights so the existing test_abstract mock with (9,) weights for (9,1) mappings keeps passing; there is no numba/JAX twin of this function; no log-det code touched.)

New test	unpatched main	with fix
`test_mapper_util.py::…scipy_and_jax_delaunay_tables__agree` (JAX)	fails: 55/64 signals differ, max 0.1003	passes, atol 1e-12
`test_mapper_util.py::…vertex_point_via_two_triangles__identical` (NumPy)	fails	passes, hand-computed values
`test_log_det_sparse.py::…vertex_coincident_adapt_split__numpy_matches_jax`	fails: H max diff 0.740 (max\|H\| 14.6), log-det gap 0.097	passes: H to 3.6e-15, log-dets differ 2.1e-9, both within 1.7e-9 of a 50-digit mpmath reference (asserted <1e-6)

pytest test_autoarray/inversion: 935 passed, 4 xfailed; rest of test_autoarray: 1099 passed; black clean. No existing expectation changed. Behaviour note for the PR: evidence values change on both backends wherever data points fall on Delaunay vertices/edges (Overlay / grid-subsample image meshes); workspace parity scripts with pinned evidences may need updating — the smoke leg will show which.

Next: once #626 merges and frees the claim, /start_library on this task and apply the patch (session 3e3dea38 scratchpad patch_627.patch, temporary) onto feature/adaptsplit-logdet-backend-parity, then ship under the gate.

Parked at the review leg (--auto): no PR opened

This --auto run parked at the review leg, so no PR was opened. Branch feature/adaptsplit-logdet-backend-parity is pushed at d78794e3: the drafted patch, plus Closes #627 and the Fable co-author trailer. It is based on 7cc71238, which includes #630.

Gate legs

Tests: pass.
- PyAutoArray: 2036 passed, 4 xfailed.
- Downstream, run against the worktree autoarray: PyAutoGalaxy 1357 passed; PyAutoLens 832 passed, 1 xfailed.
- No existing expectation changed.
Smoke: pass. pyauto-heart smoke autogalaxy autolens autolens_test euclid howtolens ran 151 scripts and notebooks, with 0 failures.
Review: FINDINGS. This is the reason for the park. The independent reader found no bug in the Delaunay parity fix itself. These checks pass:
- JAX grad is finite and matches finite differences to 8e-10.
- jit and vmap both work.
- The control test confirms the bug: without the mask the signals differ by 0.1003; with it they differ by 1e-16.
- NaN weights stay loud.
- Pixels with no counted entries still go through the existing safe denominator.
Heart: RED. The current reason is release validation FAILED (stage integrate). That RED is covered by the development override above.

The finding (I verified it): the change reaches well beyond Delaunay meshes, which the commit message does not say.

KNNBarycentric builds its weights in barycentric_weights_from_3_nearest (autoarray/inversion/mesh/interpolator/knn.py:314). It clips negative barycentric weights to exactly 0. That happens for every query outside its 3-nearest triangle, so whole regions are affected, not just the rare points that sit exactly on a vertex.
Under this patch those vertices no longer count toward pixel_counts. For a typical KNNBarycentric fit, that changes the pixel signals, the AdaptSplit H matrix and the evidence.
It also adds a new jump in the likelihood. Before the patch, the count was fixed for a given nearest-3 set. Now it drops wherever a clipped weight reaches exactly 0.
- The reviewer measured this: moving one weight from 1e-9 to 0 moved a test scalar from 6.2358 to 6.4889.
- The gradient with respect to that weight is exactly 0 at w=0.
- As lens parameters move points across a clip boundary, the evidence jumps.
Rectangular bilinear meshes are affected only in exact-coincidence cases.

A merge would therefore commit to two behaviour decisions, not one. Both need a human choice; the autonomy gate does not let an --auto run make them:

(a) Accept as is. Zero-weight entries count for no interpolator. The PR then has to say plainly that KNNBarycentric evidence changes broadly and gains these jumps.
(b) Limit the rule to the Delaunay case. Count only non-zero weights for Delaunay tables only, or more generally drop only the duplicate zero-weight entries that arise from the choice of simplex. Leave KNN clipping semantics as they were. This needs a small code change and one more test.
(c) Fix the root instead. Make the triangle choice deterministic across scipy and JAX for vertex- and edge-coincident points (the alternative rejected earlier). The count semantics stay as they are.

My recommendation is (b): it keeps the parity fix and avoids introducing a discontinuity in a mesh used for gradient-based fits. The choice is yours.

Witness status: only the synthetic part has been checked on this branch.

The committed synthetic ill-conditioned AdaptSplit test passes:
- cond(H) is about 2.2e9.
- H parity is 3.6e-15.
- The two log-dets differ by 2.1e-9.
- Both are within 1.7e-9 of a 50-digit mpmath reference.
The Aris 125×125 figure (2468.446318 on both backends, Δ4e-8) comes from the earlier simulated fix in the investigation. The Aris cache was not re-run on this branch: the scratchpad dump was deleted to free disk.

To resume: give a decision (a, b or c), then /start_library adaptsplit-logdet-backend-parity (the worktree and branch are kept), then /ship_library.

Human decision: option (b), un-parked

The live human chose option (b) via AskUserQuestion in session 3e3dea38: "Delaunay only (Recommended by review)".

Zero-weight mapping entries will be left out of the adaptive pixel-signal counts for Delaunay tables only. KNNBarycentric and rectangular meshes keep today's counting. The task resumes on feature/adaptsplit-logdet-backend-parity with --auto, under the same development override. It stops at PR-open.

fix: exclude zero-weight Delaunay mappings from adaptive pixel signals

Library PR Created

PR: #631, head f052e381, label pending-release.

This --auto run stops at PR-open. Merging is a human /prm.

Scope: Delaunay only, as the human chose: "Delaunay only (Recommended by review)".

adaptive_pixel_signals_from(..., exclude_zero_weight_mappings=False) gets a new keyword.
exclude_zero_weight_mappings is a new interpolator class attribute. It is True for Delaunay and DelaunayNN, and False for KNN, KNNBarycentric and rectangular meshes.
I amended the parked commit d78794e3 into f052e381 and pushed the unmerged feature branch with --force-with-lease.

API changes: no signature break. One keyword and one class attribute are added. Evidence values change, the same way on both backends, but only for Delaunay and DelaunayNN meshes that have data points exactly on vertices or edges. The independent reviewer confirmed KNN, KNNBarycentric and rectangular are bit-identical to main.

Gate

Tests: PyAutoArray 2039 passed, 4 xfailed. PyAutoGalaxy 1357 passed. PyAutoLens 832 passed, 1 xfailed.
Smoke: 151/151.
Review: CLEAN, from an independent Opus reader.
Heart: RED (release validation FAILED (stage integrate)), under the recorded development override.

Witness: PASS on the synthetic ill-conditioned AdaptSplit test:

H parity is 3.6e-15.
The two log-dets differ by 2.1e-9.
Both are within 1.7e-9 of an mpmath reference.

The Aris cache was not re-run on this branch. The investigation's simulated fix gave 2468.446318 on both backends.

Workspace impact: recommendation (iii), no workspace change needed.

No workspace, test or pipeline script calls adaptive_pixel_signals_from directly.
The behaviour change only alters numbers for Delaunay meshes with vertex- or edge-coincident data, such as Overlay or grid-subsample image meshes.
The smoke subsets above passed against the branch.
Any script that pins an exact evidence value for such a mesh would shift, but none failed. No /start_workspace is needed.

Campaign: at close-out, copy the before/after witness into PyAutoPulse#62, task interferometer_large_mask_aris.

Shipped

Merged in #631 (merge 0377eea; commits f052e38 fix + da5f6ce test guard). Pending release.

Root cause. The two backends built different AdaptSplit H matrices. adaptive_pixel_signals_from counted zero-weight mapping entries, and scipy and the JAX Delaunay walk choose different triangles for data points lying exactly on a vertex or edge, so those zero-weight entries landed on different pixels. Every log-det factorization agreed to 2e-8 on the same H.

Fix. A new exclude_zero_weight_mappings keyword, driven by an interpolator class attribute: True on InterpolatorDelaunay / InterpolatorDelaunayNN, False on the base interpolator and KNN. The scope is Delaunay only, per the human's decision (b) above: KNNBarycentric clips negative weights to exactly 0 across whole regions, so excluding zeros there would change whole-region signals and add a discontinuity.

Witness. PASS on a synthetic ill-conditioned AdaptSplit problem: H parity 3.6e-15, log-dets 2.1e-9 apart, and both within 1.7e-9 of a 50-digit reference. Aris was not re-run.

Behaviour change. Evidence values change on Delaunay / DelaunayNN meshes where data points sit exactly on mesh vertices or edges. KNN, KNNBarycentric and rectangular meshes are bit-identical.

Out of scope (unfiled). The absolute 1e-8 jitter in the split regularization builders makes the log-det round-off-dominated once max(coeff)^4 ≳ 1e6.

Gate: tests PyAutoArray 2039 passed / 4 xfailed, PyAutoGalaxy 1357, PyAutoLens 832 / 1 xfailed; smoke 151/151; independent review CLEAN on f052e38; Heart RED under the recorded development override, merged by human /prm. The first CI run went red on the no-JAX leg (the new KNN test lacked importorskip("jax")). da5f6ce fixed it before merge.

Campaign: PyAutoPulse#62 (interferometer_large_mask_aris). Next on the PyAutoArray claim: #628.

Navigation Menu

Repository navigation

fix: AdaptSplit regularization log-det parity between numba and JAX #627

Overview

Plan

Affected Repositories

Branch Survey

Implementation Steps

Key Files

Evidence

Heart RED development override

Original Prompt

AdaptSplit regularization log-det differs between numba and JAX backends

Root cause found (read-only investigation, 2026-10-15) — plan amended

Amended plan

Patch drafted (not yet applied — PyAutoArray claim held by #626)

Parked at the review leg (--auto): no PR opened

Human decision: option (b), un-parked

Library PR Created

Shipped

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

fix: AdaptSplit regularization log-det parity between numba and JAX #627

Description

Overview

Plan

Affected Repositories

Branch Survey

Implementation Steps

Key Files

Evidence

Heart RED development override

Original Prompt

AdaptSplit regularization log-det differs between numba and JAX backends

Activity

Jammy2211 commented on Oct 15, 2026

Jammy2211 commented on Oct 15, 2026

Root cause found (read-only investigation, 2026-10-15) — plan amended

Amended plan

Jammy2211 commented on Oct 15, 2026

Patch drafted (not yet applied — PyAutoArray claim held by #626)

Jammy2211 commented on Oct 15, 2026

Parked at the review leg (--auto): no PR opened

Jammy2211 commented on Oct 15, 2026

Human decision: option (b), un-parked

Jammy2211 commented on Oct 15, 2026

Library PR Created

Jammy2211 commented on Oct 15, 2026

Shipped

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions