diff --git a/autogalaxy/aggregator/interferometer/interferometer.py b/autogalaxy/aggregator/interferometer/interferometer.py index d4772ac4..c47bb88e 100644 --- a/autogalaxy/aggregator/interferometer/interferometer.py +++ b/autogalaxy/aggregator/interferometer/interferometer.py @@ -73,6 +73,11 @@ def _sparse_terms_from(hdu_list, real_space_mask: aa.Mask2D) -> aa.SparseTerms: The `(y, x)` phase centre is reloaded from the same scalars / header cards; a file written before it was persisted (a shorter `SPARSE_TERMS_SCALARS` array and no `PHCENTY` / `PHCENTX` cards) reloads it as `None` (not recorded), as does a `NaN` entry. + + The `oversample` fine grids are read from the `PRECISION_OPERATOR_FINE` / `DIRTY_IMAGE_FINE` HDUs, with + `oversample` / `oversample_pad` from the scalars, when the file has them (they are written only on + request, see `interferometer_hdu_list_from(include_fine_grids=True)`). A file without them -- written + without opting in, or before they existed -- reloads all four as `None`. """ header = hdu_list[0].header @@ -108,6 +113,27 @@ def header_value(field): else (float(phase_centre_y), float(phase_centre_x)) ) + oversample = header_value("oversample") + oversample_pad = header_value("oversample_pad") + + if ( + _has_hdu(hdu_list, "precision_operator_fine") + and _has_hdu(hdu_list, "dirty_image_fine") + and oversample is not None + ): + fine_grids = dict( + precision_operator_fine=np.asarray( + _hdu(hdu_list, "precision_operator_fine").data, dtype="float64" + ), + dirty_image_fine=np.asarray( + _hdu(hdu_list, "dirty_image_fine").data, dtype="float64" + ), + oversample=int(oversample), + oversample_pad=None if oversample_pad is None else float(oversample_pad), + ) + else: + fine_grids = {} + return aa.SparseTerms( nufft_precision_operator=np.asarray( _hdu(hdu_list, "nufft_precision_operator").data, dtype="float64" @@ -128,6 +154,7 @@ def header_value(field): None if transformer_class_name is None else str(transformer_class_name) ), phase_centre=phase_centre, + **fine_grids, ) diff --git a/autogalaxy/interferometer/model/analysis.py b/autogalaxy/interferometer/model/analysis.py index 95425822..e577e7d5 100644 --- a/autogalaxy/interferometer/model/analysis.py +++ b/autogalaxy/interferometer/model/analysis.py @@ -9,6 +9,7 @@ It also handles adapt images, visualization, and result wrapping into `ResultInterferometer`. """ + import logging import numpy as np from typing import Optional @@ -46,15 +47,19 @@ "transformer_class_name": "TRNSFRMR", "phase_centre_y": "PHCENTY", "phase_centre_x": "PHCENTX", + "oversample": "OVERSAMP", + "oversample_pad": "OVERSPAD", } # 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` and the -# `(y, x)` phase centre (arcsec) are `NaN` when not recorded. Entries are only ever appended, so a file -# written before an entry existed has a shorter array and the loader treats the missing entries as not -# recorded. +# header cards under `SPARSE_TERMS_HEADER_KEYS` are kept as human-readable copies only. `eps`, the +# `(y, x)` phase centre (arcsec) and the `oversample` factor / pad of the fine grids are `NaN` when not +# recorded (the oversampling is recorded only when the fine grids themselves are written, so a reloaded +# `SparseTerms` never claims an `oversample` without its grids). Entries are only ever appended, so a +# file written before an entry existed has a shorter array and the loader treats the missing entries as +# not recorded. SPARSE_TERMS_SCALARS_ORDER = ( "sum_weights", "data_term", @@ -63,10 +68,14 @@ "eps", "phase_centre_y", "phase_centre_x", + "oversample", + "oversample_pad", ) -def interferometer_hdu_list_from(dataset: aa.Interferometer): +def interferometer_hdu_list_from( + dataset: aa.Interferometer, include_fine_grids: bool = False +): """ Returns the `HDUList` written to `dataset.fits` by `AnalysisInterferometer.save_attributes`, from which the aggregator (`autogalaxy.aggregator.interferometer.interferometer._interferometer_from`) reloads the @@ -86,10 +95,22 @@ def interferometer_hdu_list_from(dataset: aa.Interferometer): 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. + Terms accumulated with `oversample=q` (`Interferometer.from_stream(..., oversample=q)`) also carry the + fine precision-operator and dirty-image grids. These are written -- as HDUs `PRECISION_OPERATOR_FINE` + and `DIRTY_IMAGE_FINE` after `SPARSE_TERMS_SCALARS`, with `oversample` / `oversample_pad` in the scalars + and the `OVERSAMP` / `OVERSPAD` cards -- **only** when `include_fine_grids=True`. They are not written + by default because `dataset.fits` is written into every search's output by `save_attributes` and the + grids grow as `q^2`: ~840 MB per file at a 400 x 400 grid and `q = 8`. Without them the file is the + same as for terms without fine grids (the oversampling scalars are `NaN`), and reloads with no fine + grids; rebuild them from the visibilities, or write this HDU list yourself with + `include_fine_grids=True`, when they are needed. + Parameters ---------- dataset The interferometer dataset being fitted. + include_fine_grids + If `True` and the dataset's `SparseTerms` carry `oversample` fine grids, write them too. """ mask = dataset.real_space_mask @@ -143,6 +164,26 @@ def interferometer_hdu_list_from(dataset: aa.Interferometer): (np.nan, np.nan) if terms.phase_centre is None else terms.phase_centre ) + write_fine_grids = include_fine_grids and ( + terms.precision_operator_fine is not None and terms.dirty_image_fine is not None + ) + + if write_fine_grids: + header_dict[SPARSE_TERMS_HEADER_KEYS["oversample"]] = int(terms.oversample) + header_dict[SPARSE_TERMS_HEADER_KEYS["oversample_pad"]] = float( + terms.oversample_pad + ) + oversample = (float(terms.oversample), float(terms.oversample_pad)) + fine_values = [ + np.asarray(terms.precision_operator_fine, dtype="float64"), + np.asarray(terms.dirty_image_fine, dtype="float64"), + ] + fine_ext_names = ["precision_operator_fine", "dirty_image_fine"] + else: + oversample = (np.nan, np.nan) + fine_values = [] + fine_ext_names = [] + return hdu_list_for_output_from( values_list=[ mask.astype("float"), @@ -158,9 +199,12 @@ def interferometer_hdu_list_from(dataset: aa.Interferometer): np.nan if terms.eps is None else float(terms.eps), float(phase_centre[0]), float(phase_centre[1]), + float(oversample[0]), + float(oversample[1]), ], dtype="float64", ), + *fine_values, ], ext_name_list=[ "mask", @@ -168,6 +212,7 @@ def interferometer_hdu_list_from(dataset: aa.Interferometer): "dirty_image", "dirty_beam", "sparse_terms_scalars", + *fine_ext_names, ], header_dict=header_dict, ) diff --git a/test_autogalaxy/interferometer/model/test_analysis_interferometer.py b/test_autogalaxy/interferometer/model/test_analysis_interferometer.py index 4a5ba8fb..49ac5129 100644 --- a/test_autogalaxy/interferometer/model/test_analysis_interferometer.py +++ b/test_autogalaxy/interferometer/model/test_analysis_interferometer.py @@ -1,4 +1,3 @@ - import autofit as af import autogalaxy as ag @@ -122,6 +121,8 @@ def test__interferometer_hdu_list_from__array_free_dataset__writes_sparse_terms( terms.n_vis, terms.eps, *terms.phase_centre, + np.nan, + np.nan, ], ) np.testing.assert_array_equal( @@ -164,9 +165,7 @@ def test__interferometer_hdu_list_from__in_memory_dataset__unchanged_layout( 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[2].data, interferometer_7.noise_map.in_array) np.testing.assert_array_equal(hdu_list[3].data, interferometer_7.uv_wavelengths) @@ -186,9 +185,7 @@ def test__save_attributes__array_free_dataset__aggregator_round_trip( dataset = _array_free_dataset_from(interferometer_7) - paths = af.DirectoryPaths( - name="array_free_round_trip", path_prefix=str(tmp_path) - ) + 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) @@ -206,8 +203,7 @@ def test__save_attributes__array_free_dataset__aggregator_round_trip( assert dataset_reloaded.sparse_operator is not None assert ( - dataset_reloaded.sparse_operator.data_term - == dataset.sparse_operator.data_term + dataset_reloaded.sparse_operator.data_term == dataset.sparse_operator.data_term ) assert ( dataset_reloaded.sparse_operator.noise_normalization @@ -393,9 +389,13 @@ def test__sparse_terms_from__no_scalars_hdu__falls_back_to_header_cards( 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"]) + 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) + terms_reloaded = _sparse_terms_from( + hdu_list, real_space_mask=dataset.real_space_mask + ) for name in ( "sum_weights", @@ -486,7 +486,7 @@ def test__sparse_terms_from__phase_centre_unrecorded_or_predating_file__reloads_ ) hdu_list = interferometer_hdu_list_from(dataset=dataset_unrecorded) - assert np.isnan(hdu_list["SPARSE_TERMS_SCALARS"].data[-2:]).all() + assert np.isnan(hdu_list["SPARSE_TERMS_SCALARS"].data[5:7]).all() assert "PHCENTY" not in hdu_list[0].header assert ( @@ -523,3 +523,175 @@ def test__sparse_terms_from__phase_centre_unrecorded_or_predating_file__reloads_ assert terms_reloaded.phase_centre is None for name in ("sum_weights", "data_term", "noise_normalization", "n_vis", "eps"): assert getattr(terms_reloaded, name) == getattr(terms, name), name + + +def _fine_array_free_dataset_from(dataset, oversample=2): + 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=aa.TransformerNUFFT, + oversample=oversample, + ) + + +def test__interferometer_hdu_list_from__fine_grids__written_only_when_opted_in( + interferometer_7, +): + """ + The `oversample` fine grids are hundreds of MB at realistic sizes (~840 MB at 400 pixels, q = 8) and + `dataset.fits` is written into every search's output, so they are written only with + `include_fine_grids=True`; by default the layout and scalars are those of terms without fine grids. + """ + import numpy as np + import pytest + + pytest.importorskip("nufftax") + + from autogalaxy.interferometer.model.analysis import interferometer_hdu_list_from + + dataset = _fine_array_free_dataset_from(interferometer_7) + terms = dataset.sparse_terms + + assert terms.oversample == 2 + + 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", + ] + assert np.isnan(hdu_list["SPARSE_TERMS_SCALARS"].data[7:9]).all() + assert "OVERSAMP" not in hdu_list[0].header + + hdu_list = interferometer_hdu_list_from(dataset=dataset, include_fine_grids=True) + + assert [hdu.name for hdu in hdu_list] == [ + "MASK", + "NUFFT_PRECISION_OPERATOR", + "DIRTY_IMAGE", + "DIRTY_BEAM", + "SPARSE_TERMS_SCALARS", + "PRECISION_OPERATOR_FINE", + "DIRTY_IMAGE_FINE", + ] + np.testing.assert_array_equal( + hdu_list["PRECISION_OPERATOR_FINE"].data, terms.precision_operator_fine + ) + np.testing.assert_array_equal( + hdu_list["DIRTY_IMAGE_FINE"].data, terms.dirty_image_fine + ) + np.testing.assert_array_equal(hdu_list["SPARSE_TERMS_SCALARS"].data[7:9], [2, 0.25]) + assert hdu_list[0].header["OVERSAMP"] == 2 + + # Terms without fine grids: opting in writes nothing extra. + hdu_list = interferometer_hdu_list_from( + dataset=_array_free_dataset_from(interferometer_7), include_fine_grids=True + ) + + assert "PRECISION_OPERATOR_FINE" not in [hdu.name for hdu in hdu_list] + assert "OVERSAMP" not in hdu_list[0].header + + +def test__sparse_terms_from__fine_grids_round_trip_through_disk( + interferometer_7, tmp_path +): + import numpy as np + import pytest + + pytest.importorskip("nufftax") + + 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 = _fine_array_free_dataset_from(interferometer_7) + terms = dataset.sparse_terms + + file_path = tmp_path / "dataset.fits" + interferometer_hdu_list_from(dataset=dataset, include_fine_grids=True).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 + ) + + np.testing.assert_array_equal( + terms_reloaded.precision_operator_fine, terms.precision_operator_fine + ) + np.testing.assert_array_equal( + terms_reloaded.dirty_image_fine, terms.dirty_image_fine + ) + assert terms_reloaded.oversample == 2 + assert isinstance(terms_reloaded.oversample, int) + assert terms_reloaded.oversample_pad == 0.25 + + # Reloaded fine terms still refuse to be summed with terms of another q. + other = _fine_array_free_dataset_from(interferometer_7, oversample=4).sparse_terms + + with pytest.raises(aa.exc.InversionException, match="oversample"): + terms_reloaded + other + + # Written without opting in: reloads with no fine grids and no oversample. + file_path = tmp_path / "dataset_default.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 + ) + + assert terms_reloaded.precision_operator_fine is None + assert terms_reloaded.dirty_image_fine is None + assert terms_reloaded.oversample is None + assert terms_reloaded.oversample_pad is None + + +def test__sparse_terms_from__file_predating_fine_grids__reloads_none(interferometer_7): + """ + A file written before the fine grids existed -- a 7-entry `SPARSE_TERMS_SCALARS` array, no `OVERSAMP` + card and no fine-grid HDUs -- reloads with `None` fine grids and oversampling, other terms intact. + """ + 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_list[0]] + + [ + ( + fits.ImageHDU( + data=hdu.data[:7], header=hdu.header, name="SPARSE_TERMS_SCALARS" + ) + if hdu.name == "SPARSE_TERMS_SCALARS" + else hdu + ) + for hdu in hdu_list[1:] + ] + ) + + terms_reloaded = _sparse_terms_from( + hdu_list, real_space_mask=dataset.real_space_mask + ) + + assert terms_reloaded.precision_operator_fine is None + assert terms_reloaded.dirty_image_fine is None + assert terms_reloaded.oversample is None + assert terms_reloaded.oversample_pad is None + 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