diff --git a/src/zedprofiler/featurization/granularity.py b/src/zedprofiler/featurization/granularity.py index 0735b6f..5daf1d4 100644 --- a/src/zedprofiler/featurization/granularity.py +++ b/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)) diff --git a/tests/benchmarking.py b/tests/benchmarking.py index a26eb58..00305ef 100644 --- a/tests/benchmarking.py +++ b/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", diff --git a/tests/featurization/test_granularity.py b/tests/featurization/test_granularity.py index acb0769..764d547 100644 --- a/tests/featurization/test_granularity.py +++ b/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()