Speed up granularity ~2x by only upsampling labeled voxels

src/zedprofiler/featurization/granularity.py

@@ -74,6 +74,35 @@ def _subsample_3d(
     return scipy.ndimage.map_coordinates(data, (k, i, j), order=order)
 
 
+def _labeled_voxel_positions(
+    masked_labels: numpy.ndarray,
+) -> tuple[tuple[numpy.ndarray, ...], numpy.ndarray]:
+    """Return coordinates and label ids for every voxel in a labeled object.
+
+    ``scipy.ndimage.mean(image, masked_labels, label_range)`` only ever reads
+    voxels where ``masked_labels`` is nonzero; everything else is discarded.
+    Precomputing just those voxel coordinates (and the label id at each one)
+    lets a per-scale loop upsample/scan only what will actually be used,
+    instead of the whole image, without changing the result.
+
+    Parameters
+    ----------
+    masked_labels : numpy.ndarray
+        Label image with 0 marking background/unlabeled voxels.
+
+    Returns
+    -------
+    tuple[tuple[numpy.ndarray, ...], numpy.ndarray]
+        ``(coords, label_ids)`` where ``coords`` is a per-axis tuple of index
+        arrays (as returned by ``numpy.nonzero``) and ``label_ids`` is the
+        label value at each of those coordinates, i.e.
+        ``masked_labels[coords]``.
+
+    """
+    coords = numpy.nonzero(masked_labels)
+    return coords, masked_labels[coords]
+
+
 def _upsample_3d(
     data: numpy.ndarray,
     subsampled_shape: numpy.ndarray,
@@ -381,20 +410,34 @@ def compute_granularity(  # noqa: C901, PLR0912, PLR0913, PLR0915
             f"Spectrum length: {granular_spectrum_length}",
         )
 
-    # Precompute the upsample coordinate grids once before the spectrum loop.
-    # They depend only on the fixed subsampled/original shapes (not on the
-    # per-scale ``rec``), so rebuilding three full-resolution float arrays via
-    # ``numpy.mgrid`` on every scale (as ``_upsample_3d`` does internally) is
-    # wasted work. Reuse the same coordinate tuple for every
-    # ``map_coordinates`` call below. Only needed when subsampling is active
-    # and there are objects to measure.
+    # Precomputing labeled-voxel positions once (instead of upsampling/
+    # scanning the whole image every scale in the loop below) gives identical
+    # per-object means for a fraction of the work when labeled objects cover
+    # a small part of the image -- the common case. Coordinates depend only
+    # on the fixed subsampled/original shapes (not on the per-scale ``rec``),
+    # so this is computed once and reused for every ``map_coordinates`` call
+    # below. See ``_labeled_voxel_positions`` for why this is equivalent.
+    # When nobjects == 0 (no labeled objects at all) this block is skipped,
+    # leaving labeled_voxel_coords/labeled_voxel_labels at their empty
+    # defaults; upsample_coords below and the per-object branch in the
+    # spectrum loop are both also gated on nobjects > 0, so no per-object
+    # work is attempted and object_measurements stays empty. This mirrors
+    # pre-optimization behavior: the returned DataFrame has zero rows but
+    # keeps its Metadata_* columns (see
+    # test_compute_granularity_zero_objects_returns_empty_dataframe).
+    labeled_voxel_coords: tuple[numpy.ndarray, ...] = ()
+    labeled_voxel_labels = numpy.array([], dtype=original_labels.dtype)
+    if nobjects > 0:
+        labeled_voxel_coords, labeled_voxel_labels = _labeled_voxel_positions(
+            masked_labels,
+        )
+    have_labeled_voxels = labeled_voxel_labels.size > 0
+
     upsample_coords: tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray] | None = None
     if subsample_size < 1.0 and nobjects > 0:
-        k, i, j = numpy.mgrid[
-            0 : original_shape[0],
-            0 : original_shape[1],
-            0 : original_shape[2],
-        ].astype(float)
+        k = labeled_voxel_coords[0].astype(float)
+        i = labeled_voxel_coords[1].astype(float)
+        j = labeled_voxel_coords[2].astype(float)
         if original_shape[0] > 1:
             k *= float(new_shape[0] - 1) / float(original_shape[0] - 1)
         if original_shape[1] > 1:
@@ -430,23 +473,31 @@ def compute_granularity(  # noqa: C901, PLR0912, PLR0913, PLR0915
             print(f"Scale 1 - gs: {gs:.4f}, currentmean: {currentmean:.6f}")
 
         # ----------------------------------------------------------
-        # Per-object granularity: upsample rec to original shape,
-        # then compute per-label means using masked_labels.
+        # Per-object granularity: upsample rec to original shape at only the
+        # voxels that belong to a labeled object, then compute per-label
+        # means using those voxels' label ids. Equivalent to upsampling the
+        # whole image and calling scipy.ndimage.mean(rec_full, masked_labels,
+        # label_range), since that call already discards everything outside
+        # labeled_voxel_coords -- just without doing the discarded work.
         # ----------------------------------------------------------
         if nobjects > 0:
             if upsample_coords is not None:
-                rec_full = scipy.ndimage.map_coordinates(
+                rec_at_object_voxels = scipy.ndimage.map_coordinates(
                     rec,
                     upsample_coords,
                     order=1,
                 )
             else:
-                rec_full = rec
+                rec_at_object_voxels = rec[labeled_voxel_coords]
 
             # Single-pass per-object mean via scipy.ndimage.mean
-            if numpy.any(masked_labels > 0):
+            if have_labeled_voxels:
                 new_object_means = _fix_scipy_ndimage_result(
-                    scipy.ndimage.mean(rec_full, masked_labels, label_range),
+                    scipy.ndimage.mean(
+                        rec_at_object_voxels,
+                        labeled_voxel_labels,
+                        label_range,
+                    ),
                 )
             else:
                 new_object_means = numpy.zeros(len(label_range))

tests/benchmarking.py

@@ -320,6 +320,10 @@ def _load_real_world_two_object_loader() -> TwoObjectLoader:
     object_ids = [int(x) for x in np.unique(label) if x != 0]
     image_set_loader = ImageSetLoader.__new__(ImageSetLoader)
     image_set_loader.image_set_name = "real-world-dr90-c00-c90"
+    # Mirrors ImageSetLoader.image_id, which falls back to image_set_name
+    # when no identifier fields are set (see BenchmarkImageSet above). This
+    # loader bypasses __init__ via __new__, so nothing sets it otherwise.
+    image_set_loader.image_id = image_set_loader.image_set_name
     image_set_loader.image_set_dict = {
         "DNA1": tifffile.imread(
             CELLPROFILER_TUTORIAL_ROOT / "input" / f"{first_image_name}.tif",

tests/featurization/test_granularity.py

@@ -9,6 +9,7 @@
 from pydantic import BaseModel, ConfigDict, field_validator
 
 from zedprofiler.featurization.granularity import (
+    _labeled_voxel_positions,
     _subsample_3d,
     _upsample_3d,
     compute_granularity,
@@ -231,3 +232,165 @@ class Dummy:
     df = compute_granularity(Dummy(), radius=1, granular_spectrum_length=2)
     assert isinstance(df, pd.DataFrame)
     assert sorted(df["Metadata_Object_ObjectID"].tolist()) == [257, 514]
+
+
+def test_labeled_voxel_positions_matches_full_array_scan() -> None:
+    """Gathering only labeled voxels must match scanning the whole array.
+
+    This pins the core equivalence the granularity per-scale loop relies on
+    for its performance optimization: scipy.ndimage.mean(image, labels,
+    label_range) run on the full array must equal the same call restricted
+    to _labeled_voxel_positions's gathered coordinates/label ids, for
+    arbitrary (including sparse, non-contiguous) label placement.
+    """
+    rng = np.random.default_rng(0)
+    shape = (6, 10, 12)
+    labels = np.zeros(shape, dtype=int)
+    labels[1, 2, 3] = 5
+    labels[1, 2, 4] = 5
+    labels[4, 8, 9] = 9
+    image = rng.uniform(0, 100, size=shape)
+    label_range = np.array([5, 9])
+
+    coords, label_ids = _labeled_voxel_positions(labels)
+
+    full_means = scipy.ndimage.mean(image, labels, label_range)
+    gathered_means = scipy.ndimage.mean(image[coords], label_ids, label_range)
+
+    np.testing.assert_allclose(full_means, gathered_means)
+
+
+def test_labeled_voxel_positions_empty_when_no_labels() -> None:
+    """An all-background label image yields empty coordinates/labels, not a crash."""
+    labels = np.zeros((4, 4, 4), dtype=int)
+    coords, label_ids = _labeled_voxel_positions(labels)
+    assert label_ids.size == 0
+    assert all(c.size == 0 for c in coords)
+
+
+def test_sparse_upsample_matches_full_array_upsample() -> None:
+    """Upsampling only labeled voxels must match upsampling the whole array.
+
+    This pins the other equivalence the granularity per-scale loop's
+    optimization relies on (the mean-gathering side is pinned by
+    test_labeled_voxel_positions_matches_full_array_scan above): evaluating
+    scipy.ndimage.map_coordinates only at the coordinates of labeled voxels,
+    scaled into the subsampled array's coordinate space, must equal
+    upsampling the *entire* subsampled array with _upsample_3d (the
+    pre-optimization approach) and then indexing at those same voxels.
+    """
+    original_shape = (10, 14, 16)
+    subsampled_shape = np.array([5.0, 7.0, 8.0])
+    rng = np.random.default_rng(1)
+    rec = rng.uniform(0, 100, size=(5, 7, 8))
+
+    labels = np.zeros(original_shape, dtype=int)
+    labels[1, 2, 3] = 5
+    labels[1, 2, 4] = 5
+    labels[8, 12, 14] = 9
+
+    coords, _label_ids = _labeled_voxel_positions(labels)
+
+    # Reference (pre-optimization): upsample the whole subsampled array,
+    # then index at the labeled voxels.
+    full_upsampled = _upsample_3d(rec, subsampled_shape, original_shape)
+    expected = full_upsampled[coords]
+
+    # Optimized: scale only the labeled voxels' coordinates into the
+    # subsampled array's space, then map_coordinates just those points.
+    k, i, j = (c.astype(float) for c in coords)
+    if original_shape[0] > 1:
+        k *= float(subsampled_shape[0] - 1) / float(original_shape[0] - 1)
+    if original_shape[1] > 1:
+        i *= float(subsampled_shape[1] - 1) / float(original_shape[1] - 1)
+    if original_shape[2] > 1:
+        j *= float(subsampled_shape[2] - 1) / float(original_shape[2] - 1)
+    actual = scipy.ndimage.map_coordinates(rec, (k, i, j), order=1)
+
+    np.testing.assert_allclose(actual, expected)
+
+
+def test_compute_granularity_zero_objects_returns_empty_dataframe() -> None:
+    """No labeled objects (nobjects == 0) must not crash the per-scale loop.
+
+    Answers https://github.com/WayScience/ZedProfiler/pull/51#discussion_r3766497646:
+    with an all-background label image, every ``nobjects > 0`` branch in
+    compute_granularity is skipped, so no per-object measurements are ever
+    recorded. The result is a zero-row DataFrame that still carries its
+    Metadata_* columns, matching pre-optimization behavior.
+    """
+    shape = (8, 8, 8)
+    img = np.zeros(shape, dtype=float)
+    lab = np.zeros(shape, dtype=int)  # no labeled objects
+
+    class Dummy:
+        image = img
+        label_image = lab
+        object_ids: ClassVar[list[int]] = []
+        image_set_loader = type("ISL", (), {"image_set_name": "s", "image_id": "s"})()
+        compartment = "Cell"
+        channel = "Ch1"
+
+    df = compute_granularity(Dummy(), radius=1, granular_spectrum_length=3)
+    assert isinstance(df, pd.DataFrame)
+    assert len(df) == 0
+    assert "Metadata_Object_ObjectID" in df.columns
+    assert "Metadata_Imaging_ImageID" in df.columns
+    assert "Metadata_Experiment_ImageSet" in df.columns
+
+
+@pytest.mark.parametrize("shape,center", [((24, 48, 48), (12, 22, 32))])
+def test_compute_granularity_sparse_object_in_larger_image(
+    shape: tuple[int, int, int],
+    center: tuple[int, int, int],
+) -> None:
+    """A small object far from the edges of a much larger, subsampled image.
+
+    This is the scenario the labeled-voxel-only upsampling optimization
+    targets: a labeled object occupying a small fraction of the image, with
+    subsampling active (the production-default code path). Exercises a
+    shape/object-size ratio far more extreme than the small synthetic cases
+    elsewhere in this file, where the object is a large fraction of the image.
+    """
+    img, lab = make_image_and_label(shape, center)
+    imgset = ImageSetLoaderModel()
+    loader = ObjectLoaderModel(
+        image=img,
+        label_image=lab,
+        object_ids=[1],
+        image_set_loader=imgset,
+    )
+
+    granular_spectrum_length = 4
+    df = compute_granularity(
+        loader,
+        radius=2,
+        granular_spectrum_length=granular_spectrum_length,
+        subsample_size=0.5,
+        image_sample_size=0.5,
+    )
+
+    assert len(df) == 1
+    value_cols = [
+        c
+        for c in df.columns
+        if c
+        not in (
+            "Metadata_Object_ObjectID",
+            "Metadata_Imaging_ImageID",
+            "Metadata_Experiment_ImageSet",
+        )
+    ]
+    # This is an end-to-end smoke test for the sparse-object code path, not a
+    # value-equivalence check: erosion/reconstruction spectra don't have a
+    # simple closed-form expected value to assert against here. Exact
+    # numeric equivalence of the optimization itself (gathering only labeled
+    # voxels instead of scanning/upsampling the whole image) is pinned
+    # directly by test_labeled_voxel_positions_matches_full_array_scan and
+    # test_sparse_upsample_matches_full_array_upsample above. So here we only
+    # check that the pipeline produces one granularity column per requested
+    # scale (granular_spectrum_length) and that none of them are NaN/inf,
+    # which would indicate the sparse-voxel path silently dropped or
+    # corrupted a scale's measurement.
+    assert len(value_cols) == granular_spectrum_length
+    assert np.isfinite(df[value_cols].to_numpy(dtype=float)).all()