feat(interferometer): fits, save_attributes and aggregator reload on array-free datasets (streaming phase 2)

autogalaxy/aggregator/agg_util.py

@@ -99,9 +99,11 @@ def mask_header_from(fit, name="dataset"):
     """
 
     header = aa.Header(header_sci_obj=fit.value(name=name)[0].header)
+    pixel_scales_y = header.header_sci_obj[Mask2DKeys.PIXSCAY.value]
+    # Files written without a `PIXSCAX` card are assumed to have square pixels.
     pixel_scales = (
-        header.header_sci_obj[Mask2DKeys.PIXSCAY.value],
-        header.header_sci_obj[Mask2DKeys.PIXSCAY.value],
+        pixel_scales_y,
+        header.header_sci_obj.get(Mask2DKeys.PIXSCAX.value, pixel_scales_y),
     )
     origin = (
         header.header_sci_obj[Mask2DKeys.ORIGINY.value],

autogalaxy/aggregator/interferometer/interferometer.py

@@ -20,10 +20,100 @@
 from functools import partial
 from typing import List
 
+import numpy as np
+
 import autofit as af
 import autoarray as aa
 
 from autogalaxy.aggregator import agg_util
+from autogalaxy.interferometer.model.analysis import (
+    SPARSE_TERMS_HEADER_KEYS,
+    SPARSE_TERMS_SCALARS_ORDER,
+)
+
+# The positions of the HDUs of the legacy in-memory `dataset.fits` layout (mask / data /
+# noise_map / uv_wavelengths), used when a file predates EXTNAMEs being read by name.
+_LEGACY_HDU_INDEX = {"mask": 0, "data": 1, "noise_map": 2, "uv_wavelengths": 3}
+
+
+def _hdu(hdu_list, name: str):
+    """
+    Returns the HDU of `hdu_list` whose EXTNAME is `name` (case-insensitive), falling back to its position
+    in the legacy in-memory `dataset.fits` layout (`_LEGACY_HDU_INDEX`) if no HDU has that EXTNAME.
+    """
+    try:
+        return hdu_list[name.upper()]
+    except KeyError:
+        return hdu_list[_LEGACY_HDU_INDEX[name]]
+
+
+def _has_hdu(hdu_list, name: str) -> bool:
+    """
+    Returns whether `hdu_list` has an HDU whose EXTNAME is `name` (case-insensitive).
+    """
+    try:
+        hdu_list[name.upper()]
+    except KeyError:
+        return False
+
+    return True
+
+
+def _sparse_terms_from(hdu_list, real_space_mask: aa.Mask2D) -> aa.SparseTerms:
+    """
+    Returns the `SparseTerms` of an array-free interferometer dataset from the `dataset.fits` HDU list written
+    by `autogalaxy.interferometer.model.analysis.interferometer_hdu_list_from`.
+
+    The arrays are read from the HDUs named `NUFFT_PRECISION_OPERATOR`, `DIRTY_IMAGE` and `DIRTY_BEAM`, the
+    numeric scalars losslessly from the `SPARSE_TERMS_SCALARS` HDU (order `SPARSE_TERMS_SCALARS_ORDER`),
+    falling back to the primary header cards (`SPARSE_TERMS_HEADER_KEYS`, which truncate exponent-form
+    float64 values) for files written before that HDU existed, the transformer class name from the header,
+    and the mask provenance (shape, pixel scales, origin) from the reloaded real-space mask.
+    """
+    header = hdu_list[0].header
+
+    if _has_hdu(hdu_list, "sparse_terms_scalars"):
+        scalars = dict(
+            zip(
+                SPARSE_TERMS_SCALARS_ORDER,
+                np.asarray(_hdu(hdu_list, "sparse_terms_scalars").data, dtype="float64"),
+            )
+        )
+
+        def header_value(field):
+            if field in scalars:
+                value = scalars[field]
+                return None if np.isnan(value) else value
+            return header.get(SPARSE_TERMS_HEADER_KEYS[field])
+
+    else:
+
+        def header_value(field):
+            return header.get(SPARSE_TERMS_HEADER_KEYS[field])
+
+    eps = header_value("eps")
+    transformer_class_name = header_value("transformer_class_name")
+
+    return aa.SparseTerms(
+        nufft_precision_operator=np.asarray(
+            _hdu(hdu_list, "nufft_precision_operator").data, dtype="float64"
+        ),
+        dirty_image_native=np.asarray(
+            _hdu(hdu_list, "dirty_image").data, dtype="float64"
+        ),
+        dirty_beam_native=np.asarray(_hdu(hdu_list, "dirty_beam").data, dtype="float64"),
+        sum_weights=float(header_value("sum_weights")),
+        data_term=float(header_value("data_term")),
+        noise_normalization=float(header_value("noise_normalization")),
+        n_vis=int(header_value("n_vis")),
+        shape_native=tuple(real_space_mask.shape_native),
+        pixel_scales=tuple(float(value) for value in real_space_mask.pixel_scales),
+        origin=tuple(float(value) for value in real_space_mask.origin),
+        eps=None if eps is None else float(eps),
+        transformer_class_name=(
+            None if transformer_class_name is None else str(transformer_class_name)
+        ),
+    )
 
 
 def _interferometer_from(
@@ -41,6 +131,11 @@ def _interferometer_from(
     - The real space mask defining the grid of the interferometer for the FFT (`dataset/real_space_mask.fits`).
     - The settings of the `Interferometer` data structure used in the fit (`dataset/settings.json`).
 
+    A fit of an array-free dataset (built by `Interferometer.from_stream` / `from_sparse_terms`) stores its
+    `SparseTerms` in `dataset.fits` instead of the visibilities (an HDU named `NUFFT_PRECISION_OPERATOR` is
+    present). It is rebuilt via `Interferometer.from_sparse_terms`, so the returned dataset is array-free with
+    its sparse operator re-attached, and a fit of it reproduces the original `log_evidence`.
+
     Each individual attribute can be loaded from the database via the `fit.value()` method.
 
     This method combines all of these attributes and returns a `Interferometer` object, including having its
@@ -64,13 +159,24 @@ def _interferometer_from(
     for fit in fit_list:
         real_space_mask, header = agg_util.mask_header_from(fit=fit)
 
+        hdu_list = fit.value(name="dataset")
+
+        if _has_hdu(hdu_list, "nufft_precision_operator"):
+            dataset_list.append(
+                aa.Interferometer.from_sparse_terms(
+                    _sparse_terms_from(hdu_list, real_space_mask=real_space_mask),
+                    real_space_mask=real_space_mask,
+                )
+            )
+            continue
+
         data = aa.Visibilities(
-            visibilities=fit.value(name="dataset")[1].data.astype("float")
+            visibilities=_hdu(hdu_list, "data").data.astype("float")
         )
         noise_map = aa.VisibilitiesNoiseMap(
-            fit.value(name="dataset")[2].data.astype("float")
+            _hdu(hdu_list, "noise_map").data.astype("float")
         )
-        uv_wavelengths = fit.value(name="dataset")[3].data
+        uv_wavelengths = _hdu(hdu_list, "uv_wavelengths").data
 
         transformer_class = fit.value(name="transformer_class")
 
@@ -102,6 +208,11 @@ def __init__(self, aggregator: af.Aggregator):
         - The real space mask defining the grid of the interferometer for the FFT (`dataset/real_space_mask.fits`).
         - The settings of the `Interferometer` data structure used in the fit (`dataset/settings.json`).
 
+    A fit of an array-free dataset (built by `Interferometer.from_stream` / `from_sparse_terms`) stores its
+    `SparseTerms` in `dataset.fits` instead of the visibilities (an HDU named `NUFFT_PRECISION_OPERATOR` is
+    present). It is rebuilt via `Interferometer.from_sparse_terms`, so the returned dataset is array-free with
+    its sparse operator re-attached, and a fit of it reproduces the original `log_evidence`.
+
         The `aggregator` contains the path to each of these files, and they can be loaded individually. This class
         can load them all at once and create an `Interferometer` object via the `_interferometer_from` method.
 

autogalaxy/interferometer/fit_interferometer.py

@@ -231,15 +231,34 @@ def profile_image(self) -> aa.Array2D:
         return self.galaxies.image_2d_from(grid=self.grids.lp, xp=self._xp)
 
     @functools.cached_property
-    def profile_visibilities(self) -> aa.Visibilities:
+    def profile_visibilities(self) -> Optional[aa.Visibilities]:
         """
         Returns the visibilities of every light profile of every galaxy, which are computed by performing
         a Fourier transform to the sum of light profile images.
 
         If the galaxies have no ordinary (non-linear) light profile (e.g. their light is entirely an MGE of
         linear Gaussians), the image is all zeros and the Fourier transform is skipped. This is decided
         structurally, so it is safe under `jax.jit`.
