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

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.

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).