Skip to content

Commit d8bc3ba

Browse files
Jammy2211claude
andcommitted
perf(over-sample): cache the non-uniform binning divisor on the OverSampler (#507)
`OverSampler.binned_array_2d_from` is called once per numpy light-profile evaluation -- 120 times per likelihood evaluation for a 60-component MGE (the data grid and the blurring grid), and once more for every ordinary light profile. On the non-uniform branch it recomputed, every call, two quantities that depend only on the `sub_size` fixed in the constructor: - `np.bincount(self.segment_ids, ...)`, the number of sub-pixels binned into each pixel, plus the zero-guard `counts[counts == 0] = 1`; - `sub_is_uniform`, an `np.isclose` over the whole `sub_size` array, which selects the branch in the first place. Both are now `cached_property`. Measured on the HST breakdown cell's over sampler (17980 sub-pixels into 15361 pixels): the divisor cost 38.8 us per call and the uniformity check 61.3 us, against 81.5 us for the whole cached call -- so a binned evaluation was more than twice the work it needed to be. The cached divisor is marked read-only (`flags.writeable = False`) and built with `np.maximum(counts, 1)` rather than the old in-place `counts[counts == 0] = 1`, because it is now handed to every caller instead of being a fresh array per call. Values are unchanged: `np.maximum` on the `bincount` output is the same int64 array the in-place guard produced, so the binned arrays are bit-identical. `cached_property` stores on the instance `__dict__`, so the cache pickles to Nautilus worker processes with the rest of the object -- asserted by a round-trip test, since `OverSampler` is pickled once per fit. Tests (`test_autoarray/operators/over_sample/test_over_sampler.py`): `..._non_uniform_counts_are_cached_and_not_mutated` (repeated calls bit-identical, cached array unchanged and not writeable), `..._zero_count_segments_do_not_divide_by_zero` (the guard the in-place mutation used to provide), and `..._pickles_with_cached_properties`. `pytest test_autoarray`: 1326 passed (1323 before, +3). Refs #507 Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Fr6iJ5T1RDARWfxttCuGkK
1 parent 88e14bc commit d8bc3ba

2 files changed

Lines changed: 122 additions & 8 deletions

File tree

‎autoarray/operators/over_sampling/over_sampler.py‎

Lines changed: 31 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -167,13 +167,41 @@ def __init__(self, mask: Mask2D, sub_size: Union[int, Array2D]):
167167
):
168168
self.segment_ids[start:end] = seg_id
169169

170-
@property
170+
@cached_property
171171
def sub_is_uniform(self) -> bool:
172172
"""
173173
Returns True if the sub_size is uniform across all pixels in the mask.
174+
175+
The `sub_size` is fixed when the `OverSampler` is created, so this is cached: it is evaluated once per
176+
instance instead of on every call to `binned_array_2d_from` (which is called once per light profile
177+
evaluation, and therefore many hundreds of times per likelihood evaluation).
174178
"""
175179
return np.all(np.isclose(self.sub_size, self.sub_size[0]))
176180

181+
@cached_property
182+
def binned_counts(self) -> np.ndarray:
183+
"""
184+
The number of sub-pixels which are binned up into every pixel of the mask, used as the divisor when
185+
binning an over sampled array via `binned_array_2d_from` for a non-uniform `sub_size`.
186+
187+
Segments which contain no sub-pixels have their count set to 1, so that the division by this array does
188+
not raise a divide-by-zero (the corresponding sum is zero, so the binned value is zero either way).
189+
190+
This is a function of the `segment_ids` alone, which are fixed when the `OverSampler` is created, so it is
191+
computed once per instance instead of on every call to `binned_array_2d_from`.
192+
193+
The returned array is read-only, because it is shared by every call and must never be mutated in place.
194+
"""
195+
counts = np.bincount(self.segment_ids, minlength=self.mask.pixels_in_mask)
196+
197+
# Avoid division by zero (on a copy, because `np.bincount` output is not reused elsewhere but the cached
198+
# array below is handed out to every caller).
199+
200+
counts = np.maximum(counts, 1)
201+
counts.flags.writeable = False
202+
203+
return counts
204+
177205
def tree_flatten(self):
178206
return (self.mask, self.sub_size), ()
179207

@@ -263,13 +291,8 @@ def binned_array_2d_from(self, array: Array2D, xp=np) -> "Array2D":
263291
self.segment_ids, weights=array, minlength=self.mask.pixels_in_mask
264292
)
265293