+
+        On an array-free dataset (built by `Interferometer.from_stream` / `from_sparse_terms`, which has no
+        `uv_wavelengths` and so no transformer) there are no visibilities to compute: this returns `None` when
+        the galaxies have no ordinary light profile, and raises an `exc.DatasetException` when they do, because
+        subtracting a light profile's visibilities needs the visibility arrays.
         """
+        if self.dataset.transformer is None:
+            if _has_light_profile_non_linear(galaxies=self.galaxies):
+                raise aa.exc.DatasetException(
+                    "This FitInterferometer's dataset is array-free (built by from_stream / "
+                    "from_sparse_terms) and has no visibilities or transformer, so the galaxies' "
+                    "ordinary (non-linear) light profiles cannot be Fourier transformed and "
+                    "subtracted. An array-free dataset supports pixelization-only and linear-light "
+                    "fits; non-linear light profiles arrive in a later phase. Use the in-memory "
+                    "constructor (`Interferometer(data=..., noise_map=..., uv_wavelengths=..., ...)`) "
+                    "to fit them."
+                )
+
+            return None
+
         if _has_light_profile_non_linear(galaxies=self.galaxies):
             return self.dataset.transformer.visibilities_from(
                 image=self.profile_image, xp=self._xp
@@ -250,11 +269,20 @@ def profile_visibilities(self) -> aa.Visibilities:
         )
 
     @functools.cached_property
-    def profile_subtracted_visibilities(self) -> aa.Visibilities:
+    def profile_subtracted_visibilities(self) -> Optional[aa.Visibilities]:
         """
         Returns the interferometer dataset's visibilities with all transformed light profile images subtracted.
