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 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Fr6iJ5T1RDARWfxttCuGkK

Author: Jammy2211

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

autoarray/inversion/inversion/imaging_numba/sparse.py

@@ -532,16 +532,51 @@ 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:
         """
         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[

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:

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