fix: write dataset.fits in save_attributes for aggregator

autolens/imaging/model/analysis.py

@@ -17,6 +17,8 @@
 import autofit as af
 import autogalaxy as ag
 
+from autoconf.fitsable import hdu_list_for_output_from
+
 from autolens.analysis.analysis.dataset import AnalysisDataset
 from autolens.analysis.latent import LatentLens
 from autolens.imaging.model.result import ResultImaging
@@ -138,6 +140,44 @@ def fit_from(
             xp=self._xp
         )
 
+    def save_attributes(self, paths: af.DirectoryPaths):
+        """
+        Before the non-linear search begins, output the imaging ``dataset.fits``
+        to the ``files`` folder so the aggregator loaders (e.g. ``ImagingAgg``,
+        ``agg_util.mask_header_from``) can always reload the dataset via
+        ``fit.value(name="dataset")``, independently of whether the visualization
+        ``fits_dataset`` output ran. The plotter interface also writes this file
+        to the ``image`` folder for inspection, but that write is gated on
+        visualization settings and is not guaranteed for every fit.
+        """
+        super().save_attributes(paths=paths)
+
+        image_list = [
+            self.dataset.data.native_for_fits,
+            self.dataset.noise_map.native_for_fits,
+            self.dataset.psf.kernel.native_for_fits,
+            self.dataset.grids.lp.over_sample_size.native_for_fits.astype("float"),
+            self.dataset.grids.pixelization.over_sample_size.native_for_fits.astype(
+                "float"
+            ),
+        ]
+
+        paths.save_fits(
+            name="dataset",
+            fits=hdu_list_for_output_from(
+                values_list=[image_list[0].mask.astype("float")] + image_list,
+                ext_name_list=[
+                    "mask",
+                    "data",
+                    "noise_map",
+                    "psf",
+                    "over_sample_size_lp",
+                    "over_sample_size_pixelization",
+                ],
+                header_dict=self.dataset.mask.header_dict,
+            ),
+        )
+
     @staticmethod
     def _register_fit_imaging_pytrees() -> None:
         """Register every type reachable from a ``FitImaging`` return value

autolens/interferometer/model/analysis.py

@@ -17,6 +17,7 @@
 from typing import Optional
 
 from autoconf.dictable import to_dict
+from autoconf.fitsable import hdu_list_for_output_from
 
 import autofit as af
 import autoarray as aa
@@ -334,6 +335,27 @@ def save_attributes(self, paths: af.DirectoryPaths):
         """
         super().save_attributes(paths=paths)
 
+        # Output `dataset.fits` to the `files` folder so the aggregator loaders
+        # (e.g. `InterferometerAgg`, `agg_util.mask_header_from`) can always
+        # reload the dataset via `fit.value(name="dataset")`, independently of
+        # whether the visualization `fits_dataset` output ran. The plotter
+        # interface also writes this file to the `image` folder for inspection,
+        # but that write is gated on visualization settings and is not
+        # guaranteed for every fit.
+        paths.save_fits(
+            name="dataset",
+            fits=hdu_list_for_output_from(
+                values_list=[
+                    self.dataset.real_space_mask.astype("float"),
+                    self.dataset.data.in_array,
+                    self.dataset.noise_map.in_array,
+                    self.dataset.uv_wavelengths,
+                ],
+                ext_name_list=["mask", "data", "noise_map", "uv_wavelengths"],
+                header_dict=self.dataset.real_space_mask.header_dict,
+            ),
+        )
+
         paths.save_json(
             "transformer_class",
             to_dict(self.dataset.transformer.__class__),

test_autolens/analysis/analysis/test_analysis_dataset.py

@@ -63,3 +63,26 @@ def test__save_results__tracer_output_to_json(analysis_imaging_7x7):
     assert tracer.galaxies[1].redshift == 1.0
 
     os.remove(paths._files_path / "tracer.json")
+
+
+def test__save_attributes__dataset_fits_output_for_aggregator(analysis_imaging_7x7):
+    # Regression guard: `save_attributes` must always write `dataset.fits` to the
+    # `files` folder so the aggregator loaders (`ImagingAgg`,
+    # `agg_util.mask_header_from`) can reload the dataset via
+    # `fit.value(name="dataset")`, independently of whether visualization ran.
+    from astropy.io import fits
+
+    paths = af.DirectoryPaths()
+
+    analysis_imaging_7x7.save_attributes(paths=paths)
+
+    dataset_fits_path = paths._files_path / "dataset.fits"
+
+    assert dataset_fits_path.exists()
+
+    with fits.open(dataset_fits_path) as hdu_list:
+        ext_names = [hdu.name for hdu in hdu_list]
+
+    assert ext_names[:4] == ["MASK", "DATA", "NOISE_MAP", "PSF"]
+
+    os.remove(dataset_fits_path)