fix: restore unconditional dataset.fits output for aggregator (save_attributes)

autogalaxy/imaging/model/analysis.py

@@ -19,6 +19,8 @@
 import autofit as af
 import autoarray as aa
 
+from autoconf.fitsable import hdu_list_for_output_from
+
 from autogalaxy.analysis.adapt_images.adapt_images import AdaptImages
 from autogalaxy.analysis.analysis.dataset import AnalysisDataset
 from autogalaxy.cosmology.model import LensingCosmology
@@ -252,3 +254,35 @@ def save_attributes(self, paths: af.DirectoryPaths):
             visualization, and the pickled objects used by the aggregator output by this function.
         """
         super().save_attributes(paths=paths)
+
+        # Output `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.
+        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,
+            ),
+        )

autogalaxy/interferometer/model/analysis.py

@@ -14,6 +14,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
@@ -246,6 +247,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_autogalaxy/analysis/analysis/test_analysis.py

@@ -45,3 +45,26 @@ def test__save_results__galaxies_output_to_json(analysis_imaging_7x7):
     assert galaxies[0].redshift == 0.5
 
     os.remove(paths._files_path / "galaxies.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)