diff --git a/autolens/imaging/model/analysis.py b/autolens/imaging/model/analysis.py index f8fc36e84..bbc2edfe3 100644 --- a/autolens/imaging/model/analysis.py +++ b/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 diff --git a/autolens/interferometer/model/analysis.py b/autolens/interferometer/model/analysis.py index c4919557f..f22c4e527 100644 --- a/autolens/interferometer/model/analysis.py +++ b/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__), diff --git a/test_autolens/analysis/analysis/test_analysis_dataset.py b/test_autolens/analysis/analysis/test_analysis_dataset.py index d92996ccc..b402664a0 100644 --- a/test_autolens/analysis/analysis/test_analysis_dataset.py +++ b/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)