Streaming visibilities for memory efficiency

[HRSAstro]

Problem

On the sparse path (Interferometer.apply_sparse_operator + InversionInterferometerSparse) the inversion itself is independent of the number of visibilities, but the dataset and the fit are not: every visibility stays resident for the whole run, and every likelihood evaluation still does O(N_vis) work. Building the sparse operator by streaming the visibilities once, and carrying the handful of per-visibility sums the fit needs alongside it, makes both the resident memory and the per-likelihood cost independent of N_vis. We have this working in pyuvimage on top of autoarray's sparse inversion, verified against the in-memory sparse fit to 1e-8, and it would be a natural addition upstream.

What the current sparse path holds and computes

Resident for the whole fit, per visibility (as Interferometer is built today):

data (complex128, 16 B), noise_map (complex128, 16 B), uv_wavelengths (2 x float64, 16 B) — 48 B/visibility minimum, ~10 GB at 2e8 visibilities, before from_fits load temporaries and the NUFFT plan.

Per likelihood evaluation, all O(N_vis) in time and transient memory, even though the sparse inversion has already reduced the data to W~ and the dirty image:

FitInterferometer.profile_visibilities allocates an N_vis array (Visibilities.zeros when there are no light profiles) and profile_subtracted_visibilities = data - profile_visibilities allocates another (autogalaxy fit_interferometer.py).
fast_chi_squared recomputes sum(d_r^2/sigma_r^2) + sum(d_i^2/sigma_i^2) over the full arrays on every call (inversion/interferometer/abstract.py), a data-only constant.
noise_normalization recomputes sum(log(2 pi sigma^2)) over the full noise map on every call.
apply_sparse_operator itself needs the whole dataset in memory to form data.real * noise_map.real**-2 + 1j * ... and the np.allclose(noise_real, noise_imag) check.

So the sparse path removes N_vis from the linear algebra but not from the process: RSS scales with N_vis and the fit does several full passes over the visibilities per likelihood call.

Proposed Solution

What streaming looks like

Read the visibilities once in chunks (from .npz, FITS, or a memory-mapped source) and accumulate, per chunk, every quantity downstream of the data that is a sum over visibilities:

quantity | accumulated as | consumer
W~ (2Ny, 2Nx) | nufft_precision_operator_from per chunk, summed (it is sum_k w_k cos(...)) | curvature matrix
dirty image F^H W d | adjoint of d_r/sigma_r^2 + i d_i/sigma_i^2 per chunk, summed | data vector
sum d^2/sigma^2 | scalar | fast_chi_squared term 3
sum log 2 pi sigma^2 | scalar | noise_normalization
sum w, dirty beam | adjoint of w per chunk | normalised dirty/residual images

Then discard the chunk. The fit is built from InterferometerSparseOperator.from_nufft_precision_operator(W~, dirty_image) plus the two scalars; chi_squared is s^T F s - 2 s^T D + const with no visibilities touched, and the residual dirty image is dirty(data) - W~ * (M s) via the same FFT-multiply the operator already uses. Nothing per visibility survives the load.

Measured in pyuvimage (peak RSS of the accumulation from a compressed .npz, 400-pixel image, 4096-visibility chunks, fresh process each):

N_vis | peak RSS over baseline | resident if held (48 B/vis, autoarray minimum)
5e5 | 36 MB | 24 MB
1e6 | 36 MB | 48 MB
2e6 | 36 MB | 96 MB
4e6 | 36 MB | 192 MB

Flat across 8x; the in-memory path grows linearly and, on a real 2e8-sample ALMA MFS cube, is the difference between a laptop and a node

Parity: the streamed fit matches the in-memory sparse fit to 1e-8 in chi^2, model image, residual map and dirty image (W~ to 2e-16, dirty image to 5e-16). It extends directly to per-channel (cube) fits, since the MFS terms are the sum of the channel terms, and to phase-centre shifts applied chunk by chunk.

Suggested shape upstream

