perf: reuse Convolver state, cache operated-matrix dict, hoist linear-func pair loop (#496 phase 1)

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

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,13 +478,25 @@ 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[
                         linear_func_param_range_0[0] : linear_func_param_range_0[1],
                         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(

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,33 +544,38 @@ 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[
                     linear_func_param_range_0[0] : linear_func_param_range_0[1],
                     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(

autoarray/operators/convolver.py

@@ -142,16 +142,15 @@ 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:
             self.blurring_mask = self.mask.derive_mask.blurring_from(
                 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:
         """

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)