From 090046f52fb0514d142611112d6ad7fefa728c7e Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Tue, 29 Sep 2026 18:03:55 +0100 Subject: [PATCH] feat: cache sparse-path scalars and add chunked SparseTerms accumulation (#588) Phase 1 of Discussion PyAutoLabs/13. InterferometerSparseOperator carries data_term and noise_normalization; fast_chi_squared and FitInterferometer.noise_normalization read them when present (fallback unchanged, bit-identical). Add SparseTerms, sparse_terms_from_chunks and Interferometer.apply_sparse_operator_from_chunks (chunked accumulation of W~, dirty image, dirty beam and the scalars, verified against the one-shot path at 1e-15). DatasetInterface(data=None) marks unsubtracted data; the noise real/imag check is a shared helper run per chunk. Co-Authored-By: Claude Fable 5.1 --- autoarray/__init__.py | 1 + autoarray/dataset/interferometer/dataset.py | 173 ++++++-- autoarray/fit/fit_interferometer.py | 13 + .../inversion/inversion/dataset_interface.py | 8 + .../inversion/interferometer/abstract.py | 38 +- .../inversion_interferometer_util.py | 376 +++++++++++++++++- .../dataset/interferometer/test_dataset.py | 144 +++++++ test_autoarray/fit/test_fit_interferometer.py | 82 ++++ .../interferometer/test_interferometer.py | 145 +++++++ .../test_inversion_interferometer_util.py | 205 ++++++++++ 10 files changed, 1149 insertions(+), 36 deletions(-) diff --git a/autoarray/__init__.py b/autoarray/__init__.py index c96ab9ea1..0445a636f 100644 --- a/autoarray/__init__.py +++ b/autoarray/__init__.py @@ -53,6 +53,7 @@ from .inversion.inversion.interferometer.mapping import InversionInterferometerMapping from .inversion.inversion.interferometer.inversion_interferometer_util import ( InterferometerSparseOperator, + SparseTerms, ) from .inversion.linear_obj.linear_obj import LinearObj from .inversion.linear_obj.func_list import AbstractLinearObjFuncList diff --git a/autoarray/dataset/interferometer/dataset.py b/autoarray/dataset/interferometer/dataset.py index 35879a31f..c7155a045 100644 --- a/autoarray/dataset/interferometer/dataset.py +++ b/autoarray/dataset/interferometer/dataset.py @@ -35,6 +35,7 @@ def disable_jax(): from autoarray.structures.visibilities import VisibilitiesNoiseMap from autoarray import exc +from autoarray.fit import fit_util from autoarray.inversion.inversion.interferometer import ( inversion_interferometer_util, ) @@ -333,35 +334,7 @@ def apply_sparse_operator( if disable_jax(): use_jax = False - noise_map_real = np.asarray(self.noise_map.real) - noise_map_imag = np.asarray(self.noise_map.imag) - - if not np.allclose(noise_map_real, noise_map_imag): - - unequal = ~np.isclose(noise_map_real, noise_map_imag) - - denominator = np.maximum(np.abs(noise_map_real), np.abs(noise_map_imag)) - relative_difference = np.abs(noise_map_real - noise_map_imag) / np.where( - denominator == 0.0, 1.0, denominator - ) - - raise exc.DatasetException( - "The sparse operator cannot be applied to this interferometer dataset because its " - "noise-map has unequal real and imaginary sigma.\n\n" - "The sparse operator's precision operator `W~ = Re(F^H W F)` is built from the " - "real-part noise sigma only (see `psf_precision_operator_from`, which passes " - "`noise_map_real` to `nufft_precision_operator_from`). That reduction is exact only " - "when every visibility has equal real and imaginary sigma " - "(`sigma_real == sigma_imag`).\n\n" - f"This dataset has {int(np.count_nonzero(unequal))} of {noise_map_real.size} " - "visibilities where the real and imaginary sigma differ (maximum relative difference " - f"{np.max(relative_difference):.3e}), so the sparse curvature matrix would silently " - "disagree with the dense path.\n\n" - "Either equalise the real and imaginary noise sigma of every visibility, or fit " - "without calling `apply_sparse_operator()` — the dense " - "`InversionInterferometerMapping` path handles unequal real and imaginary sigmas " - "correctly." - ) + inversion_interferometer_util.check_noise_map_real_imag_equal(self.noise_map) if nufft_precision_operator is None: @@ -411,10 +384,152 @@ def apply_sparse_operator( + 1j * self.data.imag * self.noise_map.imag**-2.0, ) + # The two visibility-sum scalars of the likelihood, reduced once here so that neither + # `fast_chi_squared` (term 3) nor `FitInterferometer.noise_normalization` reduces over + # the visibility arrays on every likelihood call. The expressions are the ones those + # two reductions use, so the cached values are bit-for-bit what they would compute. + data_term = float( + np.sum(self.data.array.real**2.0 / self.noise_map.array.real**2.0) + + np.sum(self.data.array.imag**2.0 / self.noise_map.array.imag**2.0) + ) + noise_normalization = float( + fit_util.noise_normalization_complex_from(noise_map=self.noise_map.array) + ) + sparse_operator = inversion_interferometer_util.InterferometerSparseOperator.from_nufft_precision_operator( nufft_precision_operator=nufft_precision_operator, dirty_image=dirty_image.array, batch_size=batch_size, + data_term=data_term, + noise_normalization=noise_normalization, + ) + + return Interferometer( + real_space_mask=self.real_space_mask, + data=self.data, + noise_map=self.noise_map, + uv_wavelengths=self.uv_wavelengths, + transformer_class=lambda uv_wavelengths, real_space_mask: self.transformer, + sparse_operator=sparse_operator, + ) + + def apply_sparse_operator_from_chunks( + self, + chunks, + *, + batch_size: int = 128, + **accumulator_kwargs, + ): + """ + Build the sparse operator by accumulating the visibilities one chunk at a time, rather + than from this dataset's resident arrays in one shot. + + This is the chunked counterpart of `apply_sparse_operator`: every quantity the sparse + likelihood needs from the visibilities -- the NUFFT precision operator `W~`, the + noise-weighted dirty image, `sum(d^2 / sigma^2)` and `sum(log(2 pi sigma^2))` -- is a sum + over visibilities, so it is accumulated per chunk by + `inversion_interferometer_util.sparse_terms_from_chunks` into a `SparseTerms` record and + turned into an `InterferometerSparseOperator` by + `InterferometerSparseOperator.from_sparse_terms`. The result equals + `apply_sparse_operator()` to summation order (verified at `rtol=1e-12` by the test + suite), and the operator carries both scalars, so a fit on it never reduces over the + visibility arrays in its likelihood. + + Chunk contract + -------------- + `chunks` is any iterable (a list, or a generator reading from disk) of + `(uv_wavelengths, data, noise_map)` triples of `K` visibilities each, where `K` may + differ per chunk: `uv_wavelengths` a real `(K, 2)` array in wavelengths, `data` and + `noise_map` a `Visibilities`, a complex `(K,)` array or a real `(K, 2)` array of + (real, imag) columns. Every chunk must have equal real and imaginary noise sigma + (checked per chunk). The chunks must together be exactly this dataset's visibilities + (in any order) for the returned dataset to be self-consistent -- this method does not + check that. + + Memory + ------ + The accumulation holds one chunk at a time, so its peak memory is set by the chunk + size, not the dataset size. Holding the visibilities costs at least 48 B/visibility + (`data` and `noise_map` complex128, 16 B each; `uv_wavelengths` 2 x float64, 16 B), before + the transformer's own per-visibility state. Measured in pyuvimage on the same + accumulation (peak RSS over baseline, 400-pixel image, 4096-visibility chunks, fresh + process each; Discussion #13): + + ====== ====================== ========================= + N_vis peak RSS of the stream resident if held (48 B/vis) + ====== ====================== ========================= + 5e5 36 MB 24 MB + 1e6 36 MB 48 MB + 2e6 36 MB 96 MB + 4e6 36 MB 192 MB + ====== ====================== ========================= + + i.e. flat across 8x, while the in-memory path grows linearly (~10 GB at 2e8 + visibilities, a real ALMA MFS cube). + + This method returns a dataset of the same shape as `apply_sparse_operator` -- this + dataset's `data`, `noise_map`, `uv_wavelengths` and transformer are retained, because + the visualizer, `save_attributes`/aggregator and dirty-image outputs still read them. + Dropping them (an array-free `Interferometer` built straight from a chunk stream) is a + follow-up that needs its own dataset contract; this method is the seam it builds on. + + Parameters + ---------- + chunks + The `(uv_wavelengths, data, noise_map)` chunks. + batch_size + The number of source-pixel columns processed per batch by the sparse operator. + accumulator_kwargs + Passed to `sparse_terms_from_chunks` (`transformer_class`, `method`, `eps`, + `chunk_size`, `chunk_k`, `use_jax`, `show_progress`). When not given, + `transformer_class`, `eps` and `chunk_size` follow this dataset's transformer (the + same defaults `psf_precision_operator_from` takes), so the accumulated terms match + `apply_sparse_operator()`. + + Returns + ------- + Interferometer + A new `Interferometer` dataset with the accumulated `InterferometerSparseOperator` + attached. + + Raises + ------ + exc.DatasetException + If any chunk has unequal real and imaginary noise sigma. + """ + if disable_jax() and accumulator_kwargs.get("use_jax", False): + accumulator_kwargs["use_jax"] = False + + transformer = self.transformer + + if isinstance(transformer, TransformerNUFFT): + accumulator_kwargs.setdefault("eps", transformer.eps) + accumulator_kwargs.setdefault("chunk_size", transformer.chunk_size) + + if "transformer_class" not in accumulator_kwargs: + transformer_eps = transformer.eps + transformer_chunk_size = transformer.chunk_size + + def transformer_class(uv_wavelengths, real_space_mask): + return TransformerNUFFT( + uv_wavelengths=uv_wavelengths, + real_space_mask=real_space_mask, + eps=transformer_eps, + chunk_size=transformer_chunk_size, + ) + + accumulator_kwargs["transformer_class"] = transformer_class + else: + accumulator_kwargs.setdefault("transformer_class", type(transformer)) + + terms = inversion_interferometer_util.sparse_terms_from_chunks( + chunks, real_space_mask=self.real_space_mask, **accumulator_kwargs + ) + + sparse_operator = inversion_interferometer_util.InterferometerSparseOperator.from_sparse_terms( + terms, + real_space_mask=self.real_space_mask, + batch_size=batch_size, ) return Interferometer( diff --git a/autoarray/fit/fit_interferometer.py b/autoarray/fit/fit_interferometer.py index 0c53eaa5b..408689de5 100644 --- a/autoarray/fit/fit_interferometer.py +++ b/autoarray/fit/fit_interferometer.py @@ -128,7 +128,20 @@ def noise_normalization(self) -> float: Returns the noise-map normalization term of the noise-map, summing the noise_map value in every pixel as: [Noise_Term] = sum(log(2*pi*[Noise]**2.0)) + + When the dataset carries a sparse operator with a precomputed `noise_normalization` + (set by `Interferometer.apply_sparse_operator` and + `apply_sparse_operator_from_chunks`), and the noise-map fitted is the dataset's own + (not one a subclass has scaled or replaced), that scalar is returned instead of + reducing over every visibility's sigma on each likelihood call. It is computed with + the same expression, so the value is identical. """ + sparse_operator = getattr(self.dataset, "sparse_operator", None) + noise_normalization = getattr(sparse_operator, "noise_normalization", None) + + if noise_normalization is not None and self.noise_map is self.dataset.noise_map: + return noise_normalization + return fit_util.noise_normalization_complex_from( noise_map=self.noise_map.array, ) diff --git a/autoarray/inversion/inversion/dataset_interface.py b/autoarray/inversion/inversion/dataset_interface.py index d5abf0101..032bcb2b5 100644 --- a/autoarray/inversion/inversion/dataset_interface.py +++ b/autoarray/inversion/inversion/dataset_interface.py @@ -33,6 +33,14 @@ def __init__( data The array of the image data containing the signal that is fitted (in PyAutoGalaxy and PyAutoLens the recommended units are electrons per second). + + For an interferometer inversion on the sparse path this may be `None`, meaning "the raw visibilities + the `sparse_operator` was built from, unmodified". It is only valid with a `sparse_operator` that + carries a precomputed `data_term` (built by `Interferometer.apply_sparse_operator` or + `apply_sparse_operator_from_chunks`): the sparse inversion's data vector already comes from the + operator's cached dirty image, and `fast_chi_squared` then reads the operator's `data_term` instead + of reducing over the visibilities, so no visibility array is touched by the likelihood. Output-only + quantities that need the data itself (e.g. `data_subtracted_dict`) are unavailable in that case. noise_map An array describing the RMS standard deviation error in each pixel used for computing quantities like the chi-squared in a fit (in PyAutoGalaxy and PyAutoLens the recommended units are electrons per second). diff --git a/autoarray/inversion/inversion/interferometer/abstract.py b/autoarray/inversion/inversion/interferometer/abstract.py index f42f65149..e1195e381 100644 --- a/autoarray/inversion/inversion/interferometer/abstract.py +++ b/autoarray/inversion/inversion/interferometer/abstract.py @@ -82,7 +82,7 @@ def operated_mapping_matrix_list(self) -> List[np.ndarray]: operated_mapping_matrix_override = linear_obj.operated_mapping_matrix_override if operated_mapping_matrix_override is not None: - expected_shape = (self.data.shape[0], linear_obj.params) + expected_shape = (self.noise_map.shape[0], linear_obj.params) if tuple(operated_mapping_matrix_override.shape) != expected_shape: raise exc.InversionException( @@ -170,6 +170,11 @@ def fast_chi_squared(self): where `s` is the reconstruction vector, `F` is the curvature matrix, `D` is the data vector, and `d_r`/`d_i` are the real/imaginary parts of the observed visibilities. + When the dataset interface's `data` is `None` the third term is read from the scalar + `sparse_operator.data_term` cached when the operator was built, so no visibility array + is reduced over. That is only correct when the data fitted is the raw data the operator + was built from (nothing subtracted), which is the contract of passing `data=None`. + This avoids computing the full mapped reconstructed visibilities and is faster than computing `chi_squared` via the residual visibilities when many source pixels are used. """ @@ -190,10 +195,31 @@ def fast_chi_squared(self): ] ) - chi_squared_term_3 = xp.sum( - self.dataset.data.array.real**2.0 / self.dataset.noise_map.array.real**2.0 - ) + xp.sum( - self.dataset.data.array.imag**2.0 / self.dataset.noise_map.array.imag**2.0 - ) + if self.dataset.data is None: + # The interface carries no visibilities (e.g. a pixelization-only fit on the sparse + # path, where nothing was subtracted from the data), so term 3 is the scalar + # `d^T N^-1 d` the sparse operator cached when it was built from the raw data. + chi_squared_term_3 = getattr( + self.dataset.sparse_operator, "data_term", None + ) + + if chi_squared_term_3 is None: + raise exc.InversionException( + "The dataset input to this interferometer inversion has `data=None`, which " + "is only valid when its `sparse_operator` carries a precomputed `data_term` " + "(sum(d_r^2/sigma_r^2) + sum(d_i^2/sigma_i^2)) -- e.g. one built by " + "`Interferometer.apply_sparse_operator` or " + "`Interferometer.apply_sparse_operator_from_chunks`. This dataset's " + "sparse operator has none, so `fast_chi_squared` cannot be computed. " + "Pass the visibilities as `data` instead." + ) + else: + chi_squared_term_3 = xp.sum( + self.dataset.data.array.real**2.0 + / self.dataset.noise_map.array.real**2.0 + ) + xp.sum( + self.dataset.data.array.imag**2.0 + / self.dataset.noise_map.array.imag**2.0 + ) return chi_squared_term_1 + chi_squared_term_2 + chi_squared_term_3 diff --git a/autoarray/inversion/inversion/interferometer/inversion_interferometer_util.py b/autoarray/inversion/inversion/interferometer/inversion_interferometer_util.py index f447bef75..44b1ef743 100644 --- a/autoarray/inversion/inversion/interferometer/inversion_interferometer_util.py +++ b/autoarray/inversion/inversion/interferometer/inversion_interferometer_util.py @@ -4,7 +4,9 @@ import numpy as np import time from pathlib import Path -from typing import Optional +from typing import Callable, Iterable, Optional, Tuple + +from autoarray import exc from autoarray.operators.transformer import _load_nufftax, nufftax_exception @@ -27,6 +29,61 @@ def disable_jax(): logger = logging.getLogger(__name__) +def check_noise_map_real_imag_equal(noise_map) -> None: + """ + Raise a `DatasetException` unless every visibility of `noise_map` has equal real and + imaginary noise sigma (`noise_map.real == noise_map.imag`, to `np.allclose`). + + This is the precondition of the sparse operator: its precision operator + `W~ = Re(F^H W F)` is built from the real-part sigma alone (see + `nufft_precision_operator_from`), a reduction that is exact only under that equality. + It is shared by `Interferometer.apply_sparse_operator`, which checks the whole resident + noise-map, and `sparse_terms_from_chunks`, which checks every chunk as it streams past. + + Parameters + ---------- + noise_map + The complex noise-map (a `VisibilitiesNoiseMap` or a complex ndarray) to check. + + Raises + ------ + exc.DatasetException + If any visibility has unequal real and imaginary noise sigma. + """ + noise_map = getattr(noise_map, "array", noise_map) + + noise_map_real = np.asarray(np.real(noise_map)) + noise_map_imag = np.asarray(np.imag(noise_map)) + + if np.allclose(noise_map_real, noise_map_imag): + return + + unequal = ~np.isclose(noise_map_real, noise_map_imag) + + denominator = np.maximum(np.abs(noise_map_real), np.abs(noise_map_imag)) + relative_difference = np.abs(noise_map_real - noise_map_imag) / np.where( + denominator == 0.0, 1.0, denominator + ) + + raise exc.DatasetException( + "The sparse operator cannot be applied to this interferometer dataset because its " + "noise-map has unequal real and imaginary sigma.\n\n" + "The sparse operator's precision operator `W~ = Re(F^H W F)` is built from the " + "real-part noise sigma only (see `psf_precision_operator_from`, which passes " + "`noise_map_real` to `nufft_precision_operator_from`). That reduction is exact only " + "when every visibility has equal real and imaginary sigma " + "(`sigma_real == sigma_imag`).\n\n" + f"This dataset has {int(np.count_nonzero(unequal))} of {noise_map_real.size} " + "visibilities where the real and imaginary sigma differ (maximum relative difference " + f"{np.max(relative_difference):.3e}), so the sparse curvature matrix would silently " + "disagree with the dense path.\n\n" + "Either equalise the real and imaginary noise sigma of every visibility, or fit " + "without calling `apply_sparse_operator()` — the dense " + "`InversionInterferometerMapping` path handles unequal real and imaginary sigmas " + "correctly." + ) + + def data_vector_via_transformed_mapping_matrix_from( transformed_mapping_matrix: np.ndarray, visibilities: np.ndarray, @@ -883,6 +940,8 @@ class InterferometerSparseOperator: M: int batch_size: int w_dtype: np.dtype + data_term: Optional[float] = None + noise_normalization: Optional[float] = None """ Cached FFT operator state for fast interferometer curvature-matrix assembly. @@ -948,6 +1007,18 @@ class InterferometerSparseOperator: Larger batch sizes improve throughput on GPU but increase memory usage. w_dtype Floating-point dtype for weights and accumulations (e.g. float64). + data_term + The scalar `sum(d_r^2 / sigma_r^2) + sum(d_i^2 / sigma_i^2)` of the visibilities the + operator was built from, i.e. `d^T N^-1 d`. When set, `fast_chi_squared` of an + inversion whose interface `data` is `None` uses it for its third term, so the + likelihood never reduces over the visibility arrays. `None` for an operator built + without it (e.g. a directly constructed one); callers then fall back to the array + reduction. + noise_normalization + The scalar `sum(log(2 pi sigma_r^2)) + sum(log(2 pi sigma_i^2))` of the same + visibilities, equal to `fit_util.noise_normalization_complex_from(noise_map)`. When + set, `FitInterferometer.noise_normalization` reads it instead of reducing over the + noise-map every likelihood call. `None` falls back to the array reduction. Cached properties ----------------- @@ -1020,6 +1091,8 @@ def from_nufft_precision_operator( dirty_image: np.ndarray, *, batch_size: int = 128, + data_term: Optional[float] = None, + noise_normalization: Optional[float] = None, ): """ Construct an `InterferometerSparseOperator` from a curvature-preload array. @@ -1083,6 +1156,53 @@ def from_nufft_precision_operator( M=M, batch_size=int(batch_size), w_dtype=nufft_precision_operator.dtype, + data_term=None if data_term is None else float(data_term), + noise_normalization=( + None if noise_normalization is None else float(noise_normalization) + ), + ) + + @classmethod + def from_sparse_terms( + cls, + terms: "SparseTerms", + *, + real_space_mask, + batch_size: int = 128, + ) -> "InterferometerSparseOperator": + """ + Construct an `InterferometerSparseOperator` from the per-visibility sums of a + `SparseTerms` record (e.g. one accumulated chunk by chunk by + `sparse_terms_from_chunks`). + + The operator stores its dirty image in the slim (masked, 1D) representation that + `Interferometer.apply_sparse_operator` gives it, so the record's native dirty image is + slimmed through `Array2D(values=..., mask=real_space_mask)`. The two scalars + `data_term` and `noise_normalization` are carried over, so an inversion and fit built + on this operator need no visibility arrays for their likelihood terms. + + Parameters + ---------- + terms + The accumulated sums of every visibility the operator is to represent. + real_space_mask + The real-space `Mask2D` the terms were accumulated on. + batch_size + Number of source-pixel columns processed per block when assembling the curvature + matrix (see `from_nufft_precision_operator`). + """ + from autoarray.structures.arrays.uniform_2d import Array2D + + dirty_image = Array2D( + values=np.asarray(terms.dirty_image_native), mask=real_space_mask + ).slim.array + + return cls.from_nufft_precision_operator( + nufft_precision_operator=terms.nufft_precision_operator, + dirty_image=dirty_image, + batch_size=batch_size, + data_term=terms.data_term, + noise_normalization=terms.noise_normalization, ) @staticmethod @@ -1680,3 +1800,257 @@ def curvature_matrix_func_list_from( ) return curvature_weights_0.T @ operated + + +@dataclass(frozen=True) +class SparseTerms: + """ + Every quantity the sparse (w-tilde) interferometer likelihood needs from the visibilities, + and nothing per-visibility. + + Each field is a sum over visibilities, so the terms of a union of visibility sets are the + field-wise sum of the terms of its parts (`__add__`). That is what lets them be accumulated + one chunk at a time by `sparse_terms_from_chunks`, with peak memory set by the chunk size + rather than the dataset size, and what makes multi-channel (MFS) terms the sum of the + per-channel terms. + + Fields + ------ + nufft_precision_operator + The `(2Ny, 2Nx)` real preload `sum_k w_k cos(...)` with `w_k = 1 / sigma_r,k^2`, in the + FFT-wrapped offset order `nufft_precision_operator_from` returns (index 0 = zero lag). + dirty_image_native + The `(Ny, Nx)` native noise-weighted dirty image `Re(F^H (d_r / sigma_r^2 + i d_i / + sigma_i^2))`, the adjoint `Interferometer.apply_sparse_operator` stores (masked pixels + zero). + dirty_beam_native + The `(Ny, Nx)` native adjoint of the weights `w = 1 / sigma_r^2` (the naturally weighted + dirty beam at the image pixels, unnormalised). Not used by the likelihood; kept for + normalised dirty images of an array-free dataset. + sum_weights + `sum(w)`, the normalisation of the naturally weighted dirty beam/image. + data_term + `sum(d_r^2 / sigma_r^2) + sum(d_i^2 / sigma_i^2)`, i.e. `d^T N^-1 d`, the third term of + `fast_chi_squared`. + noise_normalization + `sum(log(2 pi sigma_r^2)) + sum(log(2 pi sigma_i^2))`, equal to + `fit_util.noise_normalization_complex_from(noise_map)`. + n_vis + The number of visibilities summed. + """ + + nufft_precision_operator: np.ndarray + dirty_image_native: np.ndarray + dirty_beam_native: np.ndarray + sum_weights: float + data_term: float + noise_normalization: float + n_vis: int + + def __add__(self, other: "SparseTerms") -> "SparseTerms": + if not isinstance(other, SparseTerms): + return NotImplemented + + if self.nufft_precision_operator.shape != other.nufft_precision_operator.shape: + raise ValueError( + "SparseTerms can only be added when accumulated on the same real-space mask: " + f"precision operator shapes {self.nufft_precision_operator.shape} and " + f"{other.nufft_precision_operator.shape} differ." + ) + + return SparseTerms( + nufft_precision_operator=self.nufft_precision_operator + + other.nufft_precision_operator, + dirty_image_native=self.dirty_image_native + other.dirty_image_native, + dirty_beam_native=self.dirty_beam_native + other.dirty_beam_native, + sum_weights=float(self.sum_weights + other.sum_weights), + data_term=float(self.data_term + other.data_term), + noise_normalization=float( + self.noise_normalization + other.noise_normalization + ), + n_vis=int(self.n_vis + other.n_vis), + ) + + +def _complex_visibilities_from(values) -> np.ndarray: + """ + Return `values` (a `Visibilities`/`VisibilitiesNoiseMap`, a complex ndarray of shape + `(K,)` or a real ndarray of shape `(K, 2)` of (real, imag) columns) as a 1D complex128 + ndarray, without the per-row `np.apply_along_axis` the `Visibilities` constructor uses. + """ + values = np.asarray(getattr(values, "array", values)) + + if not np.iscomplexobj(values): + if values.ndim == 2 and values.shape[1] == 2: + values = values[:, 0] + 1j * values[:, 1] + + return np.asarray(values, dtype=np.complex128).ravel() + + +def sparse_terms_from_chunks( + chunks: Iterable[Tuple[np.ndarray, np.ndarray, np.ndarray]], + *, + real_space_mask, + transformer_class: Optional[Callable] = None, + method: str = "nufft", + eps: Optional[float] = None, + chunk_size: Optional[int] = None, + chunk_k: int = 2048, + use_jax: bool = False, + show_progress: bool = False, +) -> SparseTerms: + """ + Accumulate the `SparseTerms` of an interferometer dataset one chunk of visibilities at a + time, so the full visibility arrays never need to be resident at once. + + Every term is linear in the per-visibility contributions: the precision operator is a sum + over visibilities of `w_k cos(...)` (the ifftshift and Nyquist zeroing inside the NUFFT + builder are linear, so they commute with the sum -- see + `test__nufft_precision_operator_via_nufft__chunked_matches_one_shot`), the dirty image and + beam are adjoint transforms of per-visibility values, and the scalars are plain sums. The + result therefore equals the one-shot terms of the concatenated chunks to summation order. + + Chunk contract + -------------- + `chunks` is any iterable (a list, or a generator reading from disk) of + `(uv_wavelengths, data, noise_map)` triples, where + + - `uv_wavelengths` is a real `(K, 2)` array of (u, v) baselines in wavelengths; + - `data` is the `K` complex visibilities (a `Visibilities`, a complex `(K,)` array or a + real `(K, 2)` array of (real, imag) columns); + - `noise_map` is the `K` complex noise sigmas in the same forms, with equal real and + imaginary sigma per visibility (checked per chunk; the first offending chunk raises). + + `K` may differ between chunks; empty chunks are skipped. Chunks are consumed once, in + order, and only one is referenced at a time. + + Parameters + ---------- + chunks + The `(uv_wavelengths, data, noise_map)` chunks. + real_space_mask + The real-space `Mask2D` the terms are accumulated on. + transformer_class + A callable `transformer_class(uv_wavelengths=..., real_space_mask=...)` returning a + transformer (`TransformerNUFFT`, `TransformerDFT`, or a factory closing over e.g. an + `eps`) that is built per chunk to form the dirty image and beam. `None` is + `TransformerNUFFT`. + method, eps, chunk_size, chunk_k, use_jax, show_progress + Passed to `nufft_precision_operator_from` for each chunk (`chunk_size` is the NUFFT + builder's own inner chunk, a memory ceiling within a chunk). `eps=None` is `1e-12`. + + Returns + ------- + SparseTerms + The accumulated terms, with the scalars in float64. + + Raises + ------ + exc.DatasetException + If any chunk's noise-map has unequal real and imaginary sigma. + ValueError + If `chunks` yields no visibilities. + """ + if transformer_class is None: + from autoarray.operators.transformer import TransformerNUFFT + + transformer_class = TransformerNUFFT + + from autoarray.structures.visibilities import Visibilities + + if eps is None: + eps = 1.0e-12 + + terms = None + + for uv_wavelengths, data, noise_map in chunks: + + uv_wavelengths = np.asarray( + getattr(uv_wavelengths, "array", uv_wavelengths), dtype=np.float64 + ) + + if uv_wavelengths.shape[0] == 0: + continue + + data = _complex_visibilities_from(data) + noise_map = _complex_visibilities_from(noise_map) + + if not (data.shape[0] == noise_map.shape[0] == uv_wavelengths.shape[0]): + raise ValueError( + "sparse_terms_from_chunks: a chunk's uv_wavelengths, data and noise_map must " + f"have the same length, got {uv_wavelengths.shape[0]}, {data.shape[0]} and " + f"{noise_map.shape[0]}." + ) + + check_noise_map_real_imag_equal(noise_map) + + noise_map_real = noise_map.real + noise_map_imag = noise_map.imag + + transformer = transformer_class( + uv_wavelengths=uv_wavelengths, real_space_mask=real_space_mask + ) + + # The same arguments `Interferometer.psf_precision_operator_from` builds from the + # dataset's own transformer, so the per-chunk operators sum to the dataset's. + nufft_precision_operator = np.asarray( + nufft_precision_operator_from( + noise_map_real=noise_map_real, + uv_wavelengths=uv_wavelengths, + shape_masked_pixels_2d=transformer.grid.mask.shape_native_masked_pixels, + grid_radians_2d=transformer.grid.mask.derive_grid.all_false.in_radians.native.array, + method=method, + eps=eps, + chunk_size=chunk_size, + chunk_k=chunk_k, + show_progress=show_progress, + use_jax=use_jax, + ), + dtype=np.float64, + ) + + # The same expression `Interferometer.apply_sparse_operator` forms its dirty image from. + dirty_image_native = np.asarray( + transformer.image_from( + visibilities=Visibilities( + visibilities=data.real * noise_map_real**-2.0 + + 1j * data.imag * noise_map_imag**-2.0 + ), + ).native.array, + dtype=np.float64, + ) + + weights = noise_map_real**-2.0 + + dirty_beam_native = np.asarray( + transformer.image_from( + visibilities=Visibilities(visibilities=weights.astype(np.complex128)), + ).native.array, + dtype=np.float64, + ) + + chunk_terms = SparseTerms( + nufft_precision_operator=nufft_precision_operator, + dirty_image_native=dirty_image_native, + dirty_beam_native=dirty_beam_native, + sum_weights=float(np.sum(weights)), + data_term=float( + np.sum(data.real**2.0 / noise_map_real**2.0) + + np.sum(data.imag**2.0 / noise_map_imag**2.0) + ), + noise_normalization=float( + np.sum(np.log(2 * np.pi * noise_map_real**2.0)) + + np.sum(np.log(2 * np.pi * noise_map_imag**2.0)) + ), + n_vis=int(uv_wavelengths.shape[0]), + ) + + terms = chunk_terms if terms is None else terms + chunk_terms + + if terms is None: + raise ValueError( + "sparse_terms_from_chunks: `chunks` yielded no visibilities, so there is nothing " + "to accumulate." + ) + + return terms diff --git a/test_autoarray/dataset/interferometer/test_dataset.py b/test_autoarray/dataset/interferometer/test_dataset.py index 2b1d33662..7fcb99d77 100644 --- a/test_autoarray/dataset/interferometer/test_dataset.py +++ b/test_autoarray/dataset/interferometer/test_dataset.py @@ -476,3 +476,147 @@ def test__apply_sparse_operator__method_and_nufft_kwargs_reach_the_builder(mask_ rtol=1.0e-10, atol=1.0e-10 * np.abs(operator_via_numpy[0, 0]), ) + + +def _random_interferometer(mask, transformer_class, n_visibilities=40, seed=5): + rng = np.random.default_rng(seed=seed) + data = rng.normal(size=n_visibilities) + 1j * rng.normal(size=n_visibilities) + sigma = rng.uniform(0.5, 2.0, size=n_visibilities) + + return aa.Interferometer( + data=aa.Visibilities(visibilities=data), + noise_map=aa.VisibilitiesNoiseMap(visibilities=sigma + 1j * sigma), + uv_wavelengths=rng.normal(size=(n_visibilities, 2)) * 5.0e4, + real_space_mask=mask, + transformer_class=transformer_class, + ) + + +def test__apply_sparse_operator__populates_data_term_and_noise_normalization( + mask_2d_7x7, +): + dataset = _random_interferometer(mask_2d_7x7, transformer.TransformerDFT) + + sparse_operator = dataset.apply_sparse_operator(use_jax=False).sparse_operator + + data = dataset.data.array + noise_map = dataset.noise_map.array + + assert sparse_operator.data_term == np.sum( + data.real**2.0 / noise_map.real**2.0 + ) + np.sum(data.imag**2.0 / noise_map.imag**2.0) + assert ( + sparse_operator.noise_normalization + == aa.util.fit.noise_normalization_complex_from(noise_map=noise_map) + ) + + +def test__apply_sparse_operator_from_chunks__matches_apply_sparse_operator( + interferometer_7_lop, mask_2d_7x7 +): + pytest.importorskip("nufftax") + + for dataset in ( + interferometer_7_lop, + _random_interferometer(mask_2d_7x7, transformer.TransformerNUFFT), + _random_interferometer(mask_2d_7x7, transformer.TransformerDFT), + ): + uv_wavelengths = dataset.uv_wavelengths + data = dataset.data.array + noise_map = dataset.noise_map.array + n = uv_wavelengths.shape[0] + + edges = [0, 1, n // 2, n] + chunks = [ + (uv_wavelengths[k0:k1], data[k0:k1], noise_map[k0:k1]) + for k0, k1 in zip(edges[:-1], edges[1:]) + ] + + dataset_one_shot = dataset.apply_sparse_operator() + dataset_chunked = dataset.apply_sparse_operator_from_chunks(chunks) + + one_shot = dataset_one_shot.sparse_operator + chunked = dataset_chunked.sparse_operator + + np.testing.assert_allclose( + chunked.nufft_precision_operator, + one_shot.nufft_precision_operator, + rtol=1.0e-12, + atol=1.0e-12 * np.abs(one_shot.nufft_precision_operator).max(), + ) + np.testing.assert_allclose( + chunked.dirty_image, + one_shot.dirty_image, + rtol=1.0e-12, + atol=1.0e-12 * np.abs(one_shot.dirty_image).max(), + ) + assert chunked.dirty_image.shape == one_shot.dirty_image.shape + assert chunked.data_term == pytest.approx(one_shot.data_term, rel=1.0e-12) + assert chunked.noise_normalization == pytest.approx( + one_shot.noise_normalization, rel=1.0e-12 + ) + assert chunked.batch_size == one_shot.batch_size + + assert dataset_chunked.data is dataset.data + assert dataset_chunked.transformer is dataset.transformer + + +def test__apply_sparse_operator_from_chunks__inversion_is_sparse_and_matches( + mask_2d_7x7, +): + pytest.importorskip("nufftax") + + dataset = _random_interferometer(mask_2d_7x7, transformer.TransformerNUFFT) + + chunks = [ + (dataset.uv_wavelengths[k0:k1], dataset.data[k0:k1], dataset.noise_map[k0:k1]) + for k0, k1 in ((0, 13), (13, 40)) + ] + + grid = aa.Grid2D.from_mask(mask=mask_2d_7x7, over_sample_size=1) + mesh = aa.mesh.Delaunay(pixels=9) + image_mesh_grid = aa.image_mesh.Overlay(shape=(3, 3)).image_plane_mesh_grid_from( + mask=mask_2d_7x7, adapt_data=None + ) + mapper = aa.Mapper( + interpolator=mesh.interpolator_from( + source_plane_data_grid=grid, source_plane_mesh_grid=image_mesh_grid + ), + regularization=aa.reg.Constant(coefficient=1.0), + ) + + inversion_one_shot = aa.Inversion( + dataset=dataset.apply_sparse_operator(), linear_obj_list=[mapper] + ) + inversion_chunked = aa.Inversion( + dataset=dataset.apply_sparse_operator_from_chunks(chunks), + linear_obj_list=[mapper], + ) + + assert isinstance(inversion_chunked, aa.InversionInterferometerSparse) + + assert inversion_chunked.fast_chi_squared == pytest.approx( + inversion_one_shot.fast_chi_squared, rel=1.0e-10 + ) + assert inversion_chunked.log_det_curvature_reg_matrix_term == pytest.approx( + inversion_one_shot.log_det_curvature_reg_matrix_term, rel=1.0e-10 + ) + + +def test__apply_sparse_operator_from_chunks__unequal_real_imag_noise__raises( + mask_2d_7x7, +): + pytest.importorskip("nufftax") + + dataset = _random_interferometer(mask_2d_7x7, transformer.TransformerNUFFT) + + noise_map = dataset.noise_map.array.copy() + noise_map[30] = noise_map[30].real + 1.5j * noise_map[30].real + + chunks = [ + (dataset.uv_wavelengths[k0:k1], dataset.data.array[k0:k1], noise_map[k0:k1]) + for k0, k1 in ((0, 20), (20, 40)) + ] + + with pytest.raises(aa.exc.DatasetException): + dataset.apply_sparse_operator_from_chunks(chunks) diff --git a/test_autoarray/fit/test_fit_interferometer.py b/test_autoarray/fit/test_fit_interferometer.py index 5e8f1102f..2d8f1ef39 100644 --- a/test_autoarray/fit/test_fit_interferometer.py +++ b/test_autoarray/fit/test_fit_interferometer.py @@ -308,3 +308,85 @@ def test__dirty_chi_squared_map__equals_transformer_image_from_chi_squared_map( ) assert (fit_interferometer_7.dirty_chi_squared_map == dirty_chi_squared_map).all() + + +# --------------------------------------------------------------------------- +# Tests: noise normalization cached on the sparse operator +# --------------------------------------------------------------------------- + + +def _sparse_dataset(noise_normalization): + dataset, _ = _make_dataset() + + sparse_operator = aa.InterferometerSparseOperator.from_nufft_precision_operator( + nufft_precision_operator=np.ones((4, 4)), + dirty_image=np.zeros(4), + noise_normalization=noise_normalization, + ) + + return aa.Interferometer( + data=dataset.data, + noise_map=dataset.noise_map, + uv_wavelengths=dataset.uv_wavelengths, + real_space_mask=dataset.real_space_mask, + sparse_operator=sparse_operator, + ) + + +def test__noise_normalization__reads_the_sparse_operator_scalar_when_present(): + dataset = _sparse_dataset(noise_normalization=123.0) + + fit = aa.m.MockFitInterferometer( + dataset=dataset, + use_mask_in_fit=False, + model_data=aa.Visibilities(visibilities=[1.0 + 2.0j, 3.0 + 4.0j]), + ) + + assert fit.noise_normalization == 123.0 + + +def test__noise_normalization__falls_back_to_the_noise_map_without_the_scalar(): + dataset = _sparse_dataset(noise_normalization=None) + + fit = aa.m.MockFitInterferometer( + dataset=dataset, + use_mask_in_fit=False, + model_data=aa.Visibilities(visibilities=[1.0 + 2.0j, 3.0 + 4.0j]), + ) + + assert fit.noise_normalization == pytest.approx( + 4.0 * np.log(2 * np.pi * 4.0), 1.0e-12 + ) + + +def test__noise_normalization__a_replaced_noise_map_ignores_the_scalar(): + dataset = _sparse_dataset(noise_normalization=123.0) + + noise_map = aa.VisibilitiesNoiseMap(visibilities=[1.0 + 1.0j, 1.0 + 1.0j]) + + fit = aa.m.MockFitInterferometer( + dataset=dataset, + use_mask_in_fit=False, + model_data=aa.Visibilities(visibilities=[1.0 + 2.0j, 3.0 + 4.0j]), + noise_map=noise_map, + ) + + assert fit.noise_normalization == pytest.approx(4.0 * np.log(2 * np.pi), 1.0e-12) + + +def test__noise_normalization__apply_sparse_operator_scalar_matches_array_path(): + dataset, _ = _make_dataset() + + model_data = aa.Visibilities(visibilities=[1.0 + 2.0j, 3.0 + 4.0j]) + + fit_array = aa.m.MockFitInterferometer( + dataset=dataset, use_mask_in_fit=False, model_data=model_data + ) + fit_sparse = aa.m.MockFitInterferometer( + dataset=dataset.apply_sparse_operator(), + use_mask_in_fit=False, + model_data=model_data, + ) + + assert fit_sparse.dataset.sparse_operator.noise_normalization is not None + assert fit_sparse.noise_normalization == fit_array.noise_normalization diff --git a/test_autoarray/inversion/inversion/interferometer/test_interferometer.py b/test_autoarray/inversion/inversion/interferometer/test_interferometer.py index e9af578f4..57ece9b55 100644 --- a/test_autoarray/inversion/inversion/interferometer/test_interferometer.py +++ b/test_autoarray/inversion/inversion/interferometer/test_interferometer.py @@ -1176,3 +1176,148 @@ def test__interferometer_mapping__curvature_matrix_and_data_vector_evaluated_onc assert counts == {"curvature_matrix": 1, "data_vector": 1} assert log_evidence_terms == reference + + +def _sparse_interface_setup(): + mask = aa.Mask2D.circular(shape_native=(10, 10), pixel_scales=1.0, radius=3.0) + + rng = np.random.default_rng(seed=21) + n_visibilities = 30 + data = rng.normal(size=n_visibilities) + 1j * rng.normal(size=n_visibilities) + sigma = rng.uniform(0.5, 2.0, size=n_visibilities) + + dataset = aa.Interferometer( + data=aa.Visibilities(visibilities=data), + noise_map=aa.VisibilitiesNoiseMap(visibilities=sigma + 1j * sigma), + uv_wavelengths=rng.normal(size=(n_visibilities, 2)), + real_space_mask=mask, + transformer_class=aa.TransformerDFT, + ) + + dataset_sparse = dataset.apply_sparse_operator(use_jax=False) + + mapper = _mapper_from( + mask=mask, pixels=16, shape=(4, 4), regularization=aa.reg.Constant(1.0) + ) + + return dataset_sparse, mapper + + +def _interface_from(dataset_sparse, data): + return aa.DatasetInterface( + data=data, + noise_map=dataset_sparse.noise_map, + grids=dataset_sparse.grids, + transformer=dataset_sparse.transformer, + sparse_operator=dataset_sparse.sparse_operator, + ) + + +def test__fast_chi_squared__data_none_uses_the_sparse_operator_data_term(): + dataset_sparse, mapper = _sparse_interface_setup() + + inversion_array = aa.Inversion( + dataset=_interface_from(dataset_sparse, data=dataset_sparse.data), + linear_obj_list=[mapper], + ) + inversion_scalar = aa.Inversion( + dataset=_interface_from(dataset_sparse, data=None), + linear_obj_list=[mapper], + ) + + assert isinstance(inversion_scalar, aa.InversionInterferometerSparse) + + # `apply_sparse_operator` computes `data_term` with the expression term 3 reduces with, + # so the two paths agree bit-for-bit. + assert inversion_scalar.fast_chi_squared == inversion_array.fast_chi_squared + + # And both equal the chi-squared of the residual visibilities. + residual_map = aa.util.fit.residual_map_from( + data=dataset_sparse.data, + model_data=inversion_array.mapped_reconstructed_operated_data, + ) + chi_squared = aa.util.fit.chi_squared_complex_from( + chi_squared_map=aa.util.fit.chi_squared_map_complex_from( + residual_map=residual_map, noise_map=dataset_sparse.noise_map + ) + ) + + assert inversion_scalar.fast_chi_squared == pytest.approx(chi_squared, 1.0e-8) + + +def test__fast_chi_squared__data_given__the_cached_data_term_is_ignored(): + dataset_sparse, mapper = _sparse_interface_setup() + + # Stand-in for profile-subtracted data: the interface's data differs from the data the + # operator's `data_term` was built from, so term 3 must be reduced from the interface. + subtracted = aa.Visibilities(visibilities=0.5 * dataset_sparse.data.array) + + inversion = aa.Inversion( + dataset=_interface_from(dataset_sparse, data=subtracted), + linear_obj_list=[mapper], + ) + inversion_scalar = aa.Inversion( + dataset=_interface_from(dataset_sparse, data=None), + linear_obj_list=[mapper], + ) + + noise_map = dataset_sparse.noise_map.array + + term_3 = np.sum(subtracted.array.real**2.0 / noise_map.real**2.0) + np.sum( + subtracted.array.imag**2.0 / noise_map.imag**2.0 + ) + + difference = inversion.fast_chi_squared - inversion_scalar.fast_chi_squared + + assert difference == pytest.approx( + term_3 - dataset_sparse.sparse_operator.data_term, rel=1.0e-10 + ) + + +def test__fast_chi_squared__data_none_without_a_data_term__raises(): + dataset_sparse, mapper = _sparse_interface_setup() + + operator = dataset_sparse.sparse_operator + + operator_without_scalars = ( + aa.InterferometerSparseOperator.from_nufft_precision_operator( + nufft_precision_operator=operator.nufft_precision_operator, + dirty_image=operator.dirty_image, + ) + ) + + interface = aa.DatasetInterface( + data=None, + noise_map=dataset_sparse.noise_map, + grids=dataset_sparse.grids, + transformer=dataset_sparse.transformer, + sparse_operator=operator_without_scalars, + ) + + inversion = aa.Inversion(dataset=interface, linear_obj_list=[mapper]) + + with pytest.raises(aa.exc.InversionException): + inversion.fast_chi_squared + + +def test__fast_chi_squared__data_none__jax_matches_numpy(): + pytest.importorskip("jax") + + import jax.numpy as jnp + + dataset_sparse, mapper = _sparse_interface_setup() + + inversion_np = aa.Inversion( + dataset=_interface_from(dataset_sparse, data=None), + linear_obj_list=[mapper], + xp=np, + ) + inversion_jax = aa.Inversion( + dataset=_interface_from(dataset_sparse, data=None), + linear_obj_list=[mapper], + xp=jnp, + ) + + assert float(inversion_jax.fast_chi_squared) == pytest.approx( + float(inversion_np.fast_chi_squared), rel=1.0e-7 + ) 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 8a59b8a03..bfc4dd78f 100644 --- a/test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py +++ b/test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py @@ -1125,3 +1125,208 @@ def test__interferometer_sparse_operator__numpy_branch_imports_no_jax(): subprocess.run( [sys.executable, "-c", script], check=True, env=environment, timeout=600 ) + + +def _streaming_inputs(n_visibilities=60, seed=3): + """ + A circular-masked NUFFT dataset with random data and non-uniform (but equal real and + imaginary) noise, so every per-visibility weight is distinct. + """ + 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, + ) + + return mask, uv_wavelengths, data, noise_map, dataset + + +def _chunks_from(uv_wavelengths, data, noise_map, edges): + return [ + (uv_wavelengths[k0:k1], data[k0:k1], noise_map[k0:k1]) + for k0, k1 in zip(edges[:-1], edges[1:]) + ] + + +def test__sparse_terms_from_chunks__matches_one_shot_for_one_two_three_and_uneven_chunks(): + pytest.importorskip("nufftax") + + mask, uv_wavelengths, data, noise_map, dataset = _streaming_inputs() + + operator_one_shot = np.asarray(dataset.psf_precision_operator_from()) + dirty_image_one_shot = dataset.apply_sparse_operator().sparse_operator.dirty_image + + data_term = np.sum(data.real**2 / noise_map.real**2) + np.sum( + data.imag**2 / noise_map.imag**2 + ) + noise_normalization = aa.util.fit.noise_normalization_complex_from( + noise_map=noise_map + ) + + for edges in ([0, 60], [0, 30, 60], [0, 20, 40, 60], [0, 1, 17, 59, 60]): + terms = aa.util.inversion_interferometer.sparse_terms_from_chunks( + _chunks_from(uv_wavelengths, data, noise_map, edges), + real_space_mask=mask, + ) + + np.testing.assert_allclose( + terms.nufft_precision_operator, + operator_one_shot, + rtol=1.0e-12, + atol=1.0e-12 * np.abs(operator_one_shot).max(), + ) + + dirty_image = aa.Array2D(values=terms.dirty_image_native, mask=mask).slim.array + + np.testing.assert_allclose( + dirty_image, + dirty_image_one_shot, + rtol=1.0e-12, + atol=1.0e-12 * np.abs(dirty_image_one_shot).max(), + ) + + assert terms.n_vis == 60 + assert terms.data_term == pytest.approx(data_term, rel=1.0e-12) + assert terms.noise_normalization == pytest.approx( + noise_normalization, rel=1.0e-12 + ) + assert terms.sum_weights == pytest.approx( + np.sum(noise_map.real**-2.0), rel=1.0e-12 + ) + + # The dirty beam is the adjoint of the weights, whose peak (at the image centre for + # a centred mask) is bounded by sum(w). + assert terms.dirty_beam_native.shape == mask.shape_native + assert np.max(np.abs(terms.dirty_beam_native)) <= terms.sum_weights * ( + 1.0 + 1.0e-12 + ) + + +def test__sparse_terms_from_chunks__accepts_visibilities_and_real_two_column_arrays(): + pytest.importorskip("nufftax") + + mask, uv_wavelengths, data, noise_map, _ = _streaming_inputs(n_visibilities=20) + + terms_complex = aa.util.inversion_interferometer.sparse_terms_from_chunks( + [(uv_wavelengths, data, noise_map)], real_space_mask=mask + ) + terms_structures = aa.util.inversion_interferometer.sparse_terms_from_chunks( + [ + ( + uv_wavelengths, + aa.Visibilities(visibilities=data), + aa.VisibilitiesNoiseMap(visibilities=noise_map), + ) + ], + real_space_mask=mask, + ) + terms_columns = aa.util.inversion_interferometer.sparse_terms_from_chunks( + [ + ( + uv_wavelengths, + np.stack([data.real, data.imag], axis=-1), + np.stack([noise_map.real, noise_map.imag], axis=-1), + ) + ], + real_space_mask=mask, + ) + + for terms in (terms_structures, terms_columns): + np.testing.assert_array_equal( + terms.nufft_precision_operator, terms_complex.nufft_precision_operator + ) + np.testing.assert_array_equal( + terms.dirty_image_native, terms_complex.dirty_image_native + ) + assert terms.data_term == terms_complex.data_term + assert terms.noise_normalization == terms_complex.noise_normalization + + +def test__sparse_terms_from_chunks__unequal_real_imag_noise_in_a_later_chunk__raises(): + pytest.importorskip("nufftax") + + mask, uv_wavelengths, data, noise_map, _ = _streaming_inputs() + + noise_map = noise_map.copy() + noise_map[45] = noise_map[45].real + 2.0j * noise_map[45].real + + with pytest.raises(aa.exc.DatasetException): + aa.util.inversion_interferometer.sparse_terms_from_chunks( + _chunks_from(uv_wavelengths, data, noise_map, [0, 20, 40, 60]), + real_space_mask=mask, + ) + + +def test__sparse_terms_from_chunks__no_visibilities__raises(): + mask = aa.Mask2D.circular(shape_native=(10, 10), pixel_scales=0.5, radius=2.0) + + with pytest.raises(ValueError): + aa.util.inversion_interferometer.sparse_terms_from_chunks( + [], real_space_mask=mask + ) + + +def test__sparse_terms__add__is_field_wise_sum(): + def terms(value, n_vis): + return aa.SparseTerms( + nufft_precision_operator=np.full((4, 4), value), + dirty_image_native=np.full((2, 2), 2.0 * value), + dirty_beam_native=np.full((2, 2), 3.0 * value), + sum_weights=4.0 * value, + data_term=5.0 * value, + noise_normalization=6.0 * value, + n_vis=n_vis, + ) + + total = terms(1.0, 3) + terms(10.0, 4) + + np.testing.assert_array_equal(total.nufft_precision_operator, np.full((4, 4), 11.0)) + np.testing.assert_array_equal(total.dirty_image_native, np.full((2, 2), 22.0)) + np.testing.assert_array_equal(total.dirty_beam_native, np.full((2, 2), 33.0)) + assert total.sum_weights == 44.0 + assert total.data_term == 55.0 + assert total.noise_normalization == 66.0 + assert total.n_vis == 7 + + mismatched = aa.SparseTerms( + nufft_precision_operator=np.zeros((6, 6)), + dirty_image_native=np.zeros((3, 3)), + dirty_beam_native=np.zeros((3, 3)), + sum_weights=0.0, + data_term=0.0, + noise_normalization=0.0, + n_vis=0, + ) + + with pytest.raises(ValueError): + terms(1.0, 1) + mismatched + + +def test__interferometer_sparse_operator__scalars_default_to_none_and_thread_through(): + operator = aa.InterferometerSparseOperator.from_nufft_precision_operator( + nufft_precision_operator=np.ones((4, 4)), dirty_image=np.ones(3) + ) + + assert operator.data_term is None + assert operator.noise_normalization is None + + operator = aa.InterferometerSparseOperator.from_nufft_precision_operator( + nufft_precision_operator=np.ones((4, 4)), + dirty_image=np.ones(3), + data_term=np.float64(2.5), + noise_normalization=3.5, + ) + + assert operator.data_term == 2.5 + assert type(operator.data_term) is float + assert operator.noise_normalization == 3.5