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

autoarray/__init__.py

@@ -53,6 +53,7 @@
 from .inversion.inversion.interferometer.mapping import InversionInterferometerMapping
 from .inversion.inversion.interferometer.inversion_interferometer_util import (
     InterferometerSparseOperator,
+    SparseTerms,
 )
 from .inversion.linear_obj.linear_obj import LinearObj
 from .inversion.linear_obj.func_list import AbstractLinearObjFuncList

autoarray/dataset/interferometer/dataset.py

@@ -35,6 +35,7 @@ def disable_jax():
 from autoarray.structures.visibilities import VisibilitiesNoiseMap
 
 from autoarray import exc
+from autoarray.fit import fit_util
 from autoarray.inversion.inversion.interferometer import (
     inversion_interferometer_util,
 )
@@ -333,35 +334,7 @@ def apply_sparse_operator(
         if disable_jax():
             use_jax = False
 
-        noise_map_real = np.asarray(self.noise_map.real)
-        noise_map_imag = np.asarray(self.noise_map.imag)
-
-        if not np.allclose(noise_map_real, noise_map_imag):
-
-            unequal = ~np.isclose(noise_map_real, noise_map_imag)
-
-            denominator = np.maximum(np.abs(noise_map_real), np.abs(noise_map_imag))
-            relative_difference = np.abs(noise_map_real - noise_map_imag) / np.where(
-                denominator == 0.0, 1.0, denominator
-            )
-
-            raise exc.DatasetException(
-                "The sparse operator cannot be applied to this interferometer dataset because its "
-                "noise-map has unequal real and imaginary sigma.\n\n"
-                "The sparse operator's precision operator `W~ = Re(F^H W F)` is built from the "
-                "real-part noise sigma only (see `psf_precision_operator_from`, which passes "
-                "`noise_map_real` to `nufft_precision_operator_from`). That reduction is exact only "
-                "when every visibility has equal real and imaginary sigma "
-                "(`sigma_real == sigma_imag`).\n\n"
-                f"This dataset has {int(np.count_nonzero(unequal))} of {noise_map_real.size} "
-                "visibilities where the real and imaginary sigma differ (maximum relative difference "
-                f"{np.max(relative_difference):.3e}), so the sparse curvature matrix would silently "
-                "disagree with the dense path.\n\n"
-                "Either equalise the real and imaginary noise sigma of every visibility, or fit "
-                "without calling `apply_sparse_operator()` — the dense "
-                "`InversionInterferometerMapping` path handles unequal real and imaginary sigmas "
-                "correctly."
-            )
+        inversion_interferometer_util.check_noise_map_real_imag_equal(self.noise_map)
 
         if nufft_precision_operator is None:
 
@@ -411,10 +384,152 @@ def apply_sparse_operator(
             + 1j * self.data.imag * self.noise_map.imag**-2.0,
         )
 
+        # The two visibility-sum scalars of the likelihood, reduced once here so that neither
+        # `fast_chi_squared` (term 3) nor `FitInterferometer.noise_normalization` reduces over
+        # the visibility arrays on every likelihood call. The expressions are the ones those
+        # two reductions use, so the cached values are bit-for-bit what they would compute.
+        data_term = float(
+            np.sum(self.data.array.real**2.0 / self.noise_map.array.real**2.0)
+            + np.sum(self.data.array.imag**2.0 / self.noise_map.array.imag**2.0)
+        )
+        noise_normalization = float(
+            fit_util.noise_normalization_complex_from(noise_map=self.noise_map.array)
+        )
+
         sparse_operator = inversion_interferometer_util.InterferometerSparseOperator.from_nufft_precision_operator(
             nufft_precision_operator=nufft_precision_operator,
             dirty_image=dirty_image.array,
             batch_size=batch_size,
+            data_term=data_term,
+            noise_normalization=noise_normalization,
+        )
+
+        return Interferometer(
+            real_space_mask=self.real_space_mask,
+            data=self.data,
+            noise_map=self.noise_map,
+            uv_wavelengths=self.uv_wavelengths,
+            transformer_class=lambda uv_wavelengths, real_space_mask: self.transformer,
+            sparse_operator=sparse_operator,
+        )
+
+    def apply_sparse_operator_from_chunks(
+        self,
+        chunks,
+        *,
+        batch_size: int = 128,
+        **accumulator_kwargs,
+    ):
+        """
+        Build the sparse operator by accumulating the visibilities one chunk at a time, rather
+        than from this dataset's resident arrays in one shot.
+
+        This is the chunked counterpart of `apply_sparse_operator`: every quantity the sparse
+        likelihood needs from the visibilities -- the NUFFT precision operator `W~`, the
+        noise-weighted dirty image, `sum(d^2 / sigma^2)` and `sum(log(2 pi sigma^2))` -- is a sum
+        over visibilities, so it is accumulated per chunk by
+        `inversion_interferometer_util.sparse_terms_from_chunks` into a `SparseTerms` record and
+        turned into an `InterferometerSparseOperator` by
+        `InterferometerSparseOperator.from_sparse_terms`. The result equals
+        `apply_sparse_operator()` to summation order (verified at `rtol=1e-12` by the test
+        suite), and the operator carries both scalars, so a fit on it never reduces over the
+        visibility arrays in its likelihood.
+
+        Chunk contract
+        --------------
+        `chunks` is any iterable (a list, or a generator reading from disk) of
+        `(uv_wavelengths, data, noise_map)` triples of `K` visibilities each, where `K` may
+        differ per chunk: `uv_wavelengths` a real `(K, 2)` array in wavelengths, `data` and
+        `noise_map` a `Visibilities`, a complex `(K,)` array or a real `(K, 2)` array of
+        (real, imag) columns. Every chunk must have equal real and imaginary noise sigma
+        (checked per chunk). The chunks must together be exactly this dataset's visibilities
+        (in any order) for the returned dataset to be self-consistent -- this method does not
+        check that.
+
+        Memory
+        ------
+        The accumulation holds one chunk at a time, so its peak memory is set by the chunk
+        size, not the dataset size. Holding the visibilities costs at least 48 B/visibility
+        (`data` and `noise_map` complex128, 16 B each; `uv_wavelengths` 2 x float64, 16 B), before
+        the transformer's own per-visibility state. Measured in pyuvimage on the same
+        accumulation (peak RSS over baseline, 400-pixel image, 4096-visibility chunks, fresh
+        process each; Discussion #13):
+
+        ======  ======================  =========================
+        N_vis   peak RSS of the stream  resident if held (48 B/vis)
+        ======  ======================  =========================
+        5e5     36 MB                   24 MB
+        1e6     36 MB                   48 MB
+        2e6     36 MB                   96 MB
+        4e6     36 MB                   192 MB
+        ======  ======================  =========================
+
+        i.e. flat across 8x, while the in-memory path grows linearly (~10 GB at 2e8
+        visibilities, a real ALMA MFS cube).
+
+        This method returns a dataset of the same shape as `apply_sparse_operator` -- this
+        dataset's `data`, `noise_map`, `uv_wavelengths` and transformer are retained, because
+        the visualizer, `save_attributes`/aggregator and dirty-image outputs still read them.
+        Dropping them (an array-free `Interferometer` built straight from a chunk stream) is a
+        follow-up that needs its own dataset contract; this method is the seam it builds on.
+
+        Parameters
+        ----------
+        chunks
+            The `(uv_wavelengths, data, noise_map)` chunks.
+        batch_size
+            The number of source-pixel columns processed per batch by the sparse operator.
+        accumulator_kwargs
+            Passed to `sparse_terms_from_chunks` (`transformer_class`, `method`, `eps`,
+            `chunk_size`, `chunk_k`, `use_jax`, `show_progress`). When not given,
+            `transformer_class`, `eps` and `chunk_size` follow this dataset's transformer (the
+            same defaults `psf_precision_operator_from` takes), so the accumulated terms match
+            `apply_sparse_operator()`.
+
+        Returns
+        -------
+        Interferometer
+            A new `Interferometer` dataset with the accumulated `InterferometerSparseOperator`
+            attached.
+
+        Raises
+        ------
+        exc.DatasetException
+            If any chunk has unequal real and imaginary noise sigma.
+        """
+        if disable_jax() and accumulator_kwargs.get("use_jax", False):
+            accumulator_kwargs["use_jax"] = False
+
+        transformer = self.transformer
+
+        if isinstance(transformer, TransformerNUFFT):
+            accumulator_kwargs.setdefault("eps", transformer.eps)
+            accumulator_kwargs.setdefault("chunk_size", transformer.chunk_size)
+
+            if "transformer_class" not in accumulator_kwargs:
+                transformer_eps = transformer.eps
+                transformer_chunk_size = transformer.chunk_size
+
+                def transformer_class(uv_wavelengths, real_space_mask):
+                    return TransformerNUFFT(
+                        uv_wavelengths=uv_wavelengths,
+                        real_space_mask=real_space_mask,
+                        eps=transformer_eps,
+                        chunk_size=transformer_chunk_size,
+                    )
+
+                accumulator_kwargs["transformer_class"] = transformer_class
+        else:
+            accumulator_kwargs.setdefault("transformer_class", type(transformer))
+
+        terms = inversion_interferometer_util.sparse_terms_from_chunks(
+            chunks, real_space_mask=self.real_space_mask, **accumulator_kwargs
+        )
+
+        sparse_operator = inversion_interferometer_util.InterferometerSparseOperator.from_sparse_terms(
+            terms,
+            real_space_mask=self.real_space_mask,
+            batch_size=batch_size,
         )
 
         return Interferometer(

autoarray/fit/fit_interferometer.py

@@ -128,7 +128,20 @@ def noise_normalization(self) -> float:
         Returns the noise-map normalization term of the noise-map, summing the noise_map value in every pixel as:
 
         [Noise_Term] = sum(log(2*pi*[Noise]**2.0))
+
+        When the dataset carries a sparse operator with a precomputed `noise_normalization`
+        (set by `Interferometer.apply_sparse_operator` and
+        `apply_sparse_operator_from_chunks`), and the noise-map fitted is the dataset's own
+        (not one a subclass has scaled or replaced), that scalar is returned instead of
+        reducing over every visibility's sigma on each likelihood call. It is computed with
+        the same expression, so the value is identical.
         """
+        sparse_operator = getattr(self.dataset, "sparse_operator", None)
+        noise_normalization = getattr(sparse_operator, "noise_normalization", None)
+
+        if noise_normalization is not None and self.noise_map is self.dataset.noise_map:
+            return noise_normalization
+
         return fit_util.noise_normalization_complex_from(
             noise_map=self.noise_map.array,
         )

autoarray/inversion/inversion/dataset_interface.py

@@ -33,6 +33,14 @@ def __init__(
         data
             The array of the image data containing the signal that is fitted (in PyAutoGalaxy and PyAutoLens the
             recommended units are electrons per second).
+
+            For an interferometer inversion on the sparse path this may be `None`, meaning "the raw visibilities
+            the `sparse_operator` was built from, unmodified". It is only valid with a `sparse_operator` that
+            carries a precomputed `data_term` (built by `Interferometer.apply_sparse_operator` or
+            `apply_sparse_operator_from_chunks`): the sparse inversion's data vector already comes from the
+            operator's cached dirty image, and `fast_chi_squared` then reads the operator's `data_term` instead
+            of reducing over the visibilities, so no visibility array is touched by the likelihood. Output-only
+            quantities that need the data itself (e.g. `data_subtracted_dict`) are unavailable in that case.
         noise_map
             An array describing the RMS standard deviation error in each pixel used for computing quantities like the
             chi-squared in a fit (in PyAutoGalaxy and PyAutoLens the recommended units are electrons per second).

autoarray/inversion/inversion/interferometer/abstract.py

@@ -82,7 +82,7 @@ def operated_mapping_matrix_list(self) -> List[np.ndarray]:
             operated_mapping_matrix_override = linear_obj.operated_mapping_matrix_override
 
             if operated_mapping_matrix_override is not None:
-                expected_shape = (self.data.shape[0], linear_obj.params)
+                expected_shape = (self.noise_map.shape[0], linear_obj.params)
 
                 if tuple(operated_mapping_matrix_override.shape) != expected_shape:
                     raise exc.InversionException(
@@ -170,6 +170,11 @@ def fast_chi_squared(self):
         where `s` is the reconstruction vector, `F` is the curvature matrix, `D` is the data vector,
         and `d_r`/`d_i` are the real/imaginary parts of the observed visibilities.
 
+        When the dataset interface's `data` is `None` the third term is read from the scalar
+        `sparse_operator.data_term` cached when the operator was built, so no visibility array
+        is reduced over. That is only correct when the data fitted is the raw data the operator
+        was built from (nothing subtracted), which is the contract of passing `data=None`.
+
         This avoids computing the full mapped reconstructed visibilities and is faster than computing
         `chi_squared` via the residual visibilities when many source pixels are used.
         """
@@ -190,10 +195,31 @@ def fast_chi_squared(self):
             ]
         )
 
