Repository navigation
feat: potential corrections (gravitational imaging) — port Cao's package #618
Description
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;GaussianRandomFieldperturbation profile reimplemented with plain numpy FFT (no powerbox). - Phase 3 — PyAutoLens: new
autolens/potential_correction/subpackage (mirroringautolens/weak) — δψ mesh + data pairing,FitDpsiImaging(δψ-only),FitDpsiSrcImaging(joint source+δψ with full covariance),SrcFactoryinterpolator family,af.Analysisclasses, 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)
- 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 buildersdiff_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), usingautoarray.numba_utilconventions. - New regularization classes
autoarray/inversion/regularization/curvature_mask.py+fourth_order_mask.pyimplementing Cao'sCurvatureRegularizationDpsi/FourthOrderRegularizationDpsias mask-basedAbstractRegularizationsubclasses (renamed to PyAuto conventions, e.g.CurvatureMask/FourthOrderMask); export viaautoarray/inversion/regularization/__init__.pyand theaa.regnamespace. - 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. - Do not port
covariance_reg.py/ the GPy + numba-scipy kernel builders —GaussianKernelRegularization,ExponentialKernelRegularization,MaternKernelRegularizationalready exist inautoarray/inversion/regularization/. - 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(fromdpsi_mesh.py), grid pairing + operators built from the phase-1 autoarray utils.pixelization.py:DpsiPixelization,DpsiSrcPixelizationmodel components.fit.py:FitDpsiImaging(linear δψ inversion of an image residual) andFitDpsiSrcImaging(joint source+δψ: block mapping matrix[F_src | −B·D_s·D_ψ], block-diag regularization, evidence terms), re-anchored on currentFitImaging/AdaptImages/SettingsInversionAPI (over-sampling replacesSettingsImaging(sub_size=…);Tracer(galaxies=…)replacesTracer.from_galaxies).src_factory.py:SrcFactoryABC +AnalyticSrcFactory,PixSrcFactory(Delaunay/RBF), NN variant behind the existing optionalnn_pyguard; GPy GPR variant dropped.analysis.py:DpsiInvAnalysis,DpsiSrcInvAnalysis,SrcInvAnalysis(af.Analysissubclasses);InversionException→ penalty handling per current conventions (no bareexcept).util.py: anchor-point rescaling (solve_dpsi_rescale_factor, Suyu scheme), PSF blur-matrix builder (reuse/extendaa.Convolverwhere possible), inverse-noise covariance.plot/: portvisualize.pyfigures onto the matplotlib function conventions used by the current plot API.
- Fix the ported code smells: no
except:-swallowing (per no-silent-guards rule), seed the random-solution draws, symmetrise covariance handling explicitly. - 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.
- Parity script (workspace_test follow-up, not this PR): port outputs vs original repo on the same simulated dataset.
- Docs:
docs/apientries; 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
Activity
Jammy2211 commentedon Jul 17, 2026
Smoke Test Results — 2026-07-17 (phase 1: PyAutoArray PR #390)
All 44 smoke tests passed across 6 workspaces (5 skipped per
no_runpolicy).Phase 1 shipped to PR-open: PyAutoArray#390 (
pending-release) — masked-grid derivative operators, coarse-mesh interpolation,CurvatureMask/FourthOrderMaskregularizations. Fulltest_autoarraysuite: 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_correctionsubpackage) are queued behind theslam-resume-fastpathtask's PyAutoGalaxy/PyAutoLens claims.Jammy2211 commentedon Jul 17, 2026
Phases 2 + 3 shipped — 2026-07-17
Phase 2 — PyAutoGalaxy#505 (
pending-release): newautogalaxy/profiles/mass/input/package —InputDeflections,InputPotential(operator-derived deflections/convergence on cleaned masks),GaussianRandomField(numpy-FFT, no powerbox) andLinearNDInterpolatorExt. 1000 tests passed.Phase 3 — PyAutoLens#621 (
pending-release): newautolens/potential_correction/subpackage (al.pc) — δψ mesh/pairing,FitDpsiImaging, jointFitDpsiSrcImaging, source factories,af.Analysisclasses, 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_runskips). Heart RED on the same 5 pre-existing unrelated reasons, human-acknowledged.Plan update (user-directed):
jax_opskernels refactored onto the ecosystemxpAPI (write-once,xp=npdefault,jax.numpywhen passed — NOT a parallel jnp codebase), plus theIterFitDpsiSrcImagingiterative LM engine; the tar'stests/golden/are the parity anchor.Jammy2211 commentedon Jul 17, 2026
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'sjax_opskernels refactored onto the ecosystem xp API (each kernel written once,xp=npdefault,jax.numpyfor 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 2InputPotentialat each accepted step; optional gauge constraints via KKT; Laplace evidence). Upstream's duplicated numpy/jax LM loops collapsed into one xp-parameterized loop.split_cross_fromtook 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 commentedon Jul 17, 2026
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.pcstack, scored as corr(δκ_rec, δκ_true) on the δψ mesh + peak distance from the true subhalo:subhalo_recovery.py)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 thesplit_cross_fromNaN 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.
Fixed-and-regression-tested so far (candidates for the original "unknown bug"):
split_cross_fromNaN 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):
CurvatureMasknever 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 commentedon Jul 17, 2026
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:
Conclusion — the "unknown bug" decomposes into three findings, all now closed:
split_cross_fromNaN (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).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 commentedon Jul 17, 2026
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): extendal.pcfromImagingtoInterferometer— 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). Thedense_utilxp 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).
Add visibility-space iterative LM engine (interferometer phase 2a)
Jammy2211 commentedon Jul 17, 2026
Potential-correction port arc fully complete: library + workspace + workspace_test layers all merged (~20 PRs across PyAutoArray/PyAutoGalaxy/PyAutoLens/autolens_workspace/autolens_workspace_test), imaging + interferometer (sparse w-tilde route), iterative engine certified against the one-shot at mid and full tiers, uv validation campaign results on #627. Closing.