diff --git a/autoarray/dataset/imaging/dataset.py b/autoarray/dataset/imaging/dataset.py index 47e5424d0..3ac1c7d8d 100644 --- a/autoarray/dataset/imaging/dataset.py +++ b/autoarray/dataset/imaging/dataset.py @@ -61,6 +61,36 @@ def _validate_convolve_over_sample_size( ) +_SPARSE_OPERATOR_CONVOLVE_OVER_SAMPLE_SIZE_ERROR = ( + "The sparse linear algebra formalism precomputes PSF products at " + "image resolution and is incompatible with an oversampled PSF " + "(convolve_over_sample_size > 1)." +) + + +def _warn_sparse_operator_discarded(sparse_operator, method_name: str) -> None: + """ + Log a warning when a dataset rebuild discards an attached sparse operator. + + The sparse operator precomputes PSF products of every pair of masked noise-map + values, so it is invalidated by any change to the mask or noise-map. Rebuilds + which make such changes (e.g. `apply_mask`, `apply_noise_scaling`) therefore do + not carry it, and the returned dataset falls back to the (slower) mapping matrix + formalism unless the operator is re-applied. This warning stops that fallback + from happening silently. + """ + if sparse_operator is None: + return + + logger.warning( + f"IMAGING - {method_name} changes the mask or noise-map, which invalidates the " + f"sparse operator attached to this dataset, so it has been discarded. Pixelized " + f"reconstructions will use the slower mapping matrix formalism unless the sparse " + f"operator is re-applied (via apply_sparse_operator() or " + f"apply_sparse_operator_cpu()) as the last dataset operation." + ) + + class Imaging(AbstractDataset): def __init__( self, @@ -439,6 +469,10 @@ def apply_mask(self, mask: Mask2D) -> "Imaging": values=self.over_sample_size_pixelization.native, mask=mask ) + _warn_sparse_operator_discarded( + sparse_operator=self.sparse_operator, method_name="apply_mask" + ) + dataset = Imaging( data=data, noise_map=noise_map, @@ -520,6 +554,10 @@ def apply_noise_scaling( noise_map = Array2D(values=noise_map, mask=self.data.mask) + _warn_sparse_operator_discarded( + sparse_operator=self.sparse_operator, method_name="apply_noise_scaling" + ) + dataset = Imaging( data=data, noise_map=noise_map, @@ -527,6 +565,8 @@ def apply_noise_scaling( noise_covariance_matrix=self.noise_covariance_matrix, over_sample_size_lp=self.over_sample_size_lp, over_sample_size_pixelization=self.over_sample_size_pixelization, + convolve_over_sample_size_lp=self.convolve_over_sample_size_lp, + convolve_over_sample_size_pixelization=self.convolve_over_sample_size_pixelization, check_noise_map=False, ) @@ -551,6 +591,10 @@ def apply_over_sampling( This function resets the cached properties so that the new over sampling is used in the grid and grid pixelization. + The `noise_covariance_matrix` and any attached `sparse_operator` are carried over, as neither depends on + the over sampling (the sparse operator depends only on the noise-map, PSF kernel and mask). This means + `apply_sparse_operator()` / `apply_sparse_operator_cpu()` may be called before or after this method. + Parameters ---------- over_sample_size_lp @@ -566,12 +610,14 @@ def apply_over_sampling( data=self.data, noise_map=self.noise_map, psf=self.psf, + noise_covariance_matrix=self.noise_covariance_matrix, over_sample_size_lp=over_sample_size_lp or self.over_sample_size_lp, over_sample_size_pixelization=over_sample_size_pixelization or self.over_sample_size_pixelization, convolve_over_sample_size_lp=self.convolve_over_sample_size_lp, convolve_over_sample_size_pixelization=self.convolve_over_sample_size_pixelization, check_noise_map=False, + sparse_operator=self.sparse_operator, ) return dataset @@ -620,9 +666,7 @@ def apply_sparse_operator( if self.psf is not None and self.psf.convolve_over_sample_size > 1: raise exc.DatasetException( - "The sparse linear algebra formalism precomputes PSF products at " - "image resolution and is incompatible with an oversampled PSF " - "(convolve_over_sample_size > 1)." + _SPARSE_OPERATOR_CONVOLVE_OVER_SAMPLE_SIZE_ERROR ) logger.info( @@ -645,6 +689,8 @@ def apply_sparse_operator( noise_covariance_matrix=self.noise_covariance_matrix, over_sample_size_lp=self.over_sample_size_lp, over_sample_size_pixelization=self.over_sample_size_pixelization, + convolve_over_sample_size_lp=self.convolve_over_sample_size_lp, + convolve_over_sample_size_pixelization=self.convolve_over_sample_size_pixelization, check_noise_map=False, sparse_operator=sparse_operator, ) @@ -668,6 +714,11 @@ def apply_sparse_operator_cpu( A new `Imaging` dataset with a precomputed Numba-based sparse operator attached, enabling efficient pixelized source reconstruction on CPU hardware. """ + if self.psf is not None and self.psf.convolve_over_sample_size > 1: + raise exc.DatasetException( + _SPARSE_OPERATOR_CONVOLVE_OVER_SAMPLE_SIZE_ERROR + ) + try: import numba except ModuleNotFoundError: @@ -711,6 +762,8 @@ def apply_sparse_operator_cpu( noise_covariance_matrix=self.noise_covariance_matrix, over_sample_size_lp=self.over_sample_size_lp, over_sample_size_pixelization=self.over_sample_size_pixelization, + convolve_over_sample_size_lp=self.convolve_over_sample_size_lp, + convolve_over_sample_size_pixelization=self.convolve_over_sample_size_pixelization, check_noise_map=False, sparse_operator=sparse_operator, ) diff --git a/autoarray/inversion/inversion/abstract.py b/autoarray/inversion/inversion/abstract.py index 79c27677e..6241a6a16 100644 --- a/autoarray/inversion/inversion/abstract.py +++ b/autoarray/inversion/inversion/abstract.py @@ -299,8 +299,17 @@ def operated_mapping_matrix(self) -> np.ndarray: If there are multiple linear objects, the blurred mapping matrices are stacked such that their simultaneous linear equations are solved simultaneously. + + The dataset-specific `operated_mapping_matrix_list` is itself cached, so for a single linear object its + matrix is returned directly rather than copied via `hstack`, avoiding holding two identical copies of + the (large) operated mapping matrix in memory. """ - return self._xp.hstack(self.operated_mapping_matrix_list) + operated_mapping_matrix_list = self.operated_mapping_matrix_list + + if len(operated_mapping_matrix_list) == 1: + return operated_mapping_matrix_list[0] + + return self._xp.hstack(operated_mapping_matrix_list) @property def data_vector(self) -> np.ndarray: diff --git a/autoarray/inversion/inversion/imaging/abstract.py b/autoarray/inversion/inversion/imaging/abstract.py index 61736c983..477915bfd 100644 --- a/autoarray/inversion/inversion/imaging/abstract.py +++ b/autoarray/inversion/inversion/imaging/abstract.py @@ -116,7 +116,7 @@ def mapping_matrix_list(self) -> List[np.ndarray]: """ return [linear_obj.mapping_matrix for linear_obj in self.linear_obj_list] - @property + @cached_property def operated_mapping_matrix_list(self) -> List[np.ndarray]: """ The `operated_mapping_matrix` of a linear object describes the mappings between the observed data's values and diff --git a/autoarray/inversion/inversion/interferometer/abstract.py b/autoarray/inversion/inversion/interferometer/abstract.py index e1195e381..7d912efdb 100644 --- a/autoarray/inversion/inversion/interferometer/abstract.py +++ b/autoarray/inversion/inversion/interferometer/abstract.py @@ -1,6 +1,8 @@ import numpy as np from typing import Dict, List, Optional, Union +from autonerves import cached_property + from autoarray import exc from autoarray.dataset.interferometer.dataset import Interferometer from autoarray.inversion.inversion.dataset_interface import DatasetInterface @@ -61,7 +63,7 @@ def transformer(self): def mask(self) -> Mask2D: return self.transformer.real_space_mask - @property + @cached_property def operated_mapping_matrix_list(self) -> List[np.ndarray]: """ The `operated_mapping_matrix` of a linear object describes the mappings between the observed data's values diff --git a/test_autoarray/dataset/imaging/test_dataset.py b/test_autoarray/dataset/imaging/test_dataset.py index 6d19276da..4ecc9fe81 100644 --- a/test_autoarray/dataset/imaging/test_dataset.py +++ b/test_autoarray/dataset/imaging/test_dataset.py @@ -1,4 +1,5 @@ import copy +import logging import numpy as np import pytest @@ -6,6 +7,9 @@ import autoarray as aa from autoarray import exc +from autoarray.inversion.inversion.imaging_numba.sparse import ( + InversionImagingSparseNumba, +) from pathlib import Path test_data_path = Path(Path(__file__).resolve().parent) / "files" @@ -504,3 +508,95 @@ def test__convolve_over_sample_size__sparse_operator_guard(): with pytest.raises(aa.exc.DatasetException): dataset.apply_sparse_operator() + + +def test__apply_over_sampling__keeps_sparse_operator_and_noise_covariance( + masked_imaging_7x7, masked_imaging_covariance_7x7, delaunay_mapper_9_3x3 +): + # The sparse operator depends only on the noise-map, PSF kernel and mask, so it + # stays valid when only the over sampling changes and must not be dropped. + dataset_operator_first = ( + masked_imaging_7x7.apply_sparse_operator_cpu().apply_over_sampling( + over_sample_size_lp=2 + ) + ) + + assert dataset_operator_first.sparse_operator is not None + + dataset_operator_last = masked_imaging_7x7.apply_over_sampling( + over_sample_size_lp=2 + ).apply_sparse_operator_cpu() + + inversion_operator_first = aa.Inversion( + dataset=dataset_operator_first, + linear_obj_list=[delaunay_mapper_9_3x3], + ) + inversion_operator_last = aa.Inversion( + dataset=dataset_operator_last, + linear_obj_list=[delaunay_mapper_9_3x3], + ) + + assert isinstance(inversion_operator_first, InversionImagingSparseNumba) + assert inversion_operator_first.reconstruction == pytest.approx( + inversion_operator_last.reconstruction, 1.0e-8 + ) + assert inversion_operator_first.log_det_curvature_reg_matrix_term == pytest.approx( + inversion_operator_last.log_det_curvature_reg_matrix_term, 1.0e-8 + ) + + # The noise covariance matrix is independent of over sampling and is also kept. + dataset_covariance = masked_imaging_covariance_7x7.apply_over_sampling( + over_sample_size_lp=2 + ) + + assert dataset_covariance.noise_covariance_matrix == pytest.approx( + masked_imaging_covariance_7x7.noise_covariance_matrix, 1.0e-8 + ) + + +def test__apply_mask_and_noise_scaling__warn_when_sparse_operator_discarded( + masked_imaging_7x7, caplog +): + dataset = masked_imaging_7x7.apply_sparse_operator_cpu() + + with caplog.at_level(logging.WARNING, logger="autoarray.dataset.imaging.dataset"): + masked = dataset.apply_mask(mask=masked_imaging_7x7.mask) + + assert masked.sparse_operator is None + assert "sparse operator" in caplog.text + assert "apply_mask" in caplog.text + + caplog.clear() + + with caplog.at_level(logging.WARNING, logger="autoarray.dataset.imaging.dataset"): + scaled = dataset.apply_noise_scaling(mask=masked_imaging_7x7.mask) + + assert scaled.sparse_operator is None + assert "sparse operator" in caplog.text + assert "apply_noise_scaling" in caplog.text + + # No warning when there is no operator to discard. + caplog.clear() + + with caplog.at_level(logging.WARNING, logger="autoarray.dataset.imaging.dataset"): + masked_imaging_7x7.apply_mask(mask=masked_imaging_7x7.mask) + + assert "sparse operator" not in caplog.text + + +def test__convolve_over_sample_size__sparse_operator_cpu_guard(): + data = aa.Array2D.no_mask(values=np.ones((11, 11)), pixel_scales=1.0) + noise_map = aa.Array2D.no_mask(values=np.ones((11, 11)), pixel_scales=1.0) + kernel_fine = aa.Array2D.no_mask(values=np.ones((9, 9)), pixel_scales=0.5) + psf = aa.Convolver(kernel=kernel_fine) + + dataset = aa.Imaging( + data=data, + noise_map=noise_map, + psf=psf, + over_sample_size_pixelization=2, + convolve_over_sample_size_pixelization=2, + ) + + with pytest.raises(aa.exc.DatasetException): + dataset.apply_sparse_operator_cpu() diff --git a/test_autoarray/inversion/inversion/imaging/test_imaging.py b/test_autoarray/inversion/inversion/imaging/test_imaging.py index b407d6e7b..761e854a4 100644 --- a/test_autoarray/inversion/inversion/imaging/test_imaging.py +++ b/test_autoarray/inversion/inversion/imaging/test_imaging.py @@ -267,3 +267,35 @@ def test__mapping_matrix_over_sampled_for__kxs__full_bin_reproduces_mapping_matr # s=2 rows mean-binned to image resolution also reproduce mapping_matrix. binned = m_s2.reshape(n_pix, 4, mapping_matrix.shape[1]).mean(axis=1) assert binned == pytest.approx(mapping_matrix, abs=1.0e-14) + + +def test__operated_mapping_matrix_list__psf_convolution_performed_once_per_linear_obj( + masked_imaging_7x7, delaunay_mapper_9_3x3, monkeypatch +): + # The dense (mapping) route reaches `operated_mapping_matrix_list` from the + # curvature matrix / data vector (via `operated_mapping_matrix`) and again from + # `mapped_reconstructed_data_dict`. It must be cached so the PSF convolution of + # each linear object's mapping matrix happens once per inversion. + calls = [] + + convolved_mapping_matrix_from = aa.Convolver.convolved_mapping_matrix_from + + def counted(self, *args, **kwargs): + calls.append(1) + return convolved_mapping_matrix_from(self, *args, **kwargs) + + monkeypatch.setattr(aa.Convolver, "convolved_mapping_matrix_from", counted) + + inversion = aa.Inversion( + dataset=masked_imaging_7x7, + linear_obj_list=[delaunay_mapper_9_3x3], + ) + + assert isinstance(inversion, aa.InversionImagingMapping) + + inversion.log_det_curvature_reg_matrix_term + inversion.reconstruction + inversion.mapped_reconstructed_operated_data + inversion.mapped_reconstructed_data + + assert len(calls) == 1 diff --git a/test_autoarray/inversion/inversion/interferometer/test_interferometer.py b/test_autoarray/inversion/inversion/interferometer/test_interferometer.py index 57ece9b55..423384122 100644 --- a/test_autoarray/inversion/inversion/interferometer/test_interferometer.py +++ b/test_autoarray/inversion/inversion/interferometer/test_interferometer.py @@ -1178,6 +1178,40 @@ def test__interferometer_mapping__curvature_matrix_and_data_vector_evaluated_onc assert log_evidence_terms == reference +def test__operated_mapping_matrix_list__transform_performed_once_per_linear_obj( + interferometer_7_no_fft, rectangular_mapper_7x7_3x3, monkeypatch +): + # The dense (mapping) route reaches `operated_mapping_matrix_list` from the + # curvature matrix / data vector (via `operated_mapping_matrix`) and again from + # `mapped_reconstructed_data_dict`. It must be cached so the Fourier transform of + # each linear object's mapping matrix happens once per inversion. + calls = [] + + transformer_cls = type(interferometer_7_no_fft.transformer) + transform_mapping_matrix = transformer_cls.transform_mapping_matrix + + def counted(self, *args, **kwargs): + calls.append(1) + return transform_mapping_matrix(self, *args, **kwargs) + + monkeypatch.setattr(transformer_cls, "transform_mapping_matrix", counted) + + inversion = aa.Inversion( + dataset=interferometer_7_no_fft, + linear_obj_list=[rectangular_mapper_7x7_3x3], + settings=aa.Settings(), + ) + + assert isinstance(inversion, aa.InversionInterferometerMapping) + + inversion.log_det_curvature_reg_matrix_term + inversion.reconstruction + inversion.mapped_reconstructed_operated_data + inversion.mapped_reconstructed_data + + assert len(calls) == 1 + + def _sparse_interface_setup(): mask = aa.Mask2D.circular(shape_native=(10, 10), pixel_scales=1.0, radius=3.0)