Click to expand starting prompt
Streaming visibilities for memory efficiency on the sparse interferometer path
Type: feature
Target: PyAutoArray
Repos:
- PyAutoArray
- PyAutoGalaxy
Themes:
- interferometer
- sparse-operator
- memory
Difficulty: medium
Autonomy: supervised
Priority: medium
Status: draft
Consequence: glance
Witness: a fit built from chunked visibilities (Interferometer.from_stream / apply_sparse_operator_streamed) matches the in-memory apply_sparse_operator fit to 1e-8 in chi_squared, model image, residual map and dirty image; fast_chi_squared and noise_normalization read precomputed scalars when present; a pixelization-only FitInterferometer allocates no N_vis arrays per evaluation.
Review-minutes: 10
Source: GitHub Discussion https://github.com/orgs/PyAutoLabs/discussions/13
("Streaming visibilities for memory efficiency", Ideas & Proposals, HRSAstro, 2026-09-15).
Reference implementation: https://github.com/HRSAstro/pyuvimage src/pyuvimage/streaming.py
(TermsAccumulator / accumulate_sparse_terms / stub_dataset_from_terms) and
tests/test_streaming.py::test_streamed_fit_matches_the_in_memory_sparse_fit.
Request (verbatim from the discussion)
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
-
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.
-
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.
-
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
Overview
GitHub Discussion https://github.com/orgs/PyAutoLabs/discussions/13 (HRSAstro, "Streaming visibilities for memory efficiency") shows the sparse interferometer path (
Interferometer.apply_sparse_operator+InversionInterferometerSparse) removes N_vis from the linear algebra but not from the process: the visibility arrays stay resident and every likelihood evaluation still does O(N_vis) work infast_chi_squaredterm 3,FitInterferometer.noise_normalization, and autogalaxy'sprofile_visibilities/profile_subtracted_visibilitiesallocations. This task is Phase 1: carry the two data-only scalars on the sparse operator and read them per likelihood call, skip the N_vis allocations for pixelization-only fits, and add the chunked accumulation primitive (sparse_terms_from_chunks) verified equal to the in-memory path. Phase 2 (follow-up prompt, filed at ship) is the array-freeInterferometer.from_streamdataset, which also needs a save/aggregator/visualizer contract.Plan
SparseTermsrecord of per-visibility sums (W~, dirty image, dirty beam, sum of weights, sum d^2/sigma^2, sum log 2 pi sigma^2, n_vis) with field-wise addition, and one chunked accumulatorsparse_terms_from_chunksthat produces it from(uv_wavelengths, data, noise_map)chunks with the real/imag noise check run per chunk.data_termandnoise_normalizationonInterferometerSparseOperator;apply_sparse_operatorcomputes them once;fast_chi_squaredandFitInterferometer.noise_normalizationread them when present and valid, otherwise reduce as today (identical numbers).Interferometer.apply_sparse_operator_from_chunksas the public seam for streamed construction (arrays retained for now).profile_visibilities/profile_subtracted_visibilitieswhen there is no non-linear light profile on the sparse path and passdata=Noneto the inversion interface so the scalar is used.Detailed implementation plan
Affected Repositories
Branch Survey
Suggested branch:
feature/interferometer-streaming-visibilitiesWorktree:
~/Code/PyAutoLabs-wt/interferometer-streaming-visibilities/Implementation Steps (PyAutoArray)
autoarray/inversion/inversion/interferometer/inversion_interferometer_util.py: add frozen dataclassSparseTerms(nufft_precision_operator(2Ny,2Nx),dirty_image_native(Ny,Nx),dirty_beam_native,sum_weights,data_term,noise_normalization,n_vis;__add__field-wise). Addsparse_terms_from_chunks(chunks, *, real_space_mask, transformer_class, method, eps, chunk_size, chunk_k, use_jax, show_progress): per chunk run thenp.allclose(noise.real, noise.imag)check (lifted fromdataset.py:336-364into a shared helper_check_noise_real_imag_equal), sumnufft_precision_operator_from(...)(linear: ifftshift + Nyquist zeroing commute with the sum, cf.chunked_matches_one_shot), build a per-chunk transformer and sumimage_from(d_r/s_r^2 + 1j d_i/s_i^2).nativeandimage_from(w), accumulate scalars in float64.InterferometerSparseOperator(frozen dataclass ~L858) with optionaldata_term: float | None = None,noise_normalization: float | None = None; thread throughfrom_nufft_precision_operator(..., data_term=None, noise_normalization=None); addfrom_sparse_terms(terms, *, real_space_mask, batch_size)slimming the native dirty image viaArray2D(values=..., mask=mask).slim(operator stores dirty_image slim,dataset.py:414-418).autoarray/dataset/interferometer/dataset.pyapply_sparse_operator(L229-427): compute both scalars once from the resident arrays and pass them in; keep the existing W~/dirty-image construction (NUFFT plan reuse); use the shared noise-check helper. AddInterferometer.apply_sparse_operator_from_chunks(chunks, **kwargs)returning the same shape of dataset asapply_sparse_operator.autoarray/inversion/inversion/interferometer/abstract.pyfast_chi_squared(L162-200): term 3 usessparse_operator.data_termonly when the interfacedata is None; otherwise reduce as today. Scalar is a Python float, jit-safe.autoarray/fit/fit_interferometer.pynoise_normalization(L126): preferdataset.sparse_operator.noise_normalizationwhen present, else existing call.autoarray/inversion/inversion/dataset_interface.py: allowdata=Nonewhensparse_operatorcarriesdata_term; verify by grep that sparse inversion classes never readdata(sparse.pyusessparse_dirty_image/operator only).test_inversion_interferometer_util.py(chunks 1/2/3 and uneven == whole dataset rtol 1e-12; scalars ==noise_normalization_complex_fromand direct sum; unequal re/im noise in a later chunk raises);test_dataset.py(apply_sparse_operatorpopulates scalars;apply_sparse_operator_from_chunksoninterferometer_7_lopmatchesapply_sparse_operatorin W~, dirty image, scalars);test_interferometer.py(fast_chi_squaredwithdata=None+ scalar == array path, extendtest__fast_chi_squared; scalar ignored whendatagiven);fit/test_fit_interferometer.py(noise_normalizationreads the scalar; falls back when absent).Implementation Steps (PyAutoGalaxy)
autogalaxy/interferometer/fit_interferometer.py:profile_visibilities(L186-202) returnsNonewhen no non-linear light profile anddataset.sparse_operator is not None(dense path keeps zeros);profile_subtracted_visibilities(L205-209) returnsself.datawhen profile visibilities are None;galaxies_to_inversion(L212-233) passesdata=Nonein that case;model_data(L248-264) andgalaxy_model_visibilities_dicttreat None as zeros lazily (output paths only).test__profile_visibilities__linear_light_only__zeros_without_fourier_transform(L627: sparse case asserts None, novisibilities_fromcall, noVisibilities.zerosallocation; dense case keeps shape assertion);_assert_sparse_fit_matches_denseusers (L547, L586) unchanged at rel 1e-8; light-profile case asserts interfacedatais not None. Leaveautogalaxy/operate/image.pyzeros calls (output paths). Runtest_autolens/interferometerfor impact (no edits expected).apply_sparse_operator_from_chunkswith the discussion's memory table and a Phase 2 pointer; release-notes entries in both repos. File Phase 2 promptdraft/feature/autoarray/interferometer_from_stream_array_free_dataset.md(array-freeInterferometer.from_stream, transformer-less construction, save_attributes/aggregator/visualizer contract, light-profile identity sum|d-p|^2/s^2 = data_term - 2 i_p^T d~ + i_p^T W~ i_p). Reply to Discussion Feature/phase removal #13 via /community after human approval.Verification
pytest test_autoarray/inversion/inversion/interferometer test_autoarray/dataset/interferometer test_autoarray/fit; thenpytest test_autogalaxy/interferometerandtest_autolens/interferometerwith the worktree PyAutoArray installed.interferometer_7_lopand a 1e5-visibility synthetic dataset,apply_sparse_operator_from_chunks(4096-vis chunks) vsapply_sparse_operator: W~ <= 2e-16, dirty image <= 5e-16,log_evidencerel <= 1e-8 viaag.FitInterferometerpixelization-only;tracemallocshows no N_vis-sized allocation perfigure_of_meriton the sparse pixelization-only fit.jax.jitof the sparseFitInterferometer.figure_of_meritcompiles and matches numpy.Key Files
autoarray/inversion/inversion/interferometer/inversion_interferometer_util.py—SparseTerms,sparse_terms_from_chunks, operator fieldsautoarray/dataset/interferometer/dataset.py—apply_sparse_operator,apply_sparse_operator_from_chunksautoarray/inversion/inversion/interferometer/abstract.py—fast_chi_squaredautoarray/fit/fit_interferometer.py—noise_normalizationautoarray/inversion/inversion/dataset_interface.py—data=Noneautogalaxy/interferometer/fit_interferometer.py— skip N_vis allocationssrc/pyuvimage/streaming.py,tests/test_streaming.pyOriginal Prompt
Click to expand starting prompt
Streaming visibilities for memory efficiency on the sparse interferometer path
Type: feature
Target: PyAutoArray
Repos:
Themes:
Difficulty: medium
Autonomy: supervised
Priority: medium
Status: draft
Consequence: glance
Witness: a fit built from chunked visibilities (
Interferometer.from_stream/apply_sparse_operator_streamed) matches the in-memoryapply_sparse_operatorfit to 1e-8 in chi_squared, model image, residual map and dirty image;fast_chi_squaredandnoise_normalizationread precomputed scalars when present; a pixelization-onlyFitInterferometerallocates no N_vis arrays per evaluation.Review-minutes: 10
Source: GitHub Discussion https://github.com/orgs/PyAutoLabs/discussions/13
("Streaming visibilities for memory efficiency", Ideas & Proposals, HRSAstro, 2026-09-15).
Reference implementation: https://github.com/HRSAstro/pyuvimage
src/pyuvimage/streaming.py(TermsAccumulator / accumulate_sparse_terms / stub_dataset_from_terms) and
tests/test_streaming.py::test_streamed_fit_matches_the_in_memory_sparse_fit.Request (verbatim from the discussion)
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
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.
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.
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
🤖 Generated with Claude Code