Skip to content

fix: write dataset.fits once, always, to image/ (drop fits_dataset knob) #731

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Jammy2211 merged 1 commit into from
Sep 8, 2026
Merged

fix: write dataset.fits once, always, to image/ (drop fits_dataset knob) #731

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 5 additions & 4 deletions autolens/config/visualize/plots.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,19 +3,20 @@
# For example, if `plots: fit: subplot_fit=True``, the ``subplot_fit.png`` subplot file will
# be plotted every time visualization is performed.

# There are two settings which are important for inspecting results via the dataset after a fit is complete which are:
# One setting is important for inspecting results via the dataset after a fit is complete:

# - `fits_dataset`: This outputs `dataset.fits` which the database functionality may use to reperform fits.
# -`fits_adapt_images`, This outputs `adapt_images.fits` which the database functionality may use to reperform fits.

# These can be disabled to save on hard-disk space but will lead to certain database functionality being disabled.
# It can be disabled to save on hard-disk space but will lead to certain database functionality being disabled.

# The dataset itself is always output as `dataset.fits` to the `image` folder of every fit, and is not controlled
# by any setting here.

subplot_format: [png] # Output format of all subplots, can be png, pdf or both (e.g. [png, pdf])
fits_are_zoomed: false # If true, output .fits files are zoomed in on the center of the unmasked region image, saving hard-disk space.

dataset: # Settings for plots of all datasets (e.g. Imaging, Interferometer).
subplot_dataset: true # Plot subplot containing all dataset quantities (e.g. the data, noise-map, etc.)?
fits_dataset: true # Output a .fits file containing the dataset data, noise-map and other quantities?

positions: # Settings for plots with resampling image-positions on (e.g. the image).
image_with_positions: true
Expand Down
43 changes: 23 additions & 20 deletions autolens/imaging/model/analysis.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -240,12 +240,13 @@ def fit_from(
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.
to the ``image`` folder. The write is unconditional (it is not gated on any
visualization setting) but is skipped if the file already exists, so a resumed
search does not rewrite it.

The aggregator loaders (e.g. ``ImagingAgg``, ``agg_util.mask_header_from``)
reload the dataset via ``fit.value(name="dataset")``, which scans the ``image``
folder for .fits files, so this single write is all that is required.
"""
super().save_attributes(paths=paths)

Expand All @@ -259,22 +260,24 @@ def save_attributes(self, paths: af.DirectoryPaths):
),
]

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,
),
hdu_list = 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,
)

dataset_path = paths.image_path / "dataset.fits"

if not dataset_path.exists():
hdu_list.writeto(dataset_path, overwrite=True)

@staticmethod
def _register_fit_imaging_pytrees() -> None:
"""Register every type reachable from a ``FitImaging`` return value
Expand Down
46 changes: 24 additions & 22 deletions autolens/interferometer/model/analysis.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -333,10 +333,10 @@ def save_attributes(self, paths: af.DirectoryPaths):

- The positions of the brightest pixels in the lensed source which are used to discard mass models.

The following .fits files are also output via the plotter interface:
It also outputs, to the `image` folder, the file `dataset.fits`, which contains:

- The real space mask applied to the dataset, in the `PrimaryHDU` of `dataset.fits`.
- The interferometer dataset as `dataset.fits` (data / noise-map / uv_wavelengths).
- The real space mask applied to the dataset, in the `PrimaryHDU`.
- The interferometer dataset (data / noise-map / uv_wavelengths).

It is common for these attributes to be loaded by many of the template aggregator functions given in the
`aggregator` modules. For example, when using the database tools to perform a fit, the default behaviour is for
Expand All @@ -351,27 +351,29 @@ 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,
),
hdu_list = 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,
)

# `dataset.fits` is written once per search, to the `image` folder, and is written
# unconditionally (it is not gated on any visualization setting). The write is skipped
# if the file already exists, so a resumed search does not rewrite it.
#
# The aggregator loaders (e.g. `InterferometerAgg`, `agg_util.mask_header_from`) reload
# the dataset via `fit.value(name="dataset")`, which scans the `image` folder for .fits
# files, so this single write is all that is required.
dataset_path = paths.image_path / "dataset.fits"

if not dataset_path.exists():
hdu_list.writeto(dataset_path, overwrite=True)

paths.save_json(
"transformer_class",
to_dict(self.dataset.transformer.__class__),
Expand Down
53 changes: 49 additions & 4 deletions test_autolens/analysis/analysis/test_analysis_dataset.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -99,24 +99,69 @@ def test__save_results__tracer_output_to_json(analysis_imaging_7x7):

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`,
# `image` 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.
# `fit.value(name="dataset")`. The write is not gated on visualization.
from astropy.io import fits

paths = af.DirectoryPaths()

analysis_imaging_7x7.save_attributes(paths=paths)
dataset_fits_path = paths.image_path / "dataset.fits"

if dataset_fits_path.exists():
os.remove(dataset_fits_path)

dataset_fits_path = paths._files_path / "dataset.fits"
analysis_imaging_7x7.save_attributes(paths=paths)

assert dataset_fits_path.exists()
assert not (paths._files_path / "dataset.fits").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"]

# A second call (e.g. a resumed search) must not rewrite the existing file.

dataset_fits_path.write_bytes(b"sentinel")

analysis_imaging_7x7.save_attributes(paths=paths)

assert dataset_fits_path.read_bytes() == b"sentinel"

os.remove(dataset_fits_path)


def test__save_attributes__dataset_fits_output_for_aggregator__interferometer(
analysis_interferometer_7,
):
from astropy.io import fits

paths = af.DirectoryPaths()

dataset_fits_path = paths.image_path / "dataset.fits"

if dataset_fits_path.exists():
os.remove(dataset_fits_path)

analysis_interferometer_7.save_attributes(paths=paths)

assert dataset_fits_path.exists()
assert not (paths._files_path / "dataset.fits").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", "UV_WAVELENGTHS"]

# A second call (e.g. a resumed search) must not rewrite the existing file.

dataset_fits_path.write_bytes(b"sentinel")

analysis_interferometer_7.save_attributes(paths=paths)

assert dataset_fits_path.read_bytes() == b"sentinel"

os.remove(dataset_fits_path)


Expand Down
Loading