Merge pull request #455 from PyAutoLabs/feature/numba-mge-operated-matrix-memo

perf: cross-eval memo for MGE operated mapping matrices (numba CPU inversion)

Author: Jammy2211

autoarray/inversion/inversion/imaging_numba/sparse.py

@@ -1,3 +1,7 @@
+import hashlib
+import os
+import pickle
+
 import numpy as np
 from typing import Dict, List, Optional, Union
 
@@ -14,6 +18,39 @@
 
 from autoarray.inversion.inversion.imaging_numba import inversion_imaging_numba_util
 
+# Cross-evaluation memo for linear-func operated mapping matrices (the MGE
+# lens-light PSF-convolved images). A sampler builds fresh linear-func objects
+# for every likelihood evaluation, so when the light profiles are FIXED in the
+# model the identical ~60-Gaussian convolution stack is recomputed each call —
+# ~0.5 s of a ~2.4 s euclid-resolution numba CPU evaluation (autolens_profiling
+# issue #151). Entries are keyed by a fingerprint of the linear-func object's
+# full pickled state (profiles + grids + PSF), so the memo engages only when
+# the profile parameters are genuinely unchanged; any varying parameter changes
+# the key and the matrix is recomputed exactly as before. Failure modes are
+# misses, never stale hits. Scoped to this numba inversion module on purpose —
+# no other path is touched. Disable with AUTOARRAY_NUMBA_OPERATED_MEMO=0.
+_operated_mapping_matrix_memo: Dict[str, np.ndarray] = {}
+
+_OPERATED_MAPPING_MATRIX_MEMO_MAX_ENTRIES = 8
+
+
+def _operated_mapping_matrix_memo_key(linear_func) -> Optional[str]:
+    """
+    Fingerprint a linear-func object's state for the cross-evaluation memo,
+    or None if it cannot be fingerprinted (unpicklable), in which case the
+    caller falls back to the uncached computation.
+
+    Must be called before the object's own cached properties are populated:
+    a fingerprint taken afterwards would include the cached arrays and never
+    match the pre-computation fingerprint of the next evaluation (a
+    miss-every-time cache, still never a stale one).
+    """
+    try:
+        state = pickle.dumps(linear_func, protocol=pickle.HIGHEST_PROTOCOL)
+    except Exception:
+        return None
+    return hashlib.sha256(state).hexdigest()
+
 
 class InversionImagingSparseNumba(AbstractInversionImaging):
     def __init__(
@@ -63,6 +100,76 @@ def psf_weighted_data(self):
             native_index_for_slim_index=self.data.mask.derive_indexes.native_for_slim,
         )
 
+    @cached_property
+    def linear_func_operated_mapping_matrix_dict(self) -> Dict:
+        """
+        The parent property, wrapped in two caches specific to this numba CPU
+        inversion:
+
+        1. `cached_property` — the dict is built once per inversion instead of
+           on every access (this inversion reads it from several matrices).
+        2. A module-level cross-evaluation memo — when a linear func's full
+           state (light profiles + grids + PSF) fingerprints identically to a
+           previous evaluation's, its operated mapping matrix (the PSF-convolved
+           MGE image stack, ~0.5 s/eval at euclid resolution) is reused instead
+           of recomputed. Fixed-profile models hit every evaluation; models with
+           free profile parameters change the fingerprint and recompute exactly
+           as before. See the memo's module docstring for the safety argument.
+
+        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
+
+        if os.environ.get("AUTOARRAY_NUMBA_OPERATED_MEMO", "1") == "0":
+            return parent_fget(self)
+
+        linear_func_list = self.cls_list_from(cls=AbstractLinearObjFuncList)
+
+        key_list = [
+            _operated_mapping_matrix_memo_key(linear_func)
+            for linear_func in linear_func_list
+        ]
+
+        if any(key is None for key in key_list):
+            return parent_fget(self)
+
+        operated_mapping_matrix_dict = {}
+
+        for linear_func, key in zip(linear_func_list, key_list):
+            operated_mapping_matrix = _operated_mapping_matrix_memo.get(key)
+
+            if operated_mapping_matrix is None:
+                operated_mapping_matrix = linear_func.operated_mapping_matrix_override
+                if operated_mapping_matrix is None:
+                    operated_mapping_matrix = self.psf.convolved_mapping_matrix_from(
+                        mapping_matrix=self._mapping_matrix_for_convolution_from(
+                            linear_func
+                        ),
+                        mask=self.mask,
+                        xp=self._xp,
+                    )
+
+                # A copy, not a view: the memo owns its buffer outright, so
+                # marking it read-only cannot leak onto the linear func's own
+                # cached override array.
+                operated_mapping_matrix = np.array(operated_mapping_matrix)
+                operated_mapping_matrix.setflags(write=False)
+
+                while (
+                    len(_operated_mapping_matrix_memo)
+                    >= _OPERATED_MAPPING_MATRIX_MEMO_MAX_ENTRIES
+                ):
+                    _operated_mapping_matrix_memo.pop(
+                        next(iter(_operated_mapping_matrix_memo))
+                    )
+
+                _operated_mapping_matrix_memo[key] = operated_mapping_matrix
+
+            operated_mapping_matrix_dict[linear_func] = operated_mapping_matrix
+
+        return operated_mapping_matrix_dict
+
     @property
     def _data_vector_mapper(self) -> np.ndarray:
         """

test_autoarray/inversion/inversion/test_sparse_numba_operated_memo.py

