perf(radial): vectorize get_radial_zernikes via shared _zernike_scores (~2x)

src/cp_measure/core/measureobjectintensitydistribution.py

@@ -1,11 +1,11 @@
 import centrosome.cpmorphology
 import centrosome.propagate
-import centrosome.zernike
 import numpy
 import numpy.ma
 import scipy.ndimage
 import scipy.sparse
-from cp_measure.utils import masks_to_ijv
+from centrosome import zernike
+from cp_measure.utils import _zernike_scores
 from numpy.typing import NDArray
 
 """"
@@ -317,73 +317,26 @@ def get_radial_zernikes(
     """
     if labels.ndim == 3:
         return {}
-    zernike_indexes = centrosome.zernike.get_zernike_indexes(zernike_degree + 1)
+    zernike_indexes = zernike.get_zernike_indexes(zernike_degree + 1)
 
-    unique_labels = numpy.arange(1, int(labels.max()) + 1)
-    # MODIFIED: Delegate index generation to the minimum_enclosing_circle
-    # MODIFIED: We assume non-overlapping labels for now
-    # TODO Support label overlap (i.e., format in ijv)
-    # MODIFIED: Delegate indexes to minimum_enclosing_circle
-    ij, r = centrosome.cpmorphology.minimum_enclosing_circle(labels, unique_labels)
+    # Intensity-weighted moment sums via the shared helper (pixels as the per-pixel weight);
+    # radial Zernikes normalise by pixel count, not the enclosing-circle area. See _zernike_scores.
+    vr, vi, _radii, counts = _zernike_scores(labels, zernike_indexes, weight=pixels)
+
+    magnitude = numpy.sqrt(vr * vr + vi * vi) / counts[:, numpy.newaxis]
+    # CellProfiler convention: arctan2(real, imag), not textbook arctan2(imag, real).
+    # Kept for 1:1 parity — do not "fix" the argument order.
+    phase = numpy.arctan2(vr, vi)
 
     #
-    # Then compute x and y, the position of each labeled pixel
-    # within a unit circle around the object
+    # Results will be formatted in a dictionary with the following keys:
+    # Zernike{Magnitude|Phase}_{n}_{m}
+    # n - the radial moment of the Zernike
+    # m - the azimuthal moment of the Zernike
     #
-    ijv = masks_to_ijv(labels)
-
-    l_ = ijv[:, 2]  # (N,1) vector with labels
-
-    yx = (ijv[:, :2] - ij[l_ - 1, :]) / r[l_ - 1, numpy.newaxis]
-
-    z = centrosome.zernike.construct_zernike_polynomials(
-        yx[:, 1], yx[:, 0], zernike_indexes
-    )
-
-    # Filter ijv-formatted items to keep the ones inside the pixels boundary
-    ijv_mask = (ijv[:, 0] < pixels.shape[0]) & (ijv[:, 1] < pixels.shape[1])
-    # ijv_mask[ijv_mask] = pixels[ijv[ijv_mask,0], ijv[ijv_mask, 1]]
-
-    yx = yx[ijv_mask, :]
-    l_ = l_[ijv_mask]
-    z_ = z[ijv_mask, :]
-
     results: dict[str, NDArray[numpy.floating]] = {}
-    if len(l_) == 0:
-        # Cover fringe case in which all labels were filtered out
-        for mag_or_phase in ("Magnitude", "Phase"):
-            for n, m in zernike_indexes:
-                name = f"{M_CATEGORY}_Zernike{mag_or_phase}_{n}_{m}"
-                results[name] = numpy.zeros(0)
-    else:
-        # MODIFIED: Replaced sum with the updated sum_labels
-        areas = scipy.ndimage.sum_labels(
-            numpy.ones(l_.shape, int), labels=l_, index=unique_labels
-        )
-
-        #
-        # Results will be formatted in a dictionary with the following keys:
-        # Zernike{Magniture|Phase}_{n}_{m}
-        # n - the radial moment of the Zernike
-        # m - the azimuthal moment of the Zernike
-        #
-        for i, (n, m) in enumerate(zernike_indexes):
-            vr = scipy.ndimage.sum_labels(
-                pixels[ijv[:, 0], ijv[:, 1]] * z_[:, i].real,
-                labels=l_,
-                index=unique_labels,
-            )
-
-            vi = scipy.ndimage.sum_labels(
-                pixels[ijv[:, 0], ijv[:, 1]] * z[:, i].imag,
-                labels=l_,
-                index=unique_labels,
-            )
-
-            magnitude = numpy.sqrt(vr * vr + vi * vi) / areas
-            phase = numpy.arctan2(vr, vi)
-
-            results[f"{M_CATEGORY}_ZernikeMagnitude_{n}_{m}"] = magnitude
-            results[f"{M_CATEGORY}_ZernikePhase_{n}_{m}"] = phase
+    for i, (n, m) in enumerate(zernike_indexes):
+        results[f"{M_CATEGORY}_ZernikeMagnitude_{n}_{m}"] = magnitude[:, i]
+        results[f"{M_CATEGORY}_ZernikePhase_{n}_{m}"] = phase[:, i]
 
     return results