1. Interferometer.apply_sparse_operator (or a new Interferometer.from_stream(...) / apply_sparse_operator_streamed(...)) that accepts an iterable of (uv_wavelengths, data, noise_map) chunks, accumulates W~, the dirty image, sum d^2/sigma^2 and sum log 2 pi sigma^2, and returns a dataset that carries those scalars instead of the visibility arrays.

2. fast_chi_squared and noise_normalization read the precomputed scalars when present, rather than reducing the full arrays per call. This alone removes two O(N_vis) reductions from every likelihood evaluation on the current in-memory sparse path as well.

3. FitInterferometer.profile_subtracted_visibilities skipped when there are no light profiles, so a pixelization-only fit does not allocate two N_vis arrays per evaluation.

Item 2 and 3 are independent of streaming and benefit every sparse-path user today.

Reference implementation

src/pyuvimage/streaming.py in https://github.com/HRSAstro/pyuvimage — TermsAccumulator / accumulate_sparse_terms (the per-chunk sums), stub_dataset_from_terms (the Interferometer built from the operator alone), and tests/test_streaming.py::test_streamed_fit_matches_the_in_memory_sparse_fit for the parity check.

Alternatives Considered

No response

[Jammy2211]

Thanks for this — the write-up made it very easy to act on, and the pyuvimage accumulator and parity test were a big help as a reference. A first phase has now landed on main across the three libraries (not yet on PyPI; it will go out with the next release):

Items 2 and 3 of your list (every sparse-path user):

