perf(numba): one-bincount label prep for colocalization

src/cp_measure/core/numba/measurecolocalization.py

@@ -39,7 +39,7 @@
 )
 from cp_measure.core.numba._colocalization import coloc_per_object
 from cp_measure.primitives._segment_numba import flatten_pairs_grouped
-from cp_measure.primitives.segment import label_to_idx_lut
+from cp_measure.primitives.segment import labels_to_offsets
 from cp_measure.primitives.shapes import to_bzyx
 
 
@@ -52,15 +52,15 @@ def _run(masks_zyx, pixels_1_zyx, pixels_2_zyx, thr_frac, compute_rwc):
     masks = numpy.ascontiguousarray(masks_zyx)
     if not numpy.issubdtype(masks.dtype, numpy.integer):
         masks = masks.astype(numpy.intp)
-    lut, n = label_to_idx_lut(masks)
+    lut, n, offsets = labels_to_offsets(masks)
     if n == 0:
         return None
-    g1, g2, offsets = flatten_pairs_grouped(
+    g1, g2 = flatten_pairs_grouped(
         masks,
         numpy.ascontiguousarray(pixels_1_zyx, dtype=numpy.float64),
         numpy.ascontiguousarray(pixels_2_zyx, dtype=numpy.float64),
         lut,
-        n,
+        offsets,
     )
     return coloc_per_object(g1, g2, offsets, n, thr_frac, compute_rwc)
 

src/cp_measure/primitives/_segment_numba.py

@@ -229,33 +229,25 @@ def segment_quantiles(values, seg0, counts, n, mad_frac):
 
 
 @njit(cache=True)
-def flatten_pairs_grouped(masks, pix1, pix2, lut, n):
+def flatten_pairs_grouped(masks, pix1, pix2, lut, offsets):
     """Flatten two co-registered ``(Z, Y, X)`` channels to per-object blocks.
 
-    Returns ``(g1, g2, offsets)`` where object ``k`` owns the contiguous slices
+    Returns ``(g1, g2)`` where object ``k`` owns the contiguous slices
     ``g1[offsets[k] : offsets[k + 1]]`` and ``g2[...]``. Both channels are
     gathered at every ``masks > 0`` pixel in one raster scan, so ``g1`` and
     ``g2`` stay aligned. Unlike :func:`flatten_numba`, non-finite pixels are
     KEPT — this mirrors the colocalization reference's ``pixels[mask]``
-    extraction, which applies no finiteness filter. Two scans (count, then
-    scatter) give an O(M) counting-sort grouping; ``masks`` must be C-contiguous
-    integer, ``lut`` the ``label -> 0..n-1`` map from ``label_to_idx_lut``.
+    extraction, which applies no finiteness filter.
+
+    A SINGLE scatter scan: ``offsets`` (the CSR block bounds, length ``n + 1``)
+    is precomputed by :func:`cp_measure.primitives.segment.labels_to_offsets`
+    from a ``bincount``, so the count scan this kernel used to do is gone.
+    ``masks`` must be C-contiguous integer; ``lut`` the ``label -> 0..n-1`` map.
     """
     Z, Y, X = masks.shape
-    counts = np.zeros(n, np.int64)
-    for z in range(Z):
-        for y in range(Y):
-            for x in range(X):
-                L = masks[z, y, x]
-                if L > 0:
-                    counts[lut[L]] += 1
-    offsets = np.zeros(n + 1, np.int64)
-    for k in range(n):
-        offsets[k + 1] = offsets[k] + counts[k]
-    M = offsets[n]
-    g1 = np.empty(M, np.float64)
-    g2 = np.empty(M, np.float64)
-    fill = offsets[:n].copy()
+    g1 = np.empty(offsets[-1], np.float64)
+    g2 = np.empty(offsets[-1], np.float64)
+    fill = offsets[:-1].copy()
     for z in range(Z):
         for y in range(Y):
             for x in range(X):
@@ -266,4 +258,4 @@ def flatten_pairs_grouped(masks, pix1, pix2, lut, n):
                     g1[p] = pix1[z, y, x]
                     g2[p] = pix2[z, y, x]
                     fill[k] = p + 1
-    return g1, g2, offsets
+    return g1, g2

src/cp_measure/primitives/segment.py

@@ -37,3 +37,30 @@ def label_to_idx_lut(
     lut = numpy.full(max_label + 1, -1, dtype=numpy.int64)
     lut[labels] = numpy.arange(n, dtype=numpy.int64)
     return lut, n
+
+
+def labels_to_offsets(
+    masks: NDArray[numpy.integer],
+) -> tuple[NDArray[numpy.int64], int, NDArray[numpy.int64]]:
+    """Build ``(lut, n, offsets)`` from a single ``np.bincount`` pass.
+
+    Like :func:`label_to_idx_lut` but also returns the CSR ``offsets`` (length
+    ``n + 1``) giving each object's pixel block in the grouped flat order, so the
+    grouped flatten can scatter in a SINGLE scan instead of count-then-scatter.
+
+    One ``bincount`` over the raster (~2x faster here than ``find_objects`` plus a
+    separate count scan) gives per-label pixel counts directly; ``cumsum`` of the
+    present-label counts is ``offsets``. ``lut`` and segment ordering (ascending
+    present labels) are identical to :func:`label_to_idx_lut`.
+    """
+    if not numpy.issubdtype(masks.dtype, numpy.integer):
+        masks = masks.astype(numpy.intp, copy=False)
+    counts = numpy.bincount(masks.ravel())
+    present = numpy.nonzero(counts[1:])[0] + 1  # positive labels actually present
+    n = int(present.size)
+    max_label = int(present[-1]) if n else 0
+    lut = numpy.full(max_label + 1, -1, dtype=numpy.int64)
+    lut[present] = numpy.arange(n, dtype=numpy.int64)
+    offsets = numpy.zeros(n + 1, dtype=numpy.int64)
+    offsets[1:] = numpy.cumsum(counts[present])
+    return lut, n, offsets

test/test_coloc_kernels.py

@@ -12,14 +12,14 @@
 @requires_numba
 def test_flatten_pairs_grouped_blocks():
     from cp_measure.primitives._segment_numba import flatten_pairs_grouped
-    from cp_measure.primitives.segment import label_to_idx_lut
+    from cp_measure.primitives.segment import labels_to_offsets
 
     masks = np.array([[0, 1, 1], [2, 2, 0], [2, 0, 1]], np.int64)
     p1 = np.arange(9, dtype=np.float64).reshape(3, 3)
     p2 = (np.arange(9, dtype=np.float64) * 10).reshape(3, 3)
-    lut, n = label_to_idx_lut(masks[np.newaxis])
-    g1, g2, offsets = flatten_pairs_grouped(
-        masks[np.newaxis], p1[np.newaxis], p2[np.newaxis], lut, n
+    lut, n, offsets = labels_to_offsets(masks[np.newaxis])
+    g1, g2 = flatten_pairs_grouped(
+        masks[np.newaxis], p1[np.newaxis], p2[np.newaxis], lut, offsets
     )
 
     assert n == 2
@@ -35,19 +35,33 @@ def test_flatten_pairs_grouped_blocks():
 def test_flatten_pairs_grouped_keeps_nonfinite():
     """Reference extracts pixels[mask] with no finiteness filter — match it."""
     from cp_measure.primitives._segment_numba import flatten_pairs_grouped
-    from cp_measure.primitives.segment import label_to_idx_lut
+    from cp_measure.primitives.segment import labels_to_offsets
 
     masks = np.array([[1, 1, 1]], np.int64)
     p1 = np.array([[1.0, np.nan, 3.0]])
     p2 = np.array([[1.0, 2.0, np.inf]])
-    lut, n = label_to_idx_lut(masks[np.newaxis])
-    g1, g2, offsets = flatten_pairs_grouped(
-        masks[np.newaxis], p1[np.newaxis], p2[np.newaxis], lut, n
+    lut, n, offsets = labels_to_offsets(masks[np.newaxis])
+    g1, g2 = flatten_pairs_grouped(
+        masks[np.newaxis], p1[np.newaxis], p2[np.newaxis], lut, offsets
     )
     assert offsets[-1] == 3  # all three pixels kept
     assert np.isnan(g1).sum() == 1 and np.isinf(g2).sum() == 1
 
 
+@pytest.mark.parametrize("labels", [[0, 1, 2, 3], [0, 2, 5, 5, 2], [0, 0, 0]])
+def test_labels_to_offsets_agrees_with_lut(labels):
+    """bincount-based (lut, n) == find_objects-based; offsets are the CSR counts."""
+    from cp_measure.primitives.segment import label_to_idx_lut, labels_to_offsets
+
+    masks = np.array(labels, np.int64).reshape(1, 1, -1)
+    lut_ref, n_ref = label_to_idx_lut(masks)
+    lut, n, offsets = labels_to_offsets(masks)
+    assert n == n_ref
+    np.testing.assert_array_equal(lut, lut_ref)
+    assert offsets[0] == 0 and offsets[-1] == int((masks > 0).sum())
+    np.testing.assert_array_equal(np.diff(offsets), np.bincount(lut[masks[masks > 0]]))
+
+
 @requires_numba
 @pytest.mark.parametrize("seed", [0, 1, 2])
 def test_dense_rank_matches_scipy(seed):