Cucim-based extend mask

pyproject.toml

@@ -7,7 +7,7 @@ packages = ["spatiomic"]
 
 [project]
 name = "spatiomic"
-version = "0.9.1"
+version = "0.9.2"
 description = "A python toolbox for spatial omics analysis."
 requires-python = ">=3.11"
 license = { file = "LICENSE" }

spatiomic/dimension/_som.py

@@ -90,6 +90,7 @@ def __init__(
 
     def set_estimators(
         self,
+        initialize_som: bool = True,
     ) -> None:
         """Set the XPySOM and nearest neighbor finder estimators."""
         # Check whether we can use cupy to work on the GPU
@@ -98,29 +99,30 @@ def set_estimators(
         )
 
         # Initialise XPySOM instance
-        try:
-            self.som = XPySom(
-                self.node_count[0],
-                self.node_count[1],
-                self.dimension_count,
-                activation_distance=self.distance_metric,
-                neighborhood_function=self.neighborhood,
-                learning_rate=self.learning_rate_initial,
-                learning_rateN=self.learning_rate_final,
-                sigma=self.sigma_initial,
-                sigmaN=self.sigma_final,
-                xp=self.xp,
-                n_parallel=self.parallel_count,
-                random_seed=self.seed,
-            )
-        except ValueError as excp:
-            if self.distance_metric == "correlation":
-                raise ValueError(
-                    "Using XPySOM with Pearson correlation requires a custom implementation. "
-                    "You can install it via `pip install git+https://github.com/complextissue/xpysom`."
-                ) from excp
-            else:
-                raise excp
+        if initialize_som:
+            try:
+                self.som = XPySom(
+                    self.node_count[0],
+                    self.node_count[1],
+                    self.dimension_count,
+                    activation_distance=self.distance_metric,
+                    neighborhood_function=self.neighborhood,
+                    learning_rate=self.learning_rate_initial,
+                    learning_rateN=self.learning_rate_final,
+                    sigma=self.sigma_initial,
+                    sigmaN=self.sigma_final,
+                    xp=self.xp,
+                    n_parallel=self.parallel_count,
+                    random_seed=self.seed,
+                )
+            except ValueError as excp:
+                if self.distance_metric == "correlation":
+                    raise ValueError(
+                        "Using XPySOM with Pearson correlation requires a custom implementation. "
+                        "You can install it via `pip install git+https://github.com/complextissue/xpysom`."
+                    ) from excp
+                else:
+                    raise excp
 
         # Create the nearest neighbor finder
         self.neighbor_estimator = get_neighbor_finder(
@@ -366,13 +368,16 @@ def load(
         Args:
             save_path (str): The path where to load the SOM and its configuration from.
         """
-        # load and set the som and load the class config
+        # Load and set the som and load the class config
         with open(save_path, "rb") as infile:
             config, self.som = pickle.load(infile)  # nosec
 
-        # initialise an XPySOM object with the data and set the class variables
+        # Initialise an XPySOM object with the data and set the class variables
         self.set_config(**config)
 
+        # Re-initialize estimators without re-initializing the SOM
+        self.set_estimators(initialize_som=False)
+
     def save(
         self,
         save_path: str,

spatiomic/process/__init__.py

@@ -12,11 +12,14 @@
 from ._register import Register as register
 from ._zscore import ZScore as zscore
 
+standardize = zscore
+
 __all__ = [
     "arcsinh",
     "clip",
     "log1p",
     "normalize",
     "register",
+    "standardize",
     "zscore",
 ]

spatiomic/segment/_extend_mask.py

@@ -1,6 +1,6 @@
 """Mask extension implementation for spatiomic."""
 
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, cast
 
 import numpy as np
 import numpy.typing as npt
@@ -17,53 +17,21 @@ def extend_mask(
 ) -> npt.NDArray[np.integer[Any]]:
     """Extend segmentation masks by dilating them up to halfway to the nearest neighboring mask.
 
-    This function dilates each mask region by a specified number of pixels, but stops dilation
-    at the halfway point to the nearest neighboring mask to ensure fair distribution of space
-    between adjacent regions.
+    Uses Voronoi tessellation via distance transform for O(1) complexity regardless of label count.
+    Each background pixel is assigned to its nearest mask, constrained by dilation_pixels.
 
     Args:
-        masks: Input segmentation masks where each unique integer represents a different segment.
-            Background pixels should have a consistent label (default 0).
-        dilation_pixels: Maximum number of pixels to dilate each mask. The actual dilation
-            may be less if masks meet halfway. Must be positive.
+        masks: 2D input segmentation masks where each unique integer represents a different segment.
+        dilation_pixels: Maximum number of pixels to dilate each mask. Must be positive.
         background_label: The label value representing background pixels. Defaults to 0.
-        use_gpu: Whether to use GPU acceleration with CuPy. Defaults to False.
+        use_gpu: Whether to use GPU acceleration with CuPy/cuCIM. Defaults to False.
 
     Returns:
-        Extended masks with the same shape and dtype as input, where each mask has been
-        dilated up to the halfway point to neighboring masks.
+        Extended masks with the same shape and dtype as input.
 
     Raises:
         ValueError: If dilation_pixels is not positive or if input is not 2D.
-
-    Example:
-        Basic usage for extending cell masks:
-
-        ```python
-        import spatiomic as so
-        import numpy as np
-
-        # Create example segmentation masks
-        masks = np.array([[0, 0, 0, 0, 0], [0, 1, 0, 2, 0], [0, 0, 0, 0, 0], [0, 3, 0, 4, 0], [0, 0, 0, 0, 0]])
-
-        # Extend masks by 1 pixel
-        extended_masks = so.segment.extend_mask(masks, dilation_pixels=1)
-
-        # Each mask will expand towards neighboring masks but stop halfway
-        ```
-
-    Note:
-        - The algorithm uses distance transforms for efficient computation
-        - Memory usage scales with image size and number of unique labels
-        - GPU acceleration is available when CuPy is installed
     """
-    if TYPE_CHECKING or not use_gpu:
-        xp = np
-        ndimage_pkg = ndimage
-    else:
-        xp = import_package("cupy", alternative=np)
-        ndimage_pkg = import_package("cupyx.scipy.ndimage", alternative=ndimage)
-
     if dilation_pixels <= 0:
         msg = f"dilation_pixels must be positive, got {dilation_pixels}"
         raise ValueError(msg)
@@ -72,79 +40,45 @@ def extend_mask(
         msg = f"Input masks must be 2D, got {masks.ndim}D"
         raise ValueError(msg)
 
-    # Convert to GPU array if using CuPy
-    masks_xp = xp.asarray(masks)
+    original_dtype = masks.dtype
 
-    # Get unique labels excluding background
-    unique_labels = xp.unique(masks_xp)
-    unique_labels = unique_labels[unique_labels != background_label]
-
-    if len(unique_labels) == 0:
-        return masks  # No masks to extend
+    if TYPE_CHECKING or not use_gpu:
+        xp = np
+        cucim_morphology = None
+    else:
+        xp = import_package("cupy", alternative=np)
+        cucim_morphology = import_package("cucim.core.operations.morphology", alternative=None)
+        if cucim_morphology is None:
+            use_gpu = False
+            xp = np
 
-    # Create arrays for efficient vectorized computation
-    extended_masks = masks_xp.copy()
+    masks_xp = xp.asarray(masks)
     background_mask = masks_xp == background_label
 
     if not xp.any(background_mask):
-        # No background pixels to extend into
-        if xp.__name__ == "cupy" and hasattr(extended_masks, "get"):
-            return extended_masks.get()  # type: ignore[no-any-return]
-        return np.asarray(extended_masks)  # type: ignore[no-any-return]
-
-    # For efficient computation, we'll use distance transforms
-    # Create distance maps for each label
-    label_distances = xp.full((len(unique_labels), *masks_xp.shape), xp.inf, dtype=xp.float32)
-
-    for idx, label in enumerate(unique_labels):
-        # Create binary mask for this label
-        label_mask = masks_xp == label
-
-        if xp.__name__ == "cupy":
-            # Convert to numpy for scipy operations
-            label_mask_np = label_mask.get() if hasattr(label_mask, "get") else np.asarray(label_mask)
-            dist_transform = ndimage.distance_transform_edt(~label_mask_np)
-            label_distances[idx] = xp.asarray(dist_transform)
-        else:
-            dist_transform = ndimage_pkg.distance_transform_edt(~label_mask)
-            label_distances[idx] = dist_transform
-
-    # For each background pixel, find closest two labels and their distances
-    background_indices = xp.where(background_mask)
-
-    if len(background_indices[0]) > 0:
-        # Vectorized computation for all background pixels
-        distances_at_bg = label_distances[:, background_indices[0], background_indices[1]]
-
-        # Find two closest labels for each background pixel
-        sorted_indices = xp.argsort(distances_at_bg, axis=0)
-        closest_distances = xp.take_along_axis(distances_at_bg, sorted_indices, axis=0)
-        closest_labels = unique_labels[sorted_indices]
-
-        # Calculate conditions for assignment
-        closest_dist = closest_distances[0]
-        within_dilation = closest_dist <= dilation_pixels
-
-        # For pixels with multiple nearby labels, check halfway condition
-        has_second_neighbor = len(unique_labels) > 1
-        if has_second_neighbor:
-            second_closest_dist = closest_distances[1] if len(unique_labels) > 1 else xp.inf
-            halfway_dist = (closest_dist + second_closest_dist) / 2.0
-            within_halfway = closest_dist <= halfway_dist
-            assignment_condition = within_dilation & within_halfway
-        else:
-            assignment_condition = within_dilation
-
-        # Assign pixels to closest labels where conditions are met
-        valid_assignments = xp.where(assignment_condition)[0]
-        if len(valid_assignments) > 0:
-            assigned_rows = background_indices[0][valid_assignments]
-            assigned_cols = background_indices[1][valid_assignments]
-            assigned_labels = closest_labels[0][valid_assignments]
-            extended_masks[assigned_rows, assigned_cols] = assigned_labels
-
-    # Convert back to numpy if we used CuPy
-    if xp.__name__ == "cupy" and hasattr(extended_masks, "get"):
-        return extended_masks.get()  # type: ignore[no-any-return]
-
-    return np.asarray(extended_masks)  # type: ignore[no-any-return]
+        if use_gpu:
+            return cast(npt.NDArray[np.integer[Any]], masks_xp.get().astype(original_dtype))  # type: ignore[attr-defined]
+        return masks.copy()
+
+    foreground_mask = ~background_mask
+
+    if not xp.any(foreground_mask):
+        if use_gpu:
+            return cast(npt.NDArray[np.integer[Any]], masks_xp.get().astype(original_dtype))  # type: ignore[attr-defined]
+        return masks.copy()
+
+    if use_gpu and cucim_morphology is not None:
+        distances, indices = cucim_morphology.distance_transform_edt(background_mask, return_indices=True)
+        extended = masks_xp[indices[0], indices[1]]
+        extended = xp.where(distances <= dilation_pixels, extended, background_label)  # type: ignore[union-attr]
+        extended = xp.where(foreground_mask, masks_xp, extended)  # type: ignore[union-attr]
+        del distances, indices, background_mask, foreground_mask, masks_xp
+        xp.get_default_memory_pool().free_all_blocks()  # type: ignore[union-attr]
+        return cast(npt.NDArray[np.integer[Any]], extended.get().astype(original_dtype))
+
+    distances, indices = ndimage.distance_transform_edt(background_mask, return_indices=True)
+    extended = masks[indices[0], indices[1]]
+    extended = np.where(distances <= dilation_pixels, extended, background_label)
+    extended = np.where(foreground_mask, masks, extended)
+
+    return cast(npt.NDArray[np.integer[Any]], extended.astype(original_dtype))

test/dimension/test_som.py

@@ -79,27 +79,27 @@ def test_som_cpu(example_data: NDArray) -> None:
 
         # TODO: add test cases for flattening and returning distances
 
-    # test quantization error calculation
-    quantization_error = data_som.get_quantization_error(example_data)
-    assert isinstance(quantization_error, float)
+        # test quantization error calculation
+        quantization_error = data_som.get_quantization_error(example_data)
+        assert isinstance(quantization_error, float)
 
-    if distance_metric in ["correlation", "cosine"]:
-        assert quantization_error >= 0.0 and quantization_error <= 1.0, (
-            f"Quantization error out of bounds for {distance_metric} distance metric: {quantization_error}"
-        )
+        if distance_metric in ["correlation", "cosine"]:
+            assert quantization_error >= 0.0 and quantization_error <= 1.0, (
+                f"Quantization error out of bounds for {distance_metric} distance metric: {quantization_error}"
+            )
 
-    # test saving and loading
-    temp_file_name = f"{uuid4()}.p"
-    temp_file_name = os.path.join(os.path.dirname(os.path.realpath(__file__)), temp_file_name)
-    data_som.save(save_path=temp_file_name)
+        # test saving and loading
+        temp_file_name = f"{uuid4()}.p"
+        temp_file_name = os.path.join(os.path.dirname(os.path.realpath(__file__)), temp_file_name)
+        data_som.save(save_path=temp_file_name)
 
-    assert os.path.isfile(temp_file_name)
+        assert os.path.isfile(temp_file_name)
 
-    new_som = so.dimension.som()
-    new_som.load(save_path=temp_file_name)
+        new_som = so.dimension.som()
+        new_som.load(save_path=temp_file_name)
 
-    assert data_som.get_config() == new_som.get_config()
-    assert np.all(data_som.get_nodes() == new_som.get_nodes())
+        assert data_som.get_config() == new_som.get_config()
+        assert np.all(data_som.get_nodes() == new_som.get_nodes())
 
-    # remove the temp file
-    os.remove(temp_file_name)
+        # remove the temp file
+        os.remove(temp_file_name)