diff --git a/autoarray/operators/over_sampling/over_sample_util.py b/autoarray/operators/over_sampling/over_sample_util.py index 87d79d84..dab85769 100644 --- a/autoarray/operators/over_sampling/over_sample_util.py +++ b/autoarray/operators/over_sampling/over_sample_util.py @@ -1,4 +1,5 @@ from __future__ import annotations +from functools import lru_cache import numpy as np from typing import TYPE_CHECKING, Union from typing import List, Tuple @@ -154,11 +155,25 @@ def convolve_bin_segment_ids_from( An integer array of shape [total evaluation samples] of segment ids into the uniform fine grid (``total_unmasked_pixels * s**2`` segments). """ - from autoarray import exc - s = int(convolve_over_sample_size) sub_size = np.asarray(sub_size).astype("int") + return _convolve_bin_segment_ids_cached(sub_size.tobytes(), sub_size.shape[0], s) + + +@lru_cache(maxsize=16) +def _convolve_bin_segment_ids_cached( + sub_size_bytes: bytes, n_pixels: int, s: int +) -> np.ndarray: + """ + Memoized body of `convolve_bin_segment_ids_from`. The segment ids are static per + (grid, s) pair but the partial bin runs once per likelihood evaluation in a fit, + so the per-pixel construction loop is cached on the sizes' bytes. + """ + from autoarray import exc + + sub_size = np.frombuffer(sub_size_bytes, dtype="int").reshape(n_pixels) + if np.any(sub_size % s != 0): raise exc.GridException( f"Every over_sample_size entry must be divisible by " @@ -175,6 +190,7 @@ def convolve_bin_segment_ids_from( segment_ids[offset : offset + n * n] = p * s**2 + (rows // k) * s + (cols // k) offset += n * n + segment_ids.setflags(write=False) return segment_ids