diff --git a/autoarray/dataset/abstract/dataset.py b/autoarray/dataset/abstract/dataset.py index 2113f2f93..fd58d8eb8 100644 --- a/autoarray/dataset/abstract/dataset.py +++ b/autoarray/dataset/abstract/dataset.py @@ -94,7 +94,12 @@ def __init__( self.noise_covariance_matrix = noise_covariance_matrix - if noise_map is None: + # An array-free dataset (e.g. `Interferometer.from_stream`) carries neither data nor a + # noise-map: everything its likelihood needs was reduced into a sparse operator + # upstream, so there is no diagonal to take and nothing to check. + if noise_map is None and data is None and noise_covariance_matrix is None: + pass + elif noise_map is None: try: noise_map = Array2D.no_mask( values=np.diag(noise_covariance_matrix), @@ -174,6 +179,10 @@ def shape_slim(self): """ The 1D size of the dataset data array after masking, i.e. the number of unmasked pixels. """ + # An array-free dataset has no data array; the mask's slim size is not a substitute + # (for an interferometer the data is visibilities, not image pixels). + if self.data is None: + return None return self.data.shape_slim @property diff --git a/autoarray/dataset/interferometer/dataset.py b/autoarray/dataset/interferometer/dataset.py index e81abd65a..df057bbab 100644 --- a/autoarray/dataset/interferometer/dataset.py +++ b/autoarray/dataset/interferometer/dataset.py @@ -53,6 +53,9 @@ def __init__( transformer_class=TransformerNUFFT, sparse_operator: Optional[InterferometerSparseOperator] = None, raise_error_dft_visibilities_limit: bool = True, + sparse_terms: Optional[ + "inversion_interferometer_util.SparseTerms" + ] = None, ): """ An interferometer dataset, containing the visibilities data, noise-map, real-space msk, Fourier transformer and @@ -106,6 +109,20 @@ def __init__( If `True`, an exception is raised if the dataset has more than 10,000 visibilities and `transformer_class=TransformerDFT`. The DFT is too slow for large datasets and `TransformerNUFFT` should be used instead. Set to `False` to suppress this check. + sparse_terms + The `SparseTerms` record the `sparse_operator` was accumulated from, if any (set by + `from_stream`, `from_sparse_terms` and `apply_sparse_operator_from_chunks`). It + carries the naturally weighted dirty image and beam and the provenance of the + accumulation. + + Array-free datasets + ------------------- + `data`, `noise_map` and `uv_wavelengths` may all be `None` with `transformer_class=None` + when a `sparse_operator` carrying its `data_term` and `noise_normalization` is given: + this is the dataset `from_stream` / `from_sparse_terms` build, on which a sparse + (w-tilde) pixelized inversion and its `log_evidence` need no visibility arrays. Its + `transformer` is `None`, and every quantity that needs the visibilities raises an + `exc.DatasetException` naming the missing input. """ self.real_space_mask = real_space_mask @@ -118,10 +135,13 @@ def __init__( self.uv_wavelengths = uv_wavelengths - self.transformer = transformer_class( - uv_wavelengths=uv_wavelengths, - real_space_mask=real_space_mask, - ) + if uv_wavelengths is not None and transformer_class is not None: + self.transformer = transformer_class( + uv_wavelengths=uv_wavelengths, + real_space_mask=real_space_mask, + ) + else: + self.transformer = None self.grids = GridsDataset( mask=self.real_space_mask, @@ -130,8 +150,9 @@ def __init__( ) self.sparse_operator = sparse_operator + self.sparse_terms = sparse_terms - if raise_error_dft_visibilities_limit: + if raise_error_dft_visibilities_limit and self.uv_wavelengths is not None: if ( self.uv_wavelengths.shape[0] > 10000 and transformer_class == TransformerDFT @@ -227,6 +248,166 @@ def from_fits( raise_error_dft_visibilities_limit=raise_error_dft_visibilities_limit, ) + @classmethod + def from_sparse_terms( + cls, + terms: "inversion_interferometer_util.SparseTerms", + real_space_mask: Mask2D, + *, + batch_size: int = 128, + ) -> "Interferometer": + """ + Build an array-free `Interferometer` from an accumulated `SparseTerms` record. + + The returned dataset carries no visibility arrays: `data`, `noise_map`, + `uv_wavelengths` and `transformer` are all `None`. It carries the + `InterferometerSparseOperator` built from `terms` (with its cached `data_term` and + `noise_normalization`) and `terms` itself (as `sparse_terms`), which is everything a + sparse (w-tilde) pixelized inversion and its `log_evidence` read, so a fit on it + equals one on the in-memory `apply_sparse_operator()` dataset. Quantities that need + the visibilities (`amplitudes`, `dirty_image`, residual maps of a fit, ...) raise an + `exc.DatasetException`; the naturally weighted `dirty_image_natural` and `dirty_beam` + are available from the terms. + + Parameters + ---------- + terms + The accumulated sums of every visibility the dataset represents (e.g. from + `inversion_interferometer_util.sparse_terms_from_chunks`). + real_space_mask + The real-space `Mask2D` the terms were accumulated on. + batch_size + The number of source-pixel columns processed per batch by the sparse operator. + """ + shape_native = getattr(terms, "shape_native", None) + + if shape_native is not None and tuple(shape_native) != tuple( + real_space_mask.shape_native + ): + raise exc.DatasetException( + f"The SparseTerms were accumulated on a real-space mask of shape_native " + f"{tuple(shape_native)} but `from_sparse_terms` was given a mask of shape_native " + f"{tuple(real_space_mask.shape_native)}." + ) + + # Exact equality of the float tuples (as `SparseTerms.__add__` uses): the recorded + # values are copied from a `Mask2D`, so the same mask reproduces them bit-for-bit. + for name in ("pixel_scales", "origin"): + value_terms = getattr(terms, name, None) + + if value_terms is None: + continue + + value_terms = tuple(float(value) for value in value_terms) + value_mask = tuple(float(value) for value in getattr(real_space_mask, name)) + + if value_terms != value_mask: + raise exc.DatasetException( + f"The SparseTerms were accumulated on a real-space mask of {name} " + f"{value_terms} but `from_sparse_terms` was given a mask of {name} " + f"{value_mask}." + ) + + sparse_operator = InterferometerSparseOperator.from_sparse_terms( + terms, + real_space_mask=real_space_mask, + batch_size=batch_size, + ) + + return cls( + data=None, + noise_map=None, + uv_wavelengths=None, + real_space_mask=real_space_mask, + transformer_class=None, + sparse_operator=sparse_operator, + sparse_terms=terms, + ) + + @classmethod + def from_stream( + cls, + chunks, + real_space_mask: Mask2D, + *, + transformer_class=TransformerNUFFT, + method: str = "nufft", + eps: Optional[float] = None, + chunk_size: Optional[int] = None, + chunk_k: int = 2048, + use_jax: bool = False, + show_progress: bool = False, + batch_size: int = 128, + ) -> "Interferometer": + """ + Build an array-free `Interferometer` by accumulating a stream of visibility chunks, + never holding the full visibility arrays. + + `chunks` is any iterable (a list, or a generator reading from disk) of + `(uv_wavelengths, data, noise_map)` triples, with the contract of + `apply_sparse_operator_from_chunks`. They are reduced one at a time by + `inversion_interferometer_util.sparse_terms_from_chunks` into a `SparseTerms` record, + which `from_sparse_terms` turns into the dataset; peak memory is set by the chunk size, + not the dataset size. + + Parameters + ---------- + chunks + The `(uv_wavelengths, data, noise_map)` chunks. + real_space_mask + The real-space `Mask2D` the terms are accumulated on. + transformer_class, method, eps, chunk_size, chunk_k, use_jax, show_progress + Passed to `sparse_terms_from_chunks`. + batch_size + The number of source-pixel columns processed per batch by the sparse operator. + + Raises + ------ + exc.DatasetException + If any chunk has unequal real and imaginary noise sigma. + """ + if disable_jax(): + use_jax = False + + terms = inversion_interferometer_util.sparse_terms_from_chunks( + chunks, + real_space_mask=real_space_mask, + transformer_class=transformer_class, + method=method, + eps=eps, + chunk_size=chunk_size, + chunk_k=chunk_k, + use_jax=use_jax, + show_progress=show_progress, + ) + + return cls.from_sparse_terms( + terms, real_space_mask=real_space_mask, batch_size=batch_size + ) + + @property + def is_array_free(self) -> bool: + """ + `True` for a dataset built by `from_stream` / `from_sparse_terms`, which carries no + visibility arrays (`data`, `noise_map`, `uv_wavelengths`) and no transformer. + """ + return self.data is None and self.uv_wavelengths is None + + def _require(self, quantity: str, *names: str): + """ + Raise a typed `exc.DatasetException` if any of the named inputs of this dataset is + `None`, i.e. this is an array-free dataset and `quantity` cannot be computed. + """ + missing = [name for name in names if getattr(self, name) is None] + + if missing: + raise exc.DatasetException( + f"This Interferometer is array-free (built by from_stream / from_sparse_terms) " + f"and has no {' / '.join(missing)}; `{quantity}` is unavailable. Use the " + f"in-memory constructor (`Interferometer(data=..., noise_map=..., " + f"uv_wavelengths=..., ...)`) if you need it." + ) + def apply_sparse_operator( self, nufft_precision_operator=None, @@ -334,6 +515,10 @@ def apply_sparse_operator( if disable_jax(): use_jax = False + self._require( + "apply_sparse_operator", "data", "noise_map", "uv_wavelengths", "transformer" + ) + inversion_interferometer_util.check_noise_map_real_imag_equal(self.noise_map) if nufft_precision_operator is None: @@ -504,6 +689,8 @@ def apply_sparse_operator_from_chunks( if disable_jax() and accumulator_kwargs.get("use_jax", False): accumulator_kwargs["use_jax"] = False + self._require("apply_sparse_operator_from_chunks", "transformer") + transformer = self.transformer if isinstance(transformer, TransformerNUFFT): @@ -543,6 +730,7 @@ def transformer_class(uv_wavelengths, real_space_mask): uv_wavelengths=self.uv_wavelengths, transformer_class=lambda uv_wavelengths, real_space_mask: self.transformer, sparse_operator=sparse_operator, + sparse_terms=terms, ) def psf_precision_operator_from( @@ -599,6 +787,10 @@ def psf_precision_operator_from( The NUFFT precision matrix of shape (total_pixels, total_pixels) where total_pixels is the number of unmasked real-space pixels. """ + self._require( + "psf_precision_operator_from", "noise_map", "uv_wavelengths", "transformer" + ) + transformer = self.transformer # The NUFFT builder and `TransformerNUFFT` spread the same visibilities onto a mode grid @@ -645,6 +837,7 @@ def amplitudes(self): The amplitudes of the complex visibilities, defined as the absolute value of each visibility: amplitude = sqrt(real^2 + imag^2). """ + self._require("amplitudes", "data") return self.data.amplitudes @property @@ -653,6 +846,7 @@ def phases(self): The phases of the complex visibilities in radians, defined as arctan(imag / real) for each visibility. """ + self._require("phases", "data") return self.data.phases @property @@ -661,6 +855,7 @@ def uv_distances(self): The radial distance of each visibility baseline from the origin of the UV-plane, in units of wavelengths. Computed as sqrt(u^2 + v^2) for each (u, v) baseline pair. """ + self._require("uv_distances", "uv_wavelengths") return np.sqrt( np.square(self.uv_wavelengths[:, 0]) + np.square(self.uv_wavelengths[:, 1]) ) @@ -674,8 +869,72 @@ def dirty_image(self): It provides a quick visual representation of the data but is convolved with the synthesized beam (the Fourier transform of the UV-plane sampling function). """ + self._require("dirty_image", "data", "transformer") return self.transformer.image_from(visibilities=self.data) + @property + def dirty_image_natural(self): + """ + The naturally weighted, normalised dirty image `Re(F^H (w d)) / sum(w)`, with + `w = 1 / sigma^2` per visibility component. + + Unlike `dirty_image` (the unweighted adjoint of the data), this is the image whose + peak for a unit point source at the phase centre is 1 -- the standard dirty image of + radio astronomy. It is available on both dataset types: an array-free dataset reads it + from its `sparse_terms` (`dirty_image_native / sum_weights`), an in-memory one computes + it from its visibilities and transformer. + """ + from autoarray.structures.arrays.uniform_2d import Array2D + + if getattr(self, "sparse_terms", None) is not None: + return Array2D( + values=np.asarray(self.sparse_terms.dirty_image_native) + / self.sparse_terms.sum_weights, + mask=self.real_space_mask, + ) + + self._require("dirty_image_natural", "data", "noise_map", "transformer") + + noise_map = self.noise_map.array + data = self.data.array + + weighted = Visibilities( + visibilities=data.real * noise_map.real**-2.0 + + 1j * data.imag * noise_map.imag**-2.0 + ) + + return self.transformer.image_from(visibilities=weighted) / float( + np.sum(noise_map.real**-2.0) + ) + + @property + def dirty_beam(self): + """ + The naturally weighted, normalised dirty beam (synthesised beam) at the image pixels, + `Re(F^H w) / sum(w)` with `w = 1 / sigma_r^2`: the response of `dirty_image_natural` + to a unit point source at the phase centre (peak 1 at the phase-centre pixel). + + An array-free dataset reads it from its `sparse_terms` + (`dirty_beam_native / sum_weights`); an in-memory one computes it from its noise-map + and transformer. + """ + from autoarray.structures.arrays.uniform_2d import Array2D + + if getattr(self, "sparse_terms", None) is not None: + return Array2D( + values=np.asarray(self.sparse_terms.dirty_beam_native) + / self.sparse_terms.sum_weights, + mask=self.real_space_mask, + ) + + self._require("dirty_beam", "noise_map", "transformer") + + weights = self.noise_map.array.real**-2.0 + + return self.transformer.image_from( + visibilities=Visibilities(visibilities=weights.astype(np.complex128)) + ) / float(np.sum(weights)) + @property def dirty_noise_map(self): """ @@ -683,6 +942,7 @@ def dirty_noise_map(self): Provides a real-space representation of the noise levels in the dirty image. """ + self._require("dirty_noise_map", "noise_map", "transformer") return self.transformer.image_from(visibilities=self.noise_map) @property @@ -691,6 +951,7 @@ def dirty_signal_to_noise_map(self): The dirty signal-to-noise map, computed as the inverse Fourier transform of the complex signal-to-noise visibility map. """ + self._require("dirty_signal_to_noise_map", "data", "noise_map", "transformer") return self.transformer.image_from(visibilities=self.signal_to_noise_map) @property @@ -705,6 +966,7 @@ def signal_to_noise_map(self): handles the complex nature of interferometric visibilities by treating the real and imaginary parts independently. """ + self._require("signal_to_noise_map", "data", "noise_map") signal_to_noise_map_real = np.divide( np.real(self.data.array), np.real(self.noise_map.array) ) diff --git a/autoarray/fit/fit_interferometer.py b/autoarray/fit/fit_interferometer.py index 408689de5..b0873b8be 100644 --- a/autoarray/fit/fit_interferometer.py +++ b/autoarray/fit/fit_interferometer.py @@ -1,3 +1,4 @@ +import functools import numpy as np from autoarray.dataset.interferometer.dataset import Interferometer @@ -8,6 +9,7 @@ from autoarray.fit.fit_dataset import FitDataset from autoarray.fit import fit_util +from autoarray import exc from autoarray import type as ty @@ -56,6 +58,26 @@ def __init__( xp=xp, ) + def _require(self, quantity: str, *names: str): + """ + Raise a typed `exc.DatasetException` when this fit's dataset is array-free (an + `Interferometer` built by `from_stream` / `from_sparse_terms`) and one of the named + inputs (`data`, `noise_map`, `transformer`) the requested `quantity` needs is `None`. + + The `log_evidence` / `figure_of_merit` of a sparse inversion fit never reaches here: + it reads `inversion.fast_chi_squared` and the sparse operator's cached + `noise_normalization`, neither of which touches the visibility arrays. + """ + missing = [name for name in names if getattr(self, name) is None] + + if missing: + raise exc.DatasetException( + f"This FitInterferometer's dataset is array-free (built by from_stream / " + f"from_sparse_terms) and has no {' / '.join(missing)}; `{quantity}` is " + f"unavailable. Only the `log_evidence` / `figure_of_merit` of a sparse " + f"inversion can be computed; use the in-memory constructor if you need it." + ) + @property def mask(self) -> np.ndarray: """ @@ -64,6 +86,7 @@ def mask(self) -> np.ndarray: Interferometer data is not spatially masked in the same way as imaging data — all visibility measurements are included in the fit — so this always returns an unmasked array. """ + self._require("mask", "data") return np.full(shape=self.data.shape, fill_value=False) @property @@ -74,7 +97,26 @@ def transformer(self) -> ty.Transformer: This is taken directly from the interferometer dataset and is used internally to compute the `dirty_*` image-space representations of the fit quantities. """ - return self.dataset.transformer + transformer = self.dataset.transformer + + if transformer is None: + raise exc.DatasetException( + "This FitInterferometer's dataset is array-free (built by from_stream / " + "from_sparse_terms) and has no transformer; `transformer` is unavailable. " + "Use the in-memory constructor if you need it." + ) + + return transformer + + @functools.cached_property + def residual_map(self): + """ + Returns the residual-map between the visibility data and model data (data - model_data). + + Raises an `exc.DatasetException` on an array-free dataset, which has no data. + """ + self._require("residual_map", "data") + return super().residual_map @property def normalized_residual_map(self) -> np.ndarray: @@ -83,6 +125,7 @@ def normalized_residual_map(self) -> np.ndarray: Normalized_Residual = (Data - Model_Data) / Noise """ + self._require("normalized_residual_map", "data", "noise_map") return fit_util.normalized_residual_map_complex_from( residual_map=self.residual_map, noise_map=self.noise_map, @@ -95,6 +138,7 @@ def chi_squared_map(self) -> np.ndarray: Chi_Squared = ((Residuals) / (Noise)) ** 2.0 = ((Data - Model)**2.0)/(Variances) """ + self._require("chi_squared_map", "data", "noise_map") return fit_util.chi_squared_map_complex_from( residual_map=self.residual_map, noise_map=self.noise_map, @@ -104,6 +148,7 @@ def chi_squared_map(self) -> np.ndarray: def signal_to_noise_map(self) -> np.ndarray: """ The signal-to-noise_map of the dataset and noise-map which are fitted.""" + self._require("signal_to_noise_map", "data", "noise_map") signal_to_noise_map_real = self.data.real / self.noise_map.real signal_to_noise_map_real[signal_to_noise_map_real < 0] = 0.0 @@ -118,6 +163,7 @@ def chi_squared(self) -> float: """ Returns the chi-squared terms of the model data's fit to an dataset, by summing the chi-squared-map. """ + self._require("chi_squared", "data", "noise_map") return fit_util.chi_squared_complex_from( chi_squared_map=self.chi_squared_map.array, ) diff --git a/autoarray/inversion/inversion/interferometer/abstract.py b/autoarray/inversion/inversion/interferometer/abstract.py index 7d912efdb..8915c6532 100644 --- a/autoarray/inversion/inversion/interferometer/abstract.py +++ b/autoarray/inversion/inversion/interferometer/abstract.py @@ -61,7 +61,26 @@ def transformer(self): @property def mask(self) -> Mask2D: - return self.transformer.real_space_mask + """ + The real-space mask the inversion reconstructs on. + + This is read from the dataset rather than its transformer, so it is defined on an + array-free `Interferometer` (built by `from_stream` / `from_sparse_terms`), whose + transformer is `None`: an `Interferometer`'s `real_space_mask`, else a transformer's + `real_space_mask`, else the dataset's own `mask` (a `DatasetInterface`'s + `grids.lp.mask`). + """ + real_space_mask = getattr(self.dataset, "real_space_mask", None) + + if real_space_mask is not None: + return real_space_mask + + transformer = getattr(self.dataset, "transformer", None) + + if transformer is not None: + return transformer.real_space_mask + + return self.dataset.mask @cached_property def operated_mapping_matrix_list(self) -> List[np.ndarray]: diff --git a/autoarray/inversion/inversion/interferometer/inversion_interferometer_util.py b/autoarray/inversion/inversion/interferometer/inversion_interferometer_util.py index 096698e36..5bf8fd0e0 100644 --- a/autoarray/inversion/inversion/interferometer/inversion_interferometer_util.py +++ b/autoarray/inversion/inversion/interferometer/inversion_interferometer_util.py @@ -1839,6 +1839,21 @@ class SparseTerms: `fit_util.noise_normalization_complex_from(noise_map)`. n_vis The number of visibilities summed. + + Provenance + ---------- + The trailing optional fields record what the terms were accumulated on, so terms built on + different geometries or accuracies cannot be silently summed (`__add__` raises), and so an + array-free `Interferometer` (built by `Interferometer.from_stream` / + `from_sparse_terms`) can say what produced it. `None` means "not recorded" and skips the + corresponding check. + + shape_native, pixel_scales, origin + The real-space mask's native shape, pixel scales and origin. + eps + The NUFFT precision the precision operator was built with. + transformer_class_name + The class name of the transformer that formed the dirty image and beam. """ nufft_precision_operator: np.ndarray @@ -1848,8 +1863,26 @@ class SparseTerms: data_term: float noise_normalization: float n_vis: int + shape_native: Optional[tuple] = None + pixel_scales: Optional[tuple] = None + origin: Optional[tuple] = None + eps: Optional[float] = None + transformer_class_name: Optional[str] = None def __add__(self, other: "SparseTerms") -> "SparseTerms": + """ + The field-wise sum of two `SparseTerms`. + + Raises `exc.InversionException` if a provenance field (`shape_native`, + `pixel_scales`, `origin`, `eps`) is recorded on both sides with different values. + The result carries, for each provenance field, the recorded value from either side + (the left one when both are recorded), so an unrecorded operand never erases the + provenance of a recorded one. + + Known limit: the mask itself is not carried, so terms accumulated on two masks with + the same `shape_native`, `pixel_scales` and `origin` but different masked pixels + cannot be distinguished by provenance and are summed without error. + """ if not isinstance(other, SparseTerms): return NotImplemented @@ -1860,6 +1893,26 @@ def __add__(self, other: "SparseTerms") -> "SparseTerms": f"{other.nufft_precision_operator.shape} differ." ) + for name in ("shape_native", "pixel_scales", "origin", "eps"): + value_self = getattr(self, name) + value_other = getattr(other, name) + + if value_self is None or value_other is None: + continue + + if isinstance(value_self, float) or isinstance(value_other, float): + differ = float(value_self) != float(value_other) + else: + differ = tuple(value_self) != tuple(value_other) + + if differ: + raise exc.InversionException( + "SparseTerms can only be added when accumulated with the same provenance: " + f"`{name}` is {value_self!r} on one and {value_other!r} on the other. " + "Terms accumulated on different real-space masks or NUFFT accuracies do " + "not describe the same operator." + ) + return SparseTerms( nufft_precision_operator=self.nufft_precision_operator + other.nufft_precision_operator, @@ -1871,9 +1924,26 @@ def __add__(self, other: "SparseTerms") -> "SparseTerms": self.noise_normalization + other.noise_normalization ), n_vis=int(self.n_vis + other.n_vis), + shape_native=_recorded(self.shape_native, other.shape_native), + pixel_scales=_recorded(self.pixel_scales, other.pixel_scales), + origin=_recorded(self.origin, other.origin), + eps=_recorded(self.eps, other.eps), + transformer_class_name=_recorded( + self.transformer_class_name, other.transformer_class_name + ), ) +def _recorded(value_left, value_right): + """ + The provenance value a sum of two `SparseTerms` carries: the recorded (non-`None`) value + of either side, the left one when both are recorded (they are equal, or `__add__` has + already raised). Taking it from either side keeps `(U + A) + B` (with `U` unrecorded) + checking `B` against `A`'s provenance. + """ + return value_left if value_left is not None else value_right + + def _complex_visibilities_from(values) -> np.ndarray: """ Return `values` (a `Visibilities`/`VisibilitiesNoiseMap`, a complex ndarray of shape @@ -2045,6 +2115,11 @@ def sparse_terms_from_chunks( + np.sum(np.log(2 * np.pi * noise_map_imag**2.0)) ), n_vis=int(uv_wavelengths.shape[0]), + shape_native=tuple(real_space_mask.shape_native), + pixel_scales=tuple(real_space_mask.pixel_scales), + origin=tuple(real_space_mask.origin), + eps=float(eps), + transformer_class_name=type(transformer).__name__, ) terms = chunk_terms if terms is None else terms + chunk_terms diff --git a/autoarray/inversion/inversion/interferometer/sparse.py b/autoarray/inversion/inversion/interferometer/sparse.py index 914905189..49fb23ba0 100644 --- a/autoarray/inversion/inversion/interferometer/sparse.py +++ b/autoarray/inversion/inversion/interferometer/sparse.py @@ -448,6 +448,16 @@ def mapped_reconstructed_operated_data_dict( The reconstruction (in the source frame) whose values are mapped to a dictionary of values for each individual mapper (in the image-plane). """ + if self.transformer is None: + raise exc.InversionException( + "This interferometer inversion's dataset is array-free (an `Interferometer` " + "built by from_stream / from_sparse_terms) and has no transformer, so the " + "reconstruction cannot be mapped to visibilities " + "(`mapped_reconstructed_operated_data_dict`). The `log_evidence` and " + "`fast_chi_squared` do not need it; use the in-memory constructor if you need " + "model visibilities." + ) + mapped_reconstructed_operated_data_dict = {} image_dict = self.mapped_reconstructed_data_dict diff --git a/test_autoarray/dataset/interferometer/test_dataset.py b/test_autoarray/dataset/interferometer/test_dataset.py index a7d805ffe..ff0482c31 100644 --- a/test_autoarray/dataset/interferometer/test_dataset.py +++ b/test_autoarray/dataset/interferometer/test_dataset.py @@ -1,3 +1,4 @@ +import dataclasses import numpy as np import autoarray as aa @@ -667,3 +668,196 @@ def test__apply_sparse_operator_from_chunks__unequal_real_imag_noise__raises( with pytest.raises(aa.exc.DatasetException): dataset.apply_sparse_operator_from_chunks(chunks) + + +def _chunks_of(dataset, edges): + return [ + ( + dataset.uv_wavelengths[k0:k1], + dataset.data.array[k0:k1], + dataset.noise_map.array[k0:k1], + ) + for k0, k1 in zip(edges[:-1], edges[1:]) + ] + + +def test__from_stream__array_free_dataset_carries_terms_and_operator(mask_2d_7x7): + pytest.importorskip("nufftax") + + dataset = _random_interferometer(mask_2d_7x7, transformer.TransformerNUFFT) + chunks = _chunks_of(dataset, [0, 13, 40]) + + dataset_stream = aa.Interferometer.from_stream(chunks, mask_2d_7x7) + + assert dataset_stream.data is None + assert dataset_stream.noise_map is None + assert dataset_stream.uv_wavelengths is None + assert dataset_stream.transformer is None + assert dataset_stream.is_array_free + assert dataset_stream.real_space_mask is mask_2d_7x7 + assert dataset_stream.mask is mask_2d_7x7 + assert dataset_stream.shape_slim is None + + terms = dataset_stream.sparse_terms + + assert terms.n_vis == 40 + assert terms.shape_native == mask_2d_7x7.shape_native + assert terms.pixel_scales == mask_2d_7x7.pixel_scales + assert terms.origin == mask_2d_7x7.origin + assert terms.eps == 1.0e-12 + assert terms.transformer_class_name == "TransformerNUFFT" + + one_shot = dataset.apply_sparse_operator().sparse_operator + + assert dataset_stream.sparse_operator.data_term == pytest.approx( + one_shot.data_term, rel=1.0e-12 + ) + assert dataset_stream.sparse_operator.noise_normalization == pytest.approx( + one_shot.noise_normalization, rel=1.0e-12 + ) + + dataset_terms = aa.Interferometer.from_sparse_terms(terms, mask_2d_7x7) + + assert dataset_terms.transformer is None + assert dataset_terms.data is None + assert dataset_terms.sparse_terms is terms + assert dataset_terms.sparse_operator.data_term == terms.data_term + + +def test__from_sparse_terms__mask_shape_mismatch__raises(mask_2d_7x7): + pytest.importorskip("nufftax") + + dataset = _random_interferometer(mask_2d_7x7, transformer.TransformerNUFFT) + terms = aa.util.inversion_interferometer.sparse_terms_from_chunks( + _chunks_of(dataset, [0, 40]), real_space_mask=mask_2d_7x7 + ) + + other_mask = aa.Mask2D.circular(shape_native=(10, 10), pixel_scales=1.0, radius=3.0) + + with pytest.raises(aa.exc.DatasetException): + aa.Interferometer.from_sparse_terms(terms, other_mask) + + +def test__from_sparse_terms__mask_pixel_scales_and_origin_mismatch__raises(mask_2d_7x7): + pytest.importorskip("nufftax") + + dataset = _random_interferometer(mask_2d_7x7, transformer.TransformerNUFFT) + terms = aa.util.inversion_interferometer.sparse_terms_from_chunks( + _chunks_of(dataset, [0, 40]), real_space_mask=mask_2d_7x7 + ) + + # Same shape_native, different pixel_scales / origin. + for name, kwargs in ( + ("pixel_scales", dict(pixel_scales=2.0 * mask_2d_7x7.pixel_scales[0])), + ( + "origin", + dict(pixel_scales=mask_2d_7x7.pixel_scales, origin=(0.5, -0.5)), + ), + ): + other_mask = aa.Mask2D(mask=np.asarray(mask_2d_7x7), **kwargs) + + with pytest.raises(aa.exc.DatasetException, match=name): + aa.Interferometer.from_sparse_terms(terms, other_mask) + + # Unrecorded provenance skips the check. + terms_unrecorded = dataclasses.replace(terms, pixel_scales=None, origin=None) + aa.Interferometer.from_sparse_terms( + terms_unrecorded, + aa.Mask2D(mask=np.asarray(mask_2d_7x7), pixel_scales=5.0, origin=(1.0, 1.0)), + ) + + +def test__array_free__array_properties_raise_typed_exception(mask_2d_7x7): + pytest.importorskip("nufftax") + + dataset = _random_interferometer(mask_2d_7x7, transformer.TransformerNUFFT) + dataset_stream = aa.Interferometer.from_stream( + _chunks_of(dataset, [0, 40]), mask_2d_7x7 + ) + + for name in ( + "amplitudes", + "phases", + "uv_distances", + "dirty_image", + "dirty_noise_map", + "signal_to_noise_map", + "dirty_signal_to_noise_map", + ): + with pytest.raises(aa.exc.DatasetException, match="array-free"): + getattr(dataset_stream, name) + + with pytest.raises(aa.exc.DatasetException, match="array-free"): + dataset_stream.psf_precision_operator_from() + + with pytest.raises(aa.exc.DatasetException, match="array-free"): + dataset_stream.apply_sparse_operator() + + with pytest.raises(aa.exc.DatasetException, match="array-free"): + dataset_stream.apply_sparse_operator_from_chunks(_chunks_of(dataset, [0, 40])) + + +def test__dirty_image_natural_and_dirty_beam__stream_and_in_memory_agree( + mask_2d_7x7, +): + pytest.importorskip("nufftax") + + for transformer_class in (transformer.TransformerNUFFT, transformer.TransformerDFT): + dataset = _random_interferometer(mask_2d_7x7, transformer_class) + + dataset_stream = aa.Interferometer.from_stream( + _chunks_of(dataset, [0, 7, 40]), + mask_2d_7x7, + transformer_class=transformer_class, + ) + + weights = dataset.noise_map.array.real**-2.0 + + # The in-memory path computes the images from its arrays (no `sparse_terms`). + assert getattr(dataset, "sparse_terms", None) is None + + for name in ("dirty_image_natural", "dirty_beam"): + in_memory = getattr(dataset, name) + streamed = getattr(dataset_stream, name) + + assert streamed.shape_native == (7, 7) + np.testing.assert_allclose( + streamed.array, + in_memory.array, + rtol=1.0e-12, + atol=1.0e-12 * np.abs(in_memory.array).max(), + err_msg=name, + ) + + # The natural dirty image is the weighted adjoint normalised by sum(w), not the + # unweighted `dirty_image`, which is unchanged. + np.testing.assert_allclose( + dataset.dirty_image_natural.array * np.sum(weights), + dataset.apply_sparse_operator().sparse_operator.dirty_image, + rtol=1.0e-10, + atol=1.0e-10 * np.abs(dataset.dirty_image_natural.array).max() * np.sum(weights), + ) + np.testing.assert_array_equal( + dataset.dirty_image.array, + dataset.transformer.image_from(visibilities=dataset.data).array, + ) + + +def test__apply_sparse_operator_from_chunks__result_carries_sparse_terms(mask_2d_7x7): + pytest.importorskip("nufftax") + + dataset = _random_interferometer(mask_2d_7x7, transformer.TransformerNUFFT) + + dataset_chunked = dataset.apply_sparse_operator_from_chunks( + _chunks_of(dataset, [0, 20, 40]) + ) + + terms = dataset_chunked.sparse_terms + + assert isinstance(terms, aa.SparseTerms) + assert terms.n_vis == 40 + assert terms.transformer_class_name == "TransformerNUFFT" + assert terms.data_term == dataset_chunked.sparse_operator.data_term + + # The in-memory `apply_sparse_operator` path records no terms. + assert dataset.apply_sparse_operator().sparse_terms is None diff --git a/test_autoarray/fit/test_fit_interferometer.py b/test_autoarray/fit/test_fit_interferometer.py index 2d8f1ef39..d4c17f103 100644 --- a/test_autoarray/fit/test_fit_interferometer.py +++ b/test_autoarray/fit/test_fit_interferometer.py @@ -390,3 +390,129 @@ def test__noise_normalization__apply_sparse_operator_scalar_matches_array_path() assert fit_sparse.dataset.sparse_operator.noise_normalization is not None assert fit_sparse.noise_normalization == fit_array.noise_normalization + + +def _array_free_fit_setup(): + from autoarray.inversion.mesh.mesh.rectangular_rtu_adapt_density import ( + overlay_grid_from, + ) + + mask = aa.Mask2D.circular(shape_native=(10, 10), pixel_scales=0.5, radius=2.0) + + rng = np.random.default_rng(seed=7) + n_visibilities = 50 + uv_wavelengths = rng.normal(size=(n_visibilities, 2)) * 1.0e5 + data = rng.normal(size=n_visibilities) + 1j * rng.normal(size=n_visibilities) + sigma = rng.uniform(0.5, 2.0, size=n_visibilities) + noise_map = sigma + 1j * sigma + + dataset = aa.Interferometer( + data=aa.Visibilities(visibilities=data), + noise_map=aa.VisibilitiesNoiseMap(visibilities=noise_map), + uv_wavelengths=uv_wavelengths, + real_space_mask=mask, + ) + + dataset_stream = aa.Interferometer.from_stream( + [ + (uv_wavelengths[k0:k1], data[k0:k1], noise_map[k0:k1]) + for k0, k1 in ((0, 20), (20, n_visibilities)) + ], + mask, + ) + + grid = aa.Grid2D.from_mask(mask=mask, over_sample_size=1) + mesh = aa.mesh.RectangularUniform(shape=(4, 4)) + mapper = aa.Mapper( + interpolator=mesh.interpolator_from( + source_plane_data_grid=grid, + source_plane_mesh_grid=aa.Grid2DIrregular( + overlay_grid_from(shape_native=(4, 4), grid=grid) + ), + adapt_data=None, + ), + regularization=aa.reg.Constant(coefficient=1.0), + ) + + return dataset.apply_sparse_operator(), dataset_stream, mapper + + +def test__fit_interferometer__array_free_dataset__log_evidence_works_and_maps_raise(): + pytest.importorskip("nufftax") + + dataset_memory, dataset_stream, mapper = _array_free_fit_setup() + + fit_stream = aa.m.MockFitInterferometer( + dataset=dataset_stream, + inversion=aa.Inversion(dataset=dataset_stream, linear_obj_list=[mapper]), + ) + fit_memory = aa.m.MockFitInterferometer( + dataset=dataset_memory, + inversion=aa.Inversion(dataset=dataset_memory, linear_obj_list=[mapper]), + ) + + assert fit_stream.log_evidence == pytest.approx(fit_memory.log_evidence, rel=1.0e-8) + assert fit_stream.figure_of_merit == pytest.approx( + fit_memory.figure_of_merit, rel=1.0e-8 + ) + assert fit_stream.noise_normalization == pytest.approx( + fit_memory.noise_normalization, rel=1.0e-12 + ) + + for name in ( + "mask", + "transformer", + "residual_map", + "normalized_residual_map", + "chi_squared_map", + "signal_to_noise_map", + "chi_squared", + "dirty_image", + "dirty_noise_map", + "dirty_residual_map", + "dirty_chi_squared_map", + ): + with pytest.raises(aa.exc.DatasetException, match="array-free"): + getattr(fit_stream, name) + + +def test__fit_interferometer__sparse_log_evidence_never_touches_visibility_maps( + monkeypatch, +): + """ + Spy: on the in-memory sparse dataset (which *has* the arrays), make every + visibility-space quantity of the fit raise, and check `log_evidence` / + `figure_of_merit` still evaluate -- i.e. the sparse likelihood does not reach them. + """ + pytest.importorskip("nufftax") + + dataset_memory, _, mapper = _array_free_fit_setup() + + touched = [] + + def spy(name): + def fget(self): + touched.append(name) + raise AssertionError(f"log_evidence touched {name}") + + return property(fget) + + for name in ( + "mask", + "transformer", + "residual_map", + "normalized_residual_map", + "chi_squared_map", + "signal_to_noise_map", + "chi_squared", + ): + monkeypatch.setattr(aa.FitInterferometer, name, spy(name)) + + fit = aa.m.MockFitInterferometer( + dataset=dataset_memory, + inversion=aa.Inversion(dataset=dataset_memory, linear_obj_list=[mapper]), + ) + + assert np.isfinite(fit.log_evidence) + assert fit.figure_of_merit == fit.log_evidence + assert touched == [] diff --git a/test_autoarray/inversion/inversion/interferometer/test_interferometer.py b/test_autoarray/inversion/inversion/interferometer/test_interferometer.py index 6044299ae..262b9b218 100644 --- a/test_autoarray/inversion/inversion/interferometer/test_interferometer.py +++ b/test_autoarray/inversion/inversion/interferometer/test_interferometer.py @@ -1367,3 +1367,171 @@ def test__fast_chi_squared__data_none__jax_matches_numpy(): assert float(inversion_jax.fast_chi_squared) == pytest.approx( float(inversion_np.fast_chi_squared), rel=1.0e-7 ) + + +def _array_free_setup(n_visibilities=60, seed=3): + """ + A NUFFT dataset with random data and non-uniform (equal real/imaginary) noise, its + in-memory `apply_sparse_operator()` counterpart, the array-free dataset streamed from + three uneven chunks of the same visibilities, and a rectangular-mesh mapper. + """ + from autoarray.inversion.mesh.mesh.rectangular_rtu_adapt_density import ( + overlay_grid_from, + ) + + mask = aa.Mask2D.circular(shape_native=(10, 10), pixel_scales=0.5, radius=2.0) + + rng = np.random.default_rng(seed=seed) + uv_wavelengths = rng.normal(size=(n_visibilities, 2)) * 1.0e5 + data = rng.normal(size=n_visibilities) + 1j * rng.normal(size=n_visibilities) + sigma = rng.uniform(0.5, 2.0, size=n_visibilities) + noise_map = sigma + 1j * sigma + + dataset = aa.Interferometer( + data=aa.Visibilities(visibilities=data), + noise_map=aa.VisibilitiesNoiseMap(visibilities=noise_map), + uv_wavelengths=uv_wavelengths, + real_space_mask=mask, + transformer_class=aa.TransformerNUFFT, + ) + + chunks = ( + (uv_wavelengths[k0:k1], data[k0:k1], noise_map[k0:k1]) + for k0, k1 in ((0, 1), (1, 25), (25, n_visibilities)) + ) + + dataset_stream = aa.Interferometer.from_stream(chunks, mask) + + grid = aa.Grid2D.from_mask(mask=mask, over_sample_size=1) + mesh = aa.mesh.RectangularUniform(shape=(4, 4)) + interpolator = mesh.interpolator_from( + source_plane_data_grid=grid, + source_plane_mesh_grid=aa.Grid2DIrregular( + overlay_grid_from(shape_native=(4, 4), grid=grid) + ), + adapt_data=None, + ) + mapper = aa.Mapper( + interpolator=interpolator, regularization=aa.reg.Constant(coefficient=1.0) + ) + + return mask, dataset.apply_sparse_operator(), dataset_stream, mapper + + +def _assert_array_free_matches_in_memory(inversion_stream, inversion_memory, rel): + assert isinstance(inversion_stream, aa.InversionInterferometerSparse) + assert isinstance(inversion_memory, aa.InversionInterferometerSparse) + + for name in ( + "fast_chi_squared", + "regularization_term", + "log_det_curvature_reg_matrix_term", + "log_det_regularization_matrix_term", + ): + assert float(getattr(inversion_stream, name)) == pytest.approx( + float(getattr(inversion_memory, name)), rel=rel + ), name + + reconstruction = np.asarray(inversion_memory.reconstruction) + + np.testing.assert_allclose( + np.asarray(inversion_stream.reconstruction), + reconstruction, + rtol=rel, + atol=rel * np.abs(reconstruction).max(), + ) + + log_evidence_stream = aa.m.MockFitInterferometer( + dataset=inversion_stream.dataset, inversion=inversion_stream + ).log_evidence + log_evidence_memory = aa.m.MockFitInterferometer( + dataset=inversion_memory.dataset, inversion=inversion_memory + ).log_evidence + + assert float(log_evidence_stream) == pytest.approx( + float(log_evidence_memory), rel=rel + ) + + +def test__array_free_dataset__sparse_inversion_matches_in_memory__numpy(): + pytest.importorskip("nufftax") + + _, dataset_memory, dataset_stream, mapper = _array_free_setup() + + _assert_array_free_matches_in_memory( + aa.Inversion(dataset=dataset_stream, linear_obj_list=[mapper]), + aa.Inversion(dataset=dataset_memory, linear_obj_list=[mapper]), + rel=1.0e-8, + ) + + +def test__array_free_dataset__sparse_inversion_matches_in_memory__jax(): + pytest.importorskip("nufftax") + pytest.importorskip("jax") + + import jax.numpy as jnp + + _, dataset_memory, dataset_stream, mapper = _array_free_setup() + + _assert_array_free_matches_in_memory( + aa.Inversion(dataset=dataset_stream, linear_obj_list=[mapper], xp=jnp), + aa.Inversion(dataset=dataset_memory, linear_obj_list=[mapper], xp=jnp), + rel=1.0e-8, + ) + + +def test__array_free_dataset__mapped_reconstructed_operated_data_raises_typed(): + pytest.importorskip("nufftax") + + mask, dataset_memory, dataset_stream, mapper = _array_free_setup() + + inversion_stream = aa.Inversion(dataset=dataset_stream, linear_obj_list=[mapper]) + + # The real-space reconstruction is still available (it needs only the mask). + image = inversion_stream.mapped_reconstructed_data_dict[mapper] + + assert image.mask is mask + + with pytest.raises(aa.exc.InversionException, match="array-free"): + inversion_stream.mapped_reconstructed_operated_data_dict + + # The in-memory dataset still maps to visibilities. + inversion_memory = aa.Inversion(dataset=dataset_memory, linear_obj_list=[mapper]) + + assert ( + inversion_memory.mapped_reconstructed_operated_data_dict[mapper].shape[0] == 60 + ) + + +def test__inversion_interferometer_mask__is_the_real_space_mask_for_every_dataset_type(): + pytest.importorskip("nufftax") + + mask, dataset_memory, dataset_stream, mapper = _array_free_setup() + + dataset_interface = aa.DatasetInterface( + data=None, + noise_map=dataset_memory.noise_map, + grids=dataset_memory.grids, + transformer=dataset_memory.transformer, + sparse_operator=dataset_memory.sparse_operator, + ) + + # A `DatasetInterface` without a transformer falls back to `grids.lp.mask`. + dataset_interface_no_transformer = aa.DatasetInterface( + data=None, + noise_map=None, + grids=dataset_memory.grids, + sparse_operator=dataset_memory.sparse_operator, + ) + + for dataset in ( + dataset_memory, + dataset_stream, + dataset_interface, + dataset_interface_no_transformer, + ): + inversion = aa.Inversion(dataset=dataset, linear_obj_list=[mapper]) + + assert isinstance(inversion, aa.InversionInterferometerSparse) + assert (inversion.mask == mask).all() + assert inversion.mask.pixel_scales == mask.pixel_scales diff --git a/test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py b/test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py index ae8b0758d..8cc511715 100644 --- a/test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py +++ b/test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py @@ -1340,3 +1340,124 @@ def test__interferometer_sparse_operator__scalars_default_to_none_and_thread_thr assert operator.data_term == 2.5 assert type(operator.data_term) is float assert operator.noise_normalization == 3.5 + + +def _terms_with_provenance(value=1.0, n_vis=1, **provenance): + return aa.SparseTerms( + nufft_precision_operator=np.full((4, 4), value), + dirty_image_native=np.full((2, 2), value), + dirty_beam_native=np.full((2, 2), value), + sum_weights=value, + data_term=value, + noise_normalization=value, + n_vis=n_vis, + **provenance, + ) + + +def test__sparse_terms__add__provenance_matching_sums_and_is_carried(): + provenance = dict( + shape_native=(2, 2), + pixel_scales=(0.5, 0.5), + origin=(0.0, 0.0), + eps=1.0e-12, + transformer_class_name="TransformerNUFFT", + ) + + total = _terms_with_provenance(1.0, 3, **provenance) + _terms_with_provenance( + 10.0, 4, **provenance + ) + + assert total.data_term == 11.0 + assert total.n_vis == 7 + for name, value in provenance.items(): + assert getattr(total, name) == value + + # `None` on either side skips that check (the recorded value is carried). + total = _terms_with_provenance(1.0, 1, **provenance) + _terms_with_provenance(1.0, 1) + + assert total.eps == 1.0e-12 + assert total.n_vis == 2 + + +def test__sparse_terms__add__provenance_mismatch_raises(): + base = dict(shape_native=(2, 2), pixel_scales=(0.5, 0.5), eps=1.0e-12) + + for name, other_value in ( + ("shape_native", (3, 3)), + ("pixel_scales", (0.25, 0.25)), + ("eps", 1.0e-6), + ): + other = dict(base, **{name: other_value}) + + with pytest.raises(aa.exc.InversionException, match=name): + _terms_with_provenance(**base) + _terms_with_provenance(**other) + + +def test__sparse_terms__add__origin_mismatch_raises(): + with pytest.raises(aa.exc.InversionException, match="origin"): + _terms_with_provenance(origin=(0.0, 0.0)) + _terms_with_provenance( + origin=(0.5, 0.0) + ) + + +def test__sparse_terms__add__unrecorded_left_operand_keeps_right_provenance(): + unknown = _terms_with_provenance() + a = _terms_with_provenance(pixel_scales=(0.5, 0.5), origin=(0.1, 0.2), eps=1.0e-9) + b = _terms_with_provenance(pixel_scales=(1.0, 1.0)) + + merged = unknown + a + + assert merged.pixel_scales == (0.5, 0.5) + assert merged.origin == (0.1, 0.2) + assert merged.eps == 1.0e-9 + + with pytest.raises(aa.exc.InversionException, match="pixel_scales"): + (unknown + a) + b + + with pytest.raises(aa.exc.InversionException, match="pixel_scales"): + (a + unknown) + b + + +def test__sparse_terms__add__unrecorded_plus_unrecorded_stays_unrecorded(): + total = _terms_with_provenance() + _terms_with_provenance() + + for name in ( + "shape_native", + "pixel_scales", + "origin", + "eps", + "transformer_class_name", + ): + assert getattr(total, name) is None + + +def test__sparse_terms_from_chunks__records_provenance(): + pytest.importorskip("nufftax") + + mask, uv_wavelengths, data, noise_map, _ = _streaming_inputs(n_visibilities=20) + + terms = aa.util.inversion_interferometer.sparse_terms_from_chunks( + _chunks_from(uv_wavelengths, data, noise_map, [0, 10, 20]), + real_space_mask=mask, + eps=1.0e-10, + ) + + assert terms.shape_native == (10, 10) + assert terms.pixel_scales == (0.5, 0.5) + assert terms.origin == (0.0, 0.0) + assert terms.eps == 1.0e-10 + assert terms.transformer_class_name == "TransformerNUFFT" + + terms_dft = aa.util.inversion_interferometer.sparse_terms_from_chunks( + _chunks_from(uv_wavelengths, data, noise_map, [0, 20]), + real_space_mask=mask, + transformer_class=aa.TransformerDFT, + ) + + assert terms_dft.transformer_class_name == "TransformerDFT" + assert terms_dft.eps == 1.0e-12 + + # Terms accumulated at different NUFFT accuracies cannot be summed. + with pytest.raises(aa.exc.InversionException): + terms + terms_dft