Skip to content

fix: unblock jit/imaging/pixelization.py and make its log-evidence honest - #127

Merged
Jammy2211 merged 4 commits into
mainfrom
feature/pixelization-eager-jit-divergence
Aug 21, 2026
Merged

fix: unblock jit/imaging/pixelization.py and make its log-evidence honest#127
Jammy2211 merged 4 commits into
mainfrom
feature/pixelization-eager-jit-divergence

Conversation

@Jammy2211

@Jammy2211 Jammy2211 commented Aug 21, 2026

Copy link
Copy Markdown
Contributor

Unblocks jax_profiling/jit/imaging/pixelization.py and settles the
eager-vs-step-by-step log-evidence divergence reported in PyAutoGalaxy#580.

Depends on PyAutoLabs/PyAutoGalaxy#581 — the script cannot run without it.

What reproducing the report found

The original prompt (2026-07-08) described two log-evidence clusters ~292k
apart, pinned behind two constants and a FIXME. That surface evidence is stale —
PR #60 rewrote the script and there is now a single constant. But the divergence
underneath was papered over rather than resolved, and the script had stopped
running entirely.

1. It did not run on main

Two independent blockers:

  • RectangularAdaptDensity was renamed to RectangularRTUAdaptDensity
    (PyAutoArray#461 f9aceea3, breaking rename, no back-compat aliases). Renamed
    here and in the gradient/ + misc/ siblings that name it. This script's own
    fiducial moves to RectangularBilinearAdaptDensity — see "The constant" below
    for why that is a restoration rather than a change.

    The siblings are relabelled by date, not blanket-renamed. The same class
    name meant rank-CDF before the 2026-07-23 consolidation and kernel-CDF after,
    so a blanket rename would relabel 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 ... the kernel-CDF meshes
    now ARE RectangularAdaptDensity"), so rows dated 2026-07-09 and earlier
    become RectangularBilinear* and the 2026-07-26 sections stay
    RectangularRTU*. The 07-09 rows describe the rank transform in their own
    prose ("rank-transform knots", "rank-reordering jumps"). 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 — leaving jax_profiling/ internally
    consistent, every script naming the mesh it historically measured. A dated-names
    warning now sits above the status table so the next reader checks the date.

  • Basis.image_2d_list_from assumed grid.mask → PyAutoGalaxy#581.

2. The evidence used the wrong matrices

compute_log_evidence took the H built in step 11, which spans the mapper
only
(1225 source pixels), while D and F are taken from the inversion and span
every linear object — 1225 source pixels + the 60 linear MGE lens-light
Gaussians = 1285. On current main that is a hard shape error at step 12:

TypeError: add got incompatible shapes for broadcasting: (1285, 1285), (1225, 1225)

Beyond the shape, the MGE columns are unregularized, so the full H is singular
(60 zero rows) and slogdet(H) is -inf — a naive rebuild returns -inf, not a
number. Inversion.log_det_* evaluates both log-determinant terms on the
reduced (mapper-only) matrices, so the script now mirrors that. After the
fix the rebuild matches FitImaging.figure_of_merit to 13 significant figures:

log_evidence (inv matrices) = 25004.71903495484
log_evidence (reference)    = 25004.71903495436

3. The remaining gap was mislabelled

The old comment blamed "cumulative floating-point differences between
JIT-compiled and eager paths"
. It is not round-off. use_positive_only_solver
defaults to True, so the fitted reconstruction is a non-negative (fnnls)
solution whose active set is chosen iteratively — at this fiducial 362 of 1285
parameters are pinned at exactly zero
, while a plain solve(F + H, D) puts
777 of them negative. A re-solve landing on a different active set gives a
genuinely different s, hence a different χ² and sᵀHs. That is a real property
of the constrained solve, not drift in the evidence arithmetic.

So the ~1.5e-3 residual gap is legitimate and expected. The comment now says
so, and the step-by-step value — previously printed and never asserted — carries
its own pin rather than being silently ignored.

The constant: not drift, a transform swap — and it is restored

The first push here re-pinned the constant to the RTU value and blamed the
1.0e-2 gap on accumulated library drift. That was inferred from a commit
message rather than measured, and it was wrong. Measured at this fiducial:

PyAutoArray mesh eager figure_of_merit
f9aceea3^ (pre-split) RectangularAdaptDensity 25004.71903495436
f9aceea3 (post-split) RectangularRTUAdaptDensity 25004.71903495436
f9aceea3 (post-split) RectangularBilinearAdaptDensity 24746.105672366088

Two things fall out. RTU is a pure rename — bit-identical across the split.
And Bilinear reproduces the constant pinned on 2026-05-11 bit-for-bit
(.hex() equal) after 3.5 months, which also rules out generic drift: had any
other library change moved this value, the match could not be exact.

The real 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 same name kernel-CDF 25004.71903495436
2026-08-21 f9aceea3 (#461) split them; Bilinear = rank-CDF 24746.105672366088

This script never chose kernel-CDF — it inherited it silently in July. Naming
RectangularBilinearAdaptDensity explicitly restores what the pin has always
meant, so no re-pin is needed and the constant keeps its full pre-#402
history. Bilinear is also #461's designated workspace default and avoids the erf
sum that dominates the numba CPU likelihood.

The comment block records this measured provenance and states the RTU value for
anyone switching the mesh back.

Historical corroboration

The restored value is not only measured today — it is what this script's own
recorded history says it should be.

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

The kernel-CDF value has never been pinned here. From #402 (2026-07-23)
onward the script was computing a number its own assertion rejected.

Two further independent lines agree:

Committed profiling artifacts from the rank-CDF era — separate runs, separate
hardware, independent of the pin. Every fp64 run recorded eager_log_evidence
bit-identical to the constant:

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

The _mp.json siblings record 24746.105678802393 — the same runs in mixed
precision, differing only in the 7th decimal, exactly as expected.

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, which is what
Bilinear routes through — same sort/argsort/cumsum body, no erf
anywhere. The kernel-CDF path is the one carrying the erf sum.

The value is independently cross-checked across two implementations of the
likelihood: the numpy/numba path (xp=np) gives 24746.105672366088 and the
JAX full pipeline (use_jax=True) gives 24746.103473124283 — agreeing to
8.9e-8 relative.

Verification

figure_of_merit (log_evidence) = 24746.105672366088
log_evidence (step-by-step)    = 24783.19775195945
log_evidence (inv matrices)    = 24746.105672366568
log_evidence (reference)       = 24746.105672366088
full log_likelihood (JIT)      = 24746.103473124283
vmap batch                     = [24746.10419047 ×3]

Assertion PASSED: inversion-matrix log_evidence matches FitImaging.log_evidence
Correctness check PASSED
Eager regression assertion PASSED
Regression assertion PASSED
Step-by-step regression assertion PASSED

Script exits 0. PyAutoGalaxy test_autogalaxy: 1113 passed.

Not in scope

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

🤖 Generated with Claude Code

https://claude.ai/code/session_01SdsQ7c6y3iuiQf2jy8gx6K

Jammy2211 and others added 2 commits August 21, 2026 18:20
…nest

The script could not run on main and, where it did run, its account of
its own numbers was wrong. Three separate things:

1. Mesh rename. `RectangularAdaptDensity` -> `RectangularRTUAdaptDensity`
   (PyAutoArray#461 f9aceea3, breaking rename, no back-compat aliases),
   in this script and in the gradient/misc siblings that name it. RTU is
   the pure rename of the kernel-CDF class, so every value is preserved.
   The Bilinear sibling is a different transform and is not adopted here.

2. Wrong matrices in the evidence. `compute_log_evidence` took the H
   built in step 11, which spans the mapper only (1225 source pixels),
   while D and F are taken from the inversion and span every linear
   object (1225 + 60 linear MGE Gaussians = 1285). On current main that
   is a hard shape error at step 12. Beyond the shape: the MGE columns
   are unregularized, so the full H is singular and its log-det is -inf.
   `Inversion.log_det_*` evaluates both log-det terms on the *reduced*
   (mapper-only) matrices, so the script now does the same. The rebuilt
   evidence matches `FitImaging.figure_of_merit` to 13 significant
   figures (25004.71903495484 vs 25004.71903495436).

3. The eager-vs-step-by-step gap was mislabelled. The comment blamed
   "cumulative floating-point differences between JIT-compiled and eager
   paths". It is not round-off: `use_positive_only_solver` is True by
   default, so the fitted reconstruction is a non-negative fnnls solution
   with an iteratively chosen active set — 344 of 1285 parameters pinned
   at exactly zero at this fiducial. A re-solve landing on a different
   active set gives a genuinely different s, hence different chi^2 and
   s^T H s. Comment corrected, and the step-by-step value — previously
   printed and never asserted — now carries its own pin.

Constant re-pinned 24746.105672366088 (v2026.5.1.4) -> 25004.71903495436
(v2026.8.17.1), a 1.0e-2 drift over ~3.5 months of library development.
The re-pin rests on a two-implementation cross-check, not one run: the
numpy/numba path (`xp=np`) and the JAX full pipeline (`use_jax=True`)
agree to 2.2e-7 relative, with vmap reproducing the JAX value exactly.

Refs PyAutoGalaxy#580 (the `Basis` fix this script also needed)

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01SdsQ7c6y3iuiQf2jy8gx6K
…stant

Correction to the previous commit, which re-pinned the constant to the RTU
(kernel-CDF) value and attributed the 1.0e-2 gap to accumulated library
drift (border PCA stabilization / kernel-CDF numpy fast path). That
attribution was inferred from a commit message, not measured, and it was
wrong. Measuring it:

  PyAutoArray f9aceea3^ + RectangularAdaptDensity      25004.71903495436
  PyAutoArray f9aceea3  + RectangularRTUAdaptDensity   25004.71903495436
  PyAutoArray f9aceea3  + RectangularBilinearAdapt...  24746.105672366088

So RTU really is a pure rename (bit-identical across the split), and
Bilinear reproduces the constant pinned on 2026-05-11 *bit-for-bit* —
`.hex()` equal, after 3.5 months. That also rules out generic drift: had
any other library change moved this value, the match could not be exact.

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

  2026-05-11  pinned; `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, RTU = kernel-CDF

The script never chose kernel-CDF; it inherited it silently in July. Naming
`RectangularBilinearAdaptDensity` explicitly restores what this pin has
always meant, so no re-pin is needed and the constant keeps its full
pre-#402 history. Bilinear is also #461's designated workspace default and
avoids the erf sum that dominates the numba CPU likelihood.

The comment block now records the measured provenance and states the RTU
value for anyone switching the mesh back. Step-by-step pin and the fnnls
active-set figures re-measured for Bilinear (362 of 1285 pinned at zero;
777 negative under an unconstrained solve).

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

Refs PyAutoGalaxy#580

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01SdsQ7c6y3iuiQf2jy8gx6K
Jammy2211 and others added 2 commits August 21, 2026 18:52
Historical corroboration for the restored constant, found by checking
whether the Bilinear value ever appeared in this script's own recorded
history. It does — and the kernel-CDF value never did:

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

So the script has never, in its whole history, pinned the kernel-CDF
value. From #402 (2026-07-23) onward it was computing a number its own
assertion rejected.

Three independent lines of evidence now agree:

1. Committed profiling artifacts from the rank-CDF era — independent of
   the pin, produced by separate runs on separate hardware. Every fp64
   run recorded `eager_log_evidence` bit-identical to the constant:
     hpc_a100_fp64.json   A100 80GB, v2026.5.1.4
     local_gpu_fp64.json  RTX 2060,  v2026.5.8.2
     local_cpu_fp64.json  CPU,       v2026.5.8.2
   The `_mp.json` siblings give 24746.105678802393 — same runs in mixed
   precision, differing only in the 7th decimal.

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`, which is
   what Bilinear routes through — same `sort`/`argsort`/`cumsum` body,
   no `erf` anywhere. The kernel-CDF path is the one with the erf sum.

3. Direct measurement across the split (previous commit): pre-split
   `RectangularAdaptDensity` == post-split RTU == 25004.71903495436,
   while Bilinear == 24746.105672366088 == the 2026-05-11 pin.

Comment-only change; no behaviour, no values.

Refs PyAutoGalaxy#580

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01SdsQ7c6y3iuiQf2jy8gx6K
Correcting my own blanket rename two commits back. I swept
`RectangularAdaptDensity` -> `RectangularRTUAdaptDensity` across the
gradient README and the misc/gradient scripts, which is the same mistake
this task exists to fix: the class name meant a *different transform*
before and after the mesh consolidation (2026-07-23), so a blanket rename
relabels historical findings as belonging to a mesh that never produced
them.

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

  rows/sections dated 2026-07-09 and earlier -> rank-CDF -> Bilinear
  sections dated 2026-07-26                  -> kernel-CDF -> RTU

The 07-09 rows describe the rank transform in their own prose ("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" — that is the
rank path, not the erf-sum kernel path. Those are now Bilinear; the
07-26 final-assessment and regularization-matrix sections stay RTU.

Same reasoning for the three scripts, all last authored 2026-06-01 and so
pre-consolidation: gradient/imaging/pixelization.py (whose ~1e-24
staircase note is a rank-era finding), misc/mapper_grad_probe.py and
misc/pixelization_sparse_cpu.py revert to Bilinear. That also leaves
jax_profiling/ internally consistent — every script now names the mesh it
historically measured.

Added a dated-names warning above the status table so the next reader
checks the date before renaming, rather than repeating this.

Docs/labels only; no behaviour, no values. Syntax-checked.

Refs PyAutoGalaxy#580

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01SdsQ7c6y3iuiQf2jy8gx6K
@Jammy2211
Jammy2211 merged commit 77ca9d3 into main Aug 21, 2026
@Jammy2211
Jammy2211 deleted the feature/pixelization-eager-jit-divergence branch August 21, 2026 23:06
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant