perf: vectorise the over-sampled grid builder; skip the mask-edge overlay under fast plots

autoarray/operators/over_sampling/over_sample_util.py

@@ -413,21 +413,51 @@ def grid_2d_slim_over_sampled_via_mask_from(
     y_pix = (cy - rows) * sy + oy
     x_pix = (cols - cx) * sx + ox
 
-    # 5) For each valid pixel, generate its sub-pixel coords
-    coords_list = []
-    for i, s in enumerate(sub_arr):
-        dy = sy / s
-        dx = sx / s
-
-        y_off = np.linspace(+sy / 2 - dy / 2, -sy / 2 + dy / 2, s)
-        x_off = np.linspace(-sx / 2 + dx / 2, +sx / 2 - dx / 2, s)
-
-        y_sub, x_sub = np.meshgrid(y_off, x_off, indexing="ij")
-
-        coords = np.stack([y_pix[i] + y_sub.ravel(), x_pix[i] + x_sub.ravel()], axis=1)
-        coords_list.append(coords)
-
-    return np.vstack(coords_list)
+    # 5) Sub-pixel offsets, one block per pixel in row-major pixel order, each
+    #    block in ``meshgrid(y_off, x_off, indexing="ij")`` (y-major) order.
+    #
+    #    Vectorised per distinct sub-size rather than looped per pixel: the
+    #    per-pixel loop built a linspace, a meshgrid and a stack for every
+    #    unmasked pixel, and on a 2000x100 CTI frame that was ~180k iterations
+    #    and ~12s per call (four calls per bypassed CTI fit, ~50s of a 90s
+    #    smoke script — PyAutoBrain /ci_speedup, 2026-09-08). The offsets for a
+    #    given sub-size are the same for every pixel, so they are built once
+    #    per distinct sub-size and broadcast onto that sub-size's pixel centres.
+    #    Non-uniform sub-sizes keep the pixel-ordered layout through a block
+    #    start offset per pixel (a cumulative sum of the block sizes).
+    centres = np.stack([y_pix, x_pix], axis=1)
+    block_sizes = sub_arr * sub_arr
+    unique_sizes = np.unique(sub_arr)
+
+    if unique_sizes.size == 1:
+        s = int(unique_sizes[0])
+        offsets = _sub_pixel_offsets_from(sy=sy, sx=sx, sub_size=s)
+        return (centres[:, None, :] + offsets[None, :, :]).reshape(-1, 2)
+
+    starts = np.concatenate([[0], np.cumsum(block_sizes)[:-1]])
+    out = np.empty((int(block_sizes.sum()), 2), dtype=float)
+    for s in unique_sizes:
+        s = int(s)
+        sel = np.nonzero(sub_arr == s)[0]
+        offsets = _sub_pixel_offsets_from(sy=sy, sx=sx, sub_size=s)
+        rows_out = starts[sel][:, None] + np.arange(s * s)[None, :]
+        out[rows_out.ravel()] = (centres[sel][:, None, :] + offsets[None, :, :]).reshape(-1, 2)
+    return out
+
+
+def _sub_pixel_offsets_from(sy: float, sx: float, sub_size: int) -> np.ndarray:
+    """
+    The (y, x) offsets of the ``sub_size * sub_size`` sub-pixels of one pixel
+    from that pixel's centre, in the y-major order
+    ``np.meshgrid(y_off, x_off, indexing="ij")`` produces — identical values
+    and order to the per-pixel construction this replaces.
+    """
+    dy = sy / sub_size
+    dx = sx / sub_size
+    y_off = np.linspace(+sy / 2 - dy / 2, -sy / 2 + dy / 2, sub_size)
+    x_off = np.linspace(-sx / 2 + dx / 2, +sx / 2 - dx / 2, sub_size)
+    y_sub, x_sub = np.meshgrid(y_off, x_off, indexing="ij")
+    return np.stack([y_sub.ravel(), x_sub.ravel()], axis=1)
 
 
 def over_sample_size_via_radial_bins_from(

autoarray/plot/array.py

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

test_autoarray/operators/over_sample/test_over_sample_util.py

@@ -393,3 +393,71 @@ def test__convolve_bin_segment_ids_from__divisibility_guard():
         util.over_sample.convolve_bin_segment_ids_from(
             sub_size=np.array([4, 3]), convolve_over_sample_size=2
         )
+
+
+def _reference_grid_2d_slim_over_sampled_via_mask_from(mask_2d, pixel_scales, sub_size, origin=(0.0, 0.0)):
+    """The per-pixel construction `grid_2d_slim_over_sampled_via_mask_from`
+    replaced on 2026-09-08 (one linspace + meshgrid + stack per unmasked
+    pixel). Kept verbatim as the oracle: the vectorised routine must return
+    exactly these values in exactly this order."""
+    H, W = mask_2d.shape
+    sy, sx = pixel_scales
+    oy, ox = origin
+    rows, cols = np.nonzero(~mask_2d)
+    sub_arr = np.asarray(sub_size)
+    sub_arr = np.full(rows.size, sub_arr, dtype=int) if sub_arr.size == 1 else sub_arr
+    valid = sub_arr > 0
+    rows, cols, sub_arr = rows[valid], cols[valid], sub_arr[valid]
+    if sub_arr.size == 0:
+        return np.empty((0, 2), dtype=float)
+    cy, cx = (H - 1) / 2.0, (W - 1) / 2.0
+    y_pix = (cy - rows) * sy + oy
+    x_pix = (cols - cx) * sx + ox
+    coords_list = []
+    for i, s in enumerate(sub_arr):
+        dy, dx = sy / s, sx / s
+        y_off = np.linspace(+sy / 2 - dy / 2, -sy / 2 + dy / 2, s)
+        x_off = np.linspace(-sx / 2 + dx / 2, +sx / 2 - dx / 2, s)
+        y_sub, x_sub = np.meshgrid(y_off, x_off, indexing="ij")
+        coords_list.append(np.stack([y_pix[i] + y_sub.ravel(), x_pix[i] + x_sub.ravel()], axis=1))
+    return np.vstack(coords_list)
+
+
+@pytest.mark.parametrize("seed", [0, 1, 2])
+def test__grid_2d_slim_over_sampled_via_mask_from__matches_per_pixel_reference(seed):
+    rng = np.random.default_rng(seed)
+    mask_2d = rng.random((7, 9)) < 0.4
+    pixel_scales = (0.7, 0.3)
+    origin = (0.25, -0.5)
+
+    # uniform sub-size (the single-broadcast branch)
+    for sub_size in (1, 2, 3):
+        grid = aa.util.over_sample.grid_2d_slim_over_sampled_via_mask_from(
+            mask_2d=mask_2d, pixel_scales=pixel_scales, sub_size=sub_size, origin=origin
+        )
+        ref = _reference_grid_2d_slim_over_sampled_via_mask_from(mask_2d, pixel_scales, sub_size, origin)
+        assert grid.shape == ref.shape
+        assert grid == pytest.approx(ref, abs=0.0)
+
+    # per-pixel sub-sizes including zeros (skipped pixels) and mixed sizes (the block-start branch)
+    n_unmasked = int((~mask_2d).sum())
+    sub_size = rng.integers(0, 5, size=n_unmasked)
+    grid = aa.util.over_sample.grid_2d_slim_over_sampled_via_mask_from(
+        mask_2d=mask_2d, pixel_scales=pixel_scales, sub_size=sub_size, origin=origin
+    )
+    ref = _reference_grid_2d_slim_over_sampled_via_mask_from(mask_2d, pixel_scales, sub_size, origin)
+    assert grid.shape == ref.shape
+    assert grid == pytest.approx(ref, abs=0.0)
+
+
+def test__grid_2d_slim_over_sampled_via_mask_from__all_masked_or_all_zero_sub_size():
+    mask_2d = np.full((3, 3), True)
+    grid = aa.util.over_sample.grid_2d_slim_over_sampled_via_mask_from(
+        mask_2d=mask_2d, pixel_scales=(1.0, 1.0), sub_size=2
+    )
+    assert grid.shape == (0, 2)
+    mask_2d = np.full((3, 3), False)
+    grid = aa.util.over_sample.grid_2d_slim_over_sampled_via_mask_from(
+        mask_2d=mask_2d, pixel_scales=(1.0, 1.0), sub_size=np.zeros(9, dtype=int)
+    )
+    assert grid.shape == (0, 2)