feat: potential corrections (gravitational imaging) — port Cao's package

[Jammy2211]

Overview

Port Xiaoyue Cao's lensing_potential_correction package — the gravitational-imaging technique (pixelised corrections δψ to the lensing potential, reconstructed jointly with the pixelised source by maximising Bayesian evidence, Matérn-family regularisation) — into the PyAuto libraries as first-class, tested source code. The external package is a wrapper pinned to autolens 2024.1.27.4; the port modernises it to the current API and layers it correctly across autoarray / autogalaxy / autolens. All docs and docstrings cite Cao et al. (2025) via potential_correction_paper and credit the origin repository.

Plan

• Phase 1 — PyAutoArray: port the generic linear algebra — masked-grid sparse differential operators (1st/2nd/4th order, numba), curvature + fourth-order regularizations for rectangular masked meshes, and the coarse-mesh→data-grid interpolation matrix builder. Reuse the existing Gaussian/Exponential/Matérn kernel regularizations (do not port covariance_reg.py).
• Phase 2 — PyAutoGalaxy: pixelised input mass profiles (InputDeflections, InputPotential) on the current decorator API; GaussianRandomField perturbation profile reimplemented with plain numpy FFT (no powerbox).
• Phase 3 — PyAutoLens: new autolens/potential_correction/ subpackage (mirroring autolens/weak) — δψ mesh + data pairing, FitDpsiImaging (δψ-only), FitDpsiSrcImaging (joint source+δψ with full covariance), SrcFactory interpolator family, af.Analysis classes, anchor-point rescaling, visualization; docs + citations.
• NumPy/numba only — the JAX port is a known follow-up PR.
• No new heavy dependencies: GPy, multiprocess, powerbox, numba-scipy stay out (GPR source interpolator dropped or lazily optional).
• Parity check of the port against the original repo's outputs on a small simulated dataset before shipping.

Detailed implementation plan

Affected Repositories

• PyAutoLens (primary — final home of the feature)
• PyAutoArray (phase 1)
• PyAutoGalaxy (phase 2)

Work Classification

Library

Branch Survey

Repository	Current Branch	Dirty?	Claimable now?
./PyAutoArray	main	clean	yes
./PyAutoGalaxy	main	clean	no — claimed by slam-resume-fastpath
./PyAutoLens	main	1 file (paper_jax/paper.md, unrelated JOSS draft)	no — claimed by slam-resume-fastpath

