diff --git a/autogalaxy/aggregator/agg_util.py b/autogalaxy/aggregator/agg_util.py index a642f9a5..4e6784f9 100644 --- a/autogalaxy/aggregator/agg_util.py +++ b/autogalaxy/aggregator/agg_util.py @@ -99,9 +99,11 @@ def mask_header_from(fit, name="dataset"): """ header = aa.Header(header_sci_obj=fit.value(name=name)[0].header) + pixel_scales_y = header.header_sci_obj[Mask2DKeys.PIXSCAY.value] + # Files written without a `PIXSCAX` card are assumed to have square pixels. pixel_scales = ( - header.header_sci_obj[Mask2DKeys.PIXSCAY.value], - header.header_sci_obj[Mask2DKeys.PIXSCAY.value], + pixel_scales_y, + header.header_sci_obj.get(Mask2DKeys.PIXSCAX.value, pixel_scales_y), ) origin = ( header.header_sci_obj[Mask2DKeys.ORIGINY.value], diff --git a/autogalaxy/aggregator/interferometer/interferometer.py b/autogalaxy/aggregator/interferometer/interferometer.py index 6888f11e..1e10050d 100644 --- a/autogalaxy/aggregator/interferometer/interferometer.py +++ b/autogalaxy/aggregator/interferometer/interferometer.py @@ -20,10 +20,100 @@ from functools import partial from typing import List +import numpy as np + import autofit as af import autoarray as aa from autogalaxy.aggregator import agg_util +from autogalaxy.interferometer.model.analysis import ( + SPARSE_TERMS_HEADER_KEYS, + SPARSE_TERMS_SCALARS_ORDER, +) + +# The positions of the HDUs of the legacy in-memory `dataset.fits` layout (mask / data / +# noise_map / uv_wavelengths), used when a file predates EXTNAMEs being read by name. +_LEGACY_HDU_INDEX = {"mask": 0, "data": 1, "noise_map": 2, "uv_wavelengths": 3} + + +def _hdu(hdu_list, name: str): + """ + Returns the HDU of `hdu_list` whose EXTNAME is `name` (case-insensitive), falling back to its position + in the legacy in-memory `dataset.fits` layout (`_LEGACY_HDU_INDEX`) if no HDU has that EXTNAME. + """ + try: + return hdu_list[name.upper()] + except KeyError: + return hdu_list[_LEGACY_HDU_INDEX[name]] + + +def _has_hdu(hdu_list, name: str) -> bool: + """ + Returns whether `hdu_list` has an HDU whose EXTNAME is `name` (case-insensitive). + """ + try: + hdu_list[name.upper()] + except KeyError: + return False + + return True + + +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 + by `autogalaxy.interferometer.model.analysis.interferometer_hdu_list_from`. + + The arrays are read from the HDUs named `NUFFT_PRECISION_OPERATOR`, `DIRTY_IMAGE` and `DIRTY_BEAM`, the + numeric scalars losslessly from the `SPARSE_TERMS_SCALARS` HDU (order `SPARSE_TERMS_SCALARS_ORDER`), + falling back to the primary header cards (`SPARSE_TERMS_HEADER_KEYS`, which truncate exponent-form + float64 values) for files written before that HDU existed, the transformer class name from the header, + and the mask provenance (shape, pixel scales, origin) from the reloaded real-space mask. + """ + header = hdu_list[0].header + + if _has_hdu(hdu_list, "sparse_terms_scalars"): + scalars = dict( + zip( + SPARSE_TERMS_SCALARS_ORDER, + np.asarray(_hdu(hdu_list, "sparse_terms_scalars").data, dtype="float64"), + ) + ) + + def header_value(field): + if field in scalars: + value = scalars[field] + return None if np.isnan(value) else value + return header.get(SPARSE_TERMS_HEADER_KEYS[field]) + + else: + + def header_value(field): + return header.get(SPARSE_TERMS_HEADER_KEYS[field]) + + eps = header_value("eps") + transformer_class_name = header_value("transformer_class_name") + + return aa.SparseTerms( + nufft_precision_operator=np.asarray( + _hdu(hdu_list, "nufft_precision_operator").data, dtype="float64" + ), + dirty_image_native=np.asarray( + _hdu(hdu_list, "dirty_image").data, dtype="float64" + ), + dirty_beam_native=np.asarray(_hdu(hdu_list, "dirty_beam").data, dtype="float64"), + sum_weights=float(header_value("sum_weights")), + data_term=float(header_value("data_term")), + noise_normalization=float(header_value("noise_normalization")), + n_vis=int(header_value("n_vis")), + shape_native=tuple(real_space_mask.shape_native), + pixel_scales=tuple(float(value) for value in real_space_mask.pixel_scales), + origin=tuple(float(value) for value in real_space_mask.origin), + eps=None if eps is None else float(eps), + transformer_class_name=( + None if transformer_class_name is None else str(transformer_class_name) + ), + ) def _interferometer_from( @@ -41,6 +131,11 @@ def _interferometer_from( - The real space mask defining the grid of the interferometer for the FFT (`dataset/real_space_mask.fits`). - The settings of the `Interferometer` data structure used in the fit (`dataset/settings.json`). + A fit of an array-free dataset (built by `Interferometer.from_stream` / `from_sparse_terms`) stores its + `SparseTerms` in `dataset.fits` instead of the visibilities (an HDU named `NUFFT_PRECISION_OPERATOR` is + 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`. + 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 @@ -64,13 +159,24 @@ def _interferometer_from( for fit in fit_list: real_space_mask, header = agg_util.mask_header_from(fit=fit) + hdu_list = fit.value(name="dataset") + + 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, + ) + ) + continue + data = aa.Visibilities( - visibilities=fit.value(name="dataset")[1].data.astype("float") + visibilities=_hdu(hdu_list, "data").data.astype("float") ) noise_map = aa.VisibilitiesNoiseMap( - fit.value(name="dataset")[2].data.astype("float") + _hdu(hdu_list, "noise_map").data.astype("float") ) - uv_wavelengths = fit.value(name="dataset")[3].data + uv_wavelengths = _hdu(hdu_list, "uv_wavelengths").data transformer_class = fit.value(name="transformer_class") @@ -102,6 +208,11 @@ def __init__(self, aggregator: af.Aggregator): - The real space mask defining the grid of the interferometer for the FFT (`dataset/real_space_mask.fits`). - The settings of the `Interferometer` data structure used in the fit (`dataset/settings.json`). + A fit of an array-free dataset (built by `Interferometer.from_stream` / `from_sparse_terms`) stores its + `SparseTerms` in `dataset.fits` instead of the visibilities (an HDU named `NUFFT_PRECISION_OPERATOR` is + 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`. + 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/fit_interferometer.py b/autogalaxy/interferometer/fit_interferometer.py index 73109dd7..9fd3da1e 100644 --- a/autogalaxy/interferometer/fit_interferometer.py +++ b/autogalaxy/interferometer/fit_interferometer.py @@ -231,7 +231,7 @@ def profile_image(self) -> aa.Array2D: return self.galaxies.image_2d_from(grid=self.grids.lp, xp=self._xp) @functools.cached_property - def profile_visibilities(self) -> aa.Visibilities: + def profile_visibilities(self) -> Optional[aa.Visibilities]: """ Returns the visibilities of every light profile of every galaxy, which are computed by performing a Fourier transform to the sum of light profile images. @@ -239,7 +239,26 @@ def profile_visibilities(self) -> aa.Visibilities: If the galaxies have no ordinary (non-linear) light profile (e.g. their light is entirely an MGE of linear Gaussians), the image is all zeros and the Fourier transform is skipped. This is decided structurally, so it is safe under `jax.jit`. + + On an array-free dataset (built by `Interferometer.from_stream` / `from_sparse_terms`, which has no + `uv_wavelengths` and so no transformer) there are no visibilities to compute: this returns `None` when + the galaxies have no ordinary light profile, and raises an `exc.DatasetException` when they do, because + subtracting a light profile's visibilities needs the visibility arrays. """ + if self.dataset.transformer is None: + if _has_light_profile_non_linear(galaxies=self.galaxies): + raise aa.exc.DatasetException( + "This FitInterferometer's dataset is array-free (built by from_stream / " + "from_sparse_terms) and has no visibilities or transformer, so the galaxies' " + "ordinary (non-linear) light profiles cannot be Fourier transformed and " + "subtracted. An array-free dataset supports pixelization-only and linear-light " + "fits; non-linear light profiles arrive in a later phase. Use the in-memory " + "constructor (`Interferometer(data=..., noise_map=..., uv_wavelengths=..., ...)`) " + "to fit them." + ) + + return None + if _has_light_profile_non_linear(galaxies=self.galaxies): return self.dataset.transformer.visibilities_from( image=self.profile_image, xp=self._xp @@ -250,11 +269,20 @@ def profile_visibilities(self) -> aa.Visibilities: ) @functools.cached_property - def profile_subtracted_visibilities(self) -> aa.Visibilities: + def profile_subtracted_visibilities(self) -> Optional[aa.Visibilities]: """ Returns the interferometer dataset's visibilities with all transformed light profile images subtracted. + + On an array-free dataset there are no visibilities, so this is `None`. `profile_visibilities` is + evaluated first, so a fit with ordinary light profiles on an array-free dataset raises rather than + silently fitting the unsubtracted sparse terms. """ - return self.data - self.profile_visibilities + profile_visibilities = self.profile_visibilities + + if self.data is None: + return None + + return self.data - profile_visibilities @property def _uses_precomputed_data_term(self) -> bool: @@ -332,10 +360,14 @@ def inversion_with_data(self) -> Optional[aa.AbstractInversion]: from the visibilities in that case, so the data it fitted are `fit.data`: this returns a shallow copy of the inversion (solved first, so it shares the reconstruction and every other cached quantity) whose dataset interface carries `fit.data`. In every other case it returns `inversion` itself. + + On an array-free dataset (built by `Interferometer.from_stream` / `from_sparse_terms`) `fit.data` is + `None`, so there are no visibilities to carry and `inversion` itself is returned; output quantities + that read the visibilities are unavailable on such a fit. """ inversion = self.inversion - if inversion is None or inversion.dataset.data is not None: + if inversion is None or inversion.dataset.data is not None or self.data is None: return inversion # Solve first, so the copy shares the reconstruction (and everything it cached) rather than repeating it. diff --git a/autogalaxy/interferometer/model/analysis.py b/autogalaxy/interferometer/model/analysis.py index f13a9a19..62cd2dc3 100644 --- a/autogalaxy/interferometer/model/analysis.py +++ b/autogalaxy/interferometer/model/analysis.py @@ -34,6 +34,124 @@ _FIT_INTERFEROMETER_PYTREES_REGISTERED = False +# The FITS header keys (at most 8 characters, so no `HIERARCH` cards are needed) under which +# `interferometer_hdu_list_from` stores the scalar fields and recorded provenance of the +# `SparseTerms` of an array-free dataset, and from which the aggregator reloads them. +SPARSE_TERMS_HEADER_KEYS = { + "sum_weights": "SUMW", + "data_term": "DATATERM", + "noise_normalization": "NOISENRM", + "n_vis": "NVIS", + "eps": "EPS", + "transformer_class_name": "TRNSFRMR", +} + +# The order of the float64 1-D array in the `SPARSE_TERMS_SCALARS` HDU, the lossless store of the numeric +# `SparseTerms` scalars. A FITS header card holds at most 20 characters of value, so astropy truncates +# exponent-form float64 values (e.g. `1.2345678901234567e+20` is written as `1.23456789012345E+20`); the +# header cards under `SPARSE_TERMS_HEADER_KEYS` are kept as human-readable copies only. `eps` is `NaN` +# when not recorded. +SPARSE_TERMS_SCALARS_ORDER = ( + "sum_weights", + "data_term", + "noise_normalization", + "n_vis", + "eps", +) + + +def interferometer_hdu_list_from(dataset: aa.Interferometer): + """ + Returns the `HDUList` written to `dataset.fits` by `AnalysisInterferometer.save_attributes`, from which + the aggregator (`autogalaxy.aggregator.interferometer.interferometer._interferometer_from`) reloads the + dataset. + + An in-memory dataset writes the real-space mask (in the `PrimaryHDU`), the visibilities, the noise-map and + the uv-wavelengths, with EXTNAMEs `MASK`, `DATA`, `NOISE_MAP` and `UV_WAVELENGTHS`. + + An array-free dataset (built by `Interferometer.from_stream` / `from_sparse_terms`) has no visibility arrays, + so it instead writes its `SparseTerms`: the real-space mask (in the `PrimaryHDU`), the NUFFT precision + operator `W~` and the native noise-weighted dirty image and dirty beam, with EXTNAMEs `MASK`, + `NUFFT_PRECISION_OPERATOR`, `DIRTY_IMAGE` and `DIRTY_BEAM`. The scalar terms and recorded provenance are + stored losslessly as a float64 1-D array in the `SPARSE_TERMS_SCALARS` HDU (order + `SPARSE_TERMS_SCALARS_ORDER`), and copied for readability to the header under `SPARSE_TERMS_HEADER_KEYS` + (`EPS` / `TRNSFRMR` only when recorded; `TRNSFRMR` is stored only in the header); the mask + shape, pixel scales and origin come from the mask itself. The file size is set by the real-space grid, not + the number of visibilities. + + Parameters + ---------- + dataset + The interferometer dataset being fitted. + """ + mask = dataset.real_space_mask + + if not dataset.is_array_free: + return hdu_list_for_output_from( + values_list=[ + mask.astype("float"), + dataset.data.in_array, + dataset.noise_map.in_array, + dataset.uv_wavelengths, + ], + ext_name_list=["mask", "data", "noise_map", "uv_wavelengths"], + header_dict=mask.header_dict, + ) + + terms = dataset.sparse_terms + + if terms is None: + raise aa.exc.DatasetException( + "An array-free Interferometer must carry the `SparseTerms` it was built from (`sparse_terms`) " + "for its dataset to be saved; build it via `Interferometer.from_stream` / `from_sparse_terms`." + ) + + header_dict = { + **mask.header_dict, + SPARSE_TERMS_HEADER_KEYS["sum_weights"]: float(terms.sum_weights), + SPARSE_TERMS_HEADER_KEYS["data_term"]: float(terms.data_term), + SPARSE_TERMS_HEADER_KEYS["noise_normalization"]: float( + terms.noise_normalization + ), + SPARSE_TERMS_HEADER_KEYS["n_vis"]: int(terms.n_vis), + } + + if terms.eps is not None: + header_dict[SPARSE_TERMS_HEADER_KEYS["eps"]] = float(terms.eps) + + if terms.transformer_class_name is not None: + header_dict[SPARSE_TERMS_HEADER_KEYS["transformer_class_name"]] = str( + terms.transformer_class_name + ) + + return hdu_list_for_output_from( + values_list=[ + mask.astype("float"), + np.asarray(terms.nufft_precision_operator, dtype="float64"), + np.asarray(terms.dirty_image_native, dtype="float64"), + np.asarray(terms.dirty_beam_native, dtype="float64"), + np.array( + [ + float(terms.sum_weights), + float(terms.data_term), + float(terms.noise_normalization), + float(terms.n_vis), + np.nan if terms.eps is None else float(terms.eps), + ], + dtype="float64", + ), + ], + ext_name_list=[ + "mask", + "nufft_precision_operator", + "dirty_image", + "dirty_beam", + "sparse_terms_scalars", + ], + header_dict=header_dict, + ) + + class AnalysisInterferometer(AnalysisDataset): Result = ResultInterferometer Visualizer = VisualizerInterferometer @@ -232,7 +350,8 @@ def save_attributes(self, paths: af.DirectoryPaths): It also outputs, to the `image` folder, the file `dataset.fits`, which contains: - The real space mask applied to the dataset, in the `PrimaryHDU`. - - The interferometer dataset (data / noise-map / uv_wavelengths). + - The interferometer dataset (data / noise-map / uv_wavelengths), or for an array-free dataset + its `SparseTerms` (see `interferometer_hdu_list_from`). 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 @@ -247,16 +366,7 @@ def save_attributes(self, paths: af.DirectoryPaths): """ super().save_attributes(paths=paths) - 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, - ) + hdu_list = interferometer_hdu_list_from(dataset=self.dataset) # `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 @@ -270,7 +380,10 @@ def save_attributes(self, paths: af.DirectoryPaths): if not dataset_path.exists(): hdu_list.writeto(dataset_path, overwrite=True) - paths.save_json( - "transformer_class", - to_dict(self.dataset.transformer.__class__), - ) + # An array-free dataset has no transformer; its transformer class name is recorded in + # the `dataset.fits` header instead. + if self.dataset.transformer is not None: + paths.save_json( + "transformer_class", + to_dict(self.dataset.transformer.__class__), + ) diff --git a/test_autogalaxy/interferometer/model/test_analysis_interferometer.py b/test_autogalaxy/interferometer/model/test_analysis_interferometer.py index 4a785d40..79fe99db 100644 --- a/test_autogalaxy/interferometer/model/test_analysis_interferometer.py +++ b/test_autogalaxy/interferometer/model/test_analysis_interferometer.py @@ -38,3 +38,343 @@ def test__fit_figure_of_merit__matches_correct_fit_given_galaxy_profiles( fit = ag.FitInterferometer(dataset=interferometer_7, galaxies=galaxies) assert fit.log_likelihood == fit_figure_of_merit + + +def _array_free_dataset_from(dataset): + 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), + ) + + +def _pixelization_only_galaxies(): + pixelization = ag.Pixelization( + mesh=ag.mesh.RectangularUniform(shape=(3, 3)), + regularization=ag.reg.Constant(coefficient=1.0), + ) + + return [ag.Galaxy(redshift=0.5, pixelization=pixelization)] + + +class _FitStub: + """ + The part of a `PyAutoFit` aggregator `Fit` the interferometer loader reads: `value(name)` returning the + `dataset.fits` HDU list written by `save_attributes` and the saved `transformer_class` json (if any). + """ + + def __init__(self, paths): + self.paths = paths + self.children = [] + + def value(self, name): + from astropy.io import fits + from autonerves.dictable import from_dict + + if name == "dataset": + return fits.open(self.paths.image_path / "dataset.fits") + + if name == "transformer_class": + path = self.paths._files_path / "transformer_class.json" + + if not path.exists(): + return None + + return from_dict(self.paths.load_json("transformer_class")) + + return None + + +def test__interferometer_hdu_list_from__array_free_dataset__writes_sparse_terms( + interferometer_7, +): + import numpy as np + + from autogalaxy.interferometer.model.analysis import ( + SPARSE_TERMS_HEADER_KEYS, + interferometer_hdu_list_from, + ) + + dataset = _array_free_dataset_from(interferometer_7) + terms = dataset.sparse_terms + + hdu_list = interferometer_hdu_list_from(dataset=dataset) + + assert [hdu.name for hdu in hdu_list] == [ + "MASK", + "NUFFT_PRECISION_OPERATOR", + "DIRTY_IMAGE", + "DIRTY_BEAM", + "SPARSE_TERMS_SCALARS", + ] + + np.testing.assert_array_equal( + hdu_list["NUFFT_PRECISION_OPERATOR"].data, terms.nufft_precision_operator + ) + np.testing.assert_array_equal( + hdu_list["SPARSE_TERMS_SCALARS"].data, + [ + terms.sum_weights, + terms.data_term, + terms.noise_normalization, + terms.n_vis, + terms.eps, + ], + ) + np.testing.assert_array_equal( + hdu_list["DIRTY_IMAGE"].data, terms.dirty_image_native + ) + np.testing.assert_array_equal(hdu_list["DIRTY_BEAM"].data, terms.dirty_beam_native) + + header = hdu_list[0].header + + for key in SPARSE_TERMS_HEADER_KEYS.values(): + assert len(key) <= 8 + + assert header["SUMW"] == terms.sum_weights + assert header["DATATERM"] == terms.data_term + assert header["NOISENRM"] == terms.noise_normalization + assert header["NVIS"] == terms.n_vis + assert header["EPS"] == terms.eps + assert header["TRNSFRMR"] == terms.transformer_class_name + assert header["PIXSCAY"] == dataset.real_space_mask.pixel_scales[0] + + +def test__interferometer_hdu_list_from__in_memory_dataset__unchanged_layout( + interferometer_7, +): + import numpy as np + + from autogalaxy.interferometer.model.analysis import interferometer_hdu_list_from + + hdu_list = interferometer_hdu_list_from(dataset=interferometer_7) + + assert [hdu.name.lower() for hdu in hdu_list] == [ + "mask", + "data", + "noise_map", + "uv_wavelengths", + ] + + np.testing.assert_array_equal( + hdu_list[0].data, interferometer_7.real_space_mask.astype("float") + ) + np.testing.assert_array_equal(hdu_list[1].data, interferometer_7.data.in_array) + 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) + + +def test__save_attributes__array_free_dataset__aggregator_round_trip( + interferometer_7, tmp_path +): + """ + `save_attributes` writes an array-free dataset's `SparseTerms` to `dataset.fits` and the aggregator loader + rebuilds it via `Interferometer.from_sparse_terms`, so a fit of the reloaded dataset reproduces the + original `log_evidence`. + """ + import pytest + + from autogalaxy.aggregator.interferometer.interferometer import ( + _interferometer_from, + ) + + dataset = _array_free_dataset_from(interferometer_7) + + paths = af.DirectoryPaths( + name="array_free_round_trip", path_prefix=str(tmp_path) + ) + + analysis = ag.AnalysisInterferometer(dataset=dataset, use_jax=False) + analysis.save_attributes(paths=paths) + + assert (paths.image_path / "dataset.fits").exists() + assert not (paths._files_path / "transformer_class.json").exists() + + dataset_list = _interferometer_from(fit=_FitStub(paths=paths)) + + assert len(dataset_list) == 1 + + dataset_reloaded = dataset_list[0] + + assert dataset_reloaded.is_array_free + assert dataset_reloaded.sparse_operator is not None + + assert ( + dataset_reloaded.sparse_operator.data_term + == dataset.sparse_operator.data_term + ) + assert ( + dataset_reloaded.sparse_operator.noise_normalization + == dataset.sparse_operator.noise_normalization + ) + + terms = dataset.sparse_terms + terms_reloaded = dataset_reloaded.sparse_terms + + for name in ( + "sum_weights", + "data_term", + "noise_normalization", + "n_vis", + "shape_native", + "pixel_scales", + "origin", + "eps", + "transformer_class_name", + ): + assert getattr(terms_reloaded, name) == getattr(terms, name), name + + 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-8) + + +def test__save_attributes__in_memory_dataset__aggregator_round_trip( + interferometer_7, tmp_path +): + import numpy as np + + from autogalaxy.aggregator.interferometer.interferometer import ( + _interferometer_from, + ) + + paths = af.DirectoryPaths(name="in_memory_round_trip", path_prefix=str(tmp_path)) + + analysis = ag.AnalysisInterferometer(dataset=interferometer_7, use_jax=False) + analysis.save_attributes(paths=paths) + + dataset_reloaded = _interferometer_from(fit=_FitStub(paths=paths))[0] + + assert not dataset_reloaded.is_array_free + assert type(dataset_reloaded.transformer) is type(interferometer_7.transformer) + + np.testing.assert_array_equal( + dataset_reloaded.data.array, interferometer_7.data.array + ) + np.testing.assert_array_equal( + dataset_reloaded.noise_map.array, interferometer_7.noise_map.array + ) + np.testing.assert_array_equal( + dataset_reloaded.uv_wavelengths, interferometer_7.uv_wavelengths + ) + + +def test__interferometer_hdu_list_from__array_free_dataset__scalars_round_trip_losslessly_through_disk( + interferometer_7, tmp_path +): + """ + A FITS header card holds at most 20 characters of value, so astropy truncates exponent-form float64 + values (e.g. `1.2345678901234567e+20` is written as `1.23456789012345E+20`). The `SparseTerms` scalars + are therefore stored losslessly in the `SPARSE_TERMS_SCALARS` HDU, which this test round-trips through + a real `writeto` / `fits.open` for 17-significant-digit values. + """ + import dataclasses + + import autoarray as aa + from astropy.io import fits + + from autogalaxy.aggregator.interferometer.interferometer import _sparse_terms_from + from autogalaxy.interferometer.model.analysis import interferometer_hdu_list_from + + dataset = _array_free_dataset_from(interferometer_7) + + terms = dataclasses.replace( + dataset.sparse_terms, + sum_weights=1.2345678901234567e20, + data_term=1.0e-9 * (1.0 + 1.0e-15), + noise_normalization=-1.2345678901234567e-20, + eps=9.876543210987654e-13, + ) + + dataset = aa.Interferometer.from_sparse_terms( + terms, real_space_mask=dataset.real_space_mask + ) + + file_path = tmp_path / "dataset.fits" + interferometer_hdu_list_from(dataset=dataset).writeto(file_path) + + with fits.open(file_path) as hdu_list: + terms_reloaded = _sparse_terms_from( + hdu_list, real_space_mask=dataset.real_space_mask + ) + + for name in ("sum_weights", "data_term", "noise_normalization", "n_vis", "eps"): + assert getattr(terms_reloaded, name) == getattr(terms, name), ( + name, + repr(getattr(terms_reloaded, name)), + repr(getattr(terms, name)), + ) + + +def test__save_attributes__array_free_dataset__non_square_pixel_scales_round_trip( + interferometer_7, tmp_path +): + """ + `agg_util.mask_header_from` reloads both axes of the mask pixel scales (`PIXSCAY`, `PIXSCAX`), so a + non-square mask, and the `SparseTerms` provenance copied from it, round-trip through `save_attributes`. + """ + import dataclasses + + import autoarray as aa + + from autogalaxy.aggregator.interferometer.interferometer import ( + _interferometer_from, + ) + + real_space_mask = aa.Mask2D( + mask=interferometer_7.real_space_mask, + pixel_scales=(0.1, 0.2), + origin=(0.3, -0.4), + ) + + terms = dataclasses.replace( + _array_free_dataset_from(interferometer_7).sparse_terms, + pixel_scales=(0.1, 0.2), + origin=(0.3, -0.4), + ) + + dataset = aa.Interferometer.from_sparse_terms( + terms, real_space_mask=real_space_mask + ) + + paths = af.DirectoryPaths(name="non_square_round_trip", path_prefix=str(tmp_path)) + + ag.AnalysisInterferometer(dataset=dataset, use_jax=False).save_attributes( + paths=paths + ) + + dataset_reloaded = _interferometer_from(fit=_FitStub(paths=paths))[0] + + assert tuple(dataset_reloaded.real_space_mask.pixel_scales) == (0.1, 0.2) + assert tuple(dataset_reloaded.real_space_mask.origin) == (0.3, -0.4) + assert dataset_reloaded.sparse_terms.pixel_scales == terms.pixel_scales + assert dataset_reloaded.sparse_terms.origin == terms.origin + + +def test__sparse_terms_from__no_scalars_hdu__falls_back_to_header_cards( + interferometer_7, +): + from astropy.io import fits + + from autogalaxy.aggregator.interferometer.interferometer import _sparse_terms_from + from autogalaxy.interferometer.model.analysis import interferometer_hdu_list_from + + dataset = _array_free_dataset_from(interferometer_7) + terms = dataset.sparse_terms + + hdu_list = interferometer_hdu_list_from(dataset=dataset) + hdu_list = fits.HDUList([hdu for hdu in hdu_list if hdu.name != "SPARSE_TERMS_SCALARS"]) + + terms_reloaded = _sparse_terms_from(hdu_list, real_space_mask=dataset.real_space_mask) + + for name in ("sum_weights", "data_term", "noise_normalization", "n_vis", "eps"): + assert getattr(terms_reloaded, name) == getattr(terms, name), name diff --git a/test_autogalaxy/interferometer/test_fit_interferometer.py b/test_autogalaxy/interferometer/test_fit_interferometer.py index db832b09..26d2cccb 100644 --- a/test_autogalaxy/interferometer/test_fit_interferometer.py +++ b/test_autogalaxy/interferometer/test_fit_interferometer.py @@ -855,3 +855,144 @@ def figure_of_merit_from(coefficient, xp): assert float(figure_of_merit_jax) == pytest.approx( figure_of_merit_numpy, rel=1.0e-8 ) + + +def _array_free_dataset_from(dataset): + """ + The array-free counterpart of `dataset` (no visibilities, uv-wavelengths or transformer), built by + streaming its visibilities through `Interferometer.from_stream` with the same transformer class, so its + `SparseTerms` equal those `apply_sparse_operator` forms from the resident arrays. + """ + 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), + ) + + +def test__fit_figure_of_merit__array_free_dataset__pixelization_only__matches_in_memory_sparse( + interferometer_7, +): + """ + A pixelization-only fit on an array-free dataset (built by `Interferometer.from_stream`) reads everything + its log evidence needs from the sparse operator, so it equals the fit on the in-memory sparse dataset. + """ + dataset_sparse = interferometer_7.apply_sparse_operator(use_jax=False) + dataset_array_free = _array_free_dataset_from(interferometer_7) + + assert dataset_array_free.is_array_free + assert dataset_array_free.transformer is None + + galaxies = _pixelization_only_galaxies() + + fit_sparse = ag.FitInterferometer(dataset=dataset_sparse, galaxies=galaxies) + fit_array_free = ag.FitInterferometer( + dataset=dataset_array_free, galaxies=galaxies + ) + + assert fit_array_free._uses_precomputed_data_term + assert fit_array_free.inversion.dataset.data is None + assert isinstance(fit_array_free.inversion, aa.InversionInterferometerSparse) + + assert fit_array_free.figure_of_merit == pytest.approx( + fit_sparse.figure_of_merit, rel=1.0e-8 + ) + assert fit_array_free.log_evidence == pytest.approx( + fit_sparse.log_evidence, rel=1.0e-8 + ) + + # There are no visibilities to compute or carry. + assert fit_array_free.profile_visibilities is None + assert fit_array_free.profile_subtracted_visibilities is None + assert fit_array_free.inversion_with_data is fit_array_free.inversion + + +def test__fit_figure_of_merit__array_free_dataset__pixelization_only__jax_jit_matches_numpy( + interferometer_7, +): + jax = pytest.importorskip("jax") + import jax.numpy as jnp + + dataset_sparse = interferometer_7.apply_sparse_operator(use_jax=False) + dataset_array_free = _array_free_dataset_from(interferometer_7) + + def figure_of_merit_from(coefficient, dataset, xp): + fit = ag.FitInterferometer( + dataset=dataset, + galaxies=_pixelization_only_galaxies(coefficient=coefficient), + xp=xp, + ) + + assert fit.inversion.dataset.data is None + + return fit.figure_of_merit + + figure_of_merit_sparse = figure_of_merit_from( + coefficient=1.0, dataset=dataset_sparse, xp=np + ) + figure_of_merit_numpy = figure_of_merit_from( + coefficient=1.0, dataset=dataset_array_free, xp=np + ) + + figure_of_merit_jax = jax.jit( + lambda c: figure_of_merit_from(c, dataset=dataset_array_free, xp=jnp) + )(1.0) + + assert figure_of_merit_numpy == pytest.approx(figure_of_merit_sparse, rel=1.0e-8) + assert float(figure_of_merit_jax) == pytest.approx( + figure_of_merit_numpy, rel=1.0e-8 + ) + + +def test__fit_figure_of_merit__array_free_dataset__light_profile__raises( + interferometer_7, +): + """ + Subtracting an ordinary light profile's visibilities needs the visibility arrays, so a fit with one on an + array-free dataset raises a typed exception rather than silently fitting the unsubtracted sparse terms. + """ + dataset_array_free = _array_free_dataset_from(interferometer_7) + + galaxies = [ + ag.Galaxy( + redshift=0.5, + bulge=ag.lp.Sersic(intensity=0.1, centre=(0.05, 0.05)), + ), + *_pixelization_only_galaxies(), + ] + + fit = ag.FitInterferometer(dataset=dataset_array_free, galaxies=galaxies) + + assert not fit._uses_precomputed_data_term + + with pytest.raises(aa.exc.DatasetException): + fit.profile_visibilities + + with pytest.raises(aa.exc.DatasetException): + ag.FitInterferometer( + dataset=dataset_array_free, galaxies=galaxies + ).figure_of_merit + + +def test__fit_figure_of_merit__array_free_dataset__linear_light_only__matches_in_memory_sparse( + interferometer_7, +): + dataset_sparse = interferometer_7.apply_sparse_operator(use_jax=False) + dataset_array_free = _array_free_dataset_from(interferometer_7) + + galaxies = [ + ag.Galaxy( + redshift=0.5, + bulge=ag.lp_linear.Gaussian(sigma=0.5, centre=(0.05, 0.05)), + ) + ] + + fit_array_free = ag.FitInterferometer( + dataset=dataset_array_free, galaxies=galaxies + ) + + assert fit_array_free.inversion.dataset.data is None + assert fit_array_free.figure_of_merit == pytest.approx( + ag.FitInterferometer(dataset=dataset_sparse, galaxies=galaxies).figure_of_merit, + rel=1.0e-8, + )