• Interferometer.apply_sparse_operator now reduces Σ d²/σ² and Σ log 2πσ² once and caches them on the operator (InterferometerSparseOperator.data_term / .noise_normalization). fast_chi_squared and FitInterferometer.noise_normalization read the scalars, so neither does an O(N_vis) pass per likelihood call any more (PyAutoArray #589).
• A pixelization-only fit (no non-linear light profile) builds its inversion with data=None, so profile_visibilities / profile_subtracted_visibilities are never formed on the likelihood path — in both autogalaxy.FitInterferometer (#637) and autolens.FitInterferometer (#757). On a 1e5-visibility NUFFT dataset the peak allocation per figure_of_merit went from 7.6 MB to 3.2 MB, and it stays at 3.2 MB at 4e5 (it was 22 MB), with the likelihood bit-identical.
• Diagnostics that need the visibilities (subplot_of_mapper etc.) take fit.inversion_with_data, which carries fit.data and shares the reconstruction.

Item 1, first half: the per-chunk sums are in as aa.SparseTerms (W̃, dirty image, dirty beam, Σw, the two scalars, n_vis; field-wise + so channel terms sum to the MFS terms), built by sparse_terms_from_chunks(chunks, real_space_mask=..., ...) over (uv_wavelengths, data, noise_map) chunks, with InterferometerSparseOperator.from_sparse_terms(terms, real_space_mask=...) and Interferometer.apply_sparse_operator_from_chunks(chunks, ...) on top. Parity against the one-shot operator: W̃ to 6e-16, dirty image to 3e-15, scalars to 3e-16 relative. This is the same accumulation as your TermsAccumulator, so if you find any divergence from pyuvimage on real data I'd like to hear about it.

Not yet done — the array-free dataset itself. apply_sparse_operator_from_chunks still returns an Interferometer that holds the arrays and builds a TransformerNUFFT, and the save/aggregator/visualizer paths still expect data, noise_map and uv_wavelengths. So resident memory is not yet independent of N_vis upstream, and your stub_dataset_from_terms is still the way to get there for now. The plan for phase 2 is an Interferometer.from_stream(chunks, real_space_mask, ...) with no transformer, a persisted SparseTerms for saving/reloading, dirty-image visualization from the terms, the light-profile identity Σ|d−p|²/σ² = data_term − 2 i_pᵀd̃ + i_pᵀW̃i_p so light-profile fits also run array-free, and per-channel cubes. No timeline promised yet; I'll post here when it's in.

Thanks again — please shout if the API shape above doesn't fit how pyuvimage wants to feed it.

[Jammy2211]

Follow-up as promised: the array-free dataset is now in on main across PyAutoArray, PyAutoGalaxy and PyAutoLens (still pending a PyPI release — I'll note here when it's tagged).

What you can do now

dataset = aa.Interferometer.from_stream(chunks, real_space_mask=mask, transformer_class=aa.TransformerNUFFT, chunk_size=65536)

chunks is any iterable of (uv_wavelengths, data, noise_map) triples (a generator reading from disk is fine). The result has no data, noise_map, uv_wavelengths or transformer — just the accumulated SparseTerms and the sparse operator with its cached scalars — and dataset.is_array_free tells you which kind you have. On it:

• pixelization and linear-light-profile fits (ag. / al.FitInterferometer, AnalysisInterferometer, jax.jit) give the same log-evidence as the in-memory sparse fit to machine precision;
• save_attributes writes the terms into dataset.fits (W̃, dirty image, dirty beam, plus a lossless float64 extension for the scalars — FITS header cards truncate float64, which bit us in review), and the aggregator rebuilds the dataset from them with the operator re-attached, so reload reproduces the evidence exactly;
• the visualizers run end to end: visibility-space panels are skipped and the dirty panels switch to natural-weighted image / beam / model / residual from the terms (dataset.dirty_image_natural, dataset.dirty_beam, fit.dirty_model_image_natural, fit.dirty_residual_map_natural), under the same filenames.

So stub_dataset_from_terms should no longer be needed for this path.

Measured, since you asked what it buys (laptop CPU, 400-px mask at 0.05"/pix, NUFFT, synthetic visibilities):

N_vis	in-memory apply_sparse_operator	from_stream, chunk 65536
1e5	2.1 GB peak	0.85 GB, 17 s
1e6	out of memory (3.1 GB single allocation in the one-shot precision build)	1.5 GB, 16 s
4e6	out of memory	1.5 GB, 52 s
1.6e7	—	1.5 GB, 174 s
5e7	—	1.8 GB, 539 s

Streaming is linear at ~11 s per million visibilities and flat in memory, so 2e8 extrapolates to roughly 40 minutes and under 2 GB on a laptop. Two practical notes from that: chunk size matters a lot (65536 is 5× faster than the 4096 in your table, because each chunk pays a fixed NUFFT setup), and there is a further 2–3× available in the accumulator from avoiding a JAX recompile per chunk, which I'll chase in a profiling campaign.

Also in since the first draft of this note

• Ordinary (non-linear) light profiles now fit array-free too, via the Σ|d−p|²/σ² = data_term − 2 i_pᵀd̃ + i_pᵀW̃i_p identity you sketched, applied to the total real-space model image so log_likelihood, log_evidence and figure_of_merit all match the in-memory fit to machine precision (1e-15 in our tests) with or without an inversion, under jax.jit. The profile visibilities are never formed.
• Per-channel cubes: MFS terms are the sum of the per-channel terms — mfs = sum(terms_per_channel) then aa.Interferometer.from_sparse_terms(mfs, real_space_mask=mask) — tested equal to a one-shot accumulation and to the in-memory MFS dataset at 1e-12.
• Phase-centre shifts chunk by chunk: from_stream(..., phase_centre=(y0, x0)) (arcsec, our (y, x) convention) multiplies each chunk by exp(+2πi(u l0 + v m0)) so a source at (y0, x0) lands at the image origin; only the dirty image and data term change, the kernel and beam are untouched, and the phase centre is recorded as provenance so shifted and unshifted terms refuse to sum. The sign is pinned by a point-source test.
• The autolens_workspace datacube example gained an array-free variant (interferometer/features/datacube/modeling_array_free.py) that streams each channel, shows the MFS sum and the phase-centre option.

Two things we found on the way that you may hit with pyuvimage-style data: TransformerNUFFT currently ignores a non-zero mask origin (dirty images land displaced relative to the DFT on offset-origin masks; filed), and the provenance cannot yet tell apart masks with the same geometry but different masked pixels — keep one mask per cube for now.

If you get a chance to point pyuvimage's real data at from_stream I'd very much like to know whether the accumulation matches your TermsAccumulator and how the wall time looks on the 2e8 cube.