+
+        On an array-free dataset there are no visibilities, so this is `None`. `profile_visibilities` is
+        evaluated first, so a fit with ordinary light profiles on an array-free dataset raises rather than
+        silently fitting the unsubtracted sparse terms.
         """
-        return self.data - self.profile_visibilities
+        profile_visibilities = self.profile_visibilities
+
+        if self.data is None:
+            return None
+
+        return self.data - profile_visibilities
 
     @property
     def _uses_precomputed_data_term(self) -> bool:
@@ -332,10 +360,14 @@ def inversion_with_data(self) -> Optional[aa.AbstractInversion]:
         from the visibilities in that case, so the data it fitted are `fit.data`: this returns a shallow copy of
         the inversion (solved first, so it shares the reconstruction and every other cached quantity) whose dataset interface carries
         `fit.data`. In every other case it returns `inversion` itself.
+
+        On an array-free dataset (built by `Interferometer.from_stream` / `from_sparse_terms`) `fit.data` is
+        `None`, so there are no visibilities to carry and `inversion` itself is returned; output quantities
+        that read the visibilities are unavailable on such a fit.
         """
         inversion = self.inversion
 
-        if inversion is None or inversion.dataset.data is not None:
+        if inversion is None or inversion.dataset.data is not None or self.data is None:
             return inversion
 
         # Solve first, so the copy shares the reconstruction (and everything it cached) rather than repeating it.