fix: Basis on irregular grids + settle pixelization eager vs step-by-step log-evidence gap

[Jammy2211]

Overview

jax_profiling/jit/imaging/pixelization.py prints three "equivalent"
log-evidence values and asserts only two of them. Reproducing the original
report (PyAutoMind bug/autolens/pixelization_eager_vs_jit_divergence.md,
filed 2026-07-08) on clean main today found the surface evidence stale, the
underlying divergence papered over rather than resolved, and the script
unable to run at all against current libraries.

This issue unblocks the script (one genuine PyAutoGalaxy library bug plus the
RectangularAdaptDensity → RectangularRTUAdaptDensity rename from #461),
then settles why the eager and step-by-step numbers differ and makes the
script's account of it true.

Plan

• Fix Basis.image_2d_list_from in PyAutoGalaxy so a basis containing linear
  light profiles can be evaluated on a Grid2DIrregular — it currently assumes
  every grid carries a mask.
• Probe the live fit to confirm the suspected cause of the divergence: the
  eager inversion solves for a non-negative reconstruction (fnnls) while
  the script's step-by-step rebuild does a plain positive/negative solve.
• Point this one script at the renamed RectangularRTUAdaptDensity mesh
  (the wider rename fallout belongs to PyAutoArray#461's follow-up campaign).
• Replace the script's incorrect "floating-point drift" comment with the
  confirmed cause, and turn the printed step-by-step value into a real
  regression assertion.
• Re-pin EXPECTED_LOG_EVIDENCE_HST, which has drifted ~1e-2 relative.

Detailed implementation plan

What reproducing it found

1. The prompt's surface evidence is stale. The script moved to
autolens_workspace_developer/jax_profiling/jit/imaging/pixelization.py and was
rewritten by PR #60 ("Re-profile rectangular + Delaunay at full production
fiducial"). EXPECTED_LOG_EVIDENCE_HST_EAGER and the FIXME no longer exist;
there is one constant (EXPECTED_LOG_EVIDENCE_HST = 24746.105672366088)
asserted against eager, full-JIT and vmap. The -1.3e9 values in the prompt
came from a superseded, pathological fiducial.

2. The divergence was papered over, not resolved. The correctness assertion
(pixelization.py:806-830) compares log_evidence_check — rebuilt from
inversion.reconstruction and inversion.curvature_matrix — against the
reference. The genuinely independent step-by-step value is printed at line 804
and never asserted. The comment at lines 761-764 blames "cumulative
floating-point differences between JIT-compiled and eager paths (especially
through ill-conditioned solves)".

That explanation is almost certainly wrong.
PyAutoArray/autoarray/config/general.yaml:5 ships
use_positive_only_solver: true, so Inversion.reconstruction
(PyAutoArray/autoarray/inversion/inversion/abstract.py:493-545) is an fnnls
non-negative solve — optionally on an edge-zeroed subset via
use_edge_zeroed_pixels / zeroed_ids_to_keep. The script's
compute_reconstruction is a plain positive/negative solve(F + H, D). Those
are two different reconstructions s, so χ² and sᵀHs legitimately differ —
a real modelling difference, not round-off.

Prompt suspects 1, 2 and 4 are falsified: FitDataset.figure_of_merit →
log_evidence (PyAutoArray/autoarray/fit/fit_dataset.py:324-366) uses exactly
the same five-term formula the script implements.

3. The script does not run on main at all.

• al.mesh.RectangularAdaptDensity (line 233) was renamed by PyAutoArray
  f9aceea3 (perf: vectorize MGE potential over components #461) → RectangularRTUAdaptDensity (pure rename, values
  unchanged) with a new RectangularBilinearAdaptDensity sibling. "Breaking
  rename — no back-compat aliases; downstream workspaces update in follow-up
  PRs." The PyAuto API gate blocks the script on this.
• With that name shimmed, the run dies at Step 2:
  AttributeError: Grid2DIrregular does not have attribute mask at
  PyAutoGalaxy/autogalaxy/profiles/basis.py:164. Basis.image_2d_list_from
  builds its zero placeholder for LightProfileLinear members as
  aa.Array2D(values=xp.zeros(...), mask=grid.mask), assuming a masked
  Grid2D. The script deliberately wraps grids as Grid2DIrregular for
  JIT-traceability, and the lens light is an MGE-60 linear basis.

Also observed with the value-preserving RTU shim: eager
figure_of_merit = 25004.719 against the pinned 24746.106 — 1.05e-2
relative, so the pinned constant has drifted too.

Affected Repositories

• PyAutoGalaxy (primary, library)
• autolens_workspace_developer (script follow-up)
• PyAutoArray — only if the Step 2 probe redirects the investigation

Branch Survey

Repository	Current Branch	Dirty?
./PyAutoArray	main	clean
./PyAutoGalaxy	main	clean
./PyAutoLens	main	clean
./autolens_workspace_developer	main	clean

worktree_check_conflict → no conflict.

Suggested branch: feature/pixelization-eager-jit-divergence

Implementation Steps

1. PyAutoGalaxy — irregular-grid-safe Basis placeholder.
   autogalaxy/profiles/basis.py, image_2d_list_from (~lines 160-165). The
   binned is False branch already returns a bare xp.zeros(...); the default
   branch must stop assuming grid.mask. Return an Array2D only when the grid
   actually carries a mask, otherwise the bare zero array. Check the sibling
   accessors in the same file (convergence_2d_from, any other grid.mask use)
   for the same assumption before settling the shape of the fix.
   Unit test in test_autogalaxy/profiles/test_basis.py: a Basis containing a
   LightProfileLinear evaluated on a Grid2DIrregular returns zeros of the
   right length instead of raising. No JAX in unit tests.

2. Confirm the solver hypothesis before touching the script. With step 1 in
   place, probe the live fit once: inversion.settings.use_positive_only_solver,
   inversion.settings.use_edge_zeroed_pixels, and
   max|inversion.reconstruction − solve(F + H, D)|. If the reconstructions
   differ materially, the verdict is "genuine difference". If they agree to
   round-off, the hypothesis is dead and the investigation reopens on
   log_det_curvature_reg_matrix_term / curvature_reg_matrix construction
   (prompt suspect 2) instead. Do not skip this — the whole write-up depends
   on it.

3. autolens_workspace_developer — make the script honest. In
   jax_profiling/jit/imaging/pixelization.py:

   • Line 233 and the line-1079 comment: RectangularAdaptDensity →
     RectangularRTUAdaptDensity (this file only).
   • Replace the false "cumulative floating-point differences" comment
     (761-764) with the confirmed cause from step 2.
   • Assert the step-by-step value: build the step-by-step reconstruction
     through the same solver the inversion uses — reuse
     autoarray.inversion.inversion.inversion_util.reconstruction_positive_only_from
     rather than hand-rolling — so log_evidence (step-by-step) becomes a real
     regression check instead of a printed number. If matching the solver is not
     practical inside the JIT-timing path, keep the direct solve for timing and
     assert the solver-matched rebuild separately, with the reason stated in the
     comment.
   • Re-pin EXPECTED_LOG_EVIDENCE_HST to the value the green run produces,
     noting the fiducial and library version, and confirm eager / full-JIT /
     vmap / step-by-step all sit inside rtol=1e-4.

Out of scope (flag, don't do)

• The ~40 other files still naming RectangularAdaptDensity /
  RectangularAdaptImage (autogalaxy_workspace, autogalaxy_workspace_test,
  HowToLens, euclid pipeline, prior config YAMLs) — that is PyAutoArray#461's
  announced follow-up campaign. File a prompt for it.
• jax_profiling/jit/imaging/delaunay.py and the interferometer siblings —
  touch only if the step-1 library fix changes their numbers.

Key Files

• PyAutoGalaxy/autogalaxy/profiles/basis.py — image_2d_list_from, the
  grid.mask assumption at line 164.
• PyAutoGalaxy/test_autogalaxy/profiles/test_basis.py — new regression test.
• autolens_workspace_developer/jax_profiling/jit/imaging/pixelization.py —
  mesh name (233), the false FP-drift comment (761-764), the step-by-step
  print (804), the assertion block (806-830), the pinned constant (1082).
• PyAutoArray/autoarray/inversion/inversion/abstract.py — reconstruction,
  the positive-only / edge-zeroed branch (493-545).
• PyAutoArray/autoarray/fit/fit_dataset.py — log_evidence /
  figure_of_merit (324-366), confirms the five-term formula.
• PyAutoArray/autoarray/config/general.yaml — use_positive_only_solver: true.

Verification

source ~/Code/PyAutoLabs-wt/pixelization-eager-jit-divergence/activate.sh
cd PyAutoGalaxy && pytest test_autogalaxy/profiles/test_basis.py -q
cd ../autolens_workspace_developer
NUMBA_CACHE_DIR=/tmp/numba_cache MPLCONFIGDIR=/tmp/matplotlib \
  python jax_profiling/jit/imaging/pixelization.py

Green means: the script runs to completion; figure_of_merit,
log_evidence (inv matrices), log_evidence (step-by-step), full-JIT and vmap
all agree with the re-pinned EXPECTED_LOG_EVIDENCE_HST at rtol=1e-4 (or the
step-by-step gap is asserted against its own solver-matched constant with the
reason recorded in the comment). Then run the Delaunay and interferometer
pixelization siblings to confirm the PyAutoGalaxy fix moved nothing there.

Original Prompt

Click to expand starting prompt

Investigate eager FitImaging.figure_of_merit vs JIT/step-by-step divergence in rectangular pixelization

Type: bug
Target: PyAutoLens
Difficulty: too-large
Autonomy: supervised
Priority: high
Status: formalised

Context

Surfaced by the eager-numpy regression assertions added in
jax_profiling/imaging/pixelization.py (PR from
PyAutoPrompt/issued/eager_numpy_regression_assertions.md).

When the rectangular pixelization script runs, it computes the same
log_evidence via three supposedly-equivalent numpy/JAX paths and gets two
different answers:

figure_of_merit (log_evidence) = -1338814172.1831784   ← eager FitImaging
log_evidence (step-by-step)    = -1338521802.3596904   ← JIT-equivalent numpy rebuild
log_evidence (inv matrices)    = -1338814172.1831782   ← numpy rebuild from eager matrices
log_evidence (reference)       = -1338814172.1831784   ← same as figure_of_merit
full log_likelihood (JIT)      = -1338521802.3596945   ← JIT full pipeline

The split is:

• Cluster A (eager): FitImaging.figure_of_merit, log_evidence (reference),
  and log_evidence (inv matrices) all agree on -1338814172.18.
• Cluster B (JIT / step-by-step): log_evidence (step-by-step) and
  full log_likelihood (JIT full pipeline) agree on -1338521802.36.

The two clusters differ by ~292k in absolute value (~0.02% relative) — well
above float64 round-off and well above the rtol=1e-4 used in the regression
assertions. The script's own internal assertion

Assertion PASSED: inversion-matrix log_evidence matches FitImaging.log_evidence

confirms the eager and "inv matrices" numpy rebuilds agree with each other.
The JIT-style rebuild (step-by-step) disagrees with that pair and instead
matches the JIT full-pipeline result.

So something is happening in FitImaging-eager that is not the same as
re-running the same math in a JIT-style numpy pipeline. Currently the
pixelization script pins two separate regression constants — the original
EXPECTED_LOG_EVIDENCE_HST for the JIT path, and a new
EXPECTED_LOG_EVIDENCE_HST_EAGER for the eager path (with a FIXME pointing
at this prompt).

Task

Identify which cluster is correct (or whether they're both reasonable results
of a numerically-legitimate-but-subtle difference in accumulation order /
regularization), then either:

1. Fix the divergent path so both clusters converge, and collapse the two
   constants back into one.
2. Document the divergence as genuine (e.g. one path includes a term the
   other legitimately omits), update code comments accordingly, and keep the
   two constants but remove the FIXME framing.

Either outcome is fine — the goal is to understand why the two numpy
computations of the same quantity disagree and to make the code's treatment
of them honest.

Where to look

The disagreement is specifically in the rectangular pixelization log-evidence
stack. Likely suspects (in rough order of plausibility):

1. FitImaging.figure_of_merit for Inversion-based fits — does it use a
   different formulation (Bayesian evidence with full regularization term)
   than the step-by-step reconstruction in the script does?

   • PyAutoLens/autolens/imaging/fit_imaging.py
   • PyAutoArray/autoarray/inversion/inversion/abstract.py::log_evidence
     and related helper properties.

2. Regularization matrix handling — the eager path may be building the
   full regularization contribution via log_det_regularization_matrix_term
   whereas the step-by-step may be using a shortcut that agrees with what
   the JIT path does.

3. Mapping matrix numerical accumulation order — check whether
   mapping_matrix and blurred_mapping_matrix are computed with different
   intermediate dtypes or summation orders in the eager vs step-by-step
   pipelines.

4. inversion.log_evidence vs fit.figure_of_merit — check whether
   FitImaging.figure_of_merit returns
   inversion.log_evidence_with_regularization or just
   inversion.log_likelihood. The script's "step-by-step" may be comparing
   against a different decomposition.

Not suspected:

• The simulator is deterministic (seeded), and imaging/delaunay.py (using
  a different mesh but same overall fit pattern) does NOT show this
  divergence — its eager figure_of_merit matches its EXPECTED_LOG_EVIDENCE_HST
  bit-for-bit. So the bug is specifically in the rectangular-pixelization
  code path, not in the general FitImaging + Inversion framework.

• The interferometer pixelization script (jax_profiling/interferometer/pixelization.py)
  also passes the eager assertion cleanly — figure_of_merit_ref ≈ EXPECTED_LOG_EVIDENCE_SMA. So it's not a universal pixelization issue,
  it's imaging + rectangular specifically.

Verification

After fixing:

source ~/Code/PyAutoLabs-wt/<task-name>/activate.sh
cd autolens_workspace_developer
python jax_profiling/imaging/pixelization.py

Expect:

• figure_of_merit (log_evidence) == log_evidence (step-by-step) to within
  float64 round-off (or a documented, justified gap).
• EXPECTED_LOG_EVIDENCE_HST_EAGER can be removed and the eager assertion
  can revert to using EXPECTED_LOG_EVIDENCE_HST alongside the JIT/vmap
  assertions.
• Delete the FIXME comment block.

Affected repos

• PyAutoLens and/or PyAutoArray (library — most likely PyAutoArray
  inversion code, possibly PyAutoLens/fit_imaging.py)
• autolens_workspace_developer (script — remove the FIXME and the second
  constant once the library fix lands)

Suggested branch

feature/pixelization-eager-jit-divergence

Notes

• Don't start by bumping either constant. The script's internal
  "inv matrices vs step-by-step" split is the debugging starting point —
  narrow down which numpy computation is authoritative, then fix the other.
• The divergence is ~0.02%, which is well within the regime where Bayesian
  model comparison decisions could flip, so treat this as a real bug until
  proven otherwise (not just a rounding artefact).

[Jammy2211]

Verdict: not a bug — plus a silent transform swap the pin was hiding

Both PRs are open:

• Library — fix: Basis linear-profile placeholder on irregular grids #581 (Basis on irregular grids). Merge this first; the script cannot run without it.
• Workspace — fix: unblock jit/imaging/pixelization.py and make its log-evidence honest autolens_workspace_developer#127 (the script).

The reported divergence

Two causes, both settled, neither a library bug:

1. The script fed the wrong matrices to its own evidence rebuild. It used the step-11 H, which spans the mapper only (1225 source pixels), while D and F come from the inversion and span every linear object (1225 + 60 linear MGE Gaussians = 1285). On current main that is a hard shape error. Beyond the shape, the MGE columns are unregularized, so the full H is singular and slogdet(H) is -inf. Inversion.log_det_* evaluates both log-det terms on the reduced (mapper-only) matrices; the script now mirrors that and matches figure_of_merit to 13 significant figures.

2. The residual step-by-step gap is the non-negative solve, not FP drift. use_positive_only_solver defaults to True, so the fitted reconstruction is an fnnls solution with an iteratively chosen active set — 362 of 1285 parameters pinned at exactly zero, where an unconstrained solve puts 777 negative. A re-solve landing on a different active set gives a genuinely different s. The old comment blamed "cumulative floating-point differences"; that is now corrected, and the step-by-step value carries its own pin instead of being printed and ignored.

Prompt suspects 1, 2 and 4 are falsified: FitDataset.figure_of_merit → log_evidence uses exactly the five-term formula the script implements.

The constant: a transform swap, not drift

Chasing the pin turned up something the task did not set out to find. My first pass re-pinned to 25004.719... and blamed accumulated library drift — inferred from a commit message rather than measured, and wrong. Measured at this fiducial:

PyAutoArray	mesh	eager figure_of_merit
f9aceea3^	RectangularAdaptDensity	25004.71903495436
f9aceea3	RectangularRTUAdaptDensity	25004.71903495436
f9aceea3	RectangularBilinearAdaptDensity	24746.105672366088

RTU is a pure rename (bit-identical across the split), and Bilinear reproduces the 2026-05-11 pin bit-for-bit after 3.5 months — which also rules out generic drift, since any other change touching this value would have broken the exact match.

The actual history is a transform swap under an unchanged class name:

When	What	Value
2026-05-11	pinned (#60); RectangularAdaptDensity = rank-CDF	24746.105672366088
2026-07-23	22b28463 (#402) gave that name kernel-CDF	→ 25004.71903495436
2026-08-21	f9aceea3 (#461) split them; Bilinear = rank-CDF	24746.105672366088

The script never chose kernel-CDF — it inherited it silently in July, and the stale pin masked the change. Naming RectangularBilinearAdaptDensity explicitly restores what the pin always meant: no re-pin, full pre-#402 history intact, and it matches #461's designated workspace default.

Verification

figure_of_merit (log_evidence) = 24746.105672366088
log_evidence (step-by-step)    = 24783.19775195945
log_evidence (inv matrices)    = 24746.105672366568
full log_likelihood (JIT)      = 24746.103473124283
vmap batch                     = [24746.10419047 x3]
all five assertions PASSED, exit 0

Cross-checked across two independent implementations: numpy/numba (xp=np) and the JAX full pipeline (use_jax=True) agree to 8.9e-8. test_autogalaxy: 1113 passed.

Not in scope

misc/pixelization_spline_*.py reference RectangularSplineAdaptDensity, which no longer exists — a pre-existing break from a different mesh campaign. The wider rename fallout (autogalaxy_workspace, HowToLens, euclid pipeline, prior config YAMLs) is #461's announced workspace follow-up, already captured in PyAutoMind/draft/feature/autoarray/rectangular_bilinear_rtu_mesh_split.md.

[Jammy2211]

Historical sanity check: the restored value is what history says it should be

Following up on the constant — checking whether the Bilinear value ever appeared in this script's own recorded history. It did, and the kernel-CDF value never did.

git log -S"24746.105672366088"   ->  aa71013, 49fb4f6  (both 2026-05-11)
git log -S"25004.71903495436"    ->  only this task's commits

So the script has never, in its entire history, pinned the kernel-CDF value. From #402 (2026-07-23) onward it was computing a number its own assertion rejected — which is consistent with it simply being broken and unrun in that window.

Two further independent lines agree.

1. Committed profiling artifacts from the rank-CDF era. These are separate runs on separate hardware, recorded independently of the pinned constant. Every fp64 run stored eager_log_evidence bit-identical to it:

artifact	device	autolens	eager_log_evidence
hpc_a100_fp64.json	A100 80GB	2026.5.1.4	24746.105672366088
local_gpu_fp64.json	RTX 2060	2026.5.8.2	24746.105672366088
local_cpu_fp64.json	CPU	2026.5.8.2	24746.105672366088

Two GPUs, a CPU, two library versions — bit-identical. The _mp.json siblings record 24746.105678802393, the same runs in mixed precision, differing only in the 7th decimal exactly as expected.

2. Code identity. The May-era create_transforms (bc00c113:autoarray/inversion/mesh/interpolator/rectangular.py:70) is line-for-line identical to today's create_transforms_rank — the function Bilinear routes through. Same sort / argsort / cumsum body, no erf anywhere; the erf sum is the kernel-CDF path.

Combined with the direct cross-split measurement (pre-split RectangularAdaptDensity == post-split RTU == 25004.71903495436; Bilinear == 24746.105672366088), the chain is closed: today's Bilinear is the same transform, producing the same bits, that this script measured in May.

[Jammy2211]

Wrap-up — both PRs ready for review

Correction I caught while closing out

My first pass blanket-renamed RectangularAdaptDensity → RectangularRTUAdaptDensity across gradient/README.md and three scripts. That was the same error this issue exists to fix: the class name meant rank-CDF before the 2026-07-23 consolidation and kernel-CDF after, so a blanket rename relabels historical findings as belonging to a mesh that never produced them.

gradient/README.md is a mixed document and says so itself at its 2026-07-26 section — "Post-consolidation (PyAutoArray#403 — the kernel-CDF meshes now ARE RectangularAdaptDensity / RectangularAdaptImage)". So findings are now relabelled by date:

• rows dated 2026-07-09 and earlier → RectangularBilinear* (rank-CDF). Their own prose confirms it: "mesh queries always coincide with the rank-transform knots", "discrete rank-reordering jumps", and the phase-2b section documents create_transforms as "a per-axis sort + jnp.interp CDF" — the rank path, not the erf-sum kernel path.
• the 2026-07-26 final-assessment and regularization-matrix sections → stay RectangularRTU* (kernel-CDF).

The three scripts (gradient/imaging/pixelization.py, misc/mapper_grad_probe.py, misc/pixelization_sparse_cpu.py) were last authored 2026-06-01, so they are pre-consolidation and become Bilinear too. jax_profiling/ is now internally consistent — every script names the mesh it historically measured. A dated-names warning sits above the status table so the next reader checks the date before renaming.

Fixed in 08d5d86.

State

PyAutoGalaxy#581 (library)	CI green — docs-build, unittest 3.12, 3.13, unittest-nojax all SUCCESS. Mergeable.
autolens_workspace_developer#127 (script)	Mergeable. That repo has no .github/workflows, so empty checks is expected, not a failure.

Merge order: #581 first — the script cannot run without the Basis fix.

What was verified

• test_autogalaxy: 1113 passed.
• jax_profiling/jit/imaging/pixelization.py: runs to completion, exit 0, all five assertions pass. figure_of_merit = 24746.105672366088 — the original 2026-05-11 constant, restored rather than re-pinned.
• Cross-checked across two independent implementations (numpy/numba vs JAX full pipeline) to 8.9e-8, plus vmap.
• Every renamed symbol confirmed to resolve and construct against the live API.
• The renamed sibling scripts (gradient/, misc/) are syntax-checked and symbol-verified, not executed — they are long-running gradient/CPU sweeps and the change to them is name-only.

Left deliberately untouched

misc/pixelization_spline_*.py reference RectangularSplineAdaptDensity, which no longer exists at all — a pre-existing break from a different mesh campaign. The wider rename fallout (autogalaxy_workspace, HowToLens, euclid pipeline, prior config YAMLs) is #461's announced workspace follow-up, captured in PyAutoMind/draft/feature/autoarray/rectangular_bilinear_rtu_mesh_split.md.

[Jammy2211]

Both PRs merged — PyAutoGalaxy#581 (d1b46b9) and autolens_workspace_developer#127 (77ca9d3).

Outcome: the reported divergence was not a library bug, but investigating it surfaced one (the Basis irregular-grid crash), a script that could not run at all, and a silent transform swap (#402) that a stale regression pin had masked for a month. The constant was restored, not re-pinned.

Completion record: PyAutoMind/complete/2026/08/pixelization-eager-jit-divergence.md. Worktree removed, branches deleted.

The two loose ends are carried into PyAutoMind/draft/feature/autoarray/rectangular_bilinear_rtu_mesh_split.md (#461's follow-up campaign): the misc/pixelization_spline_*.py scripts broken on the deleted RectangularSplineAdaptDensity, and the three renamed sibling scripts that are symbol-verified but never executed — along with the relabel-by-date trap and a caution on that campaign's _test pin regeneration.