diff --git a/autoarray/inversion/inversion/imaging/abstract.py b/autoarray/inversion/inversion/imaging/abstract.py index b774dbf63..61736c983 100644 --- a/autoarray/inversion/inversion/imaging/abstract.py +++ b/autoarray/inversion/inversion/imaging/abstract.py @@ -1,6 +1,8 @@ import numpy as np from typing import Dict, List, Union, Type +from autonerves import cached_property + from autoarray.dataset.imaging.dataset import Imaging from autoarray.inversion.inversion.dataset_interface import DatasetInterface from autoarray.inversion.linear_obj.func_list import AbstractLinearObjFuncList @@ -180,7 +182,7 @@ def linear_func_mapping_matrix_dict(self) -> Dict: return linear_func_mapping_matrix_dict - @property + @cached_property def linear_func_operated_mapping_matrix_dict(self) -> Dict: """ The `operated_mapping_matrix` of a linear object describes the mappings between the observed data's values and diff --git a/autoarray/inversion/inversion/imaging/sparse.py b/autoarray/inversion/inversion/imaging/sparse.py index 2a49801d1..52f04fde0 100644 --- a/autoarray/inversion/inversion/imaging/sparse.py +++ b/autoarray/inversion/inversion/imaging/sparse.py @@ -450,26 +450,25 @@ def _curvature_matrix_func_list_and_mapper(self) -> np.ndarray: linear_func_param_range[0] : linear_func_param_range[1], ].set(off_diag) - for index_0, linear_func_0 in enumerate(linear_func_list): + # The linear func x linear func block is symmetric, so each weighted matrix is + # formed once and only the upper triangle of blocks is computed, with the + # mirrored block set from the transpose. + weighted_vector_list = [ + self.linear_func_operated_mapping_matrix_dict[linear_func] + / self.noise_map[:, None] + for linear_func in linear_func_list + ] - linear_func_param_range_0 = linear_func_param_range_list[index_0] + for index_0 in range(len(linear_func_list)): - weighted_vector_0 = ( - self.linear_func_operated_mapping_matrix_dict[linear_func_0] - / self.noise_map[:, None] - ) + linear_func_param_range_0 = linear_func_param_range_list[index_0] - for index_1, linear_func_1 in enumerate(linear_func_list): + for index_1 in range(index_0, len(linear_func_list)): linear_func_param_range_1 = linear_func_param_range_list[index_1] - weighted_vector_1 = ( - self.linear_func_operated_mapping_matrix_dict[linear_func_1] - / self.noise_map[:, None] - ) - diag = self._xp.dot( - weighted_vector_0.T, - weighted_vector_1, + weighted_vector_list[index_0].T, + weighted_vector_list[index_1], ) if self._xp is np: @@ -479,6 +478,12 @@ def _curvature_matrix_func_list_and_mapper(self) -> np.ndarray: linear_func_param_range_1[0] : linear_func_param_range_1[1], ] = diag + if index_1 != index_0: + curvature_matrix[ + linear_func_param_range_1[0] : linear_func_param_range_1[1], + linear_func_param_range_0[0] : linear_func_param_range_0[1], + ] = diag.T + else: curvature_matrix = curvature_matrix.at[ @@ -486,6 +491,12 @@ def _curvature_matrix_func_list_and_mapper(self) -> np.ndarray: linear_func_param_range_1[0] : linear_func_param_range_1[1], ].set(diag) + if index_1 != index_0: + curvature_matrix = curvature_matrix.at[ + linear_func_param_range_1[0] : linear_func_param_range_1[1], + linear_func_param_range_0[0] : linear_func_param_range_0[1], + ].set(diag.T) + return curvature_matrix def _mapped_reconstructed_data_dict_from( diff --git a/autoarray/inversion/inversion/imaging_numba/sparse.py b/autoarray/inversion/inversion/imaging_numba/sparse.py index 856105533..b2eae4ed5 100644 --- a/autoarray/inversion/inversion/imaging_numba/sparse.py +++ b/autoarray/inversion/inversion/imaging_numba/sparse.py @@ -119,7 +119,9 @@ def linear_func_operated_mapping_matrix_dict(self) -> Dict: Memoized matrices are returned read-only; every consumer in this class copies or derives from them (`np.array(...)`, divisions), never mutates. """ - parent_fget = AbstractInversionImaging.linear_func_operated_mapping_matrix_dict.fget + parent_fget = ( + AbstractInversionImaging.linear_func_operated_mapping_matrix_dict.func + ) if os.environ.get("AUTOARRAY_NUMBA_OPERATED_MEMO", "1") == "0": return parent_fget(self) @@ -542,26 +544,25 @@ def _curvature_matrix_func_list_and_mapper(self) -> np.ndarray: linear_func_param_range[0] : linear_func_param_range[1], ] = off_diag - for index_0, linear_func_0 in enumerate(linear_func_list): + # The linear func x linear func block is symmetric, so each weighted matrix is + # formed once and only the upper triangle of blocks is computed, with the + # mirrored block set from the transpose. + weighted_vector_list = [ + self.linear_func_operated_mapping_matrix_dict[linear_func] + / self.noise_map[:, None] + for linear_func in linear_func_list + ] - linear_func_param_range_0 = linear_func_param_range_list[index_0] + for index_0 in range(len(linear_func_list)): - weighted_vector_0 = ( - self.linear_func_operated_mapping_matrix_dict[linear_func_0] - / self.noise_map[:, None] - ) + linear_func_param_range_0 = linear_func_param_range_list[index_0] - for index_1, linear_func_1 in enumerate(linear_func_list): + for index_1 in range(index_0, len(linear_func_list)): linear_func_param_range_1 = linear_func_param_range_list[index_1] - weighted_vector_1 = ( - self.linear_func_operated_mapping_matrix_dict[linear_func_1] - / self.noise_map[:, None] - ) - diag = np.dot( - weighted_vector_0.T, - weighted_vector_1, + weighted_vector_list[index_0].T, + weighted_vector_list[index_1], ) curvature_matrix[ @@ -569,6 +570,12 @@ def _curvature_matrix_func_list_and_mapper(self) -> np.ndarray: linear_func_param_range_1[0] : linear_func_param_range_1[1], ] = diag + if index_1 != index_0: + curvature_matrix[ + linear_func_param_range_1[0] : linear_func_param_range_1[1], + linear_func_param_range_0[0] : linear_func_param_range_0[1], + ] = diag.T + return curvature_matrix def _mapped_reconstructed_data_dict_from( diff --git a/autoarray/operators/convolver.py b/autoarray/operators/convolver.py index 73fbd5f74..b02806be2 100644 --- a/autoarray/operators/convolver.py +++ b/autoarray/operators/convolver.py @@ -142,6 +142,7 @@ class determines how masked real-space data are embedded into a padded array, ) self.fft_shape = fft_shape + self.source_mask = mask self.mask = mask.resized_from(self.fft_shape, pad_value=1) if blurring_mask is None: @@ -149,9 +150,7 @@ class determines how masked real-space data are embedded into a padded array, kernel_shape_native=self.kernel.shape_native ) else: - self.blurring_mask = blurring_mask.resized_from( - self.fft_shape, pad_value=1 - ) + self.blurring_mask = blurring_mask.resized_from(self.fft_shape, pad_value=1) # Set by Convolver.state_from when convolve_over_sample_size > 1: the # permutations from per-pixel sub-block ordering to the fine mask's @@ -170,6 +169,17 @@ class determines how masked real-space data are embedded into a padded array, # complex64 kernel — see that method's body for why. self.fft_kernel_c64 = self.fft_kernel.astype(np.complex64) + def is_for_mask(self, mask) -> bool: + """ + Whether this state was built from the input mask, and can therefore be reused + instead of rebuilt (its padded FFT geometry is only valid for that mask). + """ + return ( + self.source_mask.pixel_scales == mask.pixel_scales + and self.source_mask.shape_native == mask.shape_native + and np.array_equal(np.array(self.source_mask), np.array(mask)) + ) + class Convolver: def __init__( @@ -303,7 +313,9 @@ def kernel_shape_image_resolution(self) -> Tuple[int, int]: if s == 1: return self.kernel.shape_native - return tuple(2 * int(np.ceil((k // 2) / s)) + 1 for k in self.kernel.shape_native) + return tuple( + 2 * int(np.ceil((k // 2) / s)) + 1 for k in self.kernel.shape_native + ) def state_from(self, mask): @@ -314,16 +326,10 @@ def state_from(self, mask): return self._fine_state_from(mask=mask) - if ( - mask.shape_native[0] != self.kernel.shape_native[0] - or mask.shape_native[1] != self.kernel.shape_native[1] - ): - return ConvolverState(kernel=self.kernel, mask=mask) - - if self._state is None: - return ConvolverState(kernel=self.kernel, mask=mask) + if self._state is not None and self._state.is_for_mask(mask=mask): + return self._state - return self._state + return ConvolverState(kernel=self.kernel, mask=mask) def _fine_state_from(self, mask) -> ConvolverState: """ diff --git a/test_autoarray/inversion/inversion/test_curvature_matrix_func_list_blocks.py b/test_autoarray/inversion/inversion/test_curvature_matrix_func_list_blocks.py new file mode 100644 index 000000000..bd797701a --- /dev/null +++ b/test_autoarray/inversion/inversion/test_curvature_matrix_func_list_blocks.py @@ -0,0 +1,128 @@ +""" +The linear-func x linear-func block of the `curvature_matrix` in the sparse imaging +inversions is computed from the upper triangle of blocks only, with the mirrored block +set from the transpose. This asserts the result matches a brute-force full double loop +for a random, spatially varying noise map. +""" + +import numpy as np +import pytest + +from autoarray.inversion.inversion.imaging.sparse import InversionImagingSparse +from autoarray.inversion.inversion.imaging_numba.sparse import ( + InversionImagingSparseNumba, +) +from autoarray.inversion.linear_obj.func_list import AbstractLinearObjFuncList +from autoarray.inversion.mappers.abstract import Mapper + + +class FakeLinearFunc: + def __init__(self, params): + self.params = params + + +class StubMixin: + """Bypasses the real constructor: the property under test only needs the linear + func list, their param ranges, the noise map and an empty starting matrix.""" + + def __init__(self, operated_matrix_list, noise_map): + self._func_list = [ + FakeLinearFunc(params=matrix.shape[1]) for matrix in operated_matrix_list + ] + + param_range_list = [] + total_params = 0 + for linear_func in self._func_list: + param_range_list.append([total_params, total_params + linear_func.params]) + total_params += linear_func.params + + self._param_range_list = param_range_list + self._total_params = total_params + self._noise_map = noise_map + + self.linear_func_operated_mapping_matrix_dict = { + linear_func: matrix + for linear_func, matrix in zip(self._func_list, operated_matrix_list) + } + + @property + def _xp(self): + return np + + @property + def noise_map(self): + return self._noise_map + + def cls_list_from(self, cls): + if cls is Mapper: + return [] + return self._func_list + + def param_range_list_from(self, cls): + if cls is Mapper: + return [] + return self._param_range_list + + @property + def _curvature_matrix_multi_mapper(self): + return np.zeros((self._total_params, self._total_params)) + + +class StubInversionSparse(StubMixin, InversionImagingSparse): + pass + + +class StubInversionSparseNumba(StubMixin, InversionImagingSparseNumba): + pass + + +def curvature_matrix_brute_force_from(operated_matrix_list, noise_map): + param_range_list = [] + total_params = 0 + for matrix in operated_matrix_list: + param_range_list.append([total_params, total_params + matrix.shape[1]]) + total_params += matrix.shape[1] + + curvature_matrix = np.zeros((total_params, total_params)) + + for index_0, matrix_0 in enumerate(operated_matrix_list): + for index_1, matrix_1 in enumerate(operated_matrix_list): + curvature_matrix[ + param_range_list[index_0][0] : param_range_list[index_0][1], + param_range_list[index_1][0] : param_range_list[index_1][1], + ] = np.dot((matrix_0 / noise_map[:, None]).T, matrix_1 / noise_map[:, None]) + + return curvature_matrix + + +@pytest.fixture +def operated_matrix_list_and_noise_map(): + rng = np.random.default_rng(7) + + data_pixels = 37 + + operated_matrix_list = [ + rng.normal(size=(data_pixels, params)) for params in (3, 2, 4) + ] + + # Spatially varying, non-constant, non-symmetric noise map. + noise_map = 0.5 + rng.random(data_pixels) * 2.0 + + return operated_matrix_list, noise_map + + +@pytest.mark.parametrize("cls", [StubInversionSparse, StubInversionSparseNumba]) +def test__curvature_matrix_func_list_blocks__matches_brute_force( + cls, operated_matrix_list_and_noise_map +): + operated_matrix_list, noise_map = operated_matrix_list_and_noise_map + + inversion = cls(operated_matrix_list=operated_matrix_list, noise_map=noise_map) + + curvature_matrix = inversion._curvature_matrix_func_list_and_mapper + + brute_force = curvature_matrix_brute_force_from( + operated_matrix_list=operated_matrix_list, noise_map=noise_map + ) + + assert curvature_matrix == pytest.approx(brute_force, abs=1.0e-12) diff --git a/test_autoarray/operators/test_convolver.py b/test_autoarray/operators/test_convolver.py index 309ecdc83..2aee39511 100644 --- a/test_autoarray/operators/test_convolver.py +++ b/test_autoarray/operators/test_convolver.py @@ -5,6 +5,7 @@ import pytest import autoarray as aa +from autoarray.operators.convolver import ConvolverState from pathlib import Path test_data_path = Path(Path(__file__).resolve().parent) / "files" @@ -386,9 +387,7 @@ def _ground_truth_scene(over_sample_size): values=_ground_truth_gaussian(kyy, kxx, 0.8), pixel_scales=1.0 / s ) - convolver = aa.Convolver( - kernel=kernel, normalize=True, convolve_over_sample_size=s - ) + convolver = aa.Convolver(kernel=kernel, normalize=True, convolve_over_sample_size=s) grid = aa.Grid2D.from_mask(mask=mask, over_sample_size=s) @@ -549,3 +548,135 @@ def test__convolve_over_sample_size__blurring_mask_padding__delta_kernel_identit binned = values_sub.reshape(mask.pixels_in_mask, s**2).mean(axis=1) assert np.array(convolved) == pytest.approx(binned, abs=1.0e-14) + + +def test__state_from__precomputed_state_reused_for_matching_mask(): + mask = aa.Mask2D.circular(shape_native=(15, 15), pixel_scales=1.0, radius=5.0) + + kernel = aa.Array2D.no_mask( + values=np.random.default_rng(3).random((5, 3)), pixel_scales=1.0 + ) + + state = ConvolverState(kernel=kernel, mask=mask) + + convolver = aa.Convolver(kernel=kernel, state=state) + + assert convolver.state_from(mask=mask) is state + assert convolver.state_from(mask=mask) is state + + # An identical but distinct mask object also matches the cached state. + mask_copy = aa.Mask2D.circular(shape_native=(15, 15), pixel_scales=1.0, radius=5.0) + + assert convolver.state_from(mask=mask_copy) is state + + # A different mask must not reuse the cached state. + mask_other = aa.Mask2D.circular(shape_native=(15, 15), pixel_scales=1.0, radius=4.0) + + state_other = convolver.state_from(mask=mask_other) + + assert state_other is not state + assert state_other.is_for_mask(mask=mask_other) + + # A mask of the same shape but different pixel scales is also a different mask. + mask_scales = aa.Mask2D.circular( + shape_native=(15, 15), pixel_scales=2.0, radius=10.0 + ) + + assert convolver.state_from(mask=mask_scales) is not state + + +def test__precomputed_state__convolution_bit_identical_to_no_state(): + mask = aa.Mask2D.circular(shape_native=(15, 15), pixel_scales=1.0, radius=5.0) + blurring_mask = mask.derive_mask.blurring_from(kernel_shape_native=(5, 3)) + + rng = np.random.default_rng(4) + + kernel = aa.Array2D.no_mask(values=rng.random((5, 3)), pixel_scales=1.0) + + image = aa.Array2D(values=rng.random(mask.pixels_in_mask), mask=mask) + blurring_image = aa.Array2D( + values=rng.random(blurring_mask.pixels_in_mask), mask=blurring_mask + ) + + mapping_matrix = rng.random((mask.pixels_in_mask, 4)) + blurring_mapping_matrix = rng.random((blurring_mask.pixels_in_mask, 4)) + + convolver_no_state = aa.Convolver(kernel=kernel) + convolver_state = aa.Convolver( + kernel=kernel, state=ConvolverState(kernel=kernel, mask=mask) + ) + + assert np.array_equal( + np.array( + convolver_state.convolved_image_via_real_space_np_from( + image=image, blurring_image=blurring_image + ) + ), + np.array( + convolver_no_state.convolved_image_via_real_space_np_from( + image=image, blurring_image=blurring_image + ) + ), + ) + + assert np.array_equal( + convolver_state.convolved_mapping_matrix_via_real_space_np_from( + mapping_matrix=mapping_matrix, + mask=mask, + blurring_mapping_matrix=blurring_mapping_matrix, + ), + convolver_no_state.convolved_mapping_matrix_via_real_space_np_from( + mapping_matrix=mapping_matrix, + mask=mask, + blurring_mapping_matrix=blurring_mapping_matrix, + ), + ) + + +def test__convolved_mapping_matrix_via_real_space_np__matches_image_convolution_per_column(): + # The blurring mapping matrix is ordered on `mask.derive_mask.blurring_from(...)`, + # whereas the convolution scatters it via the FFT-frame `state.blurring_mask`. The + # two orderings must agree, so this compares each column against both the image + # convolution and a brute force convolution performed on the original mask frame. + from scipy.signal import convolve as scipy_convolve + + mask = aa.Mask2D.circular(shape_native=(15, 15), pixel_scales=1.0, radius=5.0) + blurring_mask = mask.derive_mask.blurring_from(kernel_shape_native=(5, 3)) + + rng = np.random.default_rng(5) + + kernel = aa.Array2D.no_mask(values=rng.random((5, 3)), pixel_scales=1.0) + convolver = aa.Convolver(kernel=kernel) + + mapping_matrix = rng.random((mask.pixels_in_mask, 3)) + blurring_mapping_matrix = rng.random((blurring_mask.pixels_in_mask, 3)) + + convolved = convolver.convolved_mapping_matrix_via_real_space_np_from( + mapping_matrix=mapping_matrix, + mask=mask, + blurring_mapping_matrix=blurring_mapping_matrix, + ) + + native = np.zeros(mask.shape_native + (3,)) + native[mask.slim_to_native_tuple] = mapping_matrix + native[blurring_mask.slim_to_native_tuple] = blurring_mapping_matrix + + brute_force = scipy_convolve(native, kernel.native.array[..., None], mode="same")[ + mask.slim_to_native_tuple + ] + + assert convolved == pytest.approx(brute_force, abs=1.0e-12) + + for index in range(3): + image = aa.Array2D(values=mapping_matrix[:, index], mask=mask) + blurring_image = aa.Array2D( + values=blurring_mapping_matrix[:, index], mask=blurring_mask + ) + + convolved_image = convolver.convolved_image_via_real_space_np_from( + image=image, blurring_image=blurring_image + ) + + assert convolved[:, index] == pytest.approx( + np.array(convolved_image), abs=1.0e-12 + )