From dedde5ac1e3bc0278b7b320484351b322c809409 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Fri, 28 Aug 2026 15:31:59 -0400 Subject: [PATCH 1/3] perf(numba): split curvature-matrix F assembly into per-block helpers (#505) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `_curvature_matrix_func_list_and_mapper` assembled all three blocks of F in a single pass, so the breakdown harness could only time F as one step. Split the two loops out into private helpers that write their block into the `curvature_matrix` they are passed and return it: - `_curvature_matrix_mapper_func_blocks_from` — the mapper x linear-func loop - `_curvature_matrix_func_func_blocks_from` — the linear-func x linear-func loop `_curvature_matrix_mapper_diag` (the mapper x mapper block) already existed. `_curvature_matrix_func_list_and_mapper` now composes the three in the same order, and the `curvature_matrix` cached property is untouched (mirror + diag add unchanged). Pure code motion — no behaviour change, bit-identical output. `pytest test_autoarray/inversion`: 392 passed. Step 0 of #505, so autolens_profiling can time each block separately. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01SqrSVGPrFcUB1vvDsoTw3n --- .../inversion/imaging_numba/sparse.py | 62 +++++++++++++++++++ 1 file changed, 62 insertions(+) diff --git a/autoarray/inversion/inversion/imaging_numba/sparse.py b/autoarray/inversion/inversion/imaging_numba/sparse.py index b2eae4ed5..13fb0e053 100644 --- a/autoarray/inversion/inversion/imaging_numba/sparse.py +++ b/autoarray/inversion/inversion/imaging_numba/sparse.py @@ -505,10 +505,50 @@ 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 _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-convolved, + noise-weighted curvature vector of a linear function (the dense sliding-window correlation in + `curvature_matrix_off_diags_via_mapper_and_linear_func_curvature_vector_from`). + + Only the `[mapper, linear_func]` blocks are written; their transposes are filled in by the + global mirror applied in `curvature_matrix`. + + 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) @@ -544,6 +584,28 @@ def _curvature_matrix_func_list_and_mapper(self) -> np.ndarray: linear_func_param_range[0] : linear_func_param_range[1], ] = off_diag + 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. From cb91b7681ef507a43a75c928fa859da0aadbee1c Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Fri, 28 Aug 2026 15:56:15 -0400 Subject: [PATCH 2/3] perf(numba): drop the redundant passes in curvature-matrix F assembly (#505) Three redundant passes over F in the numba sparse imaging inversion, none of which changed a single value: 1. A global `curvature_matrix_mirrored_from` ran over the whole (P, P) matrix after assembly. Every block was already symmetric or had a known transpose: the mapper x mapper blocks are folded and mirrored inside `curvature_matrix_via_sparse_operator_from`, and the linear-func x linear-func blocks already wrote both triangles. The mirror is removed and the two off-diagonal block writers (mapper x mapper, mapper x linear-func) now place their transpose alongside the block, so F leaves assembly symmetric. 2. `_curvature_matrix_mapper_diag` wrapped the three `mapper.unique_mappings` arrays in `np.array(...)`, copying them on every evaluation. They are already contiguous ndarrays of the dtype the numba kernel wants (see `UniqueMappings.__init__`, which casts on construction) and every other caller in this module passes them through untouched. 3. `_curvature_matrix_mapper_func_blocks_from` formed `operated_mapping_matrix / noise_map ** 2` inside the mapper loop and then copied the result again with `np.array`. The weights do not depend on the mapper, so they are formed once per linear func ahead of the loop and passed straight to the kernel. `curvature_matrix_mirrored_from` in `inversion_imaging_numba_util` had no remaining caller (in this repo or downstream) and is deleted rather than left as dead numba code. Also corrects the `curvature_matrix` docstring: it claimed the property is "not a cached property" and is overwritten in memory by the regularization add, but it is decorated `@cached_property` and `curvature_reg_matrix` adds out-of-place via `np.add`. Bit-identical output verified by capturing `inversion.curvature_matrix` before and after on the autolens_profiling breakdown fiducial: `np.array_equal` True for both euclid and hst, with the figures of merit matching the pins exactly (hst 27661.910133664103, euclid 6213.3068738858765). `pytest test_autoarray/inversion`: 392 passed. Step 1 of #505. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01Fr6iJ5T1RDARWfxttCuGkK --- .../inversion_imaging_numba_util.py | 20 ------- .../inversion/imaging_numba/sparse.py | 56 ++++++++++++------- 2 files changed, 35 insertions(+), 41 deletions(-) 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..2603d952c 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, diff --git a/autoarray/inversion/inversion/imaging_numba/sparse.py b/autoarray/inversion/inversion/imaging_numba/sparse.py index 13fb0e053..c7baaa36e 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 @@ -540,8 +543,9 @@ def _curvature_matrix_mapper_func_blocks_from( noise-weighted curvature vector of a linear function (the dense sliding-window correlation in `curvature_matrix_off_diags_via_mapper_and_linear_func_curvature_vector_from`). - Only the `[mapper, linear_func]` blocks are written; their transposes are filled in by the - global mirror applied in `curvature_matrix`. + 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 ---------- @@ -557,6 +561,16 @@ def _curvature_matrix_mapper_func_blocks_from( cls=AbstractLinearObjFuncList ) + # The noise-weighted curvature weights of a linear func do not depend on the mapper, so + # they are formed once per linear func rather than once per (mapper, linear func) pair. + curvature_weights_list = [ + 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] @@ -564,17 +578,12 @@ def _curvature_matrix_mapper_func_blocks_from( 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( 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), + curvature_weights=curvature_weights_list[func_index], mask=self.mask.array, psf_kernel=self.psf.kernel.native.array, ) @@ -584,6 +593,11 @@ def _curvature_matrix_mapper_func_blocks_from( 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( From d43e1acaab6769660dc5255f784f77cc475322fd Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Fri, 28 Aug 2026 16:16:03 -0400 Subject: [PATCH 3/3] perf(numba): FFT the mapper x linear-func block of the curvature matrix (#505) The mapper x linear-func block of `F` was 70-85 % of F on every profiled cell (0.953 s of 1.275 s at HST resolution) and did not scale with the source mesh, only with image pixels. The cost was the dense sliding window inside `curvature_matrix_off_diags_via_mapper_and_linear_func_curvature_vector_from`: for every one of the 60 linear-func columns it expanded the noise-weighted curvature weights onto the native grid and ran an `ny x nx x ky x kx` correlation in numba. That correlation is exactly a convolution with the PSF reversed along both axes, so it now runs through the existing batched FFT convolver: - `Convolver.reversed_kernel` (new, cached) is the same convolver with its kernel reversed, reusing the preloaded `ConvolverState` geometry rebuilt for the reversed kernel. Cached because a dataset's PSF outlives the `Inversion` that is rebuilt for every likelihood evaluation, so the reversed kernel's FFT geometry is built once per fit rather than once per evaluation. - `InversionImagingSparseNumba._blurred_curvature_weights_from` correlates a linear func's curvature weights through it, once per linear func rather than once per (mapper, linear func) pair, with the weights zero outside the mask (no blurring mapping matrix) exactly as the sliding window had them. - `curvature_matrix_off_diags_via_mapper_and_blurred_curvature_weights_from` (new) is the scatter half of the old kernel: the genuinely sparse, irregular accumulation onto source pixels, which stays in numba. The old dense kernel and `convolve_with_kernel_native` are kept, unused by the inversion, as the reference the new path is asserted against in the tests. The numpy convolution path is `scipy.signal.convolve(..., mode="same")`, whose `(k - 1) // 2` "same" offset matches the kernel's `k // 2` centre for both odd and even kernel widths, and which spawns no thread pool (the profiling campaign runs one process per core). Tests: - `test_inversion_imaging_util.py` asserts the FFT block equals the dense kernel (rel 1e-6) on a small masked grid, parametrized over the file's existing asymmetric, non-square `KERNELS_ODD`; a control run without the reversal fails on all four. It also asserts `reversed_kernel` is the reversed kernel and is cached. - `test_curvature_matrix_func_list_blocks.py` gains an inversion-level test of `_curvature_matrix_mapper_func_blocks_from` against the dense kernel with an asymmetric PSF, which also pins that the block's transpose is written (the global mirror having been removed in the previous commit). A control run without the reversal fails. `pytest test_autoarray`: 1296 passed (1290 + 6 new). Seconds per evaluation, `OMP_NUM_THREADS=1 AUTOARRAY_NUMBA_OPERATED_MEMO=0`, n_repeats 10 (autolens_profiling breakdown cells): | cell | eval | F | mapper x l-func | |---------------------|-------------|-------------|-----------------| | hst bilinear | 1.738 -> 0.819 | 1.275 -> 0.387 | 0.953 -> 0.0629 | | euclid bilinear | 0.444 -> 0.273 | 0.284 -> 0.0889 | 0.198 -> 0.0224 | | hst Delaunay-1250 | 1.455 -> 0.754 | 1.055 -> 0.241 | 0.899 -> 0.0612 | | hst rectangular RTU | 8.501 -> 7.058 | 1.250 -> 0.335 | 0.929 -> 0.0630 | i.e. 15x on the block, 3.3x on F and 2.1x on the whole HST evaluation. Pinned log-likelihoods PASSED on all three pinned cells; F itself agrees with the sliding-window result to 3e-18 relative, and the hst figure of merit moves from 27661.910133664103 to 27661.91013366411 (9e-18 relative, well inside the harness rtol of 1e-4). Step 2 of #505. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01Fr6iJ5T1RDARWfxttCuGkK --- .../inversion_imaging_numba_util.py | 73 ++++++++ .../inversion/imaging_numba/sparse.py | 67 +++++-- autoarray/operators/convolver.py | 48 +++++ .../imaging/test_inversion_imaging_util.py | 96 ++++++++++ .../test_curvature_matrix_func_list_blocks.py | 172 +++++++++++++++++- 5 files changed, 439 insertions(+), 17 deletions(-) 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 2603d952c..aef50fb9f 100644 --- a/autoarray/inversion/inversion/imaging_numba/inversion_imaging_numba_util.py +++ b/autoarray/inversion/inversion/imaging_numba/inversion_imaging_numba_util.py @@ -674,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, @@ -694,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 c7baaa36e..f360fe18e 100644 --- a/autoarray/inversion/inversion/imaging_numba/sparse.py +++ b/autoarray/inversion/inversion/imaging_numba/sparse.py @@ -532,6 +532,37 @@ def _curvature_matrix_func_list_and_mapper(self) -> np.ndarray: 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: @@ -539,9 +570,13 @@ def _curvature_matrix_mapper_func_blocks_from( 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-convolved, - noise-weighted curvature vector of a linear function (the dense sliding-window correlation in - `curvature_matrix_off_diags_via_mapper_and_linear_func_curvature_vector_from`). + 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 @@ -556,17 +591,23 @@ def _curvature_matrix_mapper_func_blocks_from( 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 ) - # The noise-weighted curvature weights of a linear func do not depend on the mapper, so - # they are formed once per linear func rather than once per (mapper, linear func) pair. - curvature_weights_list = [ - np.array( - self.linear_func_operated_mapping_matrix_dict[linear_func] - / self.noise_map[:, None] ** 2 + # 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 ] @@ -578,14 +619,14 @@ def _curvature_matrix_mapper_func_blocks_from( for func_index, linear_func in enumerate(linear_func_list): linear_func_param_range = linear_func_param_range_list[func_index] - 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=curvature_weights_list[func_index], - mask=self.mask.array, - psf_kernel=self.psf.kernel.native.array, + blurred_curvature_weights=blurred_curvature_weights_list[ + func_index + ], ) curvature_matrix[ 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 + )