diff --git a/src/cp_measure/core/numba/measurecolocalization.py b/src/cp_measure/core/numba/measurecolocalization.py index 68126ea..17cdeea 100644 --- a/src/cp_measure/core/numba/measurecolocalization.py +++ b/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) diff --git a/src/cp_measure/primitives/_segment_numba.py b/src/cp_measure/primitives/_segment_numba.py index 940acc3..3c884cd 100644 --- a/src/cp_measure/primitives/_segment_numba.py +++ b/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 diff --git a/src/cp_measure/primitives/segment.py b/src/cp_measure/primitives/segment.py index e5b5894..bd2cc0a 100644 --- a/src/cp_measure/primitives/segment.py +++ b/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 diff --git a/test/test_coloc_kernels.py b/test/test_coloc_kernels.py index 675bcec..0239979 100644 --- a/test/test_coloc_kernels.py +++ b/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):