[Jammy2211]

Summary

Closes #552.

The JAX-vs-NumPy log-likelihood gap of the RectangularBilinearAdaptImage + Adapt parity script (15.5 nats on the 316-pixel autogalaxy_workspace_test dataset, 5.8 nats on the old 716-pixel one) is not a mixed-precision loss. Measured on freshly simulated data: mixed vs fp64 within either backend differs by ~1e-4 nats, fnnls and jaxnnls agree to 1e-12 on identical inputs, and the real-space and FFT convolutions agree to 1e-16. The whole gap sits in create_transforms_rank (the "rank" lattice transform of the RectangularBilinear* meshes): it accumulates the adapt-image weights in argsort order, and np.argsort (quicksort, unstable) and jnp.argsort (stable) order the heavily tied coordinates of an unlensed image grid differently — 553 of 632 permutation entries, 80 sub-pixels landing in a different mesh cell. The transform was ill-defined on ties (NumPy alone was not invariant to the input ordering). Full attribution tables: #552 (comment).

Three changes:

1. create_transforms_rank, weighted branch (autoarray/inversion/mesh/interpolator/rectangular.py): every point of a tie block takes the block's final cumulative weight — the right-continuous empirical CDF, via searchsorted(side="right") - 1. The transform now depends only on the point multiset: backend- and order-independent, jit-traceable, bit-identical on untied inputs. The unweighted branch is untouched.
2. NumPy backend stays fp64 under use_mixed_precision (mapper_util.mapping_matrix_from, Convolver.mapping_matrix_native_from): the fp32 allocation is now JAX-only, so a NumPy fit is the fp64 reference the Settings docstring promises.
3. curvature_matrix_via_mapping_matrix_from weights with fp64 1/noise_map on every backend (one xp-agnostic expression). The old JAX branch never accumulated in fp32 — the blurred mapping matrix is already fp64 — it only rounded 1/σ in F inconsistently with the fp64 D.

Settings.use_mixed_precision docstring corrected to match. First use_mixed_precision unit tests in test_autoarray (12 new tests).

After the fix, on the same data: the parity assertion gap is −8e−5 nats (316 px) and −8e−4 nats (716 px); at fp64 the backends agree to 1e−11; the mapper's sub-pixel assignment is identical entry for entry. The smoke script passes.

API Changes

No symbol was added, removed or renamed. Two behaviour changes:

• Adapt-image rectangular meshes on gridded (unlensed) data give different mesh geometry, so different absolute log-likelihoods — e.g. −1259.27 → −791.96 on the 316-px smoke data (a lower χ²). Both old backend values were arbitrary tie-breaks of the same ill-posed quantity; lensed data (no exact ties) and the density-adaptive (unweighted) mesh are unchanged bit-for-bit. Any regression baseline pinned to old values on gridded adapt-image fits needs re-baselining.
• use_mixed_precision=True no longer changes anything on the NumPy backend (was: fp32 mapping matrix and native cube). NumPy mixed and NumPy fp64 fits are now bit-identical.

See full details below.

Test Plan

• python -m pytest test_autoarray — 1511 passed, 73 skipped, 0 failed (with numba installed; without it 8 pre-existing w-tilde tests raise the documented InversionException).
• 12 new tests: test_rectangular.py (tied weighted points: NumPy = JAX to 1e-12, permutation-invariant, untied points bit-identical to the old formula, block-total at a tie), test_mapper_util.py and test_convolver.py (dtype per backend under the flag), test_inversion_util.py (fp64 weighting on JAX, matches NumPy).
• autogalaxy_workspace_test/scripts/imaging/jax_likelihood/rectangular.py against this branch: NumPy −791.956, JIT −792.88, PASS at rtol=2e-2 (the 0.9-nat residual there is the script's own import jax-before-autogalaxy x64-off defect, a workspace follow-up).
• Phase A diagnostic matrix re-run on the branch (issue comment): gap 15.47 → −8e−5 nats, 5.79 → −8e−4 nats.

Full API Changes (for automation & release notes)

Removed

• None.

Added

• None.

Renamed

• None.

Changed Signature

• None.

Changed Behaviour

• autoarray.inversion.mesh.interpolator.rectangular.create_transforms_rank(traced_points, mesh_weight_map=..., xp=) — weighted branch: tied coordinates take the tie block's final cumulative weight (right-continuous empirical CDF). Untied inputs unchanged. Affects RectangularBilinearAdaptImage (and any "rank"-transform mesh with a weight map) on data whose traced source-plane coordinates tie exactly — in practice unlensed image grids.
• autoarray.inversion.mappers.mapper_util.mapping_matrix_from(..., use_mixed_precision=True, xp=np) — returns float64 (was float32). Under xp=jax.numpy still float32.
• autoarray.operators.convolver.Convolver.mapping_matrix_native_from(..., use_mixed_precision=True, xp=np) — float64 cube (was float32). JAX unchanged.
• autoarray.inversion.inversion.inversion_util.curvature_matrix_via_mapping_matrix_from(...) — under xp=jax.numpy with settings.use_mixed_precision=True, F is weighted by fp64 1/noise_map (was float32-rounded). Result dtype float64 on every backend, as before.

Migration

• No code changes needed. Regression baselines for adapt-image rectangular pixelization fits on gridded/unlensed data will shift; re-baseline against the new values (the NumPy and JAX backends now agree to fp64 round-off).

Generated by the PyAutoLabs agent workflow.

🤖 Generated with Claude Code

https://claude.ai/code/session_018bJQFXsFc4iCqpUywydex2

---

Generated by Claude Code