feat: persist SparseTerms oversample fine grids in dataset.fits (opt-in)

autogalaxy/aggregator/interferometer/interferometer.py

@@ -73,6 +73,11 @@ def _sparse_terms_from(hdu_list, real_space_mask: aa.Mask2D) -> aa.SparseTerms:
     The `(y, x)` phase centre is reloaded from the same scalars / header cards; a file written before it was
     persisted (a shorter `SPARSE_TERMS_SCALARS` array and no `PHCENTY` / `PHCENTX` cards) reloads it as `None`
     (not recorded), as does a `NaN` entry.
+
+    The `oversample` fine grids are read from the `PRECISION_OPERATOR_FINE` / `DIRTY_IMAGE_FINE` HDUs, with
+    `oversample` / `oversample_pad` from the scalars, when the file has them (they are written only on
+    request, see `interferometer_hdu_list_from(include_fine_grids=True)`). A file without them -- written
+    without opting in, or before they existed -- reloads all four as `None`.
     """
     header = hdu_list[0].header
 
@@ -108,6 +113,27 @@ def header_value(field):
         else (float(phase_centre_y), float(phase_centre_x))
     )
 
+    oversample = header_value("oversample")
+    oversample_pad = header_value("oversample_pad")
+
+    if (
+        _has_hdu(hdu_list, "precision_operator_fine")
+        and _has_hdu(hdu_list, "dirty_image_fine")
+        and oversample is not None
+    ):
+        fine_grids = dict(
+            precision_operator_fine=np.asarray(
+                _hdu(hdu_list, "precision_operator_fine").data, dtype="float64"
+            ),
+            dirty_image_fine=np.asarray(
+                _hdu(hdu_list, "dirty_image_fine").data, dtype="float64"
+            ),
+            oversample=int(oversample),
+            oversample_pad=None if oversample_pad is None else float(oversample_pad),
+        )
+    else:
+        fine_grids = {}
+
     return aa.SparseTerms(
         nufft_precision_operator=np.asarray(
             _hdu(hdu_list, "nufft_precision_operator").data, dtype="float64"
@@ -128,6 +154,7 @@ def header_value(field):
             None if transformer_class_name is None else str(transformer_class_name)
         ),
         phase_centre=phase_centre,
+        **fine_grids,
     )
 
 

autogalaxy/interferometer/model/analysis.py

@@ -9,6 +9,7 @@
 
 It also handles adapt images, visualization, and result wrapping into `ResultInterferometer`.
 """
+
 import logging
 import numpy as np
 from typing import Optional
@@ -46,15 +47,19 @@
     "transformer_class_name": "TRNSFRMR",
     "phase_centre_y": "PHCENTY",
     "phase_centre_x": "PHCENTX",
+    "oversample": "OVERSAMP",
+    "oversample_pad": "OVERSPAD",
 }
 
 # The order of the float64 1-D array in the `SPARSE_TERMS_SCALARS` HDU, the lossless store of the numeric
 # `SparseTerms` scalars. A FITS header card holds at most 20 characters of value, so astropy truncates
 # exponent-form float64 values (e.g. `1.2345678901234567e+20` is written as `1.23456789012345E+20`); the
-# header cards under `SPARSE_TERMS_HEADER_KEYS` are kept as human-readable copies only. `eps` and the
-# `(y, x)` phase centre (arcsec) are `NaN` when not recorded. Entries are only ever appended, so a file
-# written before an entry existed has a shorter array and the loader treats the missing entries as not
-# recorded.
+# header cards under `SPARSE_TERMS_HEADER_KEYS` are kept as human-readable copies only. `eps`, the
+# `(y, x)` phase centre (arcsec) and the `oversample` factor / pad of the fine grids are `NaN` when not
+# recorded (the oversampling is recorded only when the fine grids themselves are written, so a reloaded
+# `SparseTerms` never claims an `oversample` without its grids). Entries are only ever appended, so a
+# file written before an entry existed has a shorter array and the loader treats the missing entries as
+# not recorded.
 SPARSE_TERMS_SCALARS_ORDER = (
     "sum_weights",
     "data_term",
@@ -63,10 +68,14 @@
     "eps",
     "phase_centre_y",
     "phase_centre_x",
+    "oversample",
+    "oversample_pad",
 )
 
 
