refactor: vectorize k x s segment-id construction (#362 refactor exercise) (#365)

The per-pixel python loop becomes pure numpy (repeat/cumsum/divmod on the
flat sample index): cold construction 62 ms -> 9.7 ms on a 3000-pixel
adaptive mask (cached repeat calls unchanged at ~80 us). Exact same ids as
the loop on adaptive references; suite unchanged.

Behaviour-preserving; closes the library leg of the k x s refactor
exercise. The operate/image call-site pattern was reviewed and left as-is
(clean two-step seam); no PyAutoGalaxy edits needed.

Co-authored-by: Jammy2211 <JNightingale2211@gmail.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

Author: 3 people

autoarray/operators/over_sampling/over_sample_util.py

@@ -181,14 +181,22 @@ def _convolve_bin_segment_ids_cached(
             f"sizes {np.unique(sub_size[sub_size % s != 0])} are not."
         )
 
-    segment_ids = np.empty(int(np.sum(sub_size**2)), dtype="int64")
-
-    offset = 0
-    for p, n in enumerate(sub_size):
-        k = n // s
-        rows, cols = np.divmod(np.arange(n * n), n)
-        segment_ids[offset : offset + n * n] = p * s**2 + (rows // k) * s + (cols // k)
-        offset += n * n
+    # Vectorized construction: for every sample, its pixel index p, position
+    # within the pixel's n_p x n_p block, and fine cell (row // k_p) * s + col // k_p.
+    counts = sub_size**2
+    n_samples = int(np.sum(counts))
+
+    pixel_of_sample = np.repeat(np.arange(sub_size.shape[0]), counts)
+    block_starts = np.concatenate(([0], np.cumsum(counts)[:-1]))
+    within = np.arange(n_samples) - block_starts[pixel_of_sample]
+
+    n_of_sample = sub_size[pixel_of_sample]
+    k_of_sample = n_of_sample // s
+    rows, cols = np.divmod(within, n_of_sample)
+
+    segment_ids = (
+        pixel_of_sample * s**2 + (rows // k_of_sample) * s + (cols // k_of_sample)
+    )
 
     segment_ids.setflags(write=False)
     return segment_ids