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

[Jammy2211]

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