perf(sizeshape): vectorize get_zernike on foreground pixels

src/cp_measure/core/measureobjectsizeshape.py

@@ -3,11 +3,11 @@
 from numpy.typing import NDArray
 
 import centrosome.cpmorphology
-import centrosome.zernike
 import numpy
 import scipy.ndimage
 import skimage.measure
-from cp_measure.utils import masks_to_ijv
+from centrosome import zernike
+from cp_measure.utils import _zernike_scores, masks_to_ijv
 
 __doc__ = """\
 MeasureObjectSizeShape
@@ -1011,24 +1011,20 @@ def get_zernike(
     pixels: NDArray[numpy.floating] | None,
     zernike_numbers: int = 9,
 ) -> dict[str, NDArray[numpy.floating]]:
-    """Per-object Zernike shape features (2D only).
+    """Per-object Zernike shape features (2D only). ``pixels`` is unused: these are shape moments.
 
     Labels must be contiguous ``1..N`` (see :func:`cp_measure._sanitize.sanitize`).
     """
     if masks.ndim == 3:
         return {}
-    unique_indices = numpy.unique(masks)
-    unique_indices = unique_indices[unique_indices > 0]
-    indices = list(range(1, len(unique_indices) + 1))
-    labels = masks
-    zernike_numbers = centrosome.zernike.get_zernike_indexes(zernike_numbers + 1)
+    zernike_indexes = zernike.get_zernike_indexes(zernike_numbers + 1)
 
-    zf_l = centrosome.zernike.zernike(zernike_numbers, labels, indices)
-    results = {}
-    for (n, m), z in zip(zernike_numbers, zf_l.transpose()):  # type: ignore[call-overload]
-        results[f"Zernike_{n}_{m}"] = z
+    real_sums, imag_sums, radii, _counts = _zernike_scores(masks, zernike_indexes)
+    areas = numpy.pi * radii * radii
+    with numpy.errstate(divide="ignore", invalid="ignore"):
+        score = numpy.sqrt(real_sums**2 + imag_sums**2) / areas[:, numpy.newaxis]
 
