From cdd44f799b3d4d90d360bd2971f154acd67dee4c Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Wed, 8 Jul 2026 21:19:50 +0100 Subject: [PATCH] perf: memoize k x s segment-id construction (#362 follow-up) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The segment ids are static per (grid, s) pair but the partial pre-bin runs once per likelihood evaluation in a fit, and the per-pixel construction loop cost ~60 ms on a realistic 3000-pixel adaptive mask — per evaluation. An lru_cache on the sizes' bytes makes repeat calls ~80 us (x750). The returned array is marked read-only (shared across callers); the divisibility guard fires identically through the cache. Co-Authored-By: Claude Fable 5 --- .../over_sampling/over_sample_util.py | 20 +++++++++++++++++-- 1 file changed, 18 insertions(+), 2 deletions(-) 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