@@ -0,0 +1,138 @@
+"""
+The cross-evaluation memo for linear-func operated mapping matrices in the
+numba CPU sparse inversion (`imaging_numba/sparse.py`).
+
+The memo must: reuse the matrix when a fresh linear-func object fingerprints
+identically to a previous evaluation's (the fixed-MGE campaign case); recompute
+when any state differs (free profile parameters); fall back to the uncached
+parent computation when an object cannot be fingerprinted or the memo is
+disabled; and never hand out writeable buffers.
+"""
+
+import numpy as np
+import pytest
+
+from autoarray.inversion.inversion.imaging_numba import sparse as sparse_module
+from autoarray.inversion.inversion.imaging_numba.sparse import (
+    InversionImagingSparseNumba,
+    _operated_mapping_matrix_memo,
+    _operated_mapping_matrix_memo_key,
+)
+
+
+class FakeLinearFunc:
+    """Stands in for an MGE linear-func bundle: `values` plays the role of the
+    profile parameters, and computing the override is counted class-wide so
+    tests can assert whether the convolution work actually ran."""
+
+    compute_count = 0
+
+    def __init__(self, values):
+        self.values = np.array(values, dtype=float)
+
+    @property
+    def operated_mapping_matrix_override(self):
+        type(self).compute_count += 1
+        return np.outer(self.values, np.arange(1.0, 4.0))
+
+
+class UnpicklableLinearFunc(FakeLinearFunc):
+    def __init__(self, values):
+        super().__init__(values)
+        self.blocker = lambda: None  # lambdas cannot be pickled
+
+
+class StubInversion(InversionImagingSparseNumba):
+    """Bypasses the real constructor; the property under test only needs
+    `cls_list_from` (and instance-dict storage for its cached_property)."""
+
+    def __init__(self, linear_func_list):
+        self._stub_linear_func_list = list(linear_func_list)
+
+    def cls_list_from(self, cls):
+        return self._stub_linear_func_list
+
+
+@pytest.fixture(autouse=True)
+def _clean_memo():
+    _operated_mapping_matrix_memo.clear()
+    FakeLinearFunc.compute_count = 0
+    yield
+    _operated_mapping_matrix_memo.clear()
+
+
+def test__memo_key__stable_for_equal_state__distinct_for_different_state():
+    key_a = _operated_mapping_matrix_memo_key(FakeLinearFunc([1.0, 2.0]))
+    key_b = _operated_mapping_matrix_memo_key(FakeLinearFunc([1.0, 2.0]))
+    key_c = _operated_mapping_matrix_memo_key(FakeLinearFunc([1.0, 2.5]))
+
+    assert key_a == key_b
+    assert key_a != key_c
+
+
+def test__memo_key__unpicklable_state_returns_none():
+    assert _operated_mapping_matrix_memo_key(UnpicklableLinearFunc([1.0])) is None
+
+
+def test__identical_state_across_fresh_objects__computes_once():
+    func_eval_0 = FakeLinearFunc([1.0, 2.0])
+    dict_0 = StubInversion([func_eval_0]).linear_func_operated_mapping_matrix_dict
+
+    # A sampler's next evaluation builds a FRESH object with identical state.
+    func_eval_1 = FakeLinearFunc([1.0, 2.0])
+    dict_1 = StubInversion([func_eval_1]).linear_func_operated_mapping_matrix_dict
+
+    assert FakeLinearFunc.compute_count == 1
+    assert np.array_equal(dict_0[func_eval_0], dict_1[func_eval_1])
+    assert not dict_1[func_eval_1].flags.writeable
+
+
+def test__changed_state__recomputes_and_matches_uncached_result():
+    StubInversion([FakeLinearFunc([1.0, 2.0])]).linear_func_operated_mapping_matrix_dict
+
+    func_changed = FakeLinearFunc([1.0, 3.0])
+    result = StubInversion([func_changed]).linear_func_operated_mapping_matrix_dict[
+        func_changed
+    ]
+
+    assert FakeLinearFunc.compute_count == 2
+    assert np.array_equal(result, np.outer([1.0, 3.0], np.arange(1.0, 4.0)))
+
+
+def test__cached_property__single_dict_build_per_inversion():
+    inversion = StubInversion([FakeLinearFunc([1.0, 2.0])])
+
+    dict_first = inversion.linear_func_operated_mapping_matrix_dict
+    dict_second = inversion.linear_func_operated_mapping_matrix_dict
+
+    assert dict_first is dict_second
+
+
+def test__unpicklable_func__falls_back_to_uncached_parent_and_stores_nothing():
+    func = UnpicklableLinearFunc([1.0, 2.0])
+
+    result = StubInversion([func]).linear_func_operated_mapping_matrix_dict[func]
+
+    assert np.array_equal(result, np.outer([1.0, 2.0], np.arange(1.0, 4.0)))
+    assert len(_operated_mapping_matrix_memo) == 0
+
+
+def test__env_var_disables_memo(monkeypatch):
+    monkeypatch.setenv("AUTOARRAY_NUMBA_OPERATED_MEMO", "0")
+
+    func = FakeLinearFunc([1.0, 2.0])
+    result = StubInversion([func]).linear_func_operated_mapping_matrix_dict[func]
+
+    assert np.array_equal(result, np.outer([1.0, 2.0], np.arange(1.0, 4.0)))
+    assert len(_operated_mapping_matrix_memo) == 0
+
+
+def test__memo_eviction__bounded_size():
+    for value in range(sparse_module._OPERATED_MAPPING_MATRIX_MEMO_MAX_ENTRIES + 3):
+        func = FakeLinearFunc([float(value)])
+        StubInversion([func]).linear_func_operated_mapping_matrix_dict
+
+    assert (
+        len(_operated_mapping_matrix_memo)
+        == sparse_module._OPERATED_MAPPING_MATRIX_MEMO_MAX_ENTRIES
+    )