test/test_radial_zernike.py

@@ -0,0 +1,107 @@
+"""Equivalence tests for the vectorised numpy ``get_radial_zernikes``.
+
+These lock **numerical fidelity to the original centrosome implementation**, not to
+externally-known values: ``_reference`` reproduces the pre-vectorisation path, and each test
+asserts the vectorised output matches it to floating-point round-off.
+"""
+
+import numpy
+import pytest
+from centrosome import cpmorphology, zernike
+
+from cp_measure.core.measureobjectintensitydistribution import (
+    M_CATEGORY,
+    get_radial_zernikes,
+)
+from cp_measure.utils import masks_to_ijv
+
+ATOL = 1e-9  # >> the ~1e-13 summation-order round-off, << any real signal
+
+
+def _reference(labels, pixels, zernike_degree=9):
+    """Reproduce the original centrosome implementation (the pre-vectorisation path)."""
+    zidx = zernike.get_zernike_indexes(zernike_degree + 1)
+    ul = numpy.unique(labels)
+    ul = ul[ul > 0]
+    out = {}
+    if len(ul) == 0:  # same interleaved key order as get_radial_zernikes
+        for n, m in zidx:
+            out[f"{M_CATEGORY}_ZernikeMagnitude_{n}_{m}"] = numpy.zeros(0)
+            out[f"{M_CATEGORY}_ZernikePhase_{n}_{m}"] = numpy.zeros(0)
+        return out
+    ij, r = cpmorphology.minimum_enclosing_circle(labels, ul)
+    ijv = masks_to_ijv(labels)
+    pos = ijv[:, 2] - 1  # contiguous 1..N: label -> row
+    yx = (ijv[:, :2] - ij[pos]) / r[pos, numpy.newaxis]
+    z = zernike.construct_zernike_polynomials(yx[:, 1], yx[:, 0], zidx)
+    w = pixels[ijv[:, 0], ijv[:, 1]]
+    areas = numpy.bincount(pos, minlength=len(ul)).astype(float)
+    for i, (n, m) in enumerate(zidx):
+        vr = numpy.bincount(pos, weights=w * z[:, i].real, minlength=len(ul))
+        vi = numpy.bincount(pos, weights=w * z[:, i].imag, minlength=len(ul))
+        out[f"{M_CATEGORY}_ZernikeMagnitude_{n}_{m}"] = (
+            numpy.sqrt(vr * vr + vi * vi) / areas
+        )
+        out[f"{M_CATEGORY}_ZernikePhase_{n}_{m}"] = numpy.arctan2(vr, vi)
+    return out
+
+
+def _assert_matches(masks, pixels, zernike_degree=9):
+    ref = _reference(masks, pixels, zernike_degree)
+    got = get_radial_zernikes(masks, pixels, zernike_degree)
+    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-9, 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
+
+
+def _pixels(shape, seed=0):
+    return numpy.random.default_rng(seed).random(shape)
+
+
+@pytest.mark.parametrize("zernike_degree", [5, 9, 14])
+def test_radial_zernike_irregular(zernike_degree):
+    # Irregular blobs of varied size/position cover the single- and multi-object cases;
+    # sweeping the degree checks the basis truncation.
+    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, _pixels(masks.shape), zernike_degree)
+
+
+def test_radial_zernike_object_touching_edge():
+    # Frame-clipped objects are a distinct geometry the irregular test never hits (its blobs
+    # stay clear of the border): the enclosing circle sits centred near the boundary.
+    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, _pixels(masks.shape))
+
+
+def test_radial_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, _pixels(masks.shape))
+
+
+def test_radial_zernike_empty_mask():
+    # No objects -> every feature is an empty (0,) array; _assert_matches covers keys + shapes.
+    _assert_matches(numpy.zeros((40, 40), numpy.int32), _pixels((40, 40)))