diff --git a/autoarray/operators/over_sampling/over_sample_util.py b/autoarray/operators/over_sampling/over_sample_util.py index a9bc30a86..86b591f78 100644 --- a/autoarray/operators/over_sampling/over_sample_util.py +++ b/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( diff --git a/autoarray/plot/array.py b/autoarray/plot/array.py index b15466000..b576128ed 100644 --- a/autoarray/plot/array.py +++ b/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: diff --git a/test_autoarray/operators/over_sample/test_over_sample_util.py b/test_autoarray/operators/over_sample/test_over_sample_util.py index d3160c23a..30a3bb209 100644 --- a/test_autoarray/operators/over_sample/test_over_sample_util.py +++ b/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)