Skip to content

Commit 180c8a4

Browse files
author
James Nightingale
authored
Merge pull request #534 from PyAutoLabs/claude/smoke-test-perf-g8i2uq
perf: vectorise the over-sampled grid builder; skip the mask-edge overlay under fast plots
2 parents bc113fb + abee39a commit 180c8a4

3 files changed

Lines changed: 120 additions & 16 deletions

File tree

‎autoarray/operators/over_sampling/over_sample_util.py‎

Lines changed: 45 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -413,21 +413,51 @@ def grid_2d_slim_over_sampled_via_mask_from(
413413
y_pix = (cy - rows) * sy + oy
414414
x_pix = (cols - cx) * sx + ox
415415

416-
# 5) For each valid pixel, generate its sub-pixel coords
417-
coords_list = []
418-
for i, s in enumerate(sub_arr):
419-
dy = sy / s
420-
dx = sx / s
421-
422-
y_off = np.linspace(+sy / 2 - dy / 2, -sy / 2 + dy / 2, s)
423-
x_off = np.linspace(-sx / 2 + dx / 2, +sx / 2 - dx / 2, s)
424-
425-
y_sub, x_sub = np.meshgrid(y_off, x_off, indexing="ij")
426-
427-
coords = np.stack([y_pix[i] + y_sub.ravel(), x_pix[i] + x_sub.ravel()], axis=1)
428-
coords_list.append(coords)
429-
430-
return np.vstack(coords_list)
416+
# 5) Sub-pixel offsets, one block per pixel in row-major pixel order, each
417+
# block in ``meshgrid(y_off, x_off, indexing="ij")`` (y-major) order.
418+
#
419+
# Vectorised per distinct sub-size rather than looped per pixel: the
420+
# per-pixel loop built a linspace, a meshgrid and a stack for every
421+
# unmasked pixel, and on a 2000x100 CTI frame that was ~180k iterations
422+
# and ~12s per call (four calls per bypassed CTI fit, ~50s of a 90s
423+
# smoke script — PyAutoBrain /ci_speedup, 2026-09-08). The offsets for a
424+
# given sub-size are the same for every pixel, so they are built once
425+
# per distinct sub-size and broadcast onto that sub-size's pixel centres.
426+
# Non-uniform sub-sizes keep the pixel-ordered layout through a block
427+
# start offset per pixel (a cumulative sum of the block sizes).
428+
centres = np.stack([y_pix, x_pix], axis=1)
429+
block_sizes = sub_arr * sub_arr
430+
unique_sizes = np.unique(sub_arr)
431+
432+
if unique_sizes.size == 1:
433+
s = int(unique_sizes[0])
434+
offsets = _sub_pixel_offsets_from(sy=sy, sx=sx, sub_size=s)
435+
return (centres[:, None, :] + offsets[None, :, :]).reshape(-1, 2)
436+
437+
starts = np.concatenate([[0], np.cumsum(block_sizes)[:-1]])
438+
out = np.empty((int(block_sizes.sum()), 2), dtype=float)
439+
for s in unique_sizes:
440+
s = int(s)
441+
sel = np.nonzero(sub_arr == s)[0]
442+
offsets = _sub_pixel_offsets_from(sy=sy, sx=sx, sub_size=s)
443+
rows_out = starts[sel][:, None] + np.arange(s * s)[None, :]
444+
out[rows_out.ravel()] = (centres[sel][:, None, :] + offsets[None, :, :]).reshape(-1, 2)
445+
return out
446+
447+
448+
def _sub_pixel_offsets_from(sy: float, sx: float, sub_size: int) -> np.ndarray:
449+
"""
450+
The (y, x) offsets of the ``sub_size * sub_size`` sub-pixels of one pixel
451+
from that pixel's centre, in the y-major order
452+
``np.meshgrid(y_off, x_off, indexing="ij")`` produces — identical values
453+
and order to the per-pixel construction this replaces.
454+
"""
455+
dy = sy / sub_size
456+
dx = sx / sub_size
457+
y_off = np.linspace(+sy / 2 - dy / 2, -sy / 2 + dy / 2, sub_size)
458+
x_off = np.linspace(-sx / 2 + dx / 2, +sx / 2 - dx / 2, sub_size)
459+
y_sub, x_sub = np.meshgrid(y_off, x_off, indexing="ij")
460+
return np.stack([y_sub.ravel(), x_sub.ravel()], axis=1)
431461

432462

433463
def over_sample_size_via_radial_bins_from(

‎autoarray/plot/array.py‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77

88
import numpy as np
99
from autoarray.plot.utils import (
10+
_FAST_PLOTS,
1011
subplots,
1112
apply_extent,
1213
apply_labels,
@@ -143,7 +144,12 @@ def plot_array(
143144
try:
144145
if extent is None:
145146
extent = array.geometry.extent
146-
if mask is None:
147+
# The mask-edge overlay derives the edge grid of the mask on EVERY call
148+
# (~0.4s on a 2000x100 CTI frame, times ~100 on-the-fly figures per
149+
# bypassed fit). PYAUTO_FAST_PLOTS already drops the figure before it
150+
# is rasterised or saved, so under it the overlay is never seen:
151+
# skip deriving it (PyAutoBrain /ci_speedup, 2026-09-08).
152+
if mask is None and not _FAST_PLOTS:
147153
mask = auto_mask_edge(array)
148154
array = array.native.array
149155
except AttributeError:

‎test_autoarray/operators/over_sample/test_over_sample_util.py‎

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -393,3 +393,71 @@ def test__convolve_bin_segment_ids_from__divisibility_guard():
393393
util.over_sample.convolve_bin_segment_ids_from(
394394
sub_size=np.array([4, 3]), convolve_over_sample_size=2
395395
)
396+
397+
398+
def _reference_grid_2d_slim_over_sampled_via_mask_from(mask_2d, pixel_scales, sub_size, origin=(0.0, 0.0)):
399+
"""The per-pixel construction `grid_2d_slim_over_sampled_via_mask_from`
400+
replaced on 2026-09-08 (one linspace + meshgrid + stack per unmasked
401+
pixel). Kept verbatim as the oracle: the vectorised routine must return
402+
exactly these values in exactly this order."""
403+
H, W = mask_2d.shape
404+
sy, sx = pixel_scales
405+
oy, ox = origin
406+
rows, cols = np.nonzero(~mask_2d)
407+
sub_arr = np.asarray(sub_size)
408+
sub_arr = np.full(rows.size, sub_arr, dtype=int) if sub_arr.size == 1 else sub_arr
409+
valid = sub_arr > 0
410+
rows, cols, sub_arr = rows[valid], cols[valid], sub_arr[valid]
411+
if sub_arr.size == 0:
412+
return np.empty((0, 2), dtype=float)
413+
cy, cx = (H - 1) / 2.0, (W - 1) / 2.0
414+
y_pix = (cy - rows) * sy + oy
415+
x_pix = (cols - cx) * sx + ox
416+
coords_list = []
417+
for i, s in enumerate(sub_arr):
418+
dy, dx = sy / s, sx / s
419+
y_off = np.linspace(+sy / 2 - dy / 2, -sy / 2 + dy / 2, s)
420+
x_off = np.linspace(-sx / 2 + dx / 2, +sx / 2 - dx / 2, s)
421+
y_sub, x_sub = np.meshgrid(y_off, x_off, indexing="ij")
422+
coords_list.append(np.stack([y_pix[i] + y_sub.ravel(), x_pix[i] + x_sub.ravel()], axis=1))
423+
return np.vstack(coords_list)
424+
425+
426+
@pytest.mark.parametrize("seed", [0, 1, 2])
427+
def test__grid_2d_slim_over_sampled_via_mask_from__matches_per_pixel_reference(seed):
428+
rng = np.random.default_rng(seed)
429+
mask_2d = rng.random((7, 9)) < 0.4
430+
pixel_scales = (0.7, 0.3)
431+
origin = (0.25, -0.5)
432+
433+
# uniform sub-size (the single-broadcast branch)
434+
for sub_size in (1, 2, 3):
435+
grid = aa.util.over_sample.grid_2d_slim_over_sampled_via_mask_from(
436+
mask_2d=mask_2d, pixel_scales=pixel_scales, sub_size=sub_size, origin=origin
437+
)
438+
ref = _reference_grid_2d_slim_over_sampled_via_mask_from(mask_2d, pixel_scales, sub_size, origin)
439+
assert grid.shape == ref.shape
440+
assert grid == pytest.approx(ref, abs=0.0)
441+
442+
# per-pixel sub-sizes including zeros (skipped pixels) and mixed sizes (the block-start branch)
443+
n_unmasked = int((~mask_2d).sum())
444+
sub_size = rng.integers(0, 5, size=n_unmasked)
445+
grid = aa.util.over_sample.grid_2d_slim_over_sampled_via_mask_from(
446+
mask_2d=mask_2d, pixel_scales=pixel_scales, sub_size=sub_size, origin=origin
447+
)
448+
ref = _reference_grid_2d_slim_over_sampled_via_mask_from(mask_2d, pixel_scales, sub_size, origin)
449+
assert grid.shape == ref.shape
450+
assert grid == pytest.approx(ref, abs=0.0)
451+
452+
453+
def test__grid_2d_slim_over_sampled_via_mask_from__all_masked_or_all_zero_sub_size():
454+
mask_2d = np.full((3, 3), True)
455+
grid = aa.util.over_sample.grid_2d_slim_over_sampled_via_mask_from(
456+
mask_2d=mask_2d, pixel_scales=(1.0, 1.0), sub_size=2
457+
)
458+
assert grid.shape == (0, 2)
459+
mask_2d = np.full((3, 3), False)
460+
grid = aa.util.over_sample.grid_2d_slim_over_sampled_via_mask_from(
461+
mask_2d=mask_2d, pixel_scales=(1.0, 1.0), sub_size=np.zeros(9, dtype=int)
462+
)
463+
assert grid.shape == (0, 2)

0 commit comments

Comments
 (0)