feat: route MGE-only interferometer fits through the W~ sparse operator

[Jammy2211]

Overview

An MGE-only FitInterferometer always takes the dense transform_mapping_matrix path, even after apply_sparse_operator(), because factory.py:204-208 turns the sparse operator off for func-list-only inversions. On the A100 the W~ steps take 1.9–13.5 ms against 0.86–23.8 s for the dense chain, and the dense path runs out of memory at alma and above (autolens_profiling#308). This task switches that guard off. It also fixes a seam in the sparse data vector: it reads a dirty image built from the unsubtracted visibilities, so fits that also have regular light profiles get the wrong data vector.

Plan

1. PyAutoArray: let func-list-only interferometer inversions use the sparse operator. Imaging is unchanged.
2. PyAutoArray: let an inversion dataset carry a dirty image for the subtracted data. The sparse data vector uses it when it is present and falls back to the cached one otherwise.
3. PyAutoGalaxy and PyAutoLens: when the fit has regular light profiles and a sparse operator, compute the dirty image of the profile-subtracted visibilities (one adjoint NUFFT) and pass it in. Pure linear fits keep the cached image at zero extra cost.
4. Tests:
   • MGE-only sparse vs dense parity, with and without a regular light profile.
   • A class-routing assertion.
   • A JAX-vs-NumPy check.
5. Ship the three library PRs, stacked PyAutoArray → PyAutoGalaxy → PyAutoLens with the same branch name.
6. Workspace follow-on in autolens_profiling:
   • Re-run the MGE breakdown cell on the library path (CPU and A100).
   • Update the ("interferometer","mge", alma+) VRAM rows.
   • Start only after PR feat(DatasetModel): add grid_rotation_angle for multi-band rotation #312 has merged, which releases the repo claim. It can wait for this release or run against library mains.

Detailed implementation plan

Affected Repositories

• PyAutoArray (primary)
• PyAutoGalaxy
• PyAutoLens
• autolens_profiling (workspace phase, after PR feat(DatasetModel): add grid_rotation_angle for multi-band rotation #312)

Branch Survey

Repository	Current Branch	Dirty?
./PyAutoArray	main	clean
./PyAutoGalaxy	main	clean
./PyAutoLens	main	clean
./autolens_profiling	main	untracked dataset/abell_1201 (unrelated); claimed by interferometer-mge-breakdown until #312 merges

Suggested branch: feature/interferometer-mge-w-tilde-route

Context

autoarray/inversion/inversion/factory.py:204-208 turns the sparse operator off whenever every linear object is a func-list. As a result, an MGE-only FitInterferometer always takes the dense transform_mapping_matrix path, even after apply_sparse_operator().

The #308 breakdown measured what that costs:

• W~ steps vs dense chain on the A100: 1.9–13.5 ms against 0.86–23.8 s.
• W~ vs dense on CPU at alma: 28.6 ms against 36.6 s.
• The dense path runs out of A100 memory at alma and above.

A code survey found that InversionInterferometerSparse already handles zero mappers:

• data_vector is mapping_matrix.T @ dirty_image.
• _curvature_matrix_func_list_and_mapper uses zeros when there is no mapper and then adds the func–func blocks from curvature_matrix_func_list_from.
• The solver and preconditioning code already checks has(Mapper).
• mapped_reconstructed_operated_data_dict is overridden as one NUFFT per linear object.

No test asserts that MGE-only fits take the dense path.

The survey also found a correctness seam, which you chose to fix in this task. The sparse data_vector reads sparse_operator.dirty_image, which is built from the ORIGINAL visibilities (dataset/interferometer/dataset.py:409). FitInterferometer in both galaxy and lens swaps in profile_subtracted_visibilities but passes the operator through unchanged. So any sparse fit that also has regular (non-linear) light profiles gets a wrong D, and its chi-squared terms become inconsistent with each other. Mixed mapper + MGE fits already hit this today.

The imaging version of the seam stays in its own existing prompt, draft/bug/autoarray/sparse_inversion_ignores_profile_subtracted_image.md.

Detailed plan

Branch feature/interferometer-mge-w-tilde-route in PyAutoArray, PyAutoGalaxy and PyAutoLens. Worktree ~/Code/PyAutoLabs-wt/interferometer-mge-w-tilde-route/, created by /start_library. The conflict guard passes for all three.

PyAutoArray

• inversion/inversion/factory.py (interferometer branch, around lines 202-236): remove the all(isinstance(..., AbstractLinearObjFuncList)) guard, so dataset.sparse_operator is not None alone selects sparse. _use_interferometer_numba already returns False when func-lists are present (:288-291), so MGE-only goes to InversionInterferometerSparse. Leave the imaging guard at :128-160 alone and add a comment saying why it differs.
• inversion/inversion/dataset_interface.py: add an optional kwarg sparse_dirty_image=None, documented as the dirty image of data for the sparse operator. It is needed when data differs from the dataset the operator was built from. Do not reuse the name dirty_image, because Interferometer.dirty_image (dataset.py:550) is a different quantity used in plots.
• inversion/inversion/interferometer/sparse.py:92: data_vector uses getattr(self.dataset, "sparse_dirty_image", None) when it is not None, and falls back to self.dataset.sparse_operator.dirty_image otherwise. This is the only reader (the numba subclass inherits it). Update the docstring.
• Tests in test_autoarray/inversion/inversion/interferometer/test_interferometer.py, reusing _sparse_parity_setup and _assert_sparse_matches_mapping (:453-528):
  • test__interferometer_sparse_operator__func_list_only__identical_to_mapping, with a [MockLinearObjFuncList] parity check and an assertion that the class is InversionInterferometerSparse.
  • ..._sparse_dirty_image_override__used_by_data_vector, which runs DatasetInterface with subtracted data plus sparse_dirty_image and checks it against dense to 1e-10.
  • Extend the JAX-vs-NumPy test (:751) to the func-list-only case.
• Add a factory test in test_factory.py: MGE-only on a dataset with a sparse operator gives the sparse class.

PyAutoGalaxy (autogalaxy/interferometer/fit_interferometer.py, galaxies_to_inversion)

• If self.dataset.sparse_operator is not None and any galaxy has a regular light profile, pass sparse_dirty_image=transformer.image_from(visibilities=d.real*σr⁻² + 1j*d.imag*σi⁻², xp=...) with d = profile_subtracted_visibilities. The regular-light-profile test is cls_list_from(cls=LightProfile, cls_filtered=LightProfileLinear), the same check as galaxy.py:248. The formula mirrors apply_sparse_operator. Pass .array as that site does.
• The branch is decided structurally in Python, so it is JIT-safe.
• Check that transformer.image_from accepts or dispatches xp for jax. If it does not, use linearity instead: sparse_operator.dirty_image − image_from(profile_visibilities weighted). The executor verifies which applies.
• Test in test_autogalaxy/interferometer/test_fit_interferometer.py: Sersic (regular) plus a linear Gaussian on a sparse dataset against the dense dataset, comparing log_likelihood and log_evidence to 1e-8. Control: all-linear.

PyAutoLens (autolens/interferometer/fit_interferometer.py:134, tracer_to_inversion)

• Apply the same change, with the regular-profile test over self.tracer.galaxies. Factor a small shared helper into autogalaxy (for example in autogalaxy/interferometer/fit_interferometer.py or the operate utils) so the formula lives in one place, and call it from lens.
• Test in test_autolens/interferometer/test_fit_interferometer.py: lens Sersic plus source MGE, sparse vs dense, to 1e-8.

Checks before shipping

• Full unit suites for the three libraries.
• Heart smoke via /smoke_test on the interferometer MGE and pixelization scripts in autolens_workspace, and on autolens_workspace_test interferometer sparse scripts.
• Workspace impact grep for apply_sparse_operator with MGE: the autolens_workspace modeling.py:322-324 docs prompt becomes accurate once this lands, so note it on that prompt.

Workspace phase (autolens_profiling, later, /start_workspace)

• Re-run scripts/interferometer/likelihood_breakdown/mge.py on the library path: CPU sma/alma, and A100 alma/alma_high/jvla via the existing submits.
• Witness: the library-path log_likelihood matches dense within 1e-6 nats where dense runs, and jax.jit(FitInterferometer) no longer runs out of memory.
• Update the VRAM rows in scripts/misc/vram/config.py, add a results addendum to the feat: add aa.interp_2d (NumPy + JAX bilinear interpolation) #308 note, and regenerate the README.

Verification

1. Unit tests: the new parity tests pass, and the full suites for PyAutoArray, PyAutoGalaxy and PyAutoLens stay green.
2. Local script: the alma MGE-only fit with apply_sparse_operator() has inversion class InversionInterferometerSparse and a log_likelihood equal to the dense path within 1e-6 nats. With a regular Sersic added, sparse and dense still agree, which is the seam fix.
3. A100, in the workspace phase: the jitted library fit runs at alma, alma_high and jvla.

Risks

• transformer.image_from under JAX (see the PyAutoGalaxy step).
• Stacked PRs across three libraries: merge them in order, PyAutoArray first.
• log_likelihood still does one NUFFT per linear object for the model visibilities, which is O(N_vis) but fits in memory. Measure it in the workspace phase.

Original Prompt

Click to expand starting prompt

Interferometer likelihood campaign: route MGE-only interferometer fits through the W~ sparse operator

Type: feature
Target: PyAutoArray
Repos:

• PyAutoArray
• autolens_profiling
  Themes:
• interferometer
• mge
• sparse-operator
• jax-gpu
  Difficulty: medium
  Autonomy: supervised
  Priority: high
  Status: draft
  Consequence: glance
  Witness: with apply_sparse_operator() applied, an MGE-only FitInterferometer on the alma dataset takes the func-list W~ path (inversion class is the sparse/W~ interferometer inversion, not InversionInterferometerMapping), its log_likelihood matches the dense path within 1e-6 nats, and jax.jit(FitInterferometer) no longer OOMs the A100 at alma/alma_high/jvla.
  Review-minutes: 10
  Epic: interferometer-likelihood-campaign

Source: autolens_profiling/results/notes/interferometer_mge_breakdown_2026_09.md, lever 1
(autolens_profiling#308).

Why

autoarray/inversion/inversion/factory.py:202-208 sets use_sparse_operator = False when
every linear object is an AbstractLinearObjFuncList, so an MGE-only interferometer fit
always pays the O(N_vis * n) dense NUFFT path even after apply_sparse_operator(). The
func-list W~ blocks already exist and are used by mixed mapper + MGE inversions
(inversion_interferometer_util.py:1466 operated_matrix_slim_from, :1628
curvature_matrix_func_list_from).

Measured on the #308 breakdown cell's measurement-only W~ arm (fp64, nufftax 0.6.1):
dense chain from the mapping matrix vs W~ chain — A100 sma 855 ms -> 2.40 ms, alma 937 ms
-> 1.90 ms, alma_high 3.81 s -> 3.35 ms, jvla 23.62 s -> 13.53 ms; CPU alma 36.58 s ->
28.6 ms. F~ vs dense F agrees to <= 1.6e-11 rel; figure of merit to <= 7.1e-7 nats. The
library path cannot run above sma on the A100 at all (65.9 GB alma one-shot transform).

What

• Let the factory keep the sparse operator on for func-list-only interferometer inversions
  when dataset.sparse_operator is set (imaging keeps today's behaviour unless measured).
• Data vector from the cached dirty image; the fast_chi_squared identity with per-dataset
  constants, so no transformed mapping matrix is formed.
• Check the mapped visibilities / residual map path (plots, FitInterferometer attributes
  that need transformed_mapping_matrix) stays correct — likely a lazy dense fallback.
• Re-run scripts/interferometer/likelihood_breakdown/mge.py (library path) CPU + A100 and
  update the VRAM table rows ("interferometer", "mge", alma+) in
  scripts/misc/vram/config.py.

Watch

• The one-off operator build (1.1-8.2 s A100, 18.5 s CPU alma) must be done once per
  dataset, not per likelihood.
• After this lands PDIP becomes the largest W~ step on the A100 (1.46 of 2.32 ms at alma).

[Jammy2211]

Library PRs Created — 2026-09-26

• PyAutoArray: feat: route MGE-only interferometer fits through the W~ sparse operator (#575) #576 (merge first)
• PyAutoGalaxy: fix: sparse interferometer fits with ordinary light profiles use the profile-subtracted dirty image (PyAutoArray#575) PyAutoGalaxy#629
• PyAutoLens: fix: al.FitInterferometer sparse path uses the profile-subtracted dirty image (PyAutoArray#575) PyAutoLens#750

Deviations from the plan, both within its intent:

• The corrected dirty image is computed as d~ − W~ i_p (one FFT convolution, about 2 ms at alma) rather than an adjoint NUFFT (about 4 s at alma on CPU). The PR carries the diagnosis.
• The check for ordinary light profiles looks inside a Basis, so MGE-only fits pay nothing extra.

Incidental: TransformerDFT.image_from now honours xp, so it no longer crashes under jit.

Workspace follow-on is pending: re-run the autolens_profiling library-path breakdown and update the VRAM rows, after autolens_profiling#312 merges.

Heart RED override (development only)

The human authorized the AUTONOMY.md "Human override for Heart RED (development only)" for this task (PyAutoArray#575, branch feature/interferometer-mge-w-tilde-route) on 2026-09-26. When asked how to proceed with the three library PRs, they chose "Override, open PRs". The override covers commit, push and opening pending-release PRs only. Merging needs a separate /prm with every required check green. This branch does not claim to fix Heart.

Exact RED reasons at ship time:

• release validation FAILED (stage integrate)
• workspace validation not passing (4 failed, cloud#35579888156: autolens notebooks/cluster/modeling.ipynb, autolens notebooks/weak/a2744.ipynb, autolens scripts/cluster/modeling.py, +1 more)
• manifest drift: hub organism blurb (organs present) — 7 mismatch(es) vs PyAutoMind/repos.yaml
• manifest drift: organism-map blocks (generated) — 1 mismatch(es) vs PyAutoMind/repos.yaml

Branch gates passed: 24/24 targeted interferometer smoke scripts pass on the branch (autolens_workspace 11, autogalaxy_workspace 5, autolens_workspace_test 5, autogalaxy_workspace_test 3: MGE, pixelization, linear light profiles, sparse-operator JAX assertions and likelihood pins). autolens_workspace_test/.../jax_likelihood/mge.py first failed identically on the branch and on the mains because the local dataset was stale (2026-08-06, older than the 09-18 simulator change); after regenerating it, the script passes on both. Unit suites: PyAutoArray 1706 passed, PyAutoGalaxy 1240 passed, PyAutoLens 757 passed + 1 xfailed. My own diff review is in-session only; the human reviews at merge.

[Jammy2211]

Workspace PR Created — 2026-09-26

• autolens_profiling: profiling(PyAutoArray#575): interferometer MGE breakdown on the library W~ route — CPU + A100 autolens_profiling#313 (merge after feat: route MGE-only interferometer fits through the W~ sparse operator (#575) #576, PyAutoGalaxy#629, PyAutoLens#750)

The library route re-ran on the branch, both locally and on RAL. The full MGE likelihood matches dense to 0.0 nats wherever dense runs.

A100 per-call times, library path vs dense:

• alma: 2.69 ms, where dense ran out of memory (chunked dense: 940 ms).
• jvla: 16.7 ms, against 24.2 s chunked dense.

These numbers are with the zero-image NUFFT fix, now folded into PyAutoGalaxy#629 (36d1b436) and PyAutoLens#750 (ac333b17).

Details: the "Library W~ route (2026-09-26)" section of results/notes/interferometer_mge_breakdown_2026_09.md.

[Jammy2211]

Shipped

Merged 2026-09-26, library-first:

• PyAutoArray: feat: route MGE-only interferometer fits through the W~ sparse operator (#575) #576
• PyAutoGalaxy: fix: sparse interferometer fits with ordinary light profiles use the profile-subtracted dirty image (PyAutoArray#575) PyAutoGalaxy#629
• PyAutoLens: fix: al.FitInterferometer sparse path uses the profile-subtracted dirty image (PyAutoArray#575) PyAutoLens#750
• autolens_profiling: profiling(PyAutoArray#575): interferometer MGE breakdown on the library W~ route — CPU + A100 autolens_profiling#313

What changed

• MGE-only interferometer inversions now take the W~ sparse operator when apply_sparse_operator() is applied (the factory guard is removed).
• The sparse data vector takes DatasetInterface.sparse_dirty_image. Galaxy/lens fits with ordinary light profiles pass d~ − W~ i_p. This fixes a silent bug: the unsubtracted dirty image gave a 23% wrong D at sma (logL −26.77 vs −24.34 dense).
• profile_visibilities skips the zero-image NUFFT when no ordinary light is present.
• TransformerDFT.image_from now honours xp (incidental).

Results

• A100 full library pipeline: alma 2.69 ms, jvla 16.7 ms, where the dense library path OOMs (chunked dense 940 ms / 24.2 s).
• CPU alma: 34.7 ms vs 13.1 s dense.
• |ΔlogL| = 0.0 wherever dense runs; MGE + lens Sersic matches exactly.

Known limits

• CPU alma_high OOMs in the apply_sparse_operator build (10.8 GB).
• jvla vmap64 still OOMs (41 GiB; likely the W~ curvature FFT, inferred).

Libraries are merged but not yet released (pending-release).