diff --git a/autoarray/inversion/inversion/imaging_numba/inversion_imaging_numba_util.py b/autoarray/inversion/inversion/imaging_numba/inversion_imaging_numba_util.py index e0fc59c56..aef50fb9f 100644 --- a/autoarray/inversion/inversion/imaging_numba/inversion_imaging_numba_util.py +++ b/autoarray/inversion/inversion/imaging_numba/inversion_imaging_numba_util.py @@ -492,26 +492,6 @@ def curvature_matrix_with_added_to_diag_from( return curvature_matrix -@numba_util.jit() -def curvature_matrix_mirrored_from( - curvature_matrix: np.ndarray, -) -> np.ndarray: - curvature_matrix_mirrored = np.zeros( - (curvature_matrix.shape[0], curvature_matrix.shape[1]) - ) - - for i in range(curvature_matrix.shape[0]): - for j in range(curvature_matrix.shape[1]): - if curvature_matrix[i, j] != 0: - curvature_matrix_mirrored[i, j] = curvature_matrix[i, j] - curvature_matrix_mirrored[j, i] = curvature_matrix[i, j] - if curvature_matrix[j, i] != 0: - curvature_matrix_mirrored[i, j] = curvature_matrix[j, i] - curvature_matrix_mirrored[j, i] = curvature_matrix[j, i] - - return curvature_matrix_mirrored - - @numba_util.jit() def curvature_matrix_via_sparse_operator_from( psf_precision_operator: np.ndarray, @@ -694,6 +674,72 @@ def curvature_matrix_off_diags_via_sparse_operator_from( return curvature_matrix +@numba_util.jit() +def curvature_matrix_off_diags_via_mapper_and_blurred_curvature_weights_from( + data_to_pix_unique: np.ndarray, + data_weights: np.ndarray, + pix_lengths: np.ndarray, + pix_pixels: int, + blurred_curvature_weights: np.ndarray, # shape (n_unmasked, n_funcs) +) -> np.ndarray: + """ + Returns the off-diagonal terms in the curvature matrix `F` (see Warren & Dye 2003) + between a mapper object and a linear func object, from curvature weights that have + already been correlated with the PSF. + + This is the scatter half of + `curvature_matrix_off_diags_via_mapper_and_linear_func_curvature_vector_from`: that + function expands the curvature weights onto the native grid, performs a dense + sliding-window correlation with the PSF and then runs this loop. Splitting the two lets + the correlation be done once per linear func by the batched FFT convolver (which is over + an order of magnitude faster at HST resolution) while this loop, which is genuinely + sparse and irregular, stays in numba and runs once per (mapper, linear func) pair. + + For each unique mapping between a data pixel and a pixelization pixel, the PSF-correlated + curvature weights at that data pixel are multiplied by the mapping weight and accumulated + into the off-diagonal block of the curvature matrix. This accounts for sub-pixel mappings + between data pixels and pixelization pixels. + + Parameters + ---------- + data_to_pix_unique + An array that maps every data pixel index (e.g. the masked image pixel indexes in 1D) + to its unique set of pixelization pixel indexes (see `data_slim_to_pixelization_unique_from`). + data_weights + For every unique mapping between a set of data sub-pixels and a pixelization pixel, + the weight of this mapping based on the number of sub-pixels that map to the pixelization pixel. + pix_lengths + A 1D array describing how many unique pixels each data pixel maps to. Used to iterate over + `data_to_pix_unique` and `data_weights`. + pix_pixels + The total number of pixels in the pixelization that reconstructs the data. + blurred_curvature_weights + The operated values of the linear function divided by the noise-map squared and + correlated with the PSF, with shape [n_unmasked_data_pixels, n_linear_func_pixels]. + + Returns + ------- + ndarray + The off-diagonal block of the curvature matrix `F` (see Warren & Dye 2003), + with shape [pix_pixels, n_linear_func_pixels]. + """ + data_pixels = data_weights.shape[0] + n_funcs = blurred_curvature_weights.shape[1] + + off_diag = np.zeros((pix_pixels, n_funcs)) + + for data_0 in range(data_pixels): + for pix_0_index in range(pix_lengths[data_0]): + data_0_weight = data_weights[data_0, pix_0_index] + pix_0 = data_to_pix_unique[data_0, pix_0_index] + for f in range(n_funcs): + off_diag[pix_0, f] += ( + data_0_weight * blurred_curvature_weights[data_0, f] + ) + + return off_diag + + @numba_util.jit() def curvature_matrix_off_diags_via_mapper_and_linear_func_curvature_vector_from( data_to_pix_unique: np.ndarray, @@ -714,6 +760,13 @@ def curvature_matrix_off_diags_via_mapper_and_linear_func_curvature_vector_from( noise-map squared) are expanded into the native 2D image grid, convolved with the PSF kernel, and then remapped back to the 1D slim representation. + The inversion itself no longer calls this function: the correlation is performed by the + batched FFT convolver (`Convolver.reversed_kernel`) and only the scatter/accumulate loop + below runs in numba, via + `curvature_matrix_off_diags_via_mapper_and_blurred_curvature_weights_from`. This dense + sliding-window implementation is retained as the reference the FFT path is asserted + against in the unit tests. + For each unique mapping between a data pixel and a pixelization pixel, the convolved curvature weights at that data pixel are multiplied by the mapping weights and accumulated into the off-diagonal block of the curvature matrix. This accounts for diff --git a/autoarray/inversion/inversion/imaging_numba/sparse.py b/autoarray/inversion/inversion/imaging_numba/sparse.py index b2eae4ed5..f360fe18e 100644 --- a/autoarray/inversion/inversion/imaging_numba/sparse.py +++ b/autoarray/inversion/inversion/imaging_numba/sparse.py @@ -332,10 +332,14 @@ def curvature_matrix(self) -> np.ndarray: for simultaneously. In the w-tilde formalism this requires us to consider the mappings between data and every linear object, meaning that the linear alegbra has both on and off diagonal terms. - The `curvature_matrix` computed here is overwritten in memory when the regularization matrix is added to it, - because for large matrices this avoids overhead. For this reason, `curvature_matrix` is not a cached property - to ensure if we access it after computing the `curvature_reg_matrix` it is correctly recalculated in a new - array of memory. + Every block of F is written into the matrix already symmetrized: the mapper x mapper blocks are + folded and mirrored inside `curvature_matrix_via_sparse_operator_from`, and the off-diagonal + blocks (mapper x mapper, mapper x linear-func, linear-func x linear-func) are each placed + together with their transpose. A global symmetrizing pass over the assembled matrix would + therefore be a no-op, and is not run. + + `curvature_matrix` is a cached property, and `curvature_reg_matrix` adds the regularization + matrix to it out-of-place, so the cached F is never overwritten by that addition. """ if self.has(cls=AbstractLinearObjFuncList): curvature_matrix = self._curvature_matrix_func_list_and_mapper @@ -344,10 +348,6 @@ def curvature_matrix(self) -> np.ndarray: else: curvature_matrix = self._curvature_matrix_multi_mapper - curvature_matrix = inversion_imaging_numba_util.curvature_matrix_mirrored_from( - curvature_matrix=curvature_matrix, - ) - if len(self.no_regularization_index_list) > 0: curvature_matrix = ( inversion_imaging_numba_util.curvature_matrix_with_added_to_diag_from( @@ -386,11 +386,9 @@ def _curvature_matrix_mapper_diag(self) -> Optional[np.ndarray]: psf_precision_operator=self.sparse_operator.psf_precision_operator_sparse, psf_precision_indexes=self.sparse_operator.indexes, psf_precision_lengths=self.sparse_operator.lengths, - data_to_pix_unique=np.array( - mapper_i.unique_mappings.data_to_pix_unique - ), - data_weights=np.array(mapper_i.unique_mappings.data_weights), - pix_lengths=np.array(mapper_i.unique_mappings.pix_lengths), + data_to_pix_unique=mapper_i.unique_mappings.data_to_pix_unique, + data_weights=mapper_i.unique_mappings.data_weights, + pix_lengths=mapper_i.unique_mappings.pix_lengths, pix_pixels=mapper_i.params, ) @@ -493,6 +491,11 @@ def _curvature_matrix_multi_mapper(self) -> np.ndarray: mapper_param_range_j[0] : mapper_param_range_j[1], ] = off_diag + curvature_matrix[ + mapper_param_range_j[0] : mapper_param_range_j[1], + mapper_param_range_i[0] : mapper_param_range_i[1], + ] = off_diag.T + return curvature_matrix @property @@ -505,18 +508,110 @@ def _curvature_matrix_func_list_and_mapper(self) -> np.ndarray: curvature matrix given by equation (4) and the letter F. This function computes the diagonal terms of F using the sparse_operator formalism. + + The three blocks of F are assembled by separate private helpers, so that each block can be + computed (and therefore profiled) on its own: + + - the mapper x mapper block, `_curvature_matrix_mapper_diag` (via `_curvature_matrix_multi_mapper`); + - the mapper x linear-func blocks, `_curvature_matrix_mapper_func_blocks_from`; + - the linear-func x linear-func blocks, `_curvature_matrix_func_func_blocks_from`. + + The helpers write into the `curvature_matrix` they are passed and return it, so composing them + in this order is exactly the single-pass assembly they replaced. """ curvature_matrix = self._curvature_matrix_multi_mapper + curvature_matrix = self._curvature_matrix_mapper_func_blocks_from( + curvature_matrix=curvature_matrix + ) + + curvature_matrix = self._curvature_matrix_func_func_blocks_from( + curvature_matrix=curvature_matrix + ) + + return curvature_matrix + + def _blurred_curvature_weights_from( + self, curvature_weights: np.ndarray + ) -> np.ndarray: + """ + Returns a linear func's noise-weighted curvature weights correlated with the PSF, in the + mask's slim representation with shape [n_unmasked_data_pixels, n_linear_func_pixels]. + + The mapper x linear-func block of `F` requires, at every unmasked data pixel, the + sliding-window sum ``sum_dy_dx psf[dy, dx] * weights[y + dy - cy, x + dx - cx]`` -- a + *correlation* with the PSF, not a convolution. Correlating with the PSF is exactly + convolving with the PSF reversed along both axes, so this routes through the dataset + PSF's `reversed_kernel` convolver and its batched (multi-column) convolution, which is + over an order of magnitude faster at HST resolution than the dense sliding window it + replaces. + + As in the sliding-window implementation the weights are zero everywhere outside the + mask (no blurring mapping matrix is supplied), and the result is read back only at the + unmasked pixels. + + Parameters + ---------- + curvature_weights + The operated values of a linear function divided by the noise-map squared, with + shape [n_unmasked_data_pixels, n_linear_func_pixels]. + """ + return self.psf.reversed_kernel.convolved_mapping_matrix_from( + mapping_matrix=curvature_weights, + mask=self.mask, + xp=np, + ) + + def _curvature_matrix_mapper_func_blocks_from( + self, curvature_matrix: np.ndarray + ) -> np.ndarray: + """ + Writes the mapper x linear-func off-diagonal blocks of the `curvature_matrix` into the input + matrix, returning it. + + Each block contracts a mapper's unique data-to-source-pixel mappings against the PSF-correlated, + noise-weighted curvature vector of a linear function. + + The correlation is the dominant cost of F at HST resolution, so it is done once per linear + func by `_blurred_curvature_weights_from` (batched FFT convolution) rather than once per + (mapper, linear func) pair by a dense sliding window, and only the sparse scatter of the + result onto source pixels runs in numba. + + Each `[mapper, linear_func]` block is written together with its transpose into the + `[linear_func, mapper]` block, so F leaves this helper symmetric and no global mirroring + pass is required. + + Parameters + ---------- + curvature_matrix + The (total_params, total_params) curvature matrix the blocks are written into. + """ + mapper_list = self.cls_list_from(cls=Mapper) mapper_param_range_list = self.param_range_list_from(cls=Mapper) + if len(mapper_list) == 0: + return curvature_matrix + linear_func_list = self.cls_list_from(cls=AbstractLinearObjFuncList) linear_func_param_range_list = self.param_range_list_from( cls=AbstractLinearObjFuncList ) + # Neither the noise-weighted curvature weights of a linear func nor their PSF + # correlation depend on the mapper, so both are formed once per linear func rather + # than once per (mapper, linear func) pair. + blurred_curvature_weights_list = [ + self._blurred_curvature_weights_from( + curvature_weights=np.array( + self.linear_func_operated_mapping_matrix_dict[linear_func] + / self.noise_map[:, None] ** 2 + ) + ) + for linear_func in linear_func_list + ] + for i in range(len(mapper_list)): mapper = mapper_list[i] mapper_param_range = mapper_param_range_list[i] @@ -524,19 +619,14 @@ def _curvature_matrix_func_list_and_mapper(self) -> np.ndarray: for func_index, linear_func in enumerate(linear_func_list): linear_func_param_range = linear_func_param_range_list[func_index] - data_linear_func_matrix = ( - self.linear_func_operated_mapping_matrix_dict[linear_func] - / self.noise_map[:, None] ** 2 - ) - - off_diag = inversion_imaging_numba_util.curvature_matrix_off_diags_via_mapper_and_linear_func_curvature_vector_from( + off_diag = inversion_imaging_numba_util.curvature_matrix_off_diags_via_mapper_and_blurred_curvature_weights_from( data_to_pix_unique=mapper.unique_mappings.data_to_pix_unique, data_weights=mapper.unique_mappings.data_weights, pix_lengths=mapper.unique_mappings.pix_lengths, pix_pixels=mapper.params, - curvature_weights=np.array(data_linear_func_matrix), - mask=self.mask.array, - psf_kernel=self.psf.kernel.native.array, + blurred_curvature_weights=blurred_curvature_weights_list[ + func_index + ], ) curvature_matrix[ @@ -544,6 +634,33 @@ def _curvature_matrix_func_list_and_mapper(self) -> np.ndarray: linear_func_param_range[0] : linear_func_param_range[1], ] = off_diag + curvature_matrix[ + linear_func_param_range[0] : linear_func_param_range[1], + mapper_param_range[0] : mapper_param_range[1], + ] = off_diag.T + + return curvature_matrix + + def _curvature_matrix_func_func_blocks_from( + self, curvature_matrix: np.ndarray + ) -> np.ndarray: + """ + Writes the linear-func x linear-func blocks of the `curvature_matrix` into the input matrix, + returning it. + + Each block is a BLAS `dot` of two noise-weighted operated mapping matrices. + + Parameters + ---------- + curvature_matrix + The (total_params, total_params) curvature matrix the blocks are written into. + """ + + linear_func_list = self.cls_list_from(cls=AbstractLinearObjFuncList) + linear_func_param_range_list = self.param_range_list_from( + cls=AbstractLinearObjFuncList + ) + # 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. diff --git a/autoarray/operators/convolver.py b/autoarray/operators/convolver.py index b02806be2..c3b5bd72f 100644 --- a/autoarray/operators/convolver.py +++ b/autoarray/operators/convolver.py @@ -10,6 +10,7 @@ from typing import Optional, Tuple, Union import warnings +from autonerves import cached_property from autonerves import conf from autoarray.structures.arrays.uniform_2d import Array2D from autoarray.structures.grids.uniform_2d import Grid2D @@ -317,6 +318,53 @@ def kernel_shape_image_resolution(self) -> Tuple[int, int]: 2 * int(np.ceil((k // 2) / s)) + 1 for k in self.kernel.shape_native ) + @cached_property + def reversed_kernel(self) -> "Convolver": + """ + This convolver with its kernel reversed along both axes. + + Convolving with the reversed kernel is exactly *correlating* with this convolver's + kernel, because reversing both axes of one operand converts a convolution into a + correlation:: + + (x * flip(k))[i] = sum_d k[d] x[i + d - c] = correlate(x, k)[i] + + Callers whose operator is defined as a sliding-window correlation (for example the + mapper x linear-func block of the imaging curvature matrix, which sums + ``psf[dy, dx] * image[y + dy - cy, x + dx - cx]``) can therefore route through the + batched convolution machinery by convolving with this convolver instead of hand + rolling the correlation. + + The reversed convolver inherits this one's ``use_fft`` policy and + ``convolve_over_sample_size``, and reuses this convolver's ``ConvolverState`` + geometry (rebuilt for the reversed kernel, whose Fourier transform differs) when one + was preloaded, so the FFT geometry is built once rather than once per call. + + Cached, so a `Convolver` that outlives the objects using it (a dataset's PSF outlives + the `Inversion` rebuilt for every likelihood evaluation) builds its reversed kernel and + that kernel's FFT geometry only once. + """ + kernel = Array2D.no_mask( + values=np.asarray(self.kernel.native.array)[::-1, ::-1].copy(), + pixel_scales=self.kernel.pixel_scales, + origin=self.kernel.origin, + ) + + # An oversampled state carries sub-pixel permutations that only `state_from` can + # attach, so it is left to rebuild that case rather than preloading a partial state. + state = ( + ConvolverState(kernel=kernel, mask=self._state.source_mask) + if self._state is not None and self.convolve_over_sample_size == 1 + else None + ) + + return Convolver( + kernel=kernel, + state=state, + use_fft=self._use_fft, + convolve_over_sample_size=self.convolve_over_sample_size, + ) + def state_from(self, mask): if self.convolve_over_sample_size > 1: diff --git a/test_autoarray/inversion/inversion/imaging/test_inversion_imaging_util.py b/test_autoarray/inversion/inversion/imaging/test_inversion_imaging_util.py index 1eaeecede..5f9adf5fa 100644 --- a/test_autoarray/inversion/inversion/imaging/test_inversion_imaging_util.py +++ b/test_autoarray/inversion/inversion/imaging/test_inversion_imaging_util.py @@ -434,3 +434,99 @@ def test__data_vector_via_blurred_mapping_matrix_from(): ) assert (data_vector == np.array([2.0, 3.0, 1.0])).all() + + +@pytest.mark.parametrize("kernel", KERNELS_ODD, ids=KERNEL_IDS) +def test__curvature_matrix_off_diags_via_mapper_and_blurred_curvature_weights_from__matches_dense_kernel( + kernel, +): + """ + The mapper x linear-func block of `F` used to expand the curvature weights onto the + native grid, run a dense sliding-window correlation with the PSF and scatter the result, + all inside one numba kernel. The correlation now runs as a batched FFT convolution with + the PSF reversed along both axes (`Convolver.reversed_kernel`) and only the scatter stays + in numba. + + This asserts the two produce the same block. The kernels are asymmetric and non-square, + so a missing reversal (correlation computed as a convolution) or a transposed axis cannot + pass. + """ + mask = aa.Mask2D( + mask=np.array( + [ + [True, True, True, True, True, True, True], + [True, True, False, False, False, True, True], + [True, False, False, False, False, False, True], + [True, False, False, False, False, False, True], + [True, False, False, False, False, False, True], + [True, True, False, False, False, True, True], + [True, True, True, True, True, True, True], + ] + ), + pixel_scales=1.0, + ) + + data_pixels = int(mask.pixels_in_mask) + n_funcs = 3 + pix_pixels = 5 + + rng = np.random.default_rng(505) + + curvature_weights = rng.normal(size=(data_pixels, n_funcs)) + + # Every data pixel maps to one or two source pixels, with non-uniform weights. + max_lengths = 2 + pix_lengths = rng.integers(1, max_lengths + 1, size=data_pixels).astype("int") + data_to_pix_unique = rng.integers( + 0, pix_pixels, size=(data_pixels, max_lengths) + ).astype("int") + data_weights = rng.random(size=(data_pixels, max_lengths)) + + off_diag_dense = aa.util.inversion_imaging_numba.curvature_matrix_off_diags_via_mapper_and_linear_func_curvature_vector_from( + data_to_pix_unique=data_to_pix_unique, + data_weights=data_weights, + pix_lengths=pix_lengths, + pix_pixels=pix_pixels, + curvature_weights=curvature_weights, + mask=np.array(mask), + psf_kernel=kernel, + ) + + convolver = aa.Convolver( + kernel=aa.Array2D.no_mask(values=kernel, pixel_scales=1.0), + ) + + blurred_curvature_weights = convolver.reversed_kernel.convolved_mapping_matrix_from( + mapping_matrix=curvature_weights, + mask=mask, + xp=np, + ) + + off_diag_fft = aa.util.inversion_imaging_numba.curvature_matrix_off_diags_via_mapper_and_blurred_curvature_weights_from( + data_to_pix_unique=data_to_pix_unique, + data_weights=data_weights, + pix_lengths=pix_lengths, + pix_pixels=pix_pixels, + blurred_curvature_weights=blurred_curvature_weights, + ) + + assert off_diag_fft == pytest.approx(off_diag_dense, rel=1.0e-6) + + +def test__convolver_reversed_kernel__is_the_kernel_reversed_and_convolves_as_a_correlation(): + """ + `Convolver.reversed_kernel` is the same convolver with its kernel reversed along both + axes, so convolving with it correlates with the original kernel. + """ + kernel = np.arange(1.0, 16.0).reshape(3, 5) + + convolver = aa.Convolver( + kernel=aa.Array2D.no_mask(values=kernel, pixel_scales=1.0), + ) + + assert convolver.reversed_kernel.kernel.native.array == pytest.approx( + kernel[::-1, ::-1] + ) + + # Cached, so the reversed kernel and its FFT geometry are built once. + assert convolver.reversed_kernel is convolver.reversed_kernel 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 index bd797701a..585216dcb 100644 --- a/test_autoarray/inversion/inversion/test_curvature_matrix_func_list_blocks.py +++ b/test_autoarray/inversion/inversion/test_curvature_matrix_func_list_blocks.py @@ -1,20 +1,29 @@ """ -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. +The off-diagonal blocks of the `curvature_matrix` in the sparse imaging inversions, asserted +against brute-force references. + +The linear-func x linear-func block is computed from the upper triangle of blocks only, with +the mirrored block set from the transpose; the mapper x linear-func block of the numba +inversion correlates the curvature weights with the PSF via an FFT convolution with the +reversed kernel and scatters the result in numba. """ import numpy as np import pytest +import autoarray as aa + from autoarray.inversion.inversion.imaging.sparse import InversionImagingSparse +from autoarray.inversion.inversion.imaging_numba import inversion_imaging_numba_util from autoarray.inversion.inversion.imaging_numba.sparse import ( InversionImagingSparseNumba, ) from autoarray.inversion.linear_obj.func_list import AbstractLinearObjFuncList +from autoarray.inversion.linear_obj.unique_mappings import UniqueMappings from autoarray.inversion.mappers.abstract import Mapper +ASYMMETRIC_KERNEL = np.arange(1.0, 16.0).reshape(3, 5) + class FakeLinearFunc: def __init__(self, params): @@ -126,3 +135,158 @@ def test__curvature_matrix_func_list_blocks__matches_brute_force( ) assert curvature_matrix == pytest.approx(brute_force, abs=1.0e-12) + + +# The mapper x linear-func block of the numba sparse imaging inversion no longer correlates +# the curvature weights with the PSF inside its numba kernel: the correlation runs as a +# batched FFT convolution with the reversed PSF (`Convolver.reversed_kernel`) and only the +# scatter onto source pixels stays in numba. The block is also written together with its +# transpose, since the global mirroring pass over F was removed. +# +# The test below asserts both against the retained dense sliding-window kernel, with an +# asymmetric, non-square PSF so a missing reversal or a transposed axis cannot pass. + + +class FakeMapper: + def __init__(self, params, unique_mappings): + self.params = params + self.unique_mappings = unique_mappings + + +class StubInversionMapperAndFunc(InversionImagingSparseNumba): + """Bypasses the real constructor: `_curvature_matrix_mapper_func_blocks_from` only needs + the mapper and linear func lists, their param ranges, the noise map, the mask and the + PSF.""" + + def __init__(self, mapper, operated_matrix, noise_map, mask, psf): + self._mapper = mapper + self._func = FakeLinearFunc(params=operated_matrix.shape[1]) + self._noise_map = noise_map + self._mask = mask + self._psf = psf + + self.linear_func_operated_mapping_matrix_dict = {self._func: operated_matrix} + + self._total_params = mapper.params + self._func.params + + @property + def _xp(self): + return np + + @property + def total_params(self): + return self._total_params + + @property + def noise_map(self): + return self._noise_map + + @property + def mask(self): + return self._mask + + @property + def psf(self): + return self._psf + + def cls_list_from(self, cls): + if cls is Mapper: + return [self._mapper] + return [self._func] + + def param_range_list_from(self, cls): + if cls is Mapper: + return [[0, self._mapper.params]] + return [[self._mapper.params, self.total_params]] + + +@pytest.fixture +def mapper_and_func_inversion(): + mask = aa.Mask2D( + mask=np.array( + [ + [True, True, True, True, True, True, True], + [True, True, False, False, False, True, True], + [True, False, False, False, False, False, True], + [True, False, False, False, False, False, True], + [True, False, False, False, False, False, True], + [True, True, False, False, False, True, True], + [True, True, True, True, True, True, True], + ] + ), + pixel_scales=1.0, + ) + + data_pixels = int(mask.pixels_in_mask) + pix_pixels = 5 + n_funcs = 3 + + rng = np.random.default_rng(505) + + max_lengths = 2 + unique_mappings = UniqueMappings( + data_to_pix_unique=rng.integers( + 0, pix_pixels, size=(data_pixels, max_lengths) + ).astype("int"), + data_weights=rng.random(size=(data_pixels, max_lengths)), + pix_lengths=rng.integers(1, max_lengths + 1, size=data_pixels).astype("int"), + ) + + mapper = FakeMapper(params=pix_pixels, unique_mappings=unique_mappings) + + operated_matrix = rng.normal(size=(data_pixels, n_funcs)) + noise_map = 0.5 + rng.random(data_pixels) * 2.0 + + psf = aa.Convolver( + kernel=aa.Array2D.no_mask(values=ASYMMETRIC_KERNEL, pixel_scales=1.0), + ) + + return StubInversionMapperAndFunc( + mapper=mapper, + operated_matrix=operated_matrix, + noise_map=noise_map, + mask=mask, + psf=psf, + ) + + +def test__curvature_matrix_mapper_func_blocks__matches_dense_kernel_and_places_transpose( + mapper_and_func_inversion, +): + inversion = mapper_and_func_inversion + + total_params = inversion.total_params + + curvature_matrix = inversion._curvature_matrix_mapper_func_blocks_from( + curvature_matrix=np.zeros((total_params, total_params)) + ) + + mapper = inversion._mapper + curvature_weights = np.array( + list(inversion.linear_func_operated_mapping_matrix_dict.values())[0] + / inversion.noise_map[:, None] ** 2 + ) + + off_diag = inversion_imaging_numba_util.curvature_matrix_off_diags_via_mapper_and_linear_func_curvature_vector_from( + data_to_pix_unique=mapper.unique_mappings.data_to_pix_unique, + data_weights=mapper.unique_mappings.data_weights, + pix_lengths=mapper.unique_mappings.pix_lengths, + pix_pixels=mapper.params, + curvature_weights=curvature_weights, + mask=np.array(inversion.mask), + psf_kernel=ASYMMETRIC_KERNEL, + ) + + assert curvature_matrix[: mapper.params, mapper.params :] == pytest.approx( + off_diag, rel=1.0e-6 + ) + + # The global mirroring pass is gone, so the transpose must be written here. + assert curvature_matrix[mapper.params :, : mapper.params] == pytest.approx( + off_diag.T, rel=1.0e-6 + ) + + # The mapper x mapper and linear-func x linear-func blocks are not this helper's to write. + assert curvature_matrix[: mapper.params, : mapper.params] == pytest.approx( + np.zeros((mapper.params, mapper.params)), abs=1.0e-12 + )