-    return results
+    return {f"Zernike_{n}_{m}": score[:, i] for i, (n, m) in enumerate(zernike_indexes)}
 
 
 def get_feret(

src/cp_measure/utils.py

@@ -3,6 +3,8 @@
 """
 
 import numpy
+from centrosome import zernike
+from numpy.typing import NDArray
 
 
 def _ensure_np_array(value):
@@ -49,3 +51,82 @@ def labels_to_binmasks(masks: numpy.ndarray) -> numpy.ndarray:
     labels = numpy.unique(masks)
     labels = labels[labels > 0]
     return masks == labels.reshape((-1,) + (1,) * masks.ndim)
+
+
+def _zernike_scores(
+    masks: NDArray[numpy.integer],
+    zernike_indexes: NDArray[numpy.integer],
+    weight: NDArray[numpy.floating] | None = None,
+) -> tuple[
+    NDArray[numpy.floating],
+    NDArray[numpy.floating],
+    NDArray[numpy.floating],
+    NDArray[numpy.floating],
+]:
+    """Per-object real/imaginary Zernike moment sums, vectorised on foreground pixels.
+
+    Mirrors ``zernike.zernike`` without its two area-scaling costs: the basis is never
+    scattered into a full ``(H, W, K)`` complex array, and each moment is segment-summed by
+    label with one ``numpy.bincount`` instead of a whole-image ``scipy.ndimage.sum``. The Horner
+    basis evaluation is copied from ``centrosome.construct_zernike_polynomials`` (same lookup
+    table, ``r**2 > 1`` cutoff, ``z = y + i*x`` convention), so results match centrosome to
+    round-off.
+
+    Labels must be the contiguous ``1..N`` cp_measure guarantees (see
+    :mod:`cp_measure._sanitize`); the segment index is ``label - 1``.
+
+    ``weight`` (co-shaped with ``masks``, e.g. an intensity image) scales each pixel's
+    contribution; ``None`` gives the unweighted shape moments. Returns
+    ``(real_sums, imag_sums, radii, counts)`` of shapes ``(n, K)``, ``(n, K)``, ``(n,)`` and
+    ``(n,)``, ordered by ascending label, where ``radii`` is the enclosing-circle radius and
+    ``counts`` the object pixel count. ``get_zernike`` normalises by ``pi * radii**2``; the
+    intensity-weighted radial Zernikes normalise by ``counts``.
+    """
+    # Contiguous 1..N contract: segment index is label - 1, so the label->index lookup is a
+    # plain arange. masks.max() raises on a size-0 array, so guard it.
+    n = 0 if masks.size == 0 else int(masks.max())
+    k = len(zernike_indexes)
+    labels = numpy.arange(1, n + 1)
+    lut = numpy.arange(-1, n)  # lut[0] = -1 (background); lut[label] = label - 1
+    centers, radii = zernike.minimum_enclosing_circle(masks, labels)
+    radii = numpy.asarray(radii, dtype=float)
+    real_sums = numpy.zeros((n, k))
+    imag_sums = numpy.zeros((n, k))
+    if n == 0:
+        return real_sums, imag_sums, radii, numpy.zeros(0)
+
+    # Foreground pixels, their object row, and unit-disk coordinates relative to each
+    # object's enclosing circle — no full (H, W) coordinate grid is materialised.
+    seg_full = lut[masks]
+    keep = seg_full >= 0
+    rows, cols = numpy.nonzero(keep)
+    seg = seg_full[keep]
+    counts = numpy.bincount(seg, minlength=n).astype(float)
+    # Single-pixel objects have an enclosing-circle radius of 0; the resulting 0/0 yields NaN
+    # coordinates (which the r**2 > 1 cutoff later discards), matching centrosome — suppress the
+    # warning since it is expected, not a fault.
+    with numpy.errstate(invalid="ignore", divide="ignore"):
+        ym = (rows - centers[seg, 0]) / radii[seg]
+        xm = (cols - centers[seg, 1]) / radii[seg]
+
+    coeffs = zernike.construct_zernike_lookuptable(zernike_indexes)
+    r_square = xm * xm + ym * ym
+    z = ym + 1j * xm
+    w = None if weight is None else weight[keep].astype(float)
+    z_pows = {m: z**m for m in numpy.unique(zernike_indexes[:, 1]) if m}
+    for idx, (zn, zm) in enumerate(zernike_indexes):
+        s = numpy.zeros_like(xm)
+        for c in coeffs[idx, : (zn - zm) // 2 + 1]:  # Horner scheme on r**2
+            s *= r_square
+            s += c
+        s[r_square > 1] = 0
+        if w is not None:
+            s *= w
+        if zm == 0:  # purely real moment; the imaginary segment-sum is identically zero
+            real_sums[:, idx] = numpy.bincount(seg, weights=s, minlength=n)
+        else:
+            zf = s * z_pows[zm]
+            real_sums[:, idx] = numpy.bincount(seg, weights=zf.real, minlength=n)
+            imag_sums[:, idx] = numpy.bincount(seg, weights=zf.imag, minlength=n)
+
+    return real_sums, imag_sums, radii, counts

test/test_sanitize.py

@@ -63,13 +63,16 @@ def test_entry_point_sanitizes_by_default():
 
 
 def test_raw_function_does_not_sanitize():
-    # sanitize=False returns the bare implementation: it assumes 1..N and breaks
-    # on gapped IDs (here zernike indexes past the relabelled range).
+    # sanitize=False is the bare implementation: it assumes 1..N and does not relabel
+    # gapped IDs, so its output differs from the sanitized wrapper.
     px = np.random.default_rng(0).random((64, 64))
     raw = get_core_measurements(sanitize=False)["zernike"]
-    raw(_three_objects((1, 3, 2)), px)  # contiguous: fine
-    with pytest.raises(IndexError):
-        raw(_three_objects((1, 17, 5)), px)
+    sanitized = get_core_measurements(sanitize=True)["zernike"]
+    gapped = _three_objects((1, 17, 5))
+    raw_out = raw(gapped, px)
+    san_out = sanitized(gapped, px)
+    key = next(iter(raw_out))
+    assert len(raw_out[key]) != len(san_out[key])
 
 
 def test_sanitize_helper_wraps_raw_function():

test/test_zernike.py

@@ -0,0 +1,94 @@
+"""Golden + edge tests for the vectorised numpy ``get_zernike``.
+
+The vectorised implementation avoids centrosome's full (H, W, K) scatter and the per-channel
+``scipy.ndimage.sum``; it must match ``centrosome.zernike.zernike`` (called with the actual
+label values) to floating-point round-off. Existing ``test_core_measurements`` already covers
+shape / 3D-empty; here we lock the numerical result and the edge cases.
+"""
+
+import numpy
+import pytest
+from centrosome import zernike
+
+from cp_measure.core.measureobjectsizeshape import get_zernike
+from cp_measure.utils import _zernike_scores
+
+ATOL = 1e-10  # >> the ~2e-16 round-off observed, << any real signal
+
+
+def _reference(masks, zernike_numbers=9):
+    """Old centrosome path, called with the real label values as indices."""
+    uniq = numpy.unique(masks)
+    uniq = uniq[uniq > 0]
+    zidx = zernike.get_zernike_indexes(zernike_numbers + 1)
+    zf = zernike.zernike(zidx, masks, uniq)
+    return {f"Zernike_{n}_{m}": zf[:, i] for i, (n, m) in enumerate(zidx)}
+
+
+def _assert_matches(masks, zernike_numbers=9):
+    ref = _reference(masks, zernike_numbers)
+    got = get_zernike(masks, None, zernike_numbers)
+    assert list(got) == list(ref), "key set / order changed"
+    for k in ref:
+        assert got[k].shape == ref[k].shape, k
+        assert numpy.allclose(got[k], ref[k], atol=ATOL, rtol=1e-10, equal_nan=True), (
+            f"{k}: max|diff|={numpy.nanmax(numpy.abs(got[k] - ref[k]))}"
+        )
+
+
+def _generate_square_objects(size, n, gap_frac=0.75):
+    masks = numpy.zeros((size, size), numpy.int32)
+    step = size // n
+    obj = int(step * gap_frac)
+    lab = 0
+    for a in range(n):
+        for b in range(n):
+            lab += 1
+            masks[a * step : a * step + obj, b * step : b * step + obj] = lab
+    return masks
+
+
+@pytest.mark.parametrize("zernike_numbers", [5, 9, 14])
+def test_zernike_matches_centrosome_irregular(zernike_numbers):
+    # Irregular blobs of varied size/position cover the single- and multi-object cases;
+    # sweeping the zernike number checks the basis truncation across degrees.
+    rng = numpy.random.default_rng(0)
+    masks = numpy.zeros((128, 128), numpy.int32)
+    for lab, (cy, cx) in enumerate(rng.integers(20, 108, size=(6, 2)), 1):
+        yy, xx = numpy.mgrid[0:128, 0:128]
+        masks[(yy - cy) ** 2 + (xx - cx) ** 2 < rng.integers(40, 120)] = lab
+    _assert_matches(masks, zernike_numbers)
+
+
+def test_zernike_object_touching_edge():
+    masks = numpy.zeros((64, 64), numpy.int32)
+    masks[0:20, 0:20] = 1  # clipped at the top-left corner
+    masks[40:64, 40:64] = 2  # clipped at the bottom-right corner
+    _assert_matches(masks)
+
+
+def test_zernike_single_pixel_object():
+    # minimum_enclosing_circle radius -> 0, so both paths divide by zero identically.
+    masks = numpy.zeros((32, 32), numpy.int32)
+    masks[16, 16] = 1
+    masks[5:15, 5:15] = 2  # a normal object alongside the degenerate one
+    _assert_matches(masks)
+
+
+def test_zernike_empty_mask():
+    masks = numpy.zeros((40, 40), numpy.int32)
+    got = get_zernike(masks, None)
+    zidx = zernike.get_zernike_indexes(9 + 1)
+    assert list(got) == [f"Zernike_{n}_{m}" for n, m in zidx]
+    assert all(v.shape == (0,) for v in got.values())
+
+
+def test_zernike_scores_unit_weight_equals_unweighted():
+    # The weight=None path (shape moments) used by get_zernike vs an explicit unit weight,
+    # the path the intensity-weighted radial Zernikes (PR #75) take.
+    masks = _generate_square_objects(128, 3)
+    zidx = zernike.get_zernike_indexes(9 + 1)
+    r0, i0, _rad0, c0 = _zernike_scores(masks, zidx)
+    r1, i1, _rad1, c1 = _zernike_scores(masks, zidx, weight=numpy.ones(masks.shape))
+    assert numpy.allclose(r0, r1) and numpy.allclose(i0, i1)
+    assert numpy.array_equal(c0, c1)