-def interferometer_hdu_list_from(dataset: aa.Interferometer):
+def interferometer_hdu_list_from(
+    dataset: aa.Interferometer, include_fine_grids: bool = False
+):
     """
     Returns the `HDUList` written to `dataset.fits` by `AnalysisInterferometer.save_attributes`, from which
     the aggregator (`autogalaxy.aggregator.interferometer.interferometer._interferometer_from`) reloads the
@@ -86,10 +95,22 @@ def interferometer_hdu_list_from(dataset: aa.Interferometer):
     shape, pixel scales and origin come from the mask itself. The file size is set by the real-space grid, not
     the number of visibilities.
 
+    Terms accumulated with `oversample=q` (`Interferometer.from_stream(..., oversample=q)`) also carry the
+    fine precision-operator and dirty-image grids. These are written -- as HDUs `PRECISION_OPERATOR_FINE`
+    and `DIRTY_IMAGE_FINE` after `SPARSE_TERMS_SCALARS`, with `oversample` / `oversample_pad` in the scalars
+    and the `OVERSAMP` / `OVERSPAD` cards -- **only** when `include_fine_grids=True`. They are not written
+    by default because `dataset.fits` is written into every search's output by `save_attributes` and the
+    grids grow as `q^2`: ~840 MB per file at a 400 x 400 grid and `q = 8`. Without them the file is the
+    same as for terms without fine grids (the oversampling scalars are `NaN`), and reloads with no fine
+    grids; rebuild them from the visibilities, or write this HDU list yourself with
+    `include_fine_grids=True`, when they are needed.
+
     Parameters
     ----------
     dataset
         The interferometer dataset being fitted.
+    include_fine_grids
+        If `True` and the dataset's `SparseTerms` carry `oversample` fine grids, write them too.
     """
     mask = dataset.real_space_mask
 
@@ -143,6 +164,26 @@ def interferometer_hdu_list_from(dataset: aa.Interferometer):
         (np.nan, np.nan) if terms.phase_centre is None else terms.phase_centre
     )
 
+    write_fine_grids = include_fine_grids and (
+        terms.precision_operator_fine is not None and terms.dirty_image_fine is not None
+    )
+
+    if write_fine_grids:
+        header_dict[SPARSE_TERMS_HEADER_KEYS["oversample"]] = int(terms.oversample)
+        header_dict[SPARSE_TERMS_HEADER_KEYS["oversample_pad"]] = float(
+            terms.oversample_pad
+        )
+        oversample = (float(terms.oversample), float(terms.oversample_pad))
+        fine_values = [
+            np.asarray(terms.precision_operator_fine, dtype="float64"),
+            np.asarray(terms.dirty_image_fine, dtype="float64"),
+        ]
+        fine_ext_names = ["precision_operator_fine", "dirty_image_fine"]
+    else:
+        oversample = (np.nan, np.nan)
+        fine_values = []
+        fine_ext_names = []
+
     return hdu_list_for_output_from(
         values_list=[
             mask.astype("float"),
@@ -158,16 +199,20 @@ def interferometer_hdu_list_from(dataset: aa.Interferometer):
                     np.nan if terms.eps is None else float(terms.eps),
                     float(phase_centre[0]),
                     float(phase_centre[1]),
+                    float(oversample[0]),
+                    float(oversample[1]),
                 ],
                 dtype="float64",
             ),
+            *fine_values,
         ],
         ext_name_list=[
             "mask",
             "nufft_precision_operator",
             "dirty_image",
             "dirty_beam",
             "sparse_terms_scalars",
+            *fine_ext_names,
         ],
         header_dict=header_dict,
     )