feat: xp-API dense kernels + iterative LM engine (potential correction phase 4)

[Jammy2211]

Summary

Phase 4 of the potential-correction port (#618): the dense/JAX layer and the iterative solver. Stacked on #621 (base = feature/potential-correction-port; retargets to main when it merges).

Per the ecosystem convention, the JAX layer is not a parallel jnp codebase: Cao's jax_ops kernels are refactored into al.pc.dense_util — dense linear-algebra kernels written once against the xp API (xp=np default; pass xp=jax.numpy for accelerator-ready, jit-compatible execution). Verified numerically identical under both backends and jit-able (session check; unit tests stay numpy-only per workspace policy).

On top of them, al.pc.IterFitDpsiSrcImaging ports the iterative Levenberg–Marquardt engine from the author's updated 2026 tree: joint optimization of the combined state [s | δψ], re-ray-tracing the image grid through the corrected lens (macro model + an InputPotential built from the current δψ — the Phase 2 profile) at every accepted step, with optional gauge constraints (⟨δψ,1⟩ = ⟨δψ,x⟩ = ⟨δψ,y⟩ = 0 via a KKT step) and a Laplace-approximation evidence. IterDpsiSrcInvAnalysis wraps it for evidence-based hyper-parameter sampling. Upstream's two duplicated LM loops (numpy + jax) collapse into one xp-parameterized loop.

One genuine upstream bug found and fixed during the port: split_cross_from took its clipping percentile over all Voronoi cell areas including the −1 unbounded-cell sentinels, so meshes with mostly-unbounded cells produced a negative clip → NaN gradient-arm lengths → interpolator crash. The port takes the percentile over bounded cells only, with a nearest-neighbour-spacing fallback (regression-tested). Relevant to the Phase 5 bug hunt: this NaN path fired precisely inside the iterative loop's source-gradient evaluation.

API Changes

Added only — nothing removed or changed.
See full details below.

Test Plan

• 17 new tests: every dense_util kernel verified against the sparse/Phase-3 reference formulations (joint + δψ-only + fixed-curvature evidence parity, LM Hessian/gradient/step/cost hand-computed, constrained step satisfies the constraint); IterFitDpsiSrcImaging end-to-end on the real 7×7 fixture (finite optimized state, cost beats the zero state, gauge constraints satisfied to 1e-6, finite Laplace evidence); split_cross_from NaN regression.
• Full test_autolens suite run before commit (execution contract).
• Session check (not committed, per numpy-only test policy): kernels bit-identical between xp=np and xp=jax.numpy, jit-compatible.
• Phase 5 follow-up: golden-file parity vs the author's tree + reproduction of the Vegetti-style iterative reconstruction.

Full API Changes (for automation & release notes)

Added

• al.pc.dense_util (new module, all functions take xp=np):
  • as_dense(matrix, xp) — numpy/scipy-sparse/jax → dense xp array.
  • dense_block_diag_from(*blocks, xp).
  • inverse_noise_variance_from(noise_slim, xp).
  • source_gradient_matrix_dense_from(source_gradient, xp) / dpsi_gradient_matrix_dense_from(itp_mat, Hx, Hy, xp) / dpsi_mapping_matrix_from(psf, src_grad, dpsi_grad, xp) — dense forms of the Phase 3 sparse operators.
  • log_evidence_joint_dense_from(...), log_evidence_dpsi_dense_from(...) — evidence with terms decomposed; valid flag instead of raising (JIT-compatible).
  • log_evidence_from_fixed_curvature(...) — Cholesky fast path when only regularizations change between samples.
  • LM kernels: lm_cost_from, lm_hessian_and_gradient_from, solve_lm_step_from (damped; optional equality-constrained KKT), log_evidence_lm_from (Laplace evidence at an LM solution).
• al.pc.IterFitDpsiSrcImaging(masked_imaging, lens_start, dpsi_pixelization, src_pixelization, gauge_constraints=False, ..., n_iter=20, tol=1e-6) — the iterative LM engine; solve_joint_optimization(xp=np), log_evidence(s, dpsi, xp=np), show(...).
• al.pc.IterDpsiSrcInvAnalysis(...) — af.Analysis over the iterative engine (failed inversions → logged −1e30 penalty).

Migration

• None — additive only.

Generated by the PyAutoLabs agent workflow.