From f7f6d79e10868bd99584839e560c4f1c08ade018 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Wed, 21 Oct 2026 13:03:54 +0100 Subject: [PATCH] feat: persist the interferometer primary beam through save_attributes and aggregator rebuild interferometer_hdu_list_from appends the primary beam (native float64, masked pixels 0) as a trailing PRIMARY_BEAM HDU in both the in-memory and array-free layouts, so legacy positional indices are unchanged and beamless datasets write the same file as before. The aggregator detects the HDU by name and passes primary_beam= into Interferometer / from_sparse_terms, so a rebuilt dataset and fit reproduce the beamed likelihood. Old outputs without the HDU reload with primary_beam=None. Refs PyAutoLabs/PyAutoGalaxy#654 Co-Authored-By: Claude Fable 5.1 --- .../interferometer/interferometer.py | 34 ++- autogalaxy/interferometer/model/analysis.py | 46 ++- .../model/test_analysis_interferometer.py | 267 ++++++++++++++++++ 3 files changed, 343 insertions(+), 4 deletions(-) diff --git a/autogalaxy/aggregator/interferometer/interferometer.py b/autogalaxy/aggregator/interferometer/interferometer.py index c47bb88e..42f89c94 100644 --- a/autogalaxy/aggregator/interferometer/interferometer.py +++ b/autogalaxy/aggregator/interferometer/interferometer.py @@ -2,8 +2,8 @@ Aggregator interface for loading ``Interferometer`` datasets from model-fit results. After an interferometer model-fit the dataset (complex visibilities, noise-map, -uv-wavelengths, real-space mask, and transformer class) is written to a FITS file in -the output directory or SQLite database. This module reconstructs a fully-configured +uv-wavelengths, real-space mask, transformer class and, if used, primary beam) is +written to a FITS file in the output directory or SQLite database. This module reconstructs a fully-configured ``aa.Interferometer`` object from those stored files. Two public objects are provided: @@ -18,7 +18,7 @@ from __future__ import annotations from functools import partial -from typing import List +from typing import List, Optional import numpy as np @@ -59,6 +59,22 @@ def _has_hdu(hdu_list, name: str) -> bool: return True +def _primary_beam_from(hdu_list) -> Optional[np.ndarray]: + """ + Returns the primary beam stored in the `PRIMARY_BEAM` HDU of a `dataset.fits` HDU list written by + `autogalaxy.interferometer.model.analysis.interferometer_hdu_list_from`, as a float64 array of the native + shape of the real-space mask (masked pixels 0), or `None` if the file has no such HDU (the dataset had no + primary beam, or the file was written before the beam was persisted). + + The HDU is detected by name only (`_has_hdu`), never via `_hdu`, whose fallback to the legacy positional + indices would misread another HDU as the beam. + """ + if not _has_hdu(hdu_list, "primary_beam"): + return None + + return np.asarray(hdu_list["PRIMARY_BEAM"].data, dtype="float64") + + def _sparse_terms_from(hdu_list, real_space_mask: aa.Mask2D) -> aa.SparseTerms: """ Returns the `SparseTerms` of an array-free interferometer dataset from the `dataset.fits` HDU list written @@ -178,6 +194,11 @@ def _interferometer_from( present). It is rebuilt via `Interferometer.from_sparse_terms`, so the returned dataset is array-free with its sparse operator re-attached, and a fit of it reproduces the original `log_evidence`. + If the dataset had a primary beam, `dataset.fits` has a trailing `PRIMARY_BEAM` HDU, which is passed back + as `primary_beam=` to either constructor, so the reloaded dataset (its transformer or sparse operator) and + any fit of it carry the same beam and reproduce the beamed likelihood. A file without that HDU (no beam, or + written before the beam was persisted) reloads with `primary_beam=None`. + Each individual attribute can be loaded from the database via the `fit.value()` method. This method combines all of these attributes and returns a `Interferometer` object, including having its @@ -203,11 +224,14 @@ def _interferometer_from( hdu_list = fit.value(name="dataset") + primary_beam = _primary_beam_from(hdu_list) + if _has_hdu(hdu_list, "nufft_precision_operator"): dataset_list.append( aa.Interferometer.from_sparse_terms( _sparse_terms_from(hdu_list, real_space_mask=real_space_mask), real_space_mask=real_space_mask, + primary_beam=primary_beam, ) ) continue @@ -228,6 +252,7 @@ def _interferometer_from( uv_wavelengths=uv_wavelengths, real_space_mask=real_space_mask, transformer_class=transformer_class, + primary_beam=primary_beam, ) dataset_list.append(dataset) @@ -255,6 +280,9 @@ def __init__(self, aggregator: af.Aggregator): present). It is rebuilt via `Interferometer.from_sparse_terms`, so the returned dataset is array-free with its sparse operator re-attached, and a fit of it reproduces the original `log_evidence`. + A primary beam, if the dataset had one, is reloaded from the trailing `PRIMARY_BEAM` HDU and passed back + as `primary_beam=`, so the reloaded dataset carries the same beam. + The `aggregator` contains the path to each of these files, and they can be loaded individually. This class can load them all at once and create an `Interferometer` object via the `_interferometer_from` method. diff --git a/autogalaxy/interferometer/model/analysis.py b/autogalaxy/interferometer/model/analysis.py index e577e7d5..dde81b99 100644 --- a/autogalaxy/interferometer/model/analysis.py +++ b/autogalaxy/interferometer/model/analysis.py @@ -73,6 +73,23 @@ ) +def _primary_beam_native_from(dataset: aa.Interferometer) -> Optional[np.ndarray]: + """ + Returns the primary beam of `dataset` as a float64 array of the native shape of its real-space mask + (masked pixels 0), as written to the `PRIMARY_BEAM` HDU of `dataset.fits`, or `None` if the dataset + has no primary beam. + """ + if dataset.primary_beam_slim is None: + return None + + return np.asarray( + aa.Array2D( + values=dataset.primary_beam_slim, mask=dataset.real_space_mask + ).native.array, + dtype="float64", + ) + + def interferometer_hdu_list_from( dataset: aa.Interferometer, include_fine_grids: bool = False ): @@ -105,6 +122,13 @@ def interferometer_hdu_list_from( grids; rebuild them from the visibilities, or write this HDU list yourself with `include_fine_grids=True`, when they are needed. + A dataset with a primary beam (`Interferometer(primary_beam=...)`) also writes it, in both layouts, as a + trailing HDU with EXTNAME `PRIMARY_BEAM`: a float64 array of the native shape of the real-space mask with + masked pixels 0. It is the last HDU, so the positional indices of the legacy in-memory layout (0-3) are + unchanged, and a dataset without a primary beam writes no such HDU (the file is identical to one written + before the beam was persisted). The aggregator detects it by name and passes it back as `primary_beam=`, + so a fit of the reloaded dataset reproduces the beamed likelihood. + Parameters ---------- dataset @@ -114,6 +138,15 @@ def interferometer_hdu_list_from( """ mask = dataset.real_space_mask + primary_beam_native = _primary_beam_native_from(dataset) + + if primary_beam_native is None: + primary_beam_values = [] + primary_beam_ext_names = [] + else: + primary_beam_values = [primary_beam_native] + primary_beam_ext_names = ["primary_beam"] + if not dataset.is_array_free: return hdu_list_for_output_from( values_list=[ @@ -121,8 +154,15 @@ def interferometer_hdu_list_from( dataset.data.in_array, dataset.noise_map.in_array, dataset.uv_wavelengths, + *primary_beam_values, + ], + ext_name_list=[ + "mask", + "data", + "noise_map", + "uv_wavelengths", + *primary_beam_ext_names, ], - ext_name_list=["mask", "data", "noise_map", "uv_wavelengths"], header_dict=mask.header_dict, ) @@ -205,6 +245,7 @@ def interferometer_hdu_list_from( dtype="float64", ), *fine_values, + *primary_beam_values, ], ext_name_list=[ "mask", @@ -213,6 +254,7 @@ def interferometer_hdu_list_from( "dirty_beam", "sparse_terms_scalars", *fine_ext_names, + *primary_beam_ext_names, ], header_dict=header_dict, ) @@ -418,6 +460,8 @@ def save_attributes(self, paths: af.DirectoryPaths): - The real space mask applied to the dataset, in the `PrimaryHDU`. - The interferometer dataset (data / noise-map / uv_wavelengths), or for an array-free dataset its `SparseTerms` (see `interferometer_hdu_list_from`). + - The primary beam on the real-space mask (native, masked pixels 0), as a trailing `PRIMARY_BEAM` + HDU, if the dataset has one, so the aggregator-rebuilt dataset and fit carry it. 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 diff --git a/test_autogalaxy/interferometer/model/test_analysis_interferometer.py b/test_autogalaxy/interferometer/model/test_analysis_interferometer.py index 49ac5129..c6f4393b 100644 --- a/test_autogalaxy/interferometer/model/test_analysis_interferometer.py +++ b/test_autogalaxy/interferometer/model/test_analysis_interferometer.py @@ -168,6 +168,9 @@ def test__interferometer_hdu_list_from__in_memory_dataset__unchanged_layout( np.testing.assert_array_equal(hdu_list[2].data, interferometer_7.noise_map.in_array) np.testing.assert_array_equal(hdu_list[3].data, interferometer_7.uv_wavelengths) + # A dataset without a primary beam writes no `PRIMARY_BEAM` HDU, so the layout is pinned at 4 HDUs. + assert len(hdu_list) == 4 + def test__save_attributes__array_free_dataset__aggregator_round_trip( interferometer_7, tmp_path @@ -695,3 +698,267 @@ def test__sparse_terms_from__file_predating_fine_grids__reloads_none(interferome assert terms_reloaded.phase_centre == terms.phase_centre for name in ("sum_weights", "data_term", "noise_normalization", "n_vis", "eps"): assert getattr(terms_reloaded, name) == getattr(terms, name), name + + +def _primary_beam_from(dataset): + """ + A Gaussian primary beam on the real-space mask of `dataset` (peak 1 at the centre) whose FWHM is 2.2 + arcsec, spanning a few of the mask's 1 arcsec pixels so the beam runs from 1 down to ~0.3. + """ + import autoarray as aa + from autoarray.dataset.interferometer.primary_beam import ARCSEC_PER_RADIAN + + dish_diameter_m = 12.0 + fwhm_arcsec = 2.2 + frequency_hz = 1.13 * 299792458.0 / dish_diameter_m / (fwhm_arcsec / ARCSEC_PER_RADIAN) + + return aa.primary_beam_gaussian_from( + mask=dataset.real_space_mask, + frequency_hz=frequency_hz, + dish_diameter_m=dish_diameter_m, + ) + + +def _beamed_dataset_from(dataset, primary_beam): + import autoarray as aa + + return aa.Interferometer( + data=dataset.data, + noise_map=dataset.noise_map, + uv_wavelengths=dataset.uv_wavelengths, + real_space_mask=dataset.real_space_mask, + transformer_class=type(dataset.transformer), + primary_beam=primary_beam, + ) + + +def _array_free_beamed_dataset_from(dataset, primary_beam): + import autoarray as aa + + return aa.Interferometer.from_stream( + [(dataset.uv_wavelengths, dataset.data, dataset.noise_map)], + real_space_mask=dataset.real_space_mask, + transformer_class=type(dataset.transformer), + primary_beam=primary_beam, + ) + + +def _light_galaxies(): + return [ag.Galaxy(redshift=0.5, bulge=ag.lp.Sersic(intensity=0.1))] + + +def test__primary_beam__fixture_beam_spans_expected_range(interferometer_7): + import numpy as np + + beam = np.asarray(_primary_beam_from(interferometer_7)) + + assert beam.max() == 1.0 + assert 0.25 < beam.min() < 0.4 + + +def test__save_attributes__in_memory_beamed_dataset__aggregator_round_trip( + interferometer_7, tmp_path +): + """ + An in-memory dataset with a primary beam writes it as a trailing `PRIMARY_BEAM` HDU, and the aggregator + passes it back as `primary_beam=`, so the reloaded dataset carries an identical beam (on the dataset and + its transformer) and a fit of it reproduces the beamed log likelihood, which differs from the beamless one. + """ + import numpy as np + import pytest + from astropy.io import fits + + from autogalaxy.aggregator.interferometer.interferometer import ( + _interferometer_from, + ) + + dataset = _beamed_dataset_from( + interferometer_7, primary_beam=_primary_beam_from(interferometer_7) + ) + + paths = af.DirectoryPaths(name="in_memory_beamed", path_prefix=str(tmp_path)) + + ag.AnalysisInterferometer(dataset=dataset, use_jax=False).save_attributes( + paths=paths + ) + + with fits.open(paths.image_path / "dataset.fits") as hdu_list: + assert [hdu.name for hdu in hdu_list] == [ + "MASK", + "DATA", + "NOISE_MAP", + "UV_WAVELENGTHS", + "PRIMARY_BEAM", + ] + assert hdu_list["PRIMARY_BEAM"].data.shape == tuple( + dataset.real_space_mask.shape_native + ) + + dataset_reloaded = _interferometer_from(fit=_FitStub(paths=paths))[0] + + assert not dataset_reloaded.is_array_free + + np.testing.assert_array_equal( + dataset_reloaded.primary_beam_slim, dataset.primary_beam_slim + ) + np.testing.assert_array_equal( + dataset_reloaded.transformer.primary_beam_slim, + dataset.transformer.primary_beam_slim, + ) + + for galaxies in (_light_galaxies(), _pixelization_only_galaxies()): + log_likelihood = ag.FitInterferometer( + dataset=dataset, galaxies=galaxies + ).log_likelihood + log_likelihood_reloaded = ag.FitInterferometer( + dataset=dataset_reloaded, galaxies=galaxies + ).log_likelihood + log_likelihood_beamless = ag.FitInterferometer( + dataset=interferometer_7, galaxies=galaxies + ).log_likelihood + + assert log_likelihood_reloaded == pytest.approx(log_likelihood, rel=1.0e-10) + assert log_likelihood_reloaded != pytest.approx( + log_likelihood_beamless, rel=1.0e-6 + ) + + +def test__save_attributes__array_free_beamed_dataset__aggregator_round_trip( + interferometer_7, tmp_path +): + """ + An array-free dataset with a primary beam writes it as a trailing `PRIMARY_BEAM` HDU after its + `SparseTerms`, and the aggregator passes it back to `from_sparse_terms`, so the reloaded sparse operator + carries an identical beam and fits reproduce the beamed log evidence / figure of merit. + """ + import numpy as np + import pytest + from astropy.io import fits + + from autogalaxy.aggregator.interferometer.interferometer import ( + _interferometer_from, + ) + + dataset = _array_free_beamed_dataset_from( + interferometer_7, primary_beam=_primary_beam_from(interferometer_7) + ) + + paths = af.DirectoryPaths(name="array_free_beamed", path_prefix=str(tmp_path)) + + ag.AnalysisInterferometer(dataset=dataset, use_jax=False).save_attributes( + paths=paths + ) + + with fits.open(paths.image_path / "dataset.fits") as hdu_list: + assert [hdu.name for hdu in hdu_list][-1] == "PRIMARY_BEAM" + + dataset_reloaded = _interferometer_from(fit=_FitStub(paths=paths))[0] + + assert dataset_reloaded.is_array_free + + np.testing.assert_array_equal( + dataset_reloaded.primary_beam_slim, dataset.primary_beam_slim + ) + np.testing.assert_array_equal( + dataset_reloaded.sparse_operator.primary_beam_slim, + dataset.sparse_operator.primary_beam_slim, + ) + + galaxies = _pixelization_only_galaxies() + + log_evidence = ag.FitInterferometer(dataset=dataset, galaxies=galaxies).log_evidence + log_evidence_reloaded = ag.FitInterferometer( + dataset=dataset_reloaded, galaxies=galaxies + ).log_evidence + + assert log_evidence_reloaded == pytest.approx(log_evidence, rel=1.0e-10) + assert log_evidence_reloaded != pytest.approx( + ag.FitInterferometer( + dataset=_array_free_dataset_from(interferometer_7), galaxies=galaxies + ).log_evidence, + rel=1.0e-6, + ) + + galaxies = [*_light_galaxies(), *_pixelization_only_galaxies()] + + figure_of_merit = ag.FitInterferometer( + dataset=dataset, galaxies=galaxies + ).figure_of_merit + figure_of_merit_reloaded = ag.FitInterferometer( + dataset=dataset_reloaded, galaxies=galaxies + ).figure_of_merit + + assert figure_of_merit_reloaded == pytest.approx(figure_of_merit, rel=1.0e-10) + + +def test__save_attributes__beamless_datasets__no_primary_beam_hdu_and_reload_without_beam( + interferometer_7, tmp_path +): + from astropy.io import fits + + from autogalaxy.aggregator.interferometer.interferometer import ( + _interferometer_from, + ) + + for name, dataset in ( + ("in_memory", interferometer_7), + ("array_free", _array_free_dataset_from(interferometer_7)), + ): + paths = af.DirectoryPaths(name=f"beamless_{name}", path_prefix=str(tmp_path)) + + ag.AnalysisInterferometer(dataset=dataset, use_jax=False).save_attributes( + paths=paths + ) + + with fits.open(paths.image_path / "dataset.fits") as hdu_list: + assert "PRIMARY_BEAM" not in [hdu.name for hdu in hdu_list] + + dataset_reloaded = _interferometer_from(fit=_FitStub(paths=paths))[0] + + assert dataset_reloaded.primary_beam is None + assert dataset_reloaded.primary_beam_slim is None + + +def test__interferometer_from__outputs_written_before_primary_beam_persisted__reload_without_beam( + interferometer_7, tmp_path +): + """ + `dataset.fits` files written before the primary beam was persisted have no `PRIMARY_BEAM` HDU, even for a + beamed dataset. Both layouts reload with `primary_beam=None`, including a legacy in-memory file with no + EXTNAMEs at all (read by position), whose beam must not be read from any positional HDU. + """ + from astropy.io import fits + + from autogalaxy.aggregator.interferometer.interferometer import ( + _interferometer_from, + ) + + beam = _primary_beam_from(interferometer_7) + + for name, dataset, strip_ext_names in ( + ("in_memory", _beamed_dataset_from(interferometer_7, beam), False), + ("in_memory_legacy", _beamed_dataset_from(interferometer_7, beam), True), + ("array_free", _array_free_beamed_dataset_from(interferometer_7, beam), False), + ): + paths = af.DirectoryPaths(name=f"old_{name}", path_prefix=str(tmp_path)) + + ag.AnalysisInterferometer(dataset=dataset, use_jax=False).save_attributes( + paths=paths + ) + + dataset_path = paths.image_path / "dataset.fits" + + with fits.open(dataset_path) as hdu_list: + hdus = [hdu.copy() for hdu in hdu_list if hdu.name != "PRIMARY_BEAM"] + + if strip_ext_names: + for hdu in hdus: + hdu.header.remove("EXTNAME", ignore_missing=True) + + fits.HDUList(hdus).writeto(dataset_path, overwrite=True) + + dataset_reloaded = _interferometer_from(fit=_FitStub(paths=paths))[0] + + assert dataset_reloaded.primary_beam is None + assert dataset_reloaded.primary_beam_slim is None + assert dataset_reloaded.is_array_free == dataset.is_array_free