diff --git a/autoarray/inversion/inversion/imaging_numba/sparse.py b/autoarray/inversion/inversion/imaging_numba/sparse.py index a6dd293ca..856105533 100644 --- a/autoarray/inversion/inversion/imaging_numba/sparse.py +++ b/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: """ diff --git a/test_autoarray/inversion/inversion/test_sparse_numba_operated_memo.py b/test_autoarray/inversion/inversion/test_sparse_numba_operated_memo.py new file mode 100644 index 000000000..25daa4875 --- /dev/null +++ b/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 + )