feat: potential corrections for interferometer data (visibility-space al.pc)

[Jammy2211]

Overview

Extend the potential-correction (gravitational imaging) subpackage al.pc (shipped in #618; ported from Cao et al. 2025, https://github.com/caoxiaoyue/lensing_potential_correction, cite via https://github.com/caoxiaoyue/potential_correction_paper) from Imaging to Interferometer datasets, fitting in visibility space. This reaches the technique's benchmark regime: the JVAS B1938+666 detections of Powell et al. 2025 (Nat. Astron. 9, 1714) and Vegetti et al. 2026 were made with a visibility-space gravitational-imaging code.

Plan

• New visibility-space fits in al.pc: FitDpsiInterferometer (δψ-only inversion of a visibility residual) and FitDpsiSrcInterferometer (joint source+δψ), where the correction response is −T(D_s·D_ψ) via transformer.transform_mapping_matrix (already xp-parameterized) in place of the imaging PSF blur matrix B; real/imag block stacking and complex-noise handling mirror the existing interferometer inversion conventions (noise_normalization_complex_from).
• δψ mesh built from the dataset's real-space mask (cleaned), reusing PairRegularDpsiMesh unchanged; source gradients remain image-plane (ray-traced real-space grid).
• Iterative engine support: the per-step source-mapping update uses the transformed mapping matrix instead of convolved_mapping_matrix_from; IterFitDpsiSrcInterferometer (or dataset-dispatching extension of the existing engine — decided in implementation by whichever mirrors the FitImaging/FitInterferometer split most cleanly).
• Analysis classes for evidence-based hyper-parameter sampling over interferometer data.
• al.pc.dense_util untouched — its xp kernels are data-agnostic (mapping matrix + noise vector in).
• Workspace follow-up (library-first gate): autolens_workspace_test scripts — jax_likelihood_functions/interferometer/potential_correction.py (xp/jit parity on the visibility evidence) and an interferometer subhalo_recovery regression (SimulatorInterferometer + NFW subhalo), smoke-wired.
• NumPy-first; JAX rides the transformer's existing xp support. Unit tests numpy-only. Citations maintained in all new modules.

Detailed implementation plan

Affected Repositories

• PyAutoLens (primary)
• autolens_workspace_test (workspace follow-up)

Work Classification

Both (library first, workspace follows)

Branch Survey

Repository	Current Branch	Dirty?
./PyAutoLens	main	1 file (paper_jax/paper.md, unrelated JOSS draft — warning only)
./autolens_workspace_test	main	clean

Suggested branch: feature/potential-correction-interferometer
Worktree root: ~/Code/PyAutoLabs-wt/potential-correction-interferometer/

Implementation Steps

1. autolens/potential_correction/fit_interferometer.py: FitDpsiInterferometer + FitDpsiSrcInterferometer.
   • Correction response: densify D_s @ D_ψ (sparse product, then .toarray()), apply dataset.transformer.transform_mapping_matrix → complex [n_vis, n_dpsi]; stack real/imag rows to match how the interferometer inversion's operated_mapping_matrix enters the curvature (mirror autoarray/inversion/inversion/interferometer/abstract.py::operated_mapping_matrix_list exactly).
   • Source blocks from al.FitInterferometer's inversion (operated_mapping_matrix, regularization_matrix) at the smooth starting model, as the imaging fit does with FitImaging.
   • Noise: per-component inverse variance from the complex noise map; evidence noise term via aa.util.fit.noise_normalization_complex_from; χ² over real+imag residuals.
   • δψ mesh: PairRegularDpsiMesh(real_space_mask, pixel_scale, factor); δψ regularization via the existing DpsiLinearObj adapter.
2. autolens/potential_correction/iterative.py: dataset seam — replace the imaging-only psf.convolved_mapping_matrix_from update and psf_matrix property with a dispatch on dataset type (imaging → PSF matrix path; interferometer → transform_mapping_matrix path); gauge constraints and LM kernels unchanged.
3. autolens/potential_correction/analysis.py: DpsiInvInterferometerAnalysis / DpsiSrcInvInterferometerAnalysis (or dataset-generic analyses — implementation judgment).
4. Exports in autolens/potential_correction/__init__.py + docs/api page update.
5. Unit tests (numpy-only): matrix-injection evidence-term tests for the complex-noise formulation; end-to-end finite evidence on the interferometer_7 fixture family; transformed-response shape/consistency vs direct transformer.visibilities_from on a δ-function correction.
6. Workspace step (after library PR): the two autolens_workspace_test scripts + smoke_tests.txt/env_vars.yaml wiring, mirroring the imaging ones.

Key Files

• autolens/potential_correction/{fit.py,iterative.py,analysis.py,dense_util.py} — the imaging implementations being extended.
• autoarray/operators/transformer.py — transform_mapping_matrix(mapping_matrix, xp=np) (DFT + NUFFT variants).
• autoarray/inversion/inversion/interferometer/abstract.py — the operated-mapping-matrix conventions to mirror.
• autoarray/fit/fit_util.py::noise_normalization_complex_from — the visibility noise term.

Original Prompt

Click to expand starting prompt

merge, and then the follow up will be to get this running on interferometer data as currently its just imaging

Prompt file: PyAutoMind/active/potential_correction_interferometer.md

[Jammy2211]

Requirement addendum (user, 2026-07-17): the sparse-operator route is the priority

The visibility-space fits must work through the sparse operator route (dataset.apply_sparse_operator / the w-tilde-style machinery), not only the dense transform_mapping_matrix path — that is the more important route because it scales with the number of visibilities (real targets like SDP.81 have 10⁶–10⁸ visibilities; a dense transformed mapping matrix [2·n_vis, n_src+n_dpsi] is not materializable there).

Design consequence for the correction blocks: with an implicit transform T, the joint curvature blocks are computed in real space —

• dpsi block: (D_s·D_ψ)^T · (T^H C⁻¹ T) · (D_s·D_ψ) where T^H C⁻¹ T is the w-tilde kernel over the ~10⁴ real-space mask pixels (independent of n_vis);
• data vector: (D_s·D_ψ)^T · T^H C⁻¹ d (the dirty/weighted image);
• cross block with the source mirrors the existing w-tilde source curvature.

Plan ordering updated: (1) implement the sparse/w-tilde route as the primary implementation, mirroring autoarray's interferometer w-tilde inversion utilities; (2) keep the dense transform_mapping_matrix route as the small-n_vis reference used for parity tests (the two must agree on the interferometer_7-scale fixtures); (3) the evidence's log-det terms operate on the (n_src+n_dpsi)² curvature either way, so dense_util's xp kernels remain the downstream consumers unchanged.

[Jammy2211]

Phase 1 shipped — 2026-07-17

PyAutoLens#624 (pending-release): al.pc.FitDpsiSrcInterferometer + DpsiSrcInvInterferometerAnalysis.

Per the sparse-operator requirement, the w-tilde route is the primary implementation: joint curvature [f | G]ᵀ·(TᴴC⁻¹T)·[f | G] via one curvature_matrix_diag_from call over the concatenated extent-indexed COO triplets (source mapper triplets + correction-response triplets), data vector against the dirty image, χ² via a single forward NUFFT — cost scales with real-space mask pixels, independent of n_vis. Dense transform_mapping_matrix route retained behind use_sparse_operator=False as the parity reference.

Verified: sparse route ≡ dense route (curvature, data vector, evidence) on the interferometer_7 fixture; 434 tests pass. Requires dataset.apply_sparse_operator() (loud error otherwise — no silent fallback).

Phase 2 (next on this issue): the iterative LM engine's interferometer seam (transformed per-step source-mapping update) and the autolens_workspace_test layer — visibility subhalo-recovery regression (SimulatorInterferometer + NFW subhalo) and jax_likelihood_functions/interferometer/potential_correction.py (xp/jit parity incl. the sparse operator under JAX).

[Jammy2211]

Phase 2b status + open defect diagnosis (2026-07-17)

Landed on #625: dpsi_mask on both interferometer fits — the real-space mask stays filled (FFT extent) while the δψ mesh restricts to an arc-tracing sub-mask, with response rows embedded back into the full row space. One-shot synthetic recovery improves to corr 0.35 / peak 0.13″, and sparse ≡ dense parity holds throughout.

Open defect (the concrete blocker for the iterative interferometer validation script): on the 64×64 synthetic, IterFitDpsiSrcInterferometer progresses (cost 1.18e11 → 1.57e7) but converges to a poor state. Probe result: the χ² identity is exact (identity vs direct NUFFT agree to all digits), but evaluating the one-shot solution through the nonlinear re-trace gives χ² 9.6e7 — the model breaks under re-tracing. Root cause: ag.mp.InputPotential uses nearest-neighbour extrapolation outside its mesh (LinearNDInterpolatorExt), so an arc-restricted δψ extrapolates constant-nonzero across the full circular re-trace grid, producing spurious deflections in the disk interior. The imaging engine never hits this because its data mask is the arc mask.

Fix design (next session): add a zero-fill extrapolation mode to InputPotential / LinearNDInterpolatorExt (linear inside the hull, → 0 outside, optionally tapered) and use it for the correction galaxy in both iterative engines; then re-validate the interferometer subhalo-recovery script (drafted in the worktree, held uncommitted until green) and the jax_likelihood_functions/interferometer script.

[Jammy2211]

InputPotential fix wave shipped — 2026-07-17

Three PRs (all pending-release, validated end-to-end by the interferometer recovery script):

• PyAutoGalaxy#508 — extrapolate="zero" on InputPotential/InputDeflections (+ fill on LinearNDInterpolatorExt): localized corrections now vanish outside their mesh instead of nearest-smearing spurious deflections. 1002 tests.
• PyAutoLens#626 — Marquardt-scaled LM damping (μ·diag(H) instead of μ·I): visibility-weighted curvatures (~1e11) gave μ·I only ~4 usable decades before the 1e15 cap, stalling the interferometer LM ("damping exceeded"); scale-invariant damping eliminates the stall. Both iterative engines now use zero-fill correction profiles. 439 tests.
• autolens_workspace_test#177 — the validation script: sparse ≡ dense evidence, one-shot recovery corr 0.35 / peak 0.13″ with the arc dpsi_mask, iterative mechanics (cost progress, gauge to 1e-6, finite evidence). Smoke-wired.

Diagnosis trail closed out: the stall decomposed into (1) nearest-extrapolation deflection smearing [fixed], (2) scale-blind LM damping [fixed]. What remains is a validation-regime fact, not an engine defect: at smoke scale the source mesh cannot fit the visibilities to noise level (best pure-source χ²/dof ≈ 4e3), so iterative corrections absorb source-model error — iterative recovery certification needs a realistic-uv configuration (B1938-like), recorded as the future campaign alongside the remaining jax_likelihood_functions/interferometer/potential_correction.py script.

[Jammy2211]

Shipped — engineering scope complete (2026-07-17)

All six PRs merged: PyAutoLens#624 (sparse-route joint fit), #625 (iterative LM engine + dpsi_mask), PyAutoGalaxy#508 (zero-fill extrapolation), PyAutoLens#626 (Marquardt LM damping + zero-fill wiring), autolens_workspace_test#177 (validation script) and #178 (JAX likelihood script — sparse ≡ dense ≡ xp=np ≡ xp=jnp ≡ jitted). Completion record at PyAutoMind/complete/2026/07/potential-correction-interferometer.md; worktree cleaned.

The one remaining validation tier is filed as PyAutoMind/draft/research/autolens/potential_correction_realistic_uv_campaign.md: iterative δκ-recovery certification on a realistic (B1938-like / ALMA-like) uv configuration, where the source model can reach χ²/dof ≈ 1 — smoke-scale synthetics structurally cannot certify it. This issue can be closed; the campaign gets its own issue when picked up.