Skip to content

feat: cache sparse-path scalars + chunked SparseTerms accumulation (#588) - #589

Merged
Jammy2211 merged 1 commit into
mainfrom
feature/interferometer-streaming-visibilities
Sep 30, 2026
Merged

Jammy2211 merged 1 commit into
mainfrom
feature/interferometer-streaming-visibilities

Conversation

@Jammy2211

Copy link
Copy Markdown
Collaborator

Summary

Phase 1 of GitHub Discussion https://github.com/orgs/PyAutoLabs/discussions/13 ("Streaming visibilities for memory efficiency", HRSAstro). Closes #588.

On the sparse interferometer path the inversion is independent of N_vis, but every likelihood call still reduced over the full visibility arrays twice: fast_chi_squared term 3 (sum(d_r^2/sigma_r^2) + sum(d_i^2/sigma_i^2), a data-only constant) and FitInterferometer.noise_normalization (sum(log 2 pi sigma^2)). This PR:

  • caches both scalars on InterferometerSparseOperator (data_term, noise_normalization) when apply_sparse_operator builds it, and reads them per likelihood call, with the previous reductions as the fallback (results are bit-identical);
  • adds the chunked accumulation primitive the array-free dataset (Phase 2) will be built on: SparseTerms (per-visibility sums: W~, dirty image, dirty beam, sum of weights, the two scalars, n_vis) with field-wise addition, sparse_terms_from_chunks(...) over (uv_wavelengths, data, noise_map) chunks, and Interferometer.apply_sparse_operator_from_chunks(...);
  • lets DatasetInterface(data=None) express "nothing was subtracted from the data" so a pixelization-only fit (PyAutoGalaxy PR) never touches an N_vis array per evaluation;
  • lifts the real/imag noise-equality check into check_noise_map_real_imag_equal so it runs once per chunk.

Parity witness (1e5 visibilities, 616-pixel mask, 25 chunks of 4096): chunked vs one-shot W~ 5.8e-16 rel, dirty image 3.0e-15 rel, data_term 1.5e-16 rel, noise_normalization 2.7e-16 rel; fast_chi_squared with data=None vs the array path 1.5e-16 rel (exactly 0 on the same operator).

Phase 2 (array-free Interferometer.from_stream, save/aggregator/visualizer contract) is filed as a Mind follow-up. Reference implementation: https://github.com/HRSAstro/pyuvimage src/pyuvimage/streaming.py.

API Changes

Additive only. New public names: aa.SparseTerms, sparse_terms_from_chunks, check_noise_map_real_imag_equal, InterferometerSparseOperator.from_sparse_terms, Interferometer.apply_sparse_operator_from_chunks. InterferometerSparseOperator gains two optional trailing fields (data_term, noise_normalization, default None); from_nufft_precision_operator gains matching keyword args. DatasetInterface accepts data=None (only with a sparse operator carrying data_term); fast_chi_squared and FitInterferometer.noise_normalization read the cached scalars when present. No symbol removed, no default changed.
See full details below.

Test Plan

  • pytest test_autoarray — 1756 passed (jax 0.10.2 importable, jax-marked sparse tests ran)
  • Target suites (inversion/inversion/interferometer, dataset/interferometer, fit, inversion/inversion/test_factory.py) — 181 passed
  • New tests: chunked == one-shot (1/2/3/uneven chunks, rtol 1e-12); scalars == noise_normalization_complex_from / direct sum; unequal re/im noise in a later chunk raises; apply_sparse_operator populates both scalars; apply_sparse_operator_from_chunks matches apply_sparse_operator; fast_chi_squared(data=None) == array path incl. jax; noise_normalization scalar/fallback
  • Parity witness script (numbers above)
  • CI green on unittest 3.12 / 3.13 / nojax
Full API Changes (for automation & release notes)

Added

  • autoarray.SparseTerms (frozen dataclass; nufft_precision_operator, dirty_image_native, dirty_beam_native, sum_weights, data_term, noise_normalization, n_vis; __add__ field-wise) — per-visibility sums accumulated over chunks
  • autoarray.inversion.inversion.interferometer.inversion_interferometer_util.sparse_terms_from_chunks(chunks, *, real_space_mask, transformer_class=None, method="nufft", eps=None, chunk_size=None, chunk_k=2048, use_jax=False, show_progress=False) -> SparseTerms
  • autoarray.inversion.inversion.interferometer.inversion_interferometer_util.check_noise_map_real_imag_equal(noise_map) -> None
  • InterferometerSparseOperator.from_sparse_terms(terms, *, real_space_mask, batch_size=128)
  • Interferometer.apply_sparse_operator_from_chunks(chunks, *, batch_size=128, **accumulator_kwargs) -> Interferometer

Changed Signature

  • InterferometerSparseOperator(..., data_term: Optional[float] = None, noise_normalization: Optional[float] = None) — two optional trailing fields
  • InterferometerSparseOperator.from_nufft_precision_operator(nufft_precision_operator, dirty_image, *, batch_size=128, data_term=None, noise_normalization=None)
  • DatasetInterface(data=None, ...) now permitted when sparse_operator.data_term is set

Changed Behaviour

  • Interferometer.apply_sparse_operator computes data_term and noise_normalization once and stores them on the operator (same values as the per-call reductions)
  • AbstractInversionInterferometer.fast_chi_squared uses sparse_operator.data_term for term 3 when dataset.data is None; raises InversionException if neither is available; otherwise unchanged
  • FitInterferometer.noise_normalization returns sparse_operator.noise_normalization when set and self.noise_map is self.dataset.noise_map; otherwise unchanged
  • AbstractInversionInterferometer.operated_mapping_matrix_list override check reads noise_map.shape[0] instead of data.shape[0] (same value; hardening for data=None)

Migration

  • None required.

Generated by the PyAutoLabs agent workflow.

🤖 Generated with Claude Code

…ion (#588)

Phase 1 of Discussion PyAutoLabs/13. InterferometerSparseOperator carries
data_term and noise_normalization; fast_chi_squared and
FitInterferometer.noise_normalization read them when present (fallback
unchanged, bit-identical). Add SparseTerms, sparse_terms_from_chunks and
Interferometer.apply_sparse_operator_from_chunks (chunked accumulation of
W~, dirty image, dirty beam and the scalars, verified against the
one-shot path at 1e-15). DatasetInterface(data=None) marks unsubtracted
data; the noise real/imag check is a shared helper run per chunk.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
@Jammy2211 Jammy2211 added the pending-release PR queued for the next release build label Sep 29, 2026
@Jammy2211
Jammy2211 merged commit 6d986bd into main Sep 30, 2026
3 checks passed
@Jammy2211
Jammy2211 deleted the feature/interferometer-streaming-visibilities branch September 30, 2026 07:56
Jammy2211 added a commit that referenced this pull request Sep 30, 2026
Resolve the append/append conflict in test_interferometer.py by keeping
both sides: #586's operated_mapping_matrix_list transform-once test and
main's #588/#589 fast_chi_squared data_term tests. interferometer/abstract.py
auto-merged: @cached_property on operated_mapping_matrix_list (#586) sits
alongside main's noise_map-based override shape check and data=None term-3
path (#588/#589).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@Jammy2211 Jammy2211 removed the pending-release PR queued for the next release build label Oct 7, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

feat: sparse-path precomputed terms + chunked accumulation (Discussion #13)

1 participant