From b697f940f7bd0e41d7c3cb72f52c347f8008aefb Mon Sep 17 00:00:00 2001 From: Tim Treis Date: Thu, 4 Jun 2026 00:02:38 +0200 Subject: [PATCH 1/4] feat(accelerator): numba costes colocalization (all 3 modes) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per-object port of get_correlation_costes + bisection_costes/linear_costes into core/numba/_costes.py::costes_per_object, reusing the #60 grouped layout (labels_to_offsets + flatten_pairs_grouped) — no sort. All three fast_costes modes: M_FASTER (bisection), M_FAST/M_ACCURATE (linear), control flow bit-reproduced (window math, num_true recompute cache, > vs >= threshold asymmetry). The reference's dead calculate_threshold call is skipped; thr is accepted for parity but unused. Pearson-on-subset matches scipy.stats.pearsonr's order (centre, normalise each vector, accumulate, clamp). error_model="numpy" so a constant subset yields NaN (not ZeroDivisionError), matching scipy's ConstantInputWarning -> nan. Exact vs numpy on float pixels (scale=1); integer-dtype diverges by design (the reference overflows z = fi + si in uint8/uint16). bzyx via to_bzyx-twice like the other coloc features. Speedup 41.8x (1080^2, 144 obj, float). Tests: kernel control-flow vs the real reference search at scale=255 (exercises the multi-iteration path), pearson vs scipy, regression vs reference, end-to-end golden 2D/3D/batch x 3 modes. Full suite 145 passed, lint clean. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/cp_measure/bulk.py | 6 +- src/cp_measure/core/numba/__init__.py | 8 +- src/cp_measure/core/numba/_costes.py | 240 ++++++++++++++++++ .../core/numba/measurecolocalization.py | 48 +++- test/test_backend_correctness.py | 4 +- test/test_coloc_backend.py | 12 +- test/test_costes_kernels.py | 88 +++++++ 7 files changed, 396 insertions(+), 10 deletions(-) create mode 100644 src/cp_measure/core/numba/_costes.py create mode 100644 test/test_costes_kernels.py diff --git a/src/cp_measure/bulk.py b/src/cp_measure/bulk.py index 60450a8..2bf08bb 100644 --- a/src/cp_measure/bulk.py +++ b/src/cp_measure/bulk.py @@ -49,8 +49,8 @@ def _numba_registries() -> dict[str, dict[str, Callable]]: """Registries for the 'numba' accelerator. Composes the numba implementations (``intensity`` and the ``pearson`` / - ``manders_fold`` / ``rwc`` / ``overlap`` colocalization features) with the - numpy implementations of every other feature — a single global "numba" + ``manders_fold`` / ``rwc`` / ``costes`` / ``overlap`` colocalization features) + with the numpy implementations of every other feature — a single global "numba" selection still yields a full, working feature set, accelerated where a numba backend exists. This is explicit per-function composition, NOT an error-driven fallback. @@ -61,6 +61,7 @@ def _numba_registries() -> dict[str, dict[str, Callable]]: ``overlap`` to the numpy ``_CORRELATION`` for symmetry is a separate call. """ from cp_measure.core.numba import ( + get_correlation_costes as _numba_costes, get_correlation_manders_fold as _numba_manders_fold, get_correlation_overlap as _numba_overlap, get_correlation_pearson as _numba_pearson, @@ -75,6 +76,7 @@ def _numba_registries() -> dict[str, dict[str, Callable]]: "pearson": _numba_pearson, "manders_fold": _numba_manders_fold, "rwc": _numba_rwc, + "costes": _numba_costes, "overlap": _numba_overlap, }, } diff --git a/src/cp_measure/core/numba/__init__.py b/src/cp_measure/core/numba/__init__.py index 48c00d8..19d6bda 100644 --- a/src/cp_measure/core/numba/__init__.py +++ b/src/cp_measure/core/numba/__init__.py @@ -5,12 +5,13 @@ ``numba`` extra; availability is gated by ``cp_measure._detect.HAS_NUMBA``. This backend accelerates ``intensity`` and the colocalization features -``pearson``/``manders_fold``/``rwc``/``overlap``; the global "numba" accelerator -composes them with the numpy implementations of every other feature (see -``cp_measure.bulk``). +``pearson``/``manders_fold``/``rwc``/``overlap``/``costes``; the global "numba" +accelerator composes them with the numpy implementations of every other feature +(see ``cp_measure.bulk``). """ from cp_measure.core.numba.measurecolocalization import ( + get_correlation_costes, get_correlation_manders_fold, get_correlation_overlap, get_correlation_pearson, @@ -19,6 +20,7 @@ from cp_measure.core.numba.measureobjectintensity import get_intensity __all__ = [ + "get_correlation_costes", "get_correlation_manders_fold", "get_correlation_overlap", "get_correlation_pearson", diff --git a/src/cp_measure/core/numba/_costes.py b/src/cp_measure/core/numba/_costes.py new file mode 100644 index 0000000..9be0142 --- /dev/null +++ b/src/cp_measure/core/numba/_costes.py @@ -0,0 +1,240 @@ +"""Costes automated-threshold colocalization kernel (single-threaded, cached). + +Per-object port of ``measurecolocalization.get_correlation_costes_ind`` and its +``bisection_costes`` / ``linear_costes`` search. Each object's two channels are +the contiguous blocks ``g1[offsets[k]:offsets[k+1]]`` / ``g2[...]`` from +:func:`cp_measure.primitives._segment_numba.flatten_pairs_grouped` (the same +grouped layout the other coloc features use — no sort here). + +The reference's ``calculate_threshold`` call inside ``..._ind`` is dead (its +outputs are never read), so it is not ported. ``thr`` likewise has no effect on +costes. ``scale`` (``infer_scale``, dtype-keyed: float→1) is passed in host-side. + +Exactness: bit-exact vs the numpy reference on **float** pixels (the realistic +input; ``scale==1``). Integer-dtype input diverges by design — the reference +computes ``z = fi + si`` in that dtype and overflows (uint8/uint16), corrupting +the regression; the float64 kernel here does not. See ``tasks/numba_costes_plan.md``. + +Serial: a plain object loop, no ``prange``/``nogil``. +""" + +import numpy as np +from numba import njit + + +@njit(cache=True, error_model="numpy") +def _regression_ab(g1, g2, lo, hi): + """Costes orthogonal-regression line ``si ≈ a·fi + b`` over non-zero pixels. + + ``non_zero = (fi > 0) | (si > 0)``; ``a`` from the covariance recovered via + ``cov = 0.5(var(fi+si) - var(fi) - var(si))`` (all ``ddof=1``), matching the + reference exactly. + """ + cnt = 0 + s1 = 0.0 + s2 = 0.0 + for i in range(lo, hi): + if g1[i] > 0 or g2[i] > 0: + cnt += 1 + s1 += g1[i] + s2 += g2[i] + m1 = s1 / cnt + m2 = s2 / cnt + vx = 0.0 + vy = 0.0 + vz = 0.0 + for i in range(lo, hi): + if g1[i] > 0 or g2[i] > 0: + d1 = g1[i] - m1 + d2 = g2[i] - m2 + dz = (g1[i] + g2[i]) - (m1 + m2) + vx += d1 * d1 + vy += d2 * d2 + vz += dz * dz + denom_n = cnt - 1 + xvar = vx / denom_n + yvar = vy / denom_n + zvar = vz / denom_n + covar = 0.5 * (zvar - (xvar + yvar)) + yx = yvar - xvar + a = (yx + np.sqrt(yx * yx + 4.0 * covar * covar)) / (2.0 * covar) + b = m2 - a * m1 + return a, b + + +@njit(cache=True, error_model="numpy") +def _count_combt(g1, g2, lo, hi, thr_fi, thr_si): + """``count_nonzero((fi < thr_fi) | (si < thr_si))`` over the object's block.""" + cnt = 0 + for i in range(lo, hi): + if g1[i] < thr_fi or g2[i] < thr_si: + cnt += 1 + return cnt + + +@njit(cache=True, error_model="numpy") +def _pearson_combt(g1, g2, lo, hi, thr_fi, thr_si): + """Pearson r over the ``(fi < thr_fi) | (si < thr_si)`` subset. + + Mirrors ``scipy.stats.pearsonr``'s operation order — centre, normalise each + vector by its L2 norm, accumulate the normalised products, clamp to [-1, 1] — + to stay as close as possible to the value the reference branches on. Pass + ``thr_fi = thr_si = inf`` for the full-block correlation. + """ + cnt = 0 + s1 = 0.0 + s2 = 0.0 + for i in range(lo, hi): + if g1[i] < thr_fi or g2[i] < thr_si: + cnt += 1 + s1 += g1[i] + s2 += g2[i] + if cnt == 0: + return np.nan + m1 = s1 / cnt + m2 = s2 / cnt + nx = 0.0 + ny = 0.0 + for i in range(lo, hi): + if g1[i] < thr_fi or g2[i] < thr_si: + d1 = g1[i] - m1 + d2 = g2[i] - m2 + nx += d1 * d1 + ny += d2 * d2 + normx = np.sqrt(nx) + normy = np.sqrt(ny) + r = 0.0 + for i in range(lo, hi): + if g1[i] < thr_fi or g2[i] < thr_si: + r += ((g1[i] - m1) / normx) * ((g2[i] - m2) / normy) + if r > 1.0: + r = 1.0 + elif r < -1.0: + r = -1.0 + return r + + +@njit(cache=True, error_model="numpy") +def _bisection(g1, g2, lo, hi, a, b, scale): + """``bisection_costes`` (M_FASTER): narrowing-window search for the threshold.""" + left = 1.0 + right = float(scale) + mid = np.floor((right - left) / 1.2) + left + lastmid = 0.0 + valid = 1.0 + while lastmid != mid: + thr_fi = mid / scale + thr_si = a * thr_fi + b + cnt = _count_combt(g1, g2, lo, hi, thr_fi, thr_si) + if cnt <= 2: + left = mid - 1 + else: + r = _pearson_combt(g1, g2, lo, hi, thr_fi, thr_si) + if r < 0: + left = mid - 1 + elif r >= 0: + right = mid + 1 + valid = mid + # NaN (constant subset): neither bound moves; loop exits next step. + lastmid = mid + if right - left > 6: + mid = np.floor((right - left) / 1.2) + left + else: + mid = np.floor((right - left) / 2.0) + left + thr_fi_c = (valid - 1) / scale + thr_si_c = a * thr_fi_c + b + return thr_fi_c, thr_si_c + + +@njit(cache=True, error_model="numpy") +def _linear(g1, g2, lo, hi, a, b, scale, accurate): + """``linear_costes`` (M_FAST / M_ACCURATE): step the threshold down to R≈0.""" + i_step = 1.0 / scale + fi_max = -np.inf + si_max = -np.inf + for i in range(lo, hi): + if g1[i] > fi_max: + fi_max = g1[i] + if g2[i] > si_max: + si_max = g2[i] + img_max = fi_max if fi_max > si_max else si_max + cur = i_step * (np.floor(img_max / i_step) + 1) + thr_fi_c = cur + thr_si_c = a * cur + b + while cur > fi_max and (a * cur + b) > si_max: + cur -= i_step + r = 0.0 + num_true = -1 # sentinel (reference's None: forces recompute on first pass) + while cur > i_step: + thr_fi_c = cur + thr_si_c = a * cur + b + cnt = _count_combt(g1, g2, lo, hi, thr_fi_c, thr_si_c) + if cnt < 2: # reference: scipy.stats.pearsonr raises -> break + break + if cnt != num_true: + r = _pearson_combt(g1, g2, lo, hi, thr_fi_c, thr_si_c) + num_true = cnt + if r <= 0: + break + elif accurate or cur < i_step * 10: + cur -= i_step + elif r > 0.45: + cur -= i_step * 10 + elif r > 0.35: + cur -= i_step * 5 + elif r > 0.25: + cur -= i_step * 2 + else: + cur -= i_step + return thr_fi_c, thr_si_c + + +@njit(cache=True, error_model="numpy") +def costes_per_object(g1, g2, offsets, n, scale, mode): + """Costes C1/C2 per object over grouped value blocks. + + ``mode``: 0 = bisection (M_FASTER), 1 = linear M_FAST, 2 = linear M_ACCURATE. + Returns ``(C1[n], C2[n])``. An object with no pixel above the Costes threshold + yields ``0.0`` (matching the reference's fringe-case default). + """ + C1 = np.zeros(n) + C2 = np.zeros(n) + for k in range(n): + lo = offsets[k] + hi = offsets[k + 1] + if hi - lo == 0: + continue + a, b = _regression_ab(g1, g2, lo, hi) + if mode == 0: + thr_fi_c, thr_si_c = _bisection(g1, g2, lo, hi, a, b, scale) + else: + thr_fi_c, thr_si_c = _linear(g1, g2, lo, hi, a, b, scale, mode == 2) + + any_fi = False + any_si = False + sum_ge_fi = 0.0 + sum_ge_si = 0.0 + sum_fi_c = 0.0 + sum_si_c = 0.0 + n_comb = 0 + for i in range(lo, hi): + v1 = g1[i] + v2 = g2[i] + if v1 > thr_fi_c: + any_fi = True + if v2 > thr_si_c: + any_si = True + if v1 >= thr_fi_c: + sum_ge_fi += v1 + if v2 >= thr_si_c: + sum_ge_si += v2 + if v1 > thr_fi_c and v2 > thr_si_c: + n_comb += 1 + sum_fi_c += v1 + sum_si_c += v2 + if n_comb > 0: + tot_fi = sum_ge_fi if any_fi else 0.0 + tot_si = sum_ge_si if any_si else 0.0 + C1[k] = sum_fi_c / tot_fi + C2[k] = sum_si_c / tot_si + return C1, C2 diff --git a/src/cp_measure/core/numba/measurecolocalization.py b/src/cp_measure/core/numba/measurecolocalization.py index 17cdeea..1bd3ff2 100644 --- a/src/cp_measure/core/numba/measurecolocalization.py +++ b/src/cp_measure/core/numba/measurecolocalization.py @@ -23,7 +23,8 @@ — but means the two backends can differ on genuine integer images. Real (float) intensity images are unaffected. -``costes`` lives in a stacked follow-up; it is not part of this module. +``costes`` is also here (it runs a per-object iterative threshold search rather +than the fused reduction, so it has its own kernel and runner). """ import numpy @@ -31,17 +32,25 @@ from cp_measure.core.measurecolocalization import ( F_CORRELATION_FORMAT, + F_COSTES_FORMAT, F_K_FORMAT, F_MANDERS_FORMAT, F_OVERLAP_FORMAT, F_RWC_FORMAT, F_SLOPE_FORMAT, + M_ACCURATE, + M_FAST, + M_FASTER, + infer_scale, ) from cp_measure.core.numba._colocalization import coloc_per_object +from cp_measure.core.numba._costes import costes_per_object from cp_measure.primitives._segment_numba import flatten_pairs_grouped from cp_measure.primitives.segment import labels_to_offsets from cp_measure.primitives.shapes import to_bzyx +_COSTES_MODE = {M_FASTER: 0, M_FAST: 1, M_ACCURATE: 2} + def _run(masks_zyx, pixels_1_zyx, pixels_2_zyx, thr_frac, compute_rwc): """Flatten one ``(Z, Y, X)`` image triple and run the fused kernel. @@ -150,3 +159,40 @@ def build(res): } return _featurize(pixels_1, pixels_2, masks, thr, False, build) + + +def get_correlation_costes( + pixels_1: NDArray[numpy.floating], + pixels_2: NDArray[numpy.floating], + masks: NDArray[numpy.integer], + fast_costes: str = M_FASTER, + thr: int = 15, +) -> dict[str, NDArray[numpy.floating]]: + """Costes automated-threshold Manders coefficients C1/C2. + + ``thr`` is accepted for signature parity but has no effect (in the reference it + only fed the dead ``calculate_threshold`` call). ``scale`` is dtype-derived via + ``infer_scale`` on ``pixels_1``, so float input gives ``scale == 1``. + """ + mode = _COSTES_MODE[fast_costes] + masks_list, pixels_1_list, unwrap = to_bzyx(masks, pixels_1) + _, pixels_2_list, _ = to_bzyx(masks, pixels_2) + + def run(masks_zyx, p1_zyx, p2_zyx): + m = numpy.ascontiguousarray(masks_zyx) + if not numpy.issubdtype(m.dtype, numpy.integer): + m = m.astype(numpy.intp) + lut, n, offsets = labels_to_offsets(m) + if n == 0: + return {f"{F_COSTES_FORMAT}_1": _EMPTY, f"{F_COSTES_FORMAT}_2": _EMPTY} + p1 = numpy.ascontiguousarray(p1_zyx, dtype=numpy.float64) + p2 = numpy.ascontiguousarray(p2_zyx, dtype=numpy.float64) + g1, g2 = flatten_pairs_grouped(m, p1, p2, lut, offsets) + scale = float(infer_scale(numpy.asarray(p1_zyx))) + c1, c2 = costes_per_object(g1, g2, offsets, n, scale, mode) + return {f"{F_COSTES_FORMAT}_1": c1, f"{F_COSTES_FORMAT}_2": c2} + + results = [ + run(m, p1, p2) for m, p1, p2 in zip(masks_list, pixels_1_list, pixels_2_list) + ] + return unwrap(results) diff --git a/test/test_backend_correctness.py b/test/test_backend_correctness.py index a5eea1e..e11833a 100644 --- a/test/test_backend_correctness.py +++ b/test/test_backend_correctness.py @@ -71,12 +71,10 @@ def test_set_accelerator_numba_composes_with_numpy(): ) # The colocalization features route to the numba backend too. corr = cp_measure.bulk.get_correlation_measurements() - for feature in ("pearson", "manders_fold", "rwc", "overlap"): + for feature in ("pearson", "manders_fold", "rwc", "costes", "overlap"): assert corr[feature].__module__ == ( "cp_measure.core.numba.measurecolocalization" ), feature - # costes stays on numpy until its follow-up lands. - assert corr["costes"].__module__ == "cp_measure.core.measurecolocalization" # Every other feature stays on the numpy backend. assert core["sizeshape"].__module__ == "cp_measure.core.measureobjectsizeshape" assert core["texture"].__module__ == "cp_measure.core.measuretexture" diff --git a/test/test_coloc_backend.py b/test/test_coloc_backend.py index 6b80b3c..20e9259 100644 --- a/test/test_coloc_backend.py +++ b/test/test_coloc_backend.py @@ -30,7 +30,7 @@ requires_numba = pytest.mark.skipif(not HAS_NUMBA, reason="numba not installed") -FUNCS = ["pearson", "manders_fold", "rwc", "overlap"] +FUNCS = ["pearson", "manders_fold", "rwc", "overlap", "costes"] def _numba_fn(name): @@ -116,3 +116,13 @@ def test_4d_batch_matches_per_image(name): assert isinstance(got, list) and len(got) == 2 for per_image, (m, a, b) in zip(got, [(m0, a0, b0), (m1, a1, b1)]): _assert_match(per_image, getattr(ref, f"get_correlation_{name}")(a, b, m)) + + +@requires_numba +@pytest.mark.parametrize("mode", ["Faster", "Fast", "Accurate"]) +@pytest.mark.parametrize("dim", ["2d", "3d"]) +def test_costes_modes_match_numpy(mode, dim): + masks, p1, p2 = _data_2d("cont") if dim == "2d" else _data_3d("cont") + expected = ref.get_correlation_costes(p1, p2, masks, fast_costes=mode) + got = _numba_fn("costes")(p1, p2, masks, fast_costes=mode) + _assert_match(got, expected) diff --git a/test/test_costes_kernels.py b/test/test_costes_kernels.py new file mode 100644 index 0000000..3512cef --- /dev/null +++ b/test/test_costes_kernels.py @@ -0,0 +1,88 @@ +"""Unit tests for the costes numba kernels. + +The reference ``bisection_costes`` / ``linear_costes`` accept ``scale_max`` as a +parameter, so the REAL multi-iteration search can be exercised at scale=255 on +float pixels in [0,1] — without the integer dtype that would make the reference +overflow ``z = fi + si``. This is where the control flow actually runs (the +end-to-end float64 path has scale=1 and a near-trivial search). +""" + +import numpy as np +import pytest +import scipy.stats + +import cp_measure.core.measurecolocalization as ref +from cp_measure._detect import HAS_NUMBA + +requires_numba = pytest.mark.skipif(not HAS_NUMBA, reason="numba not installed") + + +def _obj(seed, n=200): + """A correlated float object in [0,1]: si ≈ fi + noise, both clipped.""" + rng = np.random.default_rng(seed) + fi = rng.random(n) + si = np.clip(0.7 * fi + 0.3 * rng.random(n), 0.0, 1.0) + return np.ascontiguousarray(fi), np.ascontiguousarray(si) + + +@requires_numba +@pytest.mark.parametrize("seed", [0, 1, 2, 3]) +def test_regression_ab_matches_reference(seed): + from cp_measure.core.numba._costes import _regression_ab + + fi, si = _obj(seed) + a, b = _regression_ab(fi, si, 0, fi.size) + # Reference computes a,b inline; reproduce its exact expression. + nz = (fi > 0) | (si > 0) + xvar = np.var(fi[nz], ddof=1) + yvar = np.var(si[nz], ddof=1) + z = fi[nz] + si[nz] + zvar = np.var(z, ddof=1) + covar = 0.5 * (zvar - (xvar + yvar)) + a_ref = ((yvar - xvar) + np.sqrt((yvar - xvar) ** 2 + 4 * covar**2)) / (2 * covar) + b_ref = si[nz].mean() - a_ref * fi[nz].mean() + np.testing.assert_allclose(a, a_ref, rtol=1e-9) + np.testing.assert_allclose(b, b_ref, rtol=1e-9) + + +@requires_numba +@pytest.mark.parametrize("seed", [0, 1, 2]) +def test_pearson_combt_matches_scipy(seed): + from cp_measure.core.numba._costes import _pearson_combt + + fi, si = _obj(seed) + thr_fi, thr_si = 0.4, 0.35 + got = _pearson_combt(fi, si, 0, fi.size, thr_fi, thr_si) + combt = (fi < thr_fi) | (si < thr_si) + r_ref, _ = scipy.stats.pearsonr(fi[combt], si[combt]) + np.testing.assert_allclose(got, r_ref, rtol=1e-9, atol=1e-12) + # full-block (inf thresholds) == pearson over everything + got_full = _pearson_combt(fi, si, 0, fi.size, np.inf, np.inf) + np.testing.assert_allclose(got_full, scipy.stats.pearsonr(fi, si)[0], rtol=1e-9) + + +@requires_numba +@pytest.mark.parametrize("seed", [0, 1, 2, 3, 4]) +def test_bisection_matches_reference(seed): + from cp_measure.core.numba._costes import _bisection, _regression_ab + + fi, si = _obj(seed) + a, b = _regression_ab(fi, si, 0, fi.size) + thr_fi, thr_si = _bisection(fi, si, 0, fi.size, a, b, 255.0) + thr_fi_ref, thr_si_ref = ref.bisection_costes(fi, si, fi, si, 255) + np.testing.assert_allclose([thr_fi, thr_si], [thr_fi_ref, thr_si_ref], rtol=1e-6) + + +@requires_numba +@pytest.mark.parametrize("mode_name", ["Fast", "Accurate"]) +@pytest.mark.parametrize("seed", [0, 1, 2, 3]) +def test_linear_matches_reference(seed, mode_name): + from cp_measure.core.numba._costes import _linear, _regression_ab + + fi, si = _obj(seed) + a, b = _regression_ab(fi, si, 0, fi.size) + thr_fi, thr_si = _linear(fi, si, 0, fi.size, a, b, 255.0, mode_name == "Accurate") + thr_fi_ref, thr_si_ref = ref.linear_costes( + fi, si, fi, si, 255, fast_costes=mode_name + ) + np.testing.assert_allclose([thr_fi, thr_si], [thr_fi_ref, thr_si_ref], rtol=1e-6) From c760cff7de1e8a89cd31c5a048fefaaf95047cb6 Mon Sep 17 00:00:00 2001 From: Tim Treis Date: Thu, 4 Jun 2026 00:24:25 +0200 Subject: [PATCH 2/4] refactor(numba): tidy costes after /simplify MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Extract _flatten_image() (mask-contiguity + labels_to_offsets + flatten_pairs_grouped), shared by _run and the costes runner — removes the duplicated per-image prep chain. - Drop the dead any_fi/any_si flags in costes_per_object: tot_* is read only when n_comb > 0, which already guarantees a pixel strictly above each threshold, so the reference's any(>thr) guard is always true there. Behaviour-preserving; 57 coloc/costes tests green, lint clean. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/cp_measure/core/numba/_costes.py | 15 +++------ .../core/numba/measurecolocalization.py | 32 ++++++++++++------- 2 files changed, 25 insertions(+), 22 deletions(-) diff --git a/src/cp_measure/core/numba/_costes.py b/src/cp_measure/core/numba/_costes.py index 9be0142..6671257 100644 --- a/src/cp_measure/core/numba/_costes.py +++ b/src/cp_measure/core/numba/_costes.py @@ -210,8 +210,6 @@ def costes_per_object(g1, g2, offsets, n, scale, mode): else: thr_fi_c, thr_si_c = _linear(g1, g2, lo, hi, a, b, scale, mode == 2) - any_fi = False - any_si = False sum_ge_fi = 0.0 sum_ge_si = 0.0 sum_fi_c = 0.0 @@ -220,10 +218,6 @@ def costes_per_object(g1, g2, offsets, n, scale, mode): for i in range(lo, hi): v1 = g1[i] v2 = g2[i] - if v1 > thr_fi_c: - any_fi = True - if v2 > thr_si_c: - any_si = True if v1 >= thr_fi_c: sum_ge_fi += v1 if v2 >= thr_si_c: @@ -232,9 +226,10 @@ def costes_per_object(g1, g2, offsets, n, scale, mode): n_comb += 1 sum_fi_c += v1 sum_si_c += v2 + # tot_* (denominators) are read only when n_comb > 0, which guarantees a + # pixel strictly above each threshold, hence sum_ge_* covers it — so the + # reference's `if any(x > thr)` guard on the totals is always true here. if n_comb > 0: - tot_fi = sum_ge_fi if any_fi else 0.0 - tot_si = sum_ge_si if any_si else 0.0 - C1[k] = sum_fi_c / tot_fi - C2[k] = sum_si_c / tot_si + C1[k] = sum_fi_c / sum_ge_fi + C2[k] = sum_si_c / sum_ge_si return C1, C2 diff --git a/src/cp_measure/core/numba/measurecolocalization.py b/src/cp_measure/core/numba/measurecolocalization.py index 1bd3ff2..1bdefe8 100644 --- a/src/cp_measure/core/numba/measurecolocalization.py +++ b/src/cp_measure/core/numba/measurecolocalization.py @@ -52,18 +52,20 @@ _COSTES_MODE = {M_FASTER: 0, M_FAST: 1, M_ACCURATE: 2} -def _run(masks_zyx, pixels_1_zyx, pixels_2_zyx, thr_frac, compute_rwc): - """Flatten one ``(Z, Y, X)`` image triple and run the fused kernel. +def _flatten_image(masks_zyx, pixels_1_zyx, pixels_2_zyx): + """Normalise one ``(Z, Y, X)`` triple to grouped per-object value blocks. - Returns the nine per-object arrays from ``coloc_per_object``, or ``None`` when - the image holds no objects (the callers then emit empty feature arrays). + Returns ``(g1, g2, offsets, n)``; an empty image (no objects) yields + ``(None, None, offsets, 0)``. Shared by every colocalization feature so the + mask-contiguity, ``labels_to_offsets`` and ``flatten_pairs_grouped`` chain + lives in one place. """ masks = numpy.ascontiguousarray(masks_zyx) if not numpy.issubdtype(masks.dtype, numpy.integer): masks = masks.astype(numpy.intp) lut, n, offsets = labels_to_offsets(masks) if n == 0: - return None + return None, None, offsets, 0 g1, g2 = flatten_pairs_grouped( masks, numpy.ascontiguousarray(pixels_1_zyx, dtype=numpy.float64), @@ -71,6 +73,18 @@ def _run(masks_zyx, pixels_1_zyx, pixels_2_zyx, thr_frac, compute_rwc): lut, offsets, ) + return g1, g2, offsets, n + + +def _run(masks_zyx, pixels_1_zyx, pixels_2_zyx, thr_frac, compute_rwc): + """Flatten one ``(Z, Y, X)`` image triple and run the fused kernel. + + Returns the nine per-object arrays from ``coloc_per_object``, or ``None`` when + the image holds no objects (the callers then emit empty feature arrays). + """ + g1, g2, offsets, n = _flatten_image(masks_zyx, pixels_1_zyx, pixels_2_zyx) + if n == 0: + return None return coloc_per_object(g1, g2, offsets, n, thr_frac, compute_rwc) @@ -179,15 +193,9 @@ def get_correlation_costes( _, pixels_2_list, _ = to_bzyx(masks, pixels_2) def run(masks_zyx, p1_zyx, p2_zyx): - m = numpy.ascontiguousarray(masks_zyx) - if not numpy.issubdtype(m.dtype, numpy.integer): - m = m.astype(numpy.intp) - lut, n, offsets = labels_to_offsets(m) + g1, g2, offsets, n = _flatten_image(masks_zyx, p1_zyx, p2_zyx) if n == 0: return {f"{F_COSTES_FORMAT}_1": _EMPTY, f"{F_COSTES_FORMAT}_2": _EMPTY} - p1 = numpy.ascontiguousarray(p1_zyx, dtype=numpy.float64) - p2 = numpy.ascontiguousarray(p2_zyx, dtype=numpy.float64) - g1, g2 = flatten_pairs_grouped(m, p1, p2, lut, offsets) scale = float(infer_scale(numpy.asarray(p1_zyx))) c1, c2 = costes_per_object(g1, g2, offsets, n, scale, mode) return {f"{F_COSTES_FORMAT}_1": c1, f"{F_COSTES_FORMAT}_2": c2} From 8b4ee9e6e5b1745450865d4fe25b4013e053d751 Mon Sep 17 00:00:00 2001 From: Tim Treis Date: Thu, 4 Jun 2026 06:47:42 +0200 Subject: [PATCH 3/4] perf(numba): fuse costes bisection's count + pearson passes bisection_costes called _count_combt then _pearson_combt for each visited threshold, but _pearson_combt's first pass recomputes the exact same subset count. Fuse them into _count_pearson_combt -> (cnt, r): one count pass, and the Pearson passes only when cnt > 2 (else r is nan and unused). Bit-identical to the previous kernel (same predicate, same accumulation order; verified array_equal across correlated AND anti-correlated objects, i.e. both search directions). ~9% faster (23.0 -> 20.9 ms, 1080^2/144 obj). 31 costes tests green. (linear/accurate modes keep _count_combt + the num_true cache.) Co-Authored-By: Claude Opus 4.8 (1M context) --- src/cp_measure/core/numba/_costes.py | 47 ++++++++++++++++++++++++++-- 1 file changed, 45 insertions(+), 2 deletions(-) diff --git a/src/cp_measure/core/numba/_costes.py b/src/cp_measure/core/numba/_costes.py index 6671257..c3bb0ba 100644 --- a/src/cp_measure/core/numba/_costes.py +++ b/src/cp_measure/core/numba/_costes.py @@ -114,6 +114,50 @@ def _pearson_combt(g1, g2, lo, hi, thr_fi, thr_si): return r +@njit(cache=True, error_model="numpy") +def _count_pearson_combt(g1, g2, lo, hi, thr_fi, thr_si): + """``(cnt, r)`` in one shot: the subset count AND its Pearson r. + + The standalone ``_count_combt`` recomputes exactly the count that + ``_pearson_combt``'s first pass already produces, so the bisection paid two + count passes per visited threshold. This fuses them: pass 1 yields ``cnt`` + (+ sums); ``r`` is computed (passes 2-3, identical to ``_pearson_combt``) only + when ``cnt > 2`` (else returned as ``nan`` and unused). Bit-identical to the + separate calls (same predicate, same accumulation order). + """ + cnt = 0 + s1 = 0.0 + s2 = 0.0 + for i in range(lo, hi): + if g1[i] < thr_fi or g2[i] < thr_si: + cnt += 1 + s1 += g1[i] + s2 += g2[i] + if cnt <= 2: + return cnt, np.nan + m1 = s1 / cnt + m2 = s2 / cnt + nx = 0.0 + ny = 0.0 + for i in range(lo, hi): + if g1[i] < thr_fi or g2[i] < thr_si: + d1 = g1[i] - m1 + d2 = g2[i] - m2 + nx += d1 * d1 + ny += d2 * d2 + normx = np.sqrt(nx) + normy = np.sqrt(ny) + r = 0.0 + for i in range(lo, hi): + if g1[i] < thr_fi or g2[i] < thr_si: + r += ((g1[i] - m1) / normx) * ((g2[i] - m2) / normy) + if r > 1.0: + r = 1.0 + elif r < -1.0: + r = -1.0 + return cnt, r + + @njit(cache=True, error_model="numpy") def _bisection(g1, g2, lo, hi, a, b, scale): """``bisection_costes`` (M_FASTER): narrowing-window search for the threshold.""" @@ -125,11 +169,10 @@ def _bisection(g1, g2, lo, hi, a, b, scale): while lastmid != mid: thr_fi = mid / scale thr_si = a * thr_fi + b - cnt = _count_combt(g1, g2, lo, hi, thr_fi, thr_si) + cnt, r = _count_pearson_combt(g1, g2, lo, hi, thr_fi, thr_si) if cnt <= 2: left = mid - 1 else: - r = _pearson_combt(g1, g2, lo, hi, thr_fi, thr_si) if r < 0: left = mid - 1 elif r >= 0: From b0fe736f724150d67b6d2a20dfc7cc41e33de2fa Mon Sep 17 00:00:00 2001 From: Tim Treis Date: Sun, 7 Jun 2026 01:11:21 +0200 Subject: [PATCH 4/4] =?UTF-8?q?perf(numba/coloc):=20fused=20get=5Fcorrelat?= =?UTF-8?q?ion=5Fall=20=E2=80=94=20one=20kernel=20run=20for=20many=20featu?= =?UTF-8?q?res?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The five public get_correlation_* functions each re-ran the whole fused coloc_per_object kernel (+ flatten), so computing several coloc features paid ~5x redundant work and the numba backend lost to merged numpy on manders/overlap/rwc at large. Add get_correlation_all(p1, p2, masks, features=None): one flatten + one kernel pass returns the requested coloc groups (None = all). The cheap block (Pearson+slope, Manders, Overlap, K) is one pass; RWC's rank sort and the Costes kernel are gated to the requested set. Stateless — fusion happens by requesting the set in ONE call, not via any cache. It's the efficient entry point for any caller (not only featurize). The five single-feature functions become thin gated wrappers over it (single source; each now computes only its tier, so even one direct call is minimal). The numba correlation registry KEEPS per-group keys so the featurizer's per-group selection (_collect_correlation_features) keeps working — an earlier single-entry registry broke featurize with KeyError 'pearson'. large 1080^2/142obj: all-coloc 38ms (fused) vs 103ms (5 separate) vs 140ms (numpy) = 2.7x / 3.7x. Tests: subset/gating, bit-identity vs the wrappers, empty/batch/3D, unknown-group error, and featurize-runs-under-numba. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/cp_measure/bulk.py | 12 +- src/cp_measure/core/numba/__init__.py | 2 + .../core/numba/measurecolocalization.py | 155 ++++++++++-------- test/test_coloc_fused.py | 150 +++++++++++++++++ 4 files changed, 247 insertions(+), 72 deletions(-) create mode 100644 test/test_coloc_fused.py diff --git a/src/cp_measure/bulk.py b/src/cp_measure/bulk.py index 2bf08bb..6d00154 100644 --- a/src/cp_measure/bulk.py +++ b/src/cp_measure/bulk.py @@ -55,10 +55,16 @@ def _numba_registries() -> dict[str, dict[str, Callable]]: numba backend exists. This is explicit per-function composition, NOT an error-driven fallback. + The five correlation entries are kept per-group so the featurizer's per-group + selection (``_collect_correlation_features``) keeps working; each is now a thin, + gated wrapper over the shared ``get_correlation_all`` kernel. A caller wanting + several coloc features in ONE flatten+kernel pass (the efficient path) calls + ``cp_measure.core.numba.get_correlation_all(..., features=...)`` directly — it is + stateless, so fusion happens by requesting the set in one call, not by the + registry. + Note: ``overlap`` is not in the numpy ``_CORRELATION`` registry, so the numba - correlation registry intentionally exposes one feature the numpy one does not - (the numba ``overlap`` backend exists and is cheap to surface). Adding - ``overlap`` to the numpy ``_CORRELATION`` for symmetry is a separate call. + correlation registry intentionally exposes one feature the numpy one does not. """ from cp_measure.core.numba import ( get_correlation_costes as _numba_costes, diff --git a/src/cp_measure/core/numba/__init__.py b/src/cp_measure/core/numba/__init__.py index 19d6bda..6bc09ef 100644 --- a/src/cp_measure/core/numba/__init__.py +++ b/src/cp_measure/core/numba/__init__.py @@ -11,6 +11,7 @@ """ from cp_measure.core.numba.measurecolocalization import ( + get_correlation_all, get_correlation_costes, get_correlation_manders_fold, get_correlation_overlap, @@ -20,6 +21,7 @@ from cp_measure.core.numba.measureobjectintensity import get_intensity __all__ = [ + "get_correlation_all", "get_correlation_costes", "get_correlation_manders_fold", "get_correlation_overlap", diff --git a/src/cp_measure/core/numba/measurecolocalization.py b/src/cp_measure/core/numba/measurecolocalization.py index 1bdefe8..eb88317 100644 --- a/src/cp_measure/core/numba/measurecolocalization.py +++ b/src/cp_measure/core/numba/measurecolocalization.py @@ -27,6 +27,8 @@ than the fused reduction, so it has its own kernel and runner). """ +from collections.abc import Iterable + import numpy from numpy.typing import NDArray @@ -76,36 +78,14 @@ def _flatten_image(masks_zyx, pixels_1_zyx, pixels_2_zyx): return g1, g2, offsets, n -def _run(masks_zyx, pixels_1_zyx, pixels_2_zyx, thr_frac, compute_rwc): - """Flatten one ``(Z, Y, X)`` image triple and run the fused kernel. - - Returns the nine per-object arrays from ``coloc_per_object``, or ``None`` when - the image holds no objects (the callers then emit empty feature arrays). - """ - g1, g2, offsets, n = _flatten_image(masks_zyx, pixels_1_zyx, pixels_2_zyx) - if n == 0: - return None - return coloc_per_object(g1, g2, offsets, n, thr_frac, compute_rwc) - - -def _featurize(pixels_1, pixels_2, masks, thr, compute_rwc, build): - """Normalise the triple via ``to_bzyx`` and map each image through ``build``. - - ``build(result)`` turns the kernel tuple (or ``None`` for an empty image) into - one image's feature dict. ``unwrap`` then collapses a single image to that - dict, or returns the list for a batch — matching the numba intensity backend. - """ - frac = thr / 100.0 - masks_list, pixels_1_list, unwrap = to_bzyx(masks, pixels_1) - _, pixels_2_list, _ = to_bzyx(masks, pixels_2) - results = [ - build(_run(m, p1, p2, frac, compute_rwc)) - for m, p1, p2 in zip(masks_list, pixels_1_list, pixels_2_list) - ] - return unwrap(results) +_EMPTY = numpy.empty(0) -_EMPTY = numpy.empty(0) +# The five single-feature functions are thin, gated wrappers over ``get_correlation_all`` (single +# source of truth for the kernel + feature-key assembly). Each computes only its tier: the cheap +# functions skip RWC's rank sort and the Costes kernel; ``rwc`` adds the rank sort; ``costes`` runs +# only the iterative kernel. Calling several separately still runs the kernel per call (the library +# is stateless) — pass the set to ``get_correlation_all`` to collapse them into one pass. def get_correlation_pearson( @@ -113,13 +93,7 @@ def get_correlation_pearson( pixels_2: NDArray[numpy.floating], masks: NDArray[numpy.integer], ) -> dict[str, NDArray[numpy.floating]]: - def build(res): - if res is None: - return {F_CORRELATION_FORMAT: _EMPTY, F_SLOPE_FORMAT: _EMPTY} - corr, slope = res[0], res[1] - return {F_CORRELATION_FORMAT: corr, F_SLOPE_FORMAT: slope} - - return _featurize(pixels_1, pixels_2, masks, 15, False, build) + return get_correlation_all(pixels_1, pixels_2, masks, features=("pearson",)) def get_correlation_manders_fold( @@ -128,13 +102,9 @@ def get_correlation_manders_fold( masks: NDArray[numpy.integer], thr: int = 15, ) -> dict[str, NDArray[numpy.floating]]: - def build(res): - if res is None: - return {f"{F_MANDERS_FORMAT}_1": _EMPTY, f"{F_MANDERS_FORMAT}_2": _EMPTY} - m1, m2 = res[2], res[3] - return {f"{F_MANDERS_FORMAT}_1": m1, f"{F_MANDERS_FORMAT}_2": m2} - - return _featurize(pixels_1, pixels_2, masks, thr, False, build) + return get_correlation_all( + pixels_1, pixels_2, masks, features=("manders_fold",), thr=thr + ) def get_correlation_rwc( @@ -143,13 +113,7 @@ def get_correlation_rwc( masks: NDArray[numpy.integer], thr: int = 15, ) -> dict[str, NDArray[numpy.floating]]: - def build(res): - if res is None: - return {f"{F_RWC_FORMAT}_1": _EMPTY, f"{F_RWC_FORMAT}_2": _EMPTY} - rwc1, rwc2 = res[7], res[8] - return {f"{F_RWC_FORMAT}_1": rwc1, f"{F_RWC_FORMAT}_2": rwc2} - - return _featurize(pixels_1, pixels_2, masks, thr, True, build) + return get_correlation_all(pixels_1, pixels_2, masks, features=("rwc",), thr=thr) def get_correlation_overlap( @@ -158,21 +122,9 @@ def get_correlation_overlap( masks: NDArray[numpy.integer], thr: int = 15, ) -> dict[str, NDArray[numpy.floating]]: - def build(res): - if res is None: - return { - F_OVERLAP_FORMAT: _EMPTY, - f"{F_K_FORMAT}_1": _EMPTY, - f"{F_K_FORMAT}_2": _EMPTY, - } - overlap, k1, k2 = res[4], res[5], res[6] - return { - F_OVERLAP_FORMAT: overlap, - f"{F_K_FORMAT}_1": k1, - f"{F_K_FORMAT}_2": k2, - } - - return _featurize(pixels_1, pixels_2, masks, thr, False, build) + return get_correlation_all( + pixels_1, pixels_2, masks, features=("overlap",), thr=thr + ) def get_correlation_costes( @@ -188,17 +140,82 @@ def get_correlation_costes( only fed the dead ``calculate_threshold`` call). ``scale`` is dtype-derived via ``infer_scale`` on ``pixels_1``, so float input gives ``scale == 1``. """ + return get_correlation_all( + pixels_1, + pixels_2, + masks, + features=("costes",), + thr=thr, + fast_costes=fast_costes, + ) + + +_GROUPS = ("pearson", "manders_fold", "rwc", "overlap", "costes") +_CHEAP_GROUPS = frozenset({"pearson", "manders_fold", "rwc", "overlap"}) + + +def get_correlation_all( + pixels_1: NDArray[numpy.floating], + pixels_2: NDArray[numpy.floating], + masks: NDArray[numpy.integer], + features: Iterable[str] | None = None, + thr: int = 15, + fast_costes: str = M_FASTER, +) -> dict[str, NDArray[numpy.floating]]: + """Colocalization features from ONE flatten + ONE fused kernel pass per image. + + ``features`` selects which groups to return — any of ``pearson`` / ``manders_fold`` / ``rwc`` / + ``overlap`` / ``costes`` — or ``None`` for all. The shared ``coloc_per_object`` kernel runs once + for the cheap block (Pearson + slope, Manders, Overlap, K); the two expensive paths are gated by + the request: RWC's rank sort runs only if ``rwc`` is asked for, the Costes iterative kernel only + if ``costes`` is. This is the efficient entry point for any caller wanting several coloc features + at once (the single-feature functions delegate here). Stateless: collapsing N features into one + pass means requesting them in one call. + """ + want = set(_GROUPS) if features is None else set(features) + unknown = want - set(_GROUPS) + if unknown: + raise ValueError(f"unknown correlation feature group(s): {sorted(unknown)}") + need_cheap = bool(want & _CHEAP_GROUPS) + need_rwc = "rwc" in want + need_costes = "costes" in want + frac = thr / 100.0 mode = _COSTES_MODE[fast_costes] masks_list, pixels_1_list, unwrap = to_bzyx(masks, pixels_1) _, pixels_2_list, _ = to_bzyx(masks, pixels_2) def run(masks_zyx, p1_zyx, p2_zyx): g1, g2, offsets, n = _flatten_image(masks_zyx, p1_zyx, p2_zyx) - if n == 0: - return {f"{F_COSTES_FORMAT}_1": _EMPTY, f"{F_COSTES_FORMAT}_2": _EMPTY} - scale = float(infer_scale(numpy.asarray(p1_zyx))) - c1, c2 = costes_per_object(g1, g2, offsets, n, scale, mode) - return {f"{F_COSTES_FORMAT}_1": c1, f"{F_COSTES_FORMAT}_2": c2} + out: dict[str, NDArray[numpy.floating]] = {} + if need_cheap or need_rwc: + if n == 0: + corr = slope = m1 = m2 = overlap = k1 = k2 = rwc1 = rwc2 = _EMPTY + else: + corr, slope, m1, m2, overlap, k1, k2, rwc1, rwc2 = coloc_per_object( + g1, g2, offsets, n, frac, need_rwc + ) + if "pearson" in want: + out[F_CORRELATION_FORMAT] = corr + out[F_SLOPE_FORMAT] = slope + if "manders_fold" in want: + out[f"{F_MANDERS_FORMAT}_1"] = m1 + out[f"{F_MANDERS_FORMAT}_2"] = m2 + if "overlap" in want: + out[F_OVERLAP_FORMAT] = overlap + out[f"{F_K_FORMAT}_1"] = k1 + out[f"{F_K_FORMAT}_2"] = k2 + if need_rwc: + out[f"{F_RWC_FORMAT}_1"] = rwc1 + out[f"{F_RWC_FORMAT}_2"] = rwc2 + if need_costes: + if n == 0: + c1 = c2 = _EMPTY + else: + scale = float(infer_scale(numpy.asarray(p1_zyx))) + c1, c2 = costes_per_object(g1, g2, offsets, n, scale, mode) + out[f"{F_COSTES_FORMAT}_1"] = c1 + out[f"{F_COSTES_FORMAT}_2"] = c2 + return out results = [ run(m, p1, p2) for m, p1, p2 in zip(masks_list, pixels_1_list, pixels_2_list) diff --git a/test/test_coloc_fused.py b/test/test_coloc_fused.py new file mode 100644 index 0000000..f5d51e6 --- /dev/null +++ b/test/test_coloc_fused.py @@ -0,0 +1,150 @@ +"""The fused numba colocalization producer (`get_correlation_all`). + +`get_correlation_all(features=None)` returns every coloc feature from ONE flatten + ONE +`coloc_per_object` pass; with a `features` subset it returns exactly those groups, gating RWC's +rank sort and the Costes kernel to what was requested. The five single-feature functions are thin +gated wrappers over it. The numba correlation registry keeps per-group keys (so featurize's +per-group selection works) and `featurize` runs under the numba accelerator. +""" + +import numpy as np +import pytest +from conftest import ( + DEPTH_3D, + SIZE_2D, + SIZE_3D, + _stamp_objects_2d, + _stamp_objects_3d, + get_rng, +) + +import cp_measure.core.measurecolocalization as ref +from cp_measure._detect import HAS_NUMBA + +requires_numba = pytest.mark.skipif(not HAS_NUMBA, reason="numba not installed") + +# group -> the wrapper function name and the feature keys it produces +GROUP_KEYS = { + "pearson": (ref.F_CORRELATION_FORMAT, ref.F_SLOPE_FORMAT), + "manders_fold": (f"{ref.F_MANDERS_FORMAT}_1", f"{ref.F_MANDERS_FORMAT}_2"), + "overlap": (ref.F_OVERLAP_FORMAT, f"{ref.F_K_FORMAT}_1", f"{ref.F_K_FORMAT}_2"), + "rwc": (f"{ref.F_RWC_FORMAT}_1", f"{ref.F_RWC_FORMAT}_2"), + "costes": (f"{ref.F_COSTES_FORMAT}_1", f"{ref.F_COSTES_FORMAT}_2"), +} + + +def _nb(): + import cp_measure.core.numba.measurecolocalization as nb + + return nb + + +def _data_2d(): + masks = np.zeros((SIZE_2D, SIZE_2D), np.int32) + _stamp_objects_2d(masks, n_objects=3) + rng = get_rng() + return masks, rng.random(masks.shape), rng.random(masks.shape) + + +def _data_3d(): + masks = np.zeros((DEPTH_3D, SIZE_3D, SIZE_3D), np.int32) + _stamp_objects_3d(masks, n_objects=2) + rng = get_rng() + return masks, rng.random(masks.shape), rng.random(masks.shape) + + +def _union_separate(p1, p2, masks): + nb = _nb() + out = {} + for name in GROUP_KEYS: + out.update(getattr(nb, f"get_correlation_{name}")(p1, p2, masks)) + return out + + +@requires_numba +@pytest.mark.parametrize("data", [_data_2d, _data_3d]) +def test_fused_all_is_bit_identical_to_separate(data): + masks, p1, p2 = data() + fused = _nb().get_correlation_all(p1, p2, masks) # features=None -> all + sep = _union_separate(p1, p2, masks) + assert set(fused) == set(sep), set(fused).symmetric_difference(sep) + for key in sep: + np.testing.assert_array_equal(fused[key], sep[key], err_msg=f"feature {key!r}") + + +@requires_numba +@pytest.mark.parametrize("group", list(GROUP_KEYS)) +def test_subset_returns_only_requested_and_matches_wrapper(group): + masks, p1, p2 = _data_2d() + sub = _nb().get_correlation_all(p1, p2, masks, features=[group]) + assert set(sub) == set(GROUP_KEYS[group]) + # the single-feature wrapper is exactly this subset + wrapper = getattr(_nb(), f"get_correlation_{group}")(p1, p2, masks) + assert set(wrapper) == set(sub) + for key in sub: + np.testing.assert_array_equal(wrapper[key], sub[key], err_msg=key) + + +@requires_numba +def test_multi_subset_returns_exactly_those_groups(): + masks, p1, p2 = _data_2d() + out = _nb().get_correlation_all(p1, p2, masks, features=["pearson", "rwc"]) + expected = set(GROUP_KEYS["pearson"]) | set(GROUP_KEYS["rwc"]) + assert set(out) == expected + full = _nb().get_correlation_all(p1, p2, masks) + for key in expected: # values identical to the full run + np.testing.assert_array_equal(out[key], full[key], err_msg=key) + + +@requires_numba +def test_unknown_feature_raises(): + masks, p1, p2 = _data_2d() + with pytest.raises(ValueError, match="unknown correlation feature"): + _nb().get_correlation_all(p1, p2, masks, features=["bogus"]) + + +@requires_numba +def test_empty_mask_returns_empty_arrays(): + masks = np.zeros((SIZE_2D, SIZE_2D), np.int32) + rng = get_rng() + out = _nb().get_correlation_all( + rng.random(masks.shape), rng.random(masks.shape), masks + ) + for key, val in out.items(): + assert val.shape == (0,), key + + +@requires_numba +def test_batch_list_matches_per_image(): + imgs = [_data_2d(), _data_3d()] + masks = [m for m, _, _ in imgs] + p1 = [a for _, a, _ in imgs] + p2 = [b for _, _, b in imgs] + got = _nb().get_correlation_all(p1, p2, masks) + assert isinstance(got, list) and len(got) == 2 + for (m, a, b), per_image in zip(imgs, got): + sep = _union_separate(a, b, m) + assert set(per_image) == set(sep) + for key in sep: + np.testing.assert_array_equal(per_image[key], sep[key], err_msg=key) + + +@requires_numba +def test_registry_keeps_per_group_keys_and_featurize_runs(): + import cp_measure + from cp_measure.bulk import get_correlation_measurements + from cp_measure.featurizer import featurize + + cp_measure.set_accelerator("numba") + try: + reg = get_correlation_measurements() + assert {"pearson", "manders_fold", "rwc", "costes", "overlap"} <= set(reg) + # featurize must not KeyError under numba and must emit correlation columns + rng = get_rng() + image = rng.random((2, SIZE_2D, SIZE_2D)) + masks = np.zeros((1, SIZE_2D, SIZE_2D), np.int32) + _stamp_objects_2d(masks[0], n_objects=3) + _, columns, _ = featurize(image, masks) + assert any("Correlation" in c for c in columns) + finally: + cp_measure.set_accelerator(None)