266-
# Count number of items per segment
267-
counts = np.bincount(
268-
self.segment_ids, minlength=self.mask.pixels_in_mask
269-
)
270-
271-
# Avoid division by zero
272-
counts[counts == 0] = 1
294+
# Count number of items per segment (cached, as it depends only on the segment ids)
295+
counts = self.binned_counts
273296

274297
binned_array_2d = sums / counts
275298

‎test_autoarray/operators/over_sample/test_over_sampler.py‎

Lines changed: 91 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -107,3 +107,94 @@ def test__slim_index_for_sub_slim_index():
107107
)
108108

109109
assert (over_sampling.slim_for_sub_slim == slim_index_for_sub_slim_index_util).all()
110+
111+
112+
def test__over_sampler_binned_array_2d__non_uniform_counts_are_cached_and_not_mutated():
113+
"""
114+
The divisor used to bin a non-uniform over sampled array depends only on the `segment_ids`, so it is cached on
115+
the `OverSampler`. Repeated calls must therefore be bit-identical and must never mutate the cached array.
116+
"""
117+
mask = aa.Mask2D(
118+
mask=[[False, False], [True, True]],
119+
pixel_scales=1.0,
120+
)
121+
122+
over_sampler = aa.OverSampler(
123+
mask=mask, sub_size=aa.Array2D(values=[1, 2], mask=mask)
124+
)
125+
126+
assert not over_sampler.sub_is_uniform
127+
128+
arr = np.array([1.0, 5.0, 7.0, 10.0, 10.0])
129+
130+
binned_0 = over_sampler.binned_array_2d_from(array=arr)
131+
132+
counts = over_sampler.binned_counts
133+
134+
assert counts == pytest.approx(np.array([1, 4]), 1.0e-4)
135+
assert not counts.flags.writeable
136+
137+
binned_1 = over_sampler.binned_array_2d_from(array=arr)
138+
binned_2 = over_sampler.binned_array_2d_from(array=2.0 * arr)
139+
140+
assert np.array_equal(np.array(binned_0.slim), np.array(binned_1.slim))
141+
assert np.array_equal(np.array(binned_2.slim), 2.0 * np.array(binned_0.slim))
142+
143+
# The cached array is handed to every caller, so it must be unchanged after the calls above.
144+
145+
assert over_sampler.binned_counts is counts
146+
assert np.array_equal(counts, np.array([1, 4]))
147+
148+
149+
def test__over_sampler_binned_array_2d__zero_count_segments_do_not_divide_by_zero():
150+
"""
151+
A `sub_size` of zero produces a segment with no sub-pixels, whose count is guarded to 1 so the binned value is
152+
zero rather than a divide-by-zero NaN.
153+
"""
154+
mask = aa.Mask2D(
155+
mask=[[False, False], [True, True]],
156+
pixel_scales=1.0,
157+
)
158+
159+
with np.errstate(divide="ignore"):
160+
over_sampler = aa.OverSampler(
161+
mask=mask, sub_size=aa.Array2D(values=[0, 2], mask=mask)
162+
)
163+
164+
arr = np.array([1.0, 2.0, 3.0, 4.0])
165+
166+
binned = over_sampler.binned_array_2d_from(array=arr)
167+
168+
assert over_sampler.binned_counts == pytest.approx(np.array([1, 4]), 1.0e-4)
169+
assert binned.slim == pytest.approx(np.array([0.0, 2.5]), 1.0e-4)
170+
171+
172+
def test__over_sampler__pickles_with_cached_properties():
173+
"""
174+
`OverSampler` instances are pickled to Nautilus worker processes, so the cached divisor must survive a
175+
round trip.
176+
"""
177+
import pickle
178+
179+
mask = aa.Mask2D(
180+
mask=[[False, False], [True, True]],
181+
pixel_scales=1.0,
182+
)
183+
184+
over_sampler = aa.OverSampler(
185+
mask=mask, sub_size=aa.Array2D(values=[1, 2], mask=mask)
186+
)
187+
188+
arr = np.array([1.0, 5.0, 7.0, 10.0, 10.0])
189+
190+
binned = over_sampler.binned_array_2d_from(array=arr)
191+
192+
over_sampler_pickled = pickle.loads(pickle.dumps(over_sampler))
193+
194+
assert np.array_equal(
195+
over_sampler_pickled.binned_counts, over_sampler.binned_counts
196+
)
197+
assert np.array_equal(
198+
np.array(over_sampler_pickled.binned_array_2d_from(array=arr).slim),
199+
np.array(binned.slim),
200+
)

0 commit comments

Comments
 (0)