-        chi_squared_term_3 = xp.sum(
-            self.dataset.data.array.real**2.0 / self.dataset.noise_map.array.real**2.0
-        ) + xp.sum(
-            self.dataset.data.array.imag**2.0 / self.dataset.noise_map.array.imag**2.0
-        )
+        if self.dataset.data is None:
+            # The interface carries no visibilities (e.g. a pixelization-only fit on the sparse
+            # path, where nothing was subtracted from the data), so term 3 is the scalar
+            # `d^T N^-1 d` the sparse operator cached when it was built from the raw data.
+            chi_squared_term_3 = getattr(
+                self.dataset.sparse_operator, "data_term", None
+            )
+
+            if chi_squared_term_3 is None:
+                raise exc.InversionException(
+                    "The dataset input to this interferometer inversion has `data=None`, which "
+                    "is only valid when its `sparse_operator` carries a precomputed `data_term` "
+                    "(sum(d_r^2/sigma_r^2) + sum(d_i^2/sigma_i^2)) -- e.g. one built by "
+                    "`Interferometer.apply_sparse_operator` or "
+                    "`Interferometer.apply_sparse_operator_from_chunks`. This dataset's "
+                    "sparse operator has none, so `fast_chi_squared` cannot be computed. "
+                    "Pass the visibilities as `data` instead."
+                )
+        else:
+            chi_squared_term_3 = xp.sum(
+                self.dataset.data.array.real**2.0
+                / self.dataset.noise_map.array.real**2.0
+            ) + xp.sum(
+                self.dataset.data.array.imag**2.0
+                / self.dataset.noise_map.array.imag**2.0
+            )
 
         return chi_squared_term_1 + chi_squared_term_2 + chi_squared_term_3