Suggested branch: feature/potential-correction
Worktree root: ~/Code/PyAutoLabs-wt/potential-correction-port/ (created by /start_library)
Phase 1 (PyAutoArray) starts immediately; phases 2–3 queue behind slam-resume-fastpath (PyAutoGalaxy#502).

Implementation Steps

Phase 1 — PyAutoArray (generic linear algebra)

1. Port masked-grid sparse differential-operator builders from potential_correction/util.py (diff_1st_operator_numba_from, diff_2nd_operator_numba_from, the 2nd/4th-order regularization operator builders diff_2nd_operator_dpsi_reg_numba_func / diff_4th_operator_dpsi_reg_numba_func, clean_mask / iterative_clean_mask) into a new util module in the inversion layer (e.g. autoarray/inversion/regularization/derivative_util.py), using autoarray.numba_util conventions.
2. New regularization classes autoarray/inversion/regularization/curvature_mask.py + fourth_order_mask.py implementing Cao's CurvatureRegularizationDpsi / FourthOrderRegularizationDpsi as mask-based AbstractRegularization subclasses (renamed to PyAuto conventions, e.g. CurvatureMask / FourthOrderMask); export via autoarray/inversion/regularization/__init__.py and the aa.reg namespace.
3. Port the coarse-mesh→data interpolation matrix machinery (dpsi_mask_from, itp_box_mask_from, dpsi2data_itp_mat_from, linear_weight_from_box) as generic "paired coarse rectangular mesh" utilities.
4. Do not port covariance_reg.py / the GPy + numba-scipy kernel builders — GaussianKernelRegularization, ExponentialKernelRegularization, MaternKernelRegularization already exist in autoarray/inversion/regularization/.
5. Unit tests (numpy-only): operator matrices vs finite-difference expectations on small masks; regularization matrices symmetric/PSD; interpolation matrix partition-of-unity rows.

Phase 2 — PyAutoGalaxy (pixelised input mass profiles)
6. Port pix_mass.py → autogalaxy/profiles/mass/input/ (or alongside existing sheets): InputDeflections (from y/x deflection grids; convergence via 1st-order operators) and InputPotential (deflections/convergence via 1st/2nd-order operators), replacing old aa.grid_dec.* decorators with the current aa.decorators equivalents and LinearNDInterpolatorExt with a shared util.
7. Reimplement GaussianRandomField mass profile (from custom_mass.py) with plain numpy FFT — needed to simulate extended perturbations for tests/validation. Skip EPL/Multipole/EPLBoxyDisky (duplicates of existing profiles).
8. Unit tests: interpolated deflections/convergence round-trip against an analytic profile evaluated on a grid.

Phase 3 — PyAutoLens (the technique itself)
9. New subpackage autolens/potential_correction/:

• mesh.py: RegularDpsiMesh, PairRegularDpsiMesh (from dpsi_mesh.py), grid pairing + operators built from the phase-1 autoarray utils.
• pixelization.py: DpsiPixelization, DpsiSrcPixelization model components.
• fit.py: FitDpsiImaging (linear δψ inversion of an image residual) and FitDpsiSrcImaging (joint source+δψ: block mapping matrix [F_src | −B·D_s·D_ψ], block-diag regularization, evidence terms), re-anchored on current FitImaging / AdaptImages / SettingsInversion API (over-sampling replaces SettingsImaging(sub_size=…); Tracer(galaxies=…) replaces Tracer.from_galaxies).
• src_factory.py: SrcFactory ABC + AnalyticSrcFactory, PixSrcFactory (Delaunay/RBF), NN variant behind the existing optional nn_py guard; GPy GPR variant dropped.
• analysis.py: DpsiInvAnalysis, DpsiSrcInvAnalysis, SrcInvAnalysis (af.Analysis subclasses); InversionException → penalty handling per current conventions (no bare except).
• util.py: anchor-point rescaling (solve_dpsi_rescale_factor, Suyu scheme), PSF blur-matrix builder (reuse/extend aa.Convolver where possible), inverse-noise covariance.
• plot/: port visualize.py figures onto the matplotlib function conventions used by the current plot API.

10. Fix the ported code smells: no except:-swallowing (per no-silent-guards rule), seed the random-solution draws, symmetrise covariance handling explicitly.
11. Unit tests (numpy-only, no JAX): evidence pipeline on a tiny simulated dataset; δψ recovery of an injected SIS perturber sign/scale sanity check; anchor rescaling exactness on a plane.
12. Parity script (workspace_test follow-up, not this PR): port outputs vs original repo on the same simulated dataset.
13. Docs: docs/api entries; citation block (Cao et al. 2025 + both repo URLs) in the subpackage __init__/module docstrings and the relevant docs page; README/docs "cite us" material updated.

Key Files

• PyAutoArray/autoarray/inversion/regularization/ — existing kernel regs to reuse; new curvature/fourth-order mask regs land here.
• PyAutoArray/autoarray/numba_util.py — numba decorator conventions for the ported operators.
• PyAutoGalaxy/autogalaxy/profiles/mass/ — home for input pixelised mass profiles.
• PyAutoLens/autolens/weak/ — structural precedent for a new technique subpackage.
• PyAutoLens/autolens/interop/coolest.py — precedent for optional-dependency + citation handling.
• External source: potential_correction/{util,dpsi_mesh,dpsi_inv,dpsi_src_inv,pix_mass,covariance_reg,visualize}.py.

Citation requirement

Every new module docstring group and docs page cites:

• Cao et al. (2025), via https://github.com/caoxiaoyue/potential_correction_paper
• Origin implementation: https://github.com/caoxiaoyue/lensing_potential_correction

Original Prompt

Click to expand starting prompt

This is an implmeentation of the potential corrections, which supports PyAutoLens, can you take this and incorporate it proeprly into the PyAuto ecosystem source code: https://github.com/caoxiaoyue/lensing_potential_correction . It uses a lot of linear algebra, so maybe belongs in the autoarray inversion module, but its also very lensing specific so could go further down. For now, dont worry about JAX, we have a JAX code whichw e can bring in in a future PR, or maybe I'll just ask you to do it later. Also make sure docs and other materials cite this properly https://github.com/caoxiaoyue/potential_correction_paper.

Prompt file: PyAutoMind/active/potential_correction_port.md

[Jammy2211]

Smoke Test Results — 2026-07-17 (phase 1: PyAutoArray PR #390)

All 44 smoke tests passed across 6 workspaces (5 skipped per no_run policy).

Workspace	Passed	Failed	Skipped
autofit_workspace	7	0	1
autogalaxy_workspace	6	0	0
autolens_workspace	7	0	0
autolens_workspace_test	12	0	4
euclid_strong_lens_modeling_pipeline	6	0	0
HowToLens	6	0	0

Phase 1 shipped to PR-open: PyAutoArray#390 (pending-release) — masked-grid derivative operators, coarse-mesh interpolation, CurvatureMask/FourthOrderMask regularizations. Full test_autoarray suite: 923 passed. Parity-certified bit-identical vs the original implementation (one genuine upstream bug fixed: the nearest-box search wrap-around crash).

API changes: Added only — no workspace scripts affected. Heart was RED on 5 pre-existing unrelated reasons; ship proceeded on explicit human acknowledgement.

Phases 2 (PyAutoGalaxy InputDeflections/InputPotential + GRF) and 3 (autolens/potential_correction subpackage) are queued behind the slam-resume-fastpath task's PyAutoGalaxy/PyAutoLens claims.

[Jammy2211]

Phases 2 + 3 shipped — 2026-07-17

Phase 2 — PyAutoGalaxy#505 (pending-release): new autogalaxy/profiles/mass/input/ package — InputDeflections, InputPotential (operator-derived deflections/convergence on cleaned masks), GaussianRandomField (numpy-FFT, no powerbox) and LinearNDInterpolatorExt. 1000 tests passed.

Phase 3 — PyAutoLens#621 (pending-release): new autolens/potential_correction/ subpackage (al.pc) — δψ mesh/pairing, FitDpsiImaging, joint FitDpsiSrcImaging, source factories, af.Analysis classes, anchor rescaling, visualizations, docs page. 413 tests passed, incl. end-to-end finite evidence for both fits on the real 7×7 fixture. Ported from the author's updated 2026 tree (for_qiuhan_PT_jax.tar) with the published numpy evidence formulation.

Smoke (both PRs together): all 44 scripts passed across the 6 workspaces (5 pre-existing no_run skips). Heart RED on the same 5 pre-existing unrelated reasons, human-acknowledged.

Plan update (user-directed):

• Phase 4 (next): the JAX layer — Cao's jax_ops kernels refactored onto the ecosystem xp API (write-once, xp=np default, jax.numpy when passed — NOT a parallel jnp codebase), plus the IterFitDpsiSrcImaging iterative LM engine; the tar's tests/golden/ are the parity anchor.
• Phase 5 (new): hunt the unknown bug in the iterative potential-correction loop so the reconstruction reproduces Simona Vegetti's gravitational-imaging approach; validate against her recent results (Nat. Astron. 2025a, Nat. Astron. 2025b).
• Follow-up (workspace_test): end-to-end parity vs the original repo's golden outputs on its demo dataset.

[Jammy2211]

Phase 4 shipped — 2026-07-17

PyAutoLens#622 (pending-release, stacked on #621 — retargets to main when it merges): the dense/JAX layer + iterative solver.

• al.pc.dense_util — Cao's jax_ops kernels refactored onto the ecosystem xp API (each kernel written once, xp=np default, jax.numpy for accelerator execution; no parallel jnp codebase). Session-verified bit-identical between backends and jit-compatible; committed unit tests stay numpy-only per workspace policy.
• al.pc.IterFitDpsiSrcImaging + IterDpsiSrcInvAnalysis — the iterative LM engine (re-ray-traces through the corrected lens via the Phase 2 InputPotential at each accepted step; optional gauge constraints via KKT; Laplace evidence). Upstream's duplicated numpy/jax LM loops collapsed into one xp-parameterized loop.
• Upstream bug found + fixed: split_cross_from took its clip percentile over the −1 unbounded-Voronoi-cell sentinels → negative clip → NaN gradient arms → interpolator crash inside the iterative loop's source-gradient evaluation. Fixed (bounded-cells percentile + spacing fallback) with a regression test. Logged as a Phase 5 lead.

Tests: 430 passed (17 new). Phase 5 (reproduce the Vegetti-style reconstruction; hunt the remaining unknown bug) is next — leads so far: the NaN path above, LM jitter/logdet fragility (upstream commit 220c815), anchor/gauge degeneracy handling, and source-gradient evaluation at updated vs initial positions.

[Jammy2211]

Phase 5 — reproduction results (2026-07-17)

Reproduction of the tar's demo (200×200, SIE + 1e10 M☉ NFW at (1.81, 0.0)) on the merged al.pc stack, scored as corr(δκ_rec, δκ_true) on the δψ mesh + peak distance from the true subhalo:

Run	Hyper-params	corr	peak dist	verdict
One-shot joint fit, 200×200	Matérn52 c=2000, s=4	0.77	0.21″	✅ recovers
Iterative LM, 120×120 (subhalo_recovery.py)	Matérn52 c=2000, s=4	0.78	0.16″	✅ recovers, gauge holds
Iterative LM, 200×200	demo-2 stored values: Matérn52 c=803465.55, s=0.4497	0.18	1.82″	❌ does not localize

The failing run used the hyper-parameters stored in the tar's 2_inversion_powell.py — values tuned on the pre-fix upstream code (which had the split_cross_from NaN bug live in exactly this code path). Small scale (0.45″ ≈ 1 δψ pixel) makes the Matérn regularization near-diagonal, leaving δψ under-regularized.

Two hypotheses, discriminator in flight: rerunning the 200×200 iterative at the known-good c=2000/s=4.

• Recovers → iterative engine is sound; the "unknown bug" = the fixed NaN crash + stale hyper-parameter calibration.
• Fails → a second defect triggered by full resolution, independent of the regularization scale.

Fixed-and-regression-tested so far (candidates for the original "unknown bug"): split_cross_from NaN arms (percentile over unbounded-cell sentinels — fires inside the iterative loop's source-gradient evaluation) and the coarse-mesh box-search wrap-around KeyError.

Also empirically confirmed (walkthrough script, ws#284): CurvatureMask never localizes the test subhalo at any coefficient scanned (10–10⁴), MaternKernel ν=5/2 localizes it to one coarse pixel — Cao et al. 2025's central claim.

Benchmarks recorded for the reproduction target: Powell et al. 2025 (Nat. Astron. 9, 1714 — B1938+666, (1.13±0.04)×10⁶ M☉ at 26σ, ~10–20 source+potential Newton iterations) and Vegetti et al. 2026 (same object, (1.8±0.1)×10⁶ M☉ UD+PM, CDM tension).

[Jammy2211]

Phase 5 — RESOLVED: the iterative engine is sound; the "unknown bug" is explained (2026-07-17)

The discriminating run landed: 200×200 demo, iterative LM engine, Matérn52 c=2000 / s=4 →
corr(δκ_rec, δκ_true) = 0.73, peak 0.21″ from the true subhalo (the same position the one-shot finds), Laplace log-evidence 1.6901e4, converged in 964 s with almost no rejected steps.

Completing the table from the previous comment:

Run (200×200, iterative)	Hyper-params	corr	peak dist	log-evidence
demo-2 stored values	c=803465.55, s=0.4497	0.18	1.82″	1.6598e4
known-good values	c=2000, s=4	0.73	0.21″	1.6901e4

Conclusion — the "unknown bug" decomposes into three findings, all now closed:

1. split_cross_from NaN (percentile over the −1 unbounded-Voronoi sentinels → NaN gradient arms → crash/poisoned D_s inside the iterative loop's source-gradient evaluation) — found + fixed + regression-tested (feat: xp-API dense kernels + iterative LM engine (potential correction phase 4) #622).
2. Coarse-mesh box-search wrap-around KeyError — found + fixed + regression-tested (Feature/remove preloads #390).
3. Stale hyper-parameter calibration: the demo's stored iterative values were tuned on the pre-fix code and are miscalibrated (s=0.45″ ≈ one δψ pixel → near-diagonal regularization → the corrections soak up noise instead of the subhalo). Critically, the evidence itself discriminates correctly (1.6901e4 > 1.6598e4), so evidence-based hyper-parameter sampling (DpsiSrcInvAnalysis / IterDpsiSrcInvAnalysis) recovers the good regime automatically — no manual tuning needed, exactly as the method intends.

The Powell/Vegetti-style iterative source+potential reconstruction now reproduces end-to-end on the merged stack. Remaining (optional) follow-up: an evidence-sampled hyper-parameter run on the demo dataset to confirm the search lands at the recovering regime, and pushing toward the B1938+666-scale sensitivity benchmarks of Powell et al. 2025 / Vegetti et al. 2026.

[Jammy2211]

Shipped — task complete (2026-07-17)

All six phases merged: PyAutoArray#390, PyAutoGalaxy#505, PyAutoLens#621, PyAutoLens#622, autolens_workspace#284, autolens_workspace_test#176. Phase 5 verdict above; completion record at PyAutoMind/complete/2026/07/potential-correction-port.md; worktree cleaned.

Follow-up filed (PyAutoMind/draft/feature/autolens/potential_correction_interferometer.md): extend al.pc from Imaging to Interferometer — replace the PSF blur matrix with the transformer/NUFFT operator in the correction response, visibility-space noise/evidence terms, and the iterative engine's mapping update — bringing the technique to the VLBI/radio regime of the B1938+666 detections (Powell et al. 2025; Vegetti et al. 2026). The dense_util xp kernels are already data-agnostic, so the work concentrates in the fit-layer operators.

This issue can be closed when the follow-up is picked up (it will get its own issue via start_dev).