diff --git a/autoarray/dataset/interferometer/dataset.py b/autoarray/dataset/interferometer/dataset.py index ef1d1f1a6..a0732c373 100644 --- a/autoarray/dataset/interferometer/dataset.py +++ b/autoarray/dataset/interferometer/dataset.py @@ -339,6 +339,7 @@ def from_stream( show_progress: bool = False, batch_size: int = 128, phase_centre: Optional[Tuple[float, float]] = None, + pool_noise_map: bool = False, ) -> "Interferometer": """ Build an array-free `Interferometer` by accumulating a stream of visibility chunks, @@ -351,6 +352,19 @@ def from_stream( which `from_sparse_terms` turns into the dataset; peak memory is set by the chunk size, not the dataset size. + The sparse terms assume every visibility has equal real and imaginary noise sigma: the + precision operator and dirty beam are built from the real-part sigma alone. By default + a chunk with unequal sigmas raises. For thermal noise a difference usually comes from + the noise estimator (1-2 % when estimated by differencing adjacent visibilities); pass + `pool_noise_map=True` to pool each chunk's sigmas in quadrature, + `sigma^2 = (sigma_real^2 + sigma_imag^2) / 2` (which preserves the total variance), + or pool them yourself before handing the chunks over with + `inversion_interferometer_util.noise_map_pooled_from`. Pooling is approximate when the + sigmas genuinely differ (the sparse curvature then drops the `cos(a + b)` term of the + unequal weights; see `sparse_terms_from_chunks`), and the pooled `noise_normalization` + differs from the unpooled one, so do not compare log-evidences of a pooled and an + unpooled fit. + Parameters ---------- chunks @@ -367,11 +381,17 @@ def from_stream( `(y0, x0)` lands at the image origin; recorded as `sparse_terms.phase_centre`. `None` applies no shift (recorded as `(0.0, 0.0)`). See `sparse_terms_from_chunks`. + pool_noise_map + If `True`, pool every chunk's real and imaginary noise sigma in quadrature before + forming the terms instead of raising on unequal sigmas, logging the median and + maximum difference once (a warning when the median exceeds 25 %). See + `sparse_terms_from_chunks`. Raises ------ exc.DatasetException - If any chunk has unequal real and imaginary noise sigma. + If any chunk has unequal real and imaginary noise sigma and `pool_noise_map` is + `False`. """ if disable_jax(): use_jax = False @@ -387,6 +407,7 @@ def from_stream( use_jax=use_jax, show_progress=show_progress, phase_centre=phase_centre, + pool_noise_map=pool_noise_map, ) return cls.from_sparse_terms( @@ -427,6 +448,7 @@ def apply_sparse_operator( show_progress: bool = False, show_memory: bool = False, use_jax: bool = False, + pool_noise_map: bool = False, ): """ Precompute the NUFFT precision operator for efficient pixelized source reconstruction. @@ -493,6 +515,11 @@ def apply_sparse_operator( the harness pays 2.3-3.2 s of compile for a backend it asked to disable. (The same switch demotes the `"nufft"` builder to the NumPy brute force inside `nufft_precision_operator_from`.) + pool_noise_map + If `True`, pool the real and imaginary noise sigma of every visibility in + quadrature, `sigma^2 = (sigma_real^2 + sigma_imag^2) / 2`, before building the + operator, instead of raising on unequal sigmas (see "Precondition" below). The + returned dataset carries the pooled `noise_map`. Precondition ------------ @@ -502,19 +529,39 @@ def apply_sparse_operator( `psf_precision_operator_from`, which passes `noise_map_real` to `nufft_precision_operator_from`), a reduction that is exact only under that equality. With unequal sigmas the sparse curvature matrix silently disagrees with - the dense `InversionInterferometerMapping` path, so this method raises a + the dense `InversionInterferometerMapping` path, so by default this method raises a `DatasetException` rather than returning a wrong operator. + For thermal noise the real and imaginary parts of one visibility have the same + variance, so a measured difference usually comes from the noise estimator (1-2 % when + the noise is estimated by differencing adjacent visibilities). `pool_noise_map=True` + pools the two in quadrature (which preserves the total variance) and logs the median + and maximum difference once: an info line up to a 25 % median difference, a warning + above it, where the difference may be real. The equivalent by hand is building the + dataset with `noise_map=inversion_interferometer_util.noise_map_pooled_from(noise_map)`. + + Pooling is approximate when the sigmas genuinely differ: with per-visibility weights + `w_r`, `w_i` the exact curvature is `sum wbar cos(a - b) + dw cos(a + b)` + (`wbar = (w_r + w_i) / 2`, `dw = (w_r - w_i) / 2`) and `W~` holds only the + `cos(a - b)` part, so the dense `InversionInterferometerMapping` path (no + `apply_sparse_operator`) remains the exact one. The returned dataset carries the pooled + `noise_map`, so the dense residual and chi-squared maps and the cached sparse + `data_term` / `noise_normalization` describe the same noise model. The pooled + `noise_normalization` differs from the unpooled one, so do not compare the + log-evidences of a pooled and an unpooled fit of the same data. + Returns ------- Interferometer A new `Interferometer` dataset with the precomputed `InterferometerSparseOperator` attached, enabling efficient pixelized source reconstruction via the sparse linear algebra formalism. + With `pool_noise_map=True` its `noise_map` is the pooled one. Raises ------ exc.DatasetException - If any visibility has unequal real and imaginary noise sigma. + If any visibility has unequal real and imaginary noise sigma and `pool_noise_map` + is `False`. """ # `use_jax` now only selects between the two brute forces (the `"nufft"` builder @@ -527,6 +574,41 @@ def apply_sparse_operator( "apply_sparse_operator", "data", "noise_map", "uv_wavelengths", "transformer" ) + if pool_noise_map: + # Every term is built from the pooled dataset, and it is the pooled dataset that is + # returned, so the dense residual / chi-squared maps and the cached sparse + # `data_term` / `noise_normalization` describe the same noise model. + noise_map = np.asarray(self.noise_map.array, dtype=np.complex128) + + pooling_record = inversion_interferometer_util._NoiseMapPoolingRecord() + pooling_record.add(noise_map) + + dataset_pooled = Interferometer( + real_space_mask=self.real_space_mask, + data=self.data, + noise_map=VisibilitiesNoiseMap( + visibilities=inversion_interferometer_util.noise_map_pooled_from( + noise_map + ) + ), + uv_wavelengths=self.uv_wavelengths, + transformer_class=lambda uv_wavelengths, real_space_mask: self.transformer, + ) + + pooling_record.log() + + return dataset_pooled.apply_sparse_operator( + nufft_precision_operator=nufft_precision_operator, + batch_size=batch_size, + method=method, + eps=eps, + nufft_chunk_size=nufft_chunk_size, + chunk_k=chunk_k, + show_progress=show_progress, + show_memory=show_memory, + use_jax=use_jax, + ) + inversion_interferometer_util.check_noise_map_real_imag_equal(self.noise_map) if nufft_precision_operator is None: @@ -639,9 +721,15 @@ def apply_sparse_operator_from_chunks( 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. + (checked per chunk), the assumption the sparse operator rests on (see + `apply_sparse_operator`, "Precondition"). `pool_noise_map=True` is refused here: the + returned dataset retains this dataset's unpooled `noise_map`, which would then describe + a different noise model from the operator. Pool the dataset itself first (build it + with `noise_map=inversion_interferometer_util.noise_map_pooled_from(noise_map)` and + take the chunks from it), or use `Interferometer.from_stream(..., pool_noise_map=True)` + (array-free, no retained noise-map). 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 ------ @@ -681,10 +769,11 @@ def apply_sparse_operator_from_chunks( `chunk_size`, `chunk_k`, `use_jax`, `show_progress`). `phase_centre` is rejected: it would shift the operator's terms but not this dataset's retained `data`, so the two would describe different phase centres; use `Interferometer.from_stream` - (array-free, no retained data) to stream with a phase-centre shift. 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()`. + (array-free, no retained data) to stream with a phase-centre shift. + `pool_noise_map=True` is rejected for the same reason (see "Chunk contract"). + 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 ------- @@ -695,8 +784,8 @@ def apply_sparse_operator_from_chunks( Raises ------ exc.DatasetException - If any chunk has unequal real and imaginary noise sigma, or a `phase_centre` is - passed. + If any chunk has unequal real and imaginary noise sigma, or a `phase_centre` or + `pool_noise_map=True` is passed. """ if accumulator_kwargs.get("phase_centre") is not None: raise exc.DatasetException( @@ -708,6 +797,19 @@ def apply_sparse_operator_from_chunks( "array-free dataset with no retained visibilities." ) + if accumulator_kwargs.get("pool_noise_map", False): + raise exc.DatasetException( + "Interferometer.apply_sparse_operator_from_chunks does not take " + "`pool_noise_map=True`: it would build the sparse operator from the pooled " + "noise-map but retain this dataset's unpooled `noise_map`, so the two would " + "describe different noise models and give contradictory likelihoods. Pool the " + "dataset first -- `Interferometer(..., noise_map=" + "inversion_interferometer_util.noise_map_pooled_from(noise_map))`, with the " + "chunks taken from it -- or use `Interferometer.from_stream(..., " + "pool_noise_map=True)`, which builds an array-free dataset with no retained " + "noise-map." + ) + if disable_jax() and accumulator_kwargs.get("use_jax", False): accumulator_kwargs["use_jax"] = False diff --git a/autoarray/inversion/inversion/interferometer/inversion_interferometer_util.py b/autoarray/inversion/inversion/interferometer/inversion_interferometer_util.py index cad4afdba..d4ec227a3 100644 --- a/autoarray/inversion/inversion/interferometer/inversion_interferometer_util.py +++ b/autoarray/inversion/inversion/interferometer/inversion_interferometer_util.py @@ -42,6 +42,13 @@ def check_noise_map_real_imag_equal(noise_map) -> None: 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. + With per-visibility weights `w_r = 1 / sigma_real^2`, `w_i = 1 / sigma_imag^2` the exact + curvature is `sum_k wbar cos(a - b) + dw cos(a + b)` with `wbar = (w_r + w_i) / 2` and + `dw = (w_r - w_i) / 2`; only the `cos(a - b)` (lag-only, Toeplitz) part is representable by + `W~`, so no equal-sigma substitute is exact when the sigmas genuinely differ. When they + differ only by noise-estimator scatter, pool them in quadrature with `pool_noise_map=True` + (or `noise_map_pooled_from`), which skips this check. + Parameters ---------- noise_map @@ -79,13 +86,183 @@ def check_noise_map_real_imag_equal(noise_map) -> None: "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 " + "For thermal noise the real and imaginary parts of one visibility have the same " + "variance, so a small difference (typically 1-2 % when the noise is estimated by " + "differencing adjacent visibilities) is estimator scatter. Pool the two sigmas in " + "quadrature, sigma^2 = (sigma_real^2 + sigma_imag^2) / 2, which preserves the total " + "variance: pass `pool_noise_map=True` to `apply_sparse_operator`, `from_stream` or " + "`sparse_terms_from_chunks`, or pool the noise-map yourself with " + "`inversion_interferometer_util.noise_map_pooled_from(noise_map)` before building the " + "dataset or chunks. Pooling is an approximation when the sigmas genuinely differ: the " + "exact curvature is `sum wbar cos(a - b) + dw cos(a + b)` (`wbar`, `dw` the mean and " + "half-difference of the real and imaginary weights) and `W~` can only hold the " + "`cos(a - b)` part.\n\n" + "Alternatively fit without calling `apply_sparse_operator()` — the dense " "`InversionInterferometerMapping` path handles unequal real and imaginary sigmas " - "correctly." + "exactly." + ) + + +# The median fractional real/imag sigma difference above which pooling logs a warning rather +# than an info line: a difference this large may be real rather than noise-estimator scatter +# (the threshold pyuvimage uses, Discussion #13). +NOISE_MAP_POOLING_WARNING_FRACTION = 0.25 + + +def noise_map_pooled_from(noise_map) -> np.ndarray: + """ + Return `noise_map` with the real and imaginary sigma of every visibility pooled in + quadrature, `sigma = sqrt((sigma_real^2 + sigma_imag^2) / 2)`, as the complex array + `sigma + 1j * sigma`. + + The sparse operator assumes equal real and imaginary sigma per visibility (see + `check_noise_map_real_imag_equal`). For thermal noise the two parts of one visibility have + the same variance, so a measured difference usually comes from the noise estimator (1-2 % + when estimated by differencing adjacent visibilities); quadrature pooling preserves the + total variance `sigma_real^2 + sigma_imag^2` (and so the chi-squared expectation) and is + the better estimate of both. It differs from the arithmetic mean of the two weights + `1 / sigma^2` only at second order in the fractional asymmetry (4e-4 at 2 %). + + Visibilities whose real and imaginary sigma are already exactly equal are returned + unchanged (bit for bit), so pooling a pooled noise-map is a no-op. + + Parameters + ---------- + noise_map + The noise-map, in any form `sparse_terms_from_chunks` accepts: a + `VisibilitiesNoiseMap`, a complex `(K,)` array or a real `(K, 2)` array of + (real, imag) columns. + + Returns + ------- + np.ndarray + The pooled complex128 `(K,)` noise-map, with equal real and imaginary parts. + """ + noise_map = _complex_visibilities_from(noise_map) + + noise_map_real = noise_map.real + noise_map_imag = noise_map.imag + + sigma = np.where( + noise_map_real == noise_map_imag, + noise_map_real, + np.sqrt((noise_map_real**2.0 + noise_map_imag**2.0) / 2.0), + ) + + return sigma + 1j * sigma + + +def _noise_map_real_imag_fraction_from(noise_map) -> np.ndarray: + """ + The per-visibility fractional real/imag sigma difference `|re - im| / max(|re|, |im|)` + (0 where both are 0). + """ + noise_map = _complex_visibilities_from(noise_map) + + noise_map_real = np.abs(noise_map.real) + noise_map_imag = np.abs(noise_map.imag) + + denominator = np.maximum(noise_map_real, noise_map_imag) + + return np.abs(noise_map_real - noise_map_imag) / np.where( + denominator == 0.0, 1.0, denominator ) +def noise_map_real_imag_asymmetry_from(noise_map) -> Tuple[float, float]: + """ + Return the median and maximum fractional difference between the real and imaginary sigma + of `noise_map`, `|sigma_real - sigma_imag| / max(sigma_real, sigma_imag)` per visibility. + + The median is what `pool_noise_map=True` judges a noise-map by (so one bad baseline does + not escalate the message); the maximum is reported beside it. + + Parameters + ---------- + noise_map + The noise-map, in any form `noise_map_pooled_from` accepts. + + Returns + ------- + (float, float) + The `(median, max)` fractional asymmetry, e.g. `(0.02, 0.05)` for a 2 % median and a + 5 % maximum difference. + """ + fraction = _noise_map_real_imag_fraction_from(noise_map) + + return float(np.median(fraction)), float(np.max(fraction)) + + +class _NoiseMapPoolingRecord: + """ + Accumulates the real/imag sigma asymmetry of every noise-map pooled by one + `pool_noise_map=True` call -- one chunk at a time for `sparse_terms_from_chunks` -- so one + message is logged per call, not per chunk. + + The median is taken from a fixed histogram of the fractional asymmetry (bins of `1e-5`, so + a resolution of 0.001 %) rather than from the per-visibility values, so the record's memory + does not grow with the number of visibilities streamed; the maximum is exact. + """ + + _bin_edges = np.linspace(0.0, 1.0, 100_001) + + def __init__(self): + self.counts = np.zeros(self._bin_edges.size - 1, dtype=np.int64) + self.maximum = 0.0 + self.unequal = False + + def add(self, noise_map: np.ndarray): + """ + Record the asymmetry of a complex noise-map about to be pooled. + """ + if np.allclose(noise_map.real, noise_map.imag, atol=0.0): + return + + self.unequal = True + + fraction = _noise_map_real_imag_fraction_from(noise_map) + + self.counts += np.histogram(fraction, bins=self._bin_edges)[0] + self.maximum = max(self.maximum, float(np.max(fraction))) + + @property + def median(self) -> float: + cumulative = np.cumsum(self.counts) + index = int(np.searchsorted(cumulative, 0.5 * cumulative[-1])) + + return float(0.5 * (self._bin_edges[index] + self._bin_edges[index + 1])) + + def log(self): + """ + Log one line describing the pooling: nothing when every pooled noise-map already had + equal real and imaginary sigma (to the relative tolerance of + `check_noise_map_real_imag_equal`), an info line for a median difference up to + `NOISE_MAP_POOLING_WARNING_FRACTION`, a warning above it. + """ + if not self.unequal: + return + + median = self.median + + summary = ( + f"INTERFEROMETER - pooled sigma_re / sigma_im in quadrature, " + f"sigma^2 = (sigma_re^2 + sigma_im^2) / 2, for the sparse operator: median " + f"difference {100.0 * median:.2f} %, max {100.0 * self.maximum:.2f} %" + ) + + if median <= NOISE_MAP_POOLING_WARNING_FRACTION: + logger.info(f"{summary} (consistent with noise-estimator scatter).") + return + + logger.warning( + f"{summary}. A difference this large may be real: pooling weights the real and " + f"imaginary parts equally, which the noise-map does not, so the sparse curvature " + f"matrix drops the cos(a + b) term the unequal weights carry. The dense " + f"`InversionInterferometerMapping` path (fit without `apply_sparse_operator`) " + f"is exact for unequal real and imaginary sigmas." + ) + + def data_vector_via_transformed_mapping_matrix_from( transformed_mapping_matrix: np.ndarray, visibilities: np.ndarray, @@ -2075,6 +2252,7 @@ def sparse_terms_from_chunks( use_jax: bool = False, show_progress: bool = False, phase_centre: Optional[Tuple[float, float]] = None, + pool_noise_map: bool = False, ) -> SparseTerms: """ Accumulate the `SparseTerms` of an interferometer dataset one chunk of visibilities at a @@ -2097,6 +2275,12 @@ def sparse_terms_from_chunks( 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). + The sparse terms assume this equality: the precision operator and dirty beam are built + from the real-part sigma alone (see `check_noise_map_real_imag_equal`). With + `pool_noise_map=True` the check is replaced by quadrature pooling, + `sigma^2 = (sigma_real^2 + sigma_imag^2) / 2`, of each chunk's noise-map before any + term is formed (see "Unequal real and imaginary sigma" below); alternatively pool + before handing the chunks over, with `noise_map_pooled_from`. `K` may differ between chunks; empty chunks are skipped. Chunks are consumed once, in order, and only one is referenced at a time. @@ -2109,6 +2293,27 @@ def sparse_terms_from_chunks( / `__radd__`) is the multi-frequency-synthesis (MFS) terms, equal to one accumulation over every channel's chunks to summation order. + Unequal real and imaginary sigma + -------------------------------- + For thermal noise the real and imaginary parts of one visibility have the same variance, + so a measured difference usually comes from the noise estimator (1-2 % when the noise is + estimated by differencing adjacent visibilities). By default such a chunk raises. With + `pool_noise_map=True` each chunk's noise-map is replaced by `noise_map_pooled_from` -- + `sigma^2 = (sigma_real^2 + sigma_imag^2) / 2`, which preserves the total variance -- + before every term (precision operator, dirty image, dirty beam, `sum_weights`, + `data_term`, `noise_normalization`), so all of them describe one noise model; the result + is bit-identical to the default call on chunks pooled beforehand. One message is logged + per call, from the asymmetry accumulated over every chunk: none when every chunk already + had equal sigmas, an info line with the median and maximum difference when the median is + at most 25 %, and a warning above that, where the difference may be real. + + Pooling is an approximation when the sigmas genuinely differ: the exact curvature is + `sum wbar cos(a - b) + dw cos(a + b)` (`wbar`, `dw` the mean and half-difference of the + real and imaginary weights) and the precision operator holds only the `cos(a - b)` part; + the dense `InversionInterferometerMapping` path is exact. The pooled `noise_normalization` + differs from the unpooled one at second order in the asymmetry per visibility, so do not + compare the log-evidences of a pooled and an unpooled fit of the same data. + Phase-centre shift ------------------ With `phase_centre=(y0, x0)` (arcseconds, autoarray `(y, x)` order like a mask `origin`), @@ -2125,7 +2330,9 @@ def sparse_terms_from_chunks( accumulation. `data_term` is invariant under a unit phase only when the real and imaginary sigmas are exactly equal; `check_noise_map_real_imag_equal` accepts them equal to a relative tolerance, so with slightly unequal sigmas it differs from the unshifted value at that - tolerance (it is still the correct `data_term` of the shifted data). The shift is recorded + tolerance (it is still the correct `data_term` of the shifted data). With + `pool_noise_map=True` the sigmas are exactly equal after pooling, so `data_term` is + exactly phase-invariant. The shift is recorded as `SparseTerms.phase_centre` provenance (`(0.0, 0.0)` when no shift is applied), so terms with different phase centres -- including shifted and unshifted ones -- refuse to be summed. @@ -2148,6 +2355,10 @@ def sparse_terms_from_chunks( The `(y, x)` phase-centre shift in arcseconds applied to every chunk's visibilities before forming the dirty image (see "Phase-centre shift" above). `None` applies no shift and records `(0.0, 0.0)`. + pool_noise_map + If `True`, pool every chunk's real and imaginary sigma in quadrature before forming + any term, instead of raising on unequal sigmas (see "Unequal real and imaginary + sigma" above). `False` (the default) keeps the equal-sigma check. Returns ------- @@ -2157,7 +2368,8 @@ def sparse_terms_from_chunks( Raises ------ exc.DatasetException - If any chunk's noise-map has unequal real and imaginary sigma. + If any chunk's noise-map has unequal real and imaginary sigma and `pool_noise_map` + is `False`. ValueError If `chunks` yields no visibilities. """ @@ -2183,6 +2395,8 @@ def sparse_terms_from_chunks( m0 = phase_centre[0] * arcsec_to_rad l0 = phase_centre[1] * arcsec_to_rad + pooling_record = _NoiseMapPoolingRecord() if pool_noise_map else None + terms = None for uv_wavelengths, data, noise_map in chunks: @@ -2204,7 +2418,13 @@ def sparse_terms_from_chunks( f"{noise_map.shape[0]}." ) - check_noise_map_real_imag_equal(noise_map) + # Pooling replaces the noise-map before every term below, so the precision operator, + # dirty image, beam and both scalars all describe the same (pooled) noise model. + if pool_noise_map: + pooling_record.add(noise_map) + noise_map = noise_map_pooled_from(noise_map) + else: + check_noise_map_real_imag_equal(noise_map) noise_map_real = noise_map.real noise_map_imag = noise_map.imag @@ -2290,4 +2510,7 @@ def sparse_terms_from_chunks( "to accumulate." ) + if pool_noise_map: + pooling_record.log() + return terms diff --git a/test_autoarray/dataset/interferometer/test_dataset.py b/test_autoarray/dataset/interferometer/test_dataset.py index 345f04eb3..23e4447f7 100644 --- a/test_autoarray/dataset/interferometer/test_dataset.py +++ b/test_autoarray/dataset/interferometer/test_dataset.py @@ -1056,3 +1056,200 @@ def test__from_sparse_terms__sum_of_per_channel_terms_matches_in_memory_mfs( _delaunay_mapper(mask_2d_7x7), rel=1.0e-10, ) + + +def _asymmetric_interferometer(mask, transformer_class, fraction, n_visibilities=40): + """ + `_random_interferometer` with the imaginary sigma of every visibility set to the real + sigma times `1 - fraction` or `1 / (1 - fraction)` (random choice), so every visibility has + a fractional real/imag asymmetry of exactly `fraction` -- noise-estimator scatter at 2 %. + """ + dataset = _random_interferometer( + mask, transformer_class, n_visibilities=n_visibilities + ) + + sigma = dataset.noise_map.array.real + sign = np.random.default_rng(seed=7).choice([-1.0, 1.0], size=n_visibilities) + factor = np.where(sign > 0, 1.0 / (1.0 - fraction), 1.0 - fraction) + + return aa.Interferometer( + data=dataset.data, + noise_map=aa.VisibilitiesNoiseMap(visibilities=sigma + 1j * sigma * factor), + uv_wavelengths=dataset.uv_wavelengths, + real_space_mask=mask, + transformer_class=transformer_class, + ) + + +def _log_evidence_of(dataset, mapper): + inversion = aa.Inversion(dataset=dataset, linear_obj_list=[mapper]) + + return inversion, float( + aa.m.MockFitInterferometer(dataset=dataset, inversion=inversion).log_evidence + ) + + +def test__apply_sparse_operator__pool_noise_map__returns_pooled_dataset_matching_mapping( + mask_2d_7x7, +): + dataset = _asymmetric_interferometer( + mask_2d_7x7, transformer.TransformerDFT, fraction=0.02 + ) + + # The default still refuses unequal sigmas. + with pytest.raises(aa.exc.DatasetException, match="pool_noise_map=True"): + dataset.apply_sparse_operator(use_jax=False) + + dataset_sparse = dataset.apply_sparse_operator(use_jax=False, pool_noise_map=True) + + noise_map_pooled = aa.util.inversion_interferometer.noise_map_pooled_from( + dataset.noise_map + ) + + # The returned dataset carries the pooled noise-map, so its dense statistics describe the + # same noise model as the cached sparse scalars. + np.testing.assert_array_equal(dataset_sparse.noise_map.array, noise_map_pooled) + assert dataset_sparse.sparse_operator.noise_normalization == ( + aa.util.fit.noise_normalization_complex_from(noise_map=noise_map_pooled) + ) + + dataset_pooled = aa.Interferometer( + data=dataset.data, + noise_map=aa.VisibilitiesNoiseMap(visibilities=noise_map_pooled), + uv_wavelengths=dataset.uv_wavelengths, + real_space_mask=mask_2d_7x7, + transformer_class=transformer.TransformerDFT, + ) + + mapper = _delaunay_mapper(mask_2d_7x7) + + inversion_sparse, log_evidence_sparse = _log_evidence_of(dataset_sparse, mapper) + inversion_mapping, log_evidence_mapping = _log_evidence_of(dataset_pooled, mapper) + + assert isinstance(inversion_sparse, aa.InversionInterferometerSparse) + assert isinstance(inversion_mapping, aa.InversionInterferometerMapping) + + assert log_evidence_sparse == pytest.approx(log_evidence_mapping, rel=1.0e-8) + + +def test__apply_sparse_operator__pool_noise_map__approximation_size_at_two_percent( + mask_2d_7x7, +): + """ + Documentation pin of what pooling costs when the sigmas really are unequal: the pooled + sparse fit against the exact dense `InversionInterferometerMapping` fit of the unpooled + data, at a 2 % real/imag asymmetry on 40 visibilities. + + Measured (2026-10-07, test config): the pooled log-evidence is 9.9e-3 nats (7.7e-5 + relative) above the unpooled one. Redrawing which visibilities have the larger imaginary + sigma (sign seeds 0-7) spreads it over 0.010-0.096 nats in magnitude (at most 7.4e-4 + relative), because most of it is chi-squared and image-dependent: each visibility's real + and imaginary residuals are weighted by the pooled rather than their own sigma, an error + first order in the asymmetry and random in sign. The pooled `noise_normalization` adds a + fixed second-order shift (8.2e-3 nats here). This is why the log-evidences of a pooled and + an unpooled fit of the same data must not be compared; the bounds below only pin the order + of magnitude. + """ + dataset = _asymmetric_interferometer( + mask_2d_7x7, transformer.TransformerDFT, fraction=0.02 + ) + + mapper = _delaunay_mapper(mask_2d_7x7) + + _, log_evidence_unpooled_mapping = _log_evidence_of(dataset, mapper) + _, log_evidence_pooled_sparse = _log_evidence_of( + dataset.apply_sparse_operator(use_jax=False, pool_noise_map=True), mapper + ) + + difference = abs(log_evidence_pooled_sparse - log_evidence_unpooled_mapping) + + assert 1.0e-3 < difference < 1.0e-1 + + +def test__apply_sparse_operator__pool_noise_map__equal_sigmas_unchanged_and_silent( + mask_2d_7x7, caplog +): + dataset = _random_interferometer(mask_2d_7x7, transformer.TransformerDFT) + + with caplog.at_level("INFO"): + dataset_pooled = dataset.apply_sparse_operator( + use_jax=False, pool_noise_map=True + ) + + assert not any("pooled" in record.getMessage() for record in caplog.records) + + dataset_default = dataset.apply_sparse_operator(use_jax=False) + + np.testing.assert_array_equal( + dataset_pooled.noise_map.array, dataset.noise_map.array + ) + np.testing.assert_array_equal( + dataset_pooled.sparse_operator.dirty_image, + dataset_default.sparse_operator.dirty_image, + ) + assert ( + dataset_pooled.sparse_operator.data_term + == dataset_default.sparse_operator.data_term + ) + assert ( + dataset_pooled.sparse_operator.noise_normalization + == dataset_default.sparse_operator.noise_normalization + ) + + +def test__from_stream__pool_noise_map__matches_pre_pooled_stream(mask_2d_7x7): + dataset = _asymmetric_interferometer( + mask_2d_7x7, transformer.TransformerDFT, fraction=0.02 + ) + + kwargs = dict(transformer_class=transformer.TransformerDFT, method="numpy") + + with pytest.raises(aa.exc.DatasetException): + aa.Interferometer.from_stream( + _chunks_of(dataset, [0, 13, 40]), mask_2d_7x7, **kwargs + ) + + dataset_stream = aa.Interferometer.from_stream( + _chunks_of(dataset, [0, 13, 40]), mask_2d_7x7, pool_noise_map=True, **kwargs + ) + + noise_map_pooled = aa.util.inversion_interferometer.noise_map_pooled_from( + dataset.noise_map + ) + + chunks_pre_pooled = [ + (uv_wavelengths, data, noise_map_pooled[k0:k1]) + for (uv_wavelengths, data, _), (k0, k1) in zip( + _chunks_of(dataset, [0, 13, 40]), ((0, 13), (13, 40)) + ) + ] + + dataset_pre_pooled = aa.Interferometer.from_stream( + chunks_pre_pooled, mask_2d_7x7, **kwargs + ) + + np.testing.assert_array_equal( + dataset_stream.sparse_terms.nufft_precision_operator, + dataset_pre_pooled.sparse_terms.nufft_precision_operator, + ) + np.testing.assert_array_equal( + dataset_stream.sparse_terms.dirty_image_native, + dataset_pre_pooled.sparse_terms.dirty_image_native, + ) + assert dataset_stream.sparse_terms.data_term == ( + dataset_pre_pooled.sparse_terms.data_term + ) + assert dataset_stream.sparse_terms.noise_normalization == ( + dataset_pre_pooled.sparse_terms.noise_normalization + ) + + +def test__apply_sparse_operator_from_chunks__pool_noise_map__raises(mask_2d_7x7): + dataset = _asymmetric_interferometer( + mask_2d_7x7, transformer.TransformerDFT, fraction=0.02 + ) + + with pytest.raises(aa.exc.DatasetException, match="noise_map_pooled_from"): + dataset.apply_sparse_operator_from_chunks( + _chunks_of(dataset, [0, 40]), pool_noise_map=True, method="numpy" + ) 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 545d7cc5b..6ed07db77 100644 --- a/test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py +++ b/test_autoarray/inversion/inversion/interferometer/test_inversion_interferometer_util.py @@ -1,3 +1,5 @@ +import dataclasses +import logging import os import autoarray as aa @@ -1277,6 +1279,178 @@ def test__sparse_terms_from_chunks__unequal_real_imag_noise_in_a_later_chunk__ra ) +def test__check_noise_map_real_imag_equal__error_names_pooling(): + with pytest.raises(aa.exc.DatasetException, match="pool_noise_map=True"): + aa.util.inversion_interferometer.check_noise_map_real_imag_equal( + np.array([1.0 + 1.02j]) + ) + + +def _asymmetric_noise_map(n_visibilities, fraction, seed=11): + """ + A noise-map whose imaginary sigma is the real sigma times `1 - fraction` or + `1 / (1 - fraction)` (random choice), so every visibility has a fractional real/imag + asymmetry `|re - im| / max(re, im)` of exactly `fraction`, of either sign. + """ + rng = np.random.default_rng(seed=seed) + sigma = rng.uniform(0.5, 2.0, size=n_visibilities) + sign = rng.choice([-1.0, 1.0], size=n_visibilities) + + return sigma + 1j * sigma * np.where(sign > 0, 1.0 / (1.0 - fraction), 1.0 - fraction) + + +def test__noise_map_pooled_from__preserves_total_variance_with_equal_parts(): + noise_map = _asymmetric_noise_map(n_visibilities=30, fraction=0.02) + noise_map[4] = 1.3 + 1.3j + + pooled = aa.util.inversion_interferometer.noise_map_pooled_from(noise_map) + + assert pooled.dtype == np.complex128 + np.testing.assert_array_equal(pooled.real, pooled.imag) + np.testing.assert_allclose( + pooled.real**2 + pooled.imag**2, + noise_map.real**2 + noise_map.imag**2, + rtol=1.0e-14, + ) + + # Exactly equal sigmas are returned bit for bit, so pooling is idempotent. + assert pooled[4] == 1.3 + 1.3j + np.testing.assert_array_equal( + aa.util.inversion_interferometer.noise_map_pooled_from(pooled), pooled + ) + + # The same forms `sparse_terms_from_chunks` accepts. + np.testing.assert_array_equal( + aa.util.inversion_interferometer.noise_map_pooled_from( + aa.VisibilitiesNoiseMap(visibilities=noise_map) + ), + pooled, + ) + np.testing.assert_array_equal( + aa.util.inversion_interferometer.noise_map_pooled_from( + np.stack([noise_map.real, noise_map.imag], axis=-1) + ), + pooled, + ) + + +def test__noise_map_real_imag_asymmetry_from__two_percent_map(): + sigma = np.array([1.0, 2.0, 0.5, 1.5]) + noise_map = sigma + 1j * sigma * np.array([1.0, 0.98, 1.0, 1.0 / 0.98]) + + median, maximum = aa.util.inversion_interferometer.noise_map_real_imag_asymmetry_from( + noise_map + ) + + # Fractions 0, 2 %, 0, 2 % (the difference over the larger of the two sigmas). + assert median == pytest.approx(0.01, rel=1.0e-12) + assert maximum == pytest.approx(0.02, rel=1.0e-12) + + assert aa.util.inversion_interferometer.noise_map_real_imag_asymmetry_from( + sigma + 1j * sigma + ) == (0.0, 0.0) + + +def test__sparse_terms_from_chunks__pool_noise_map__matches_pre_pooled_chunks_bit_identically(): + """ + The witness of the pooling option: pooling inside the accumulator is the same as pooling + each chunk beforehand with `noise_map_pooled_from` and calling the default (checking) + accumulator, field by field and bit for bit; the default call on the unpooled chunks + still raises. + """ + pytest.importorskip("nufftax") + + mask, uv_wavelengths, data, _, _ = _streaming_inputs() + + noise_map = _asymmetric_noise_map(n_visibilities=60, fraction=0.02) + + edges = [0, 20, 40, 60] + + terms_pooled = aa.util.inversion_interferometer.sparse_terms_from_chunks( + _chunks_from(uv_wavelengths, data, noise_map, edges), + real_space_mask=mask, + pool_noise_map=True, + ) + + terms_pre_pooled = aa.util.inversion_interferometer.sparse_terms_from_chunks( + _chunks_from( + uv_wavelengths, + data, + aa.util.inversion_interferometer.noise_map_pooled_from(noise_map), + edges, + ), + real_space_mask=mask, + ) + + for field in dataclasses.fields(terms_pooled): + value_pooled = getattr(terms_pooled, field.name) + value_pre_pooled = getattr(terms_pre_pooled, field.name) + + if isinstance(value_pooled, np.ndarray): + np.testing.assert_array_equal(value_pooled, value_pre_pooled, field.name) + else: + assert value_pooled == value_pre_pooled, field.name + + with pytest.raises(aa.exc.DatasetException): + aa.util.inversion_interferometer.sparse_terms_from_chunks( + _chunks_from(uv_wavelengths, data, noise_map, edges), + real_space_mask=mask, + ) + + +_IIU_LOGGER = "autoarray.inversion.inversion.interferometer.inversion_interferometer_util" + + +def _pooled_terms_log(caplog, noise_map): + """ + Run `sparse_terms_from_chunks(pool_noise_map=True)` with the DFT transformer and the + brute-force builder over three chunks and return the pooling log records. The inputs come + from `_streaming_inputs`, which builds a `TransformerNUFFT`, so callers need `nufftax`. + """ + mask, uv_wavelengths, data, _, _ = _streaming_inputs() + + caplog.clear() + + with caplog.at_level(logging.INFO, logger=_IIU_LOGGER): + aa.util.inversion_interferometer.sparse_terms_from_chunks( + _chunks_from(uv_wavelengths, data, noise_map, [0, 20, 40, 60]), + real_space_mask=mask, + transformer_class=aa.TransformerDFT, + method="numpy", + pool_noise_map=True, + ) + + return [record for record in caplog.records if "pooled" in record.getMessage()] + + +def test__sparse_terms_from_chunks__pool_noise_map__logs_once_per_call(caplog): + pytest.importorskip("nufftax") + + # 2 % estimator scatter: one info line for the whole stream, not one per chunk. + records = _pooled_terms_log( + caplog, _asymmetric_noise_map(n_visibilities=60, fraction=0.02) + ) + + assert len(records) == 1 + assert records[0].levelno == logging.INFO + assert "median difference 2.00 %" in records[0].getMessage() + + # 30 %: the difference may be real, so a warning naming the exact dense path. + records = _pooled_terms_log( + caplog, _asymmetric_noise_map(n_visibilities=60, fraction=0.3) + ) + + assert len(records) == 1 + assert records[0].levelno == logging.WARNING + assert "median difference 30.00 %" in records[0].getMessage() + assert "InversionInterferometerMapping" in records[0].getMessage() + + # Already equal: pooling is a no-op and nothing is logged. + sigma = np.random.default_rng(seed=2).uniform(0.5, 2.0, size=60) + + assert _pooled_terms_log(caplog, sigma + 1j * sigma) == [] + + def test__sparse_terms_from_chunks__no_visibilities__raises(): mask = aa.Mask2D.circular(shape_native=(10, 10), pixel_scales=0.5, radius=2.0)