add tests

Author: LucaMarconato

docs/api.md

@@ -37,6 +37,7 @@ Operations on `SpatialData` objects.
     transform
     rasterize
     rasterize_bins
+    rasterize_bins_link_table_to_labels
     to_circles
     to_polygons
     aggregate

src/spatialdata/__init__.py

@@ -27,6 +27,7 @@
     "concatenate",
     "rasterize",
     "rasterize_bins",
+    "rasterize_bins_link_table_to_labels",
     "to_circles",
     "to_polygons",
     "transform",
@@ -61,7 +62,7 @@
 from spatialdata._core.operations.aggregate import aggregate
 from spatialdata._core.operations.map import map_raster, relabel_sequential
 from spatialdata._core.operations.rasterize import rasterize
-from spatialdata._core.operations.rasterize_bins import rasterize_bins
+from spatialdata._core.operations.rasterize_bins import rasterize_bins, rasterize_bins_link_table_to_labels
 from spatialdata._core.operations.transform import transform
 from spatialdata._core.operations.vectorize import to_circles, to_polygons
 from spatialdata._core.query._utils import get_bounding_box_corners

src/spatialdata/_core/operations/rasterize_bins.py

@@ -5,6 +5,7 @@
 import dask.array as da
 import numpy as np
 import pandas as pd
+from anndata import AnnData
 from dask.dataframe import DataFrame as DaskDataFrame
 from geopandas import GeoDataFrame
 from numpy.random import default_rng
@@ -21,6 +22,8 @@
 
 RNG = default_rng(0)
 
+__all__ = ["rasterize_bins", "rasterize_bins_link_table_to_labels"]
+
 
 if TYPE_CHECKING:
     from spatialdata import SpatialData
@@ -58,7 +61,8 @@ def rasterize_bins(
         If `False` this function returns a `xarray.DataArray` of shape `(c, y, x)` with dimension
         of `c` equal to the number of key(s) specified in `value_key`, or the number of var names
         in `table_name` if `value_key` is `None`.  If `True`, will return labels of shape `(y, x)`,
-        where each bin of the `bins` element will be represented as a pixel.
+        where each bin of the `bins` element will be represented as a pixel. The table by default will not be set to
+        annotate the new rasterized labels; see the Notes section for how to do this.
 
     Returns
     -------
@@ -78,6 +82,10 @@ def rasterize_bins(
 
     If `spatialdata-plot` is used to visualized the returned image, the parameter `scale='full'` needs to be passed to
     `.render_shapes()`, to disable an automatic rasterization that would confict with the rasterization performed here.
+
+    When `return_region_as_labels` is `True`, the function will return a labels layer that is not annotated by default
+    by the table. To change the annotation target of the table you can call the helper function
+    `spatialdata.rasterize_bins_link_table_to_labels()`.
     """
     element = sdata[bins]
     table = sdata.tables[table_name]
@@ -159,22 +167,12 @@ def rasterize_bins(
         transformations = {cs: to_bins.compose_with(t) for cs, t in bins_transformations.items()}
 
     if return_region_as_labels:
-        dtype = _get_uint_dtype(table.obs[instance_key].max())
-        _min_value = table.obs[instance_key].min()
-        # TODO: add a new column instead of modyfing the table inplace
-        # TODO: do not modify the index of the elements
-        if _min_value == 0:
-            logger.info(
-                f"The minimum value of the instance key column ('table.obs[{instance_key}]') has been"
-                " detected to be 0. Since the label 0 is reserved for the background, "
-                "both the instance key column in 'table.obs' "
-                f"and the index of the annotating element '{bins}' is incremented by 1."
-            )
-            table.obs[instance_key] += 1
-            element.index += 1
+        new_instance_key = _get_relabeled_column_name(instance_key)
+        table.obs[new_instance_key] = _relabel_labels(table=table, instance_key=instance_key)
+        dtype = table.obs[new_instance_key].dtype
         labels_element = np.zeros((n_rows, n_cols), dtype=dtype)
         # make labels layer that can visualy represent the cells
-        labels_element[y, x] = table.obs[instance_key].values.T
+        labels_element[y, x] = table.obs[new_instance_key].values.T
 
         return Labels2DModel.parse(data=labels_element, dims=("y", "x"), transformations=transformations)
 
@@ -242,17 +240,69 @@ def channel_rasterization(block_id: tuple[int, int, int] | None) -> ArrayLike:
     )
 
 
-def _get_uint_dtype(value: int) -> str:
+def _get_uint_dtype(labels_values_count: int) -> str:
     max_uint64 = np.iinfo(np.uint64).max
     max_uint32 = np.iinfo(np.uint32).max
     max_uint16 = np.iinfo(np.uint16).max
 
-    if max_uint16 >= value:
+    if max_uint16 >= labels_values_count:
         dtype = "uint16"
-    elif max_uint32 >= value:
+    elif max_uint32 >= labels_values_count:
         dtype = "uint32"
-    elif max_uint64 >= value:
+    elif max_uint64 >= labels_values_count:
         dtype = "uint64"
     else:
-        raise ValueError(f"Maximum cell number is {value}. Values higher than {max_uint64} are not supported.")
+        raise ValueError(
+            f"Maximum cell number is {labels_values_count}. Values higher than {max_uint64} are not supported."
+        )
     return dtype
+
+
+def _get_relabeled_column_name(column_name: str) -> str:
+    return f"relabeled_{column_name}"
+
+
+def _relabel_labels(table: AnnData, instance_key: str) -> pd.Series:
+    labels_values_count = len(table.obs[instance_key].unique())
+
+    is_not_numeric = not np.issubdtype(table.obs[instance_key].dtype, np.number)
+    zero_in_instance_key = 0 in table.obs[instance_key].values
+    has_gaps = not is_not_numeric and labels_values_count != table.obs[instance_key].max() + int(zero_in_instance_key)
+
+    relabeling_is_needed = is_not_numeric or zero_in_instance_key or has_gaps
+    if relabeling_is_needed:
+        logger.info(
+            f"The instance_key column in 'table.obs' ('table.obs[{instance_key}]') will be relabeled to ensure"
+            " a numeric data type, with a continuous range and without including the value 0 (which is reserved "
+            "for the background). The new labels will be stored in a new column named "
+            f"{_get_relabeled_column_name(instance_key)!r}."
+        )
+
+    relabeled_instance_key_column = table.obs[instance_key].astype("category").cat.codes + int(zero_in_instance_key)
+    # uses only allowed dtypes that passes our model validations, in particualr no uint8
+    dtype = _get_uint_dtype(labels_values_count=relabeled_instance_key_column.max())
+    return relabeled_instance_key_column.astype(dtype)
+
+
+def rasterize_bins_link_table_to_labels(sdata: SpatialData, table_name: str, rasterized_labels_name: str) -> None:
+    """
+    Annotates the table with the rasterized labels.
+
+    This function should be called after rasterizing the bins (with `return_regions_as_labels` is `True`) and adding
+    the rasterized labels in the spatial data object.
+
+    Parameters
+    ----------
+    sdata
+        The spatial data object containing the rasterized labels.
+    table_name
+        The name of the table to be annotated.
+    rasterized_labels_name
+        The name of the rasterized labels in the spatial data object.
+    """
+    _, region_key, instance_key = get_table_keys(sdata[table_name])
+    sdata[table_name].obs[region_key] = rasterized_labels_name
+    relabled_instance_key = _get_relabeled_column_name(instance_key)
+    sdata.set_table_annotates_spatialelement(
+        table_name=table_name, region=rasterized_labels_name, region_key=region_key, instance_key=relabled_instance_key
+    )

src/spatialdata/transformations/_utils.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import warnings
 from functools import singledispatch
 from typing import TYPE_CHECKING, Any, Optional, Union
 
@@ -8,7 +9,6 @@
 from geopandas import GeoDataFrame
 from xarray import DataArray, Dataset, DataTree
 
-from spatialdata._logging import logger
 from spatialdata._types import ArrayLike
 
 if TYPE_CHECKING:
@@ -253,10 +253,12 @@ def scale_radii(radii: ArrayLike, affine: Affine, axes: tuple[str, ...]) -> Arra
     modules = np.absolute(eigenvalues)
     if not np.allclose(modules, modules[0]):
         scale_factor = np.mean(modules)
-        logger.warning(
+        warnings.warn(
             "The vector part of the transformation matrix is not isotropic, the radius will be scaled by the average "
             f"of the modules of eigenvalues of the affine transformation matrix.\nmatrix={matrix}\n"
-            f"eigenvalues={eigenvalues}\nscale_factor={scale_factor}"
+            f"eigenvalues={eigenvalues}\nscale_factor={scale_factor}",
+            UserWarning,
+            stacklevel=2,
         )
     else:
         scale_factor = modules[0]

tests/core/operations/test_rasterize_bins.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import logging
 import re
 
 import numpy as np
@@ -12,8 +13,13 @@
 from shapely.geometry import Polygon
 
 from spatialdata._core.data_extent import are_extents_equal, get_extent
-from spatialdata._core.operations.rasterize_bins import rasterize_bins
+from spatialdata._core.operations.rasterize_bins import (
+    _relabel_labels,
+    rasterize_bins,
+    rasterize_bins_link_table_to_labels,
+)
 from spatialdata._core.spatialdata import SpatialData
+from spatialdata._logging import logger
 from spatialdata._types import ArrayLike
 from spatialdata.models.models import Image2DModel, Labels2DModel, PointsModel, ShapesModel, TableModel
 from spatialdata.transformations.transformations import Scale
@@ -80,18 +86,21 @@ def test_rasterize_bins(geometry: str, value_key: str | list[str] | None, return
     if return_region_as_labels:
         labels_name = "labels"
         sdata[labels_name] = rasterized
-        adata = sdata["table"]
-        adata.obs["region"] = labels_name
-        adata.obs["region"] = adata.obs["region"].astype("category")
-        del adata.uns[TableModel.ATTRS_KEY]
-        adata = TableModel.parse(
-            adata,
-            region=labels_name,
-            region_key="region",
-            instance_key="instance_id",
-        )
-        del sdata["table"]
-        sdata["table"] = adata
+
+        rasterize_bins_link_table_to_labels(sdata=sdata, table_name="table", rasterized_labels_name=labels_name)
+        # adata = sdata["table"]
+        # adata.obs["region"] = labels_name
+        # adata.obs["region"] = adata.obs["region"].astype("category")
+        # del adata.uns[TableModel.ATTRS_KEY]
+        # adata = TableModel.parse(
+        #     adata,
+        #     region=labels_name,
+        #     region_key="region",
+        #     instance_key="instance_id",
+        # )
+        # del sdata["table"]
+        # sdata["table"] = adata
+
         # this fails because table already annotated by labels layer
         with pytest.raises(
             ValueError,
@@ -265,3 +274,38 @@ def _get_sdata(n: int):
             row_key="row_index",
             value_key="instance_id",
         )
+
+
+def test_relabel_labels(caplog):
+    obs = DataFrame(
+        data={
+            "instance_key0": np.arange(1, 11),
+            "instance_key1": np.arange(10),
+            "instance_key2": [1, 2] + list(range(4, 12)),
+            "instance_key3": [str(i) for i in range(1, 11)],
+        }
+    )
+    adata = AnnData(X=RNG.normal(size=(10, 2)), obs=obs)
+    _relabel_labels(table=adata, instance_key="instance_key0")
+    # check logger info message
+    expected_log_message = (
+        "will be relabeled to ensure a numeric data type, with a continuous range and without including the value 0 ("
+        "which is reserved for the background). The new labels will be stored in a new column named"
+    )
+    logger.propagate = True
+    with caplog.at_level(logging.INFO):
+        _relabel_labels(table=adata, instance_key="instance_key1")
+        assert expected_log_message in caplog.text
+
+    with caplog.at_level(logging.INFO):
+        _relabel_labels(table=adata, instance_key="instance_key2")
+        assert expected_log_message in caplog.text
+
+    with caplog.at_level(logging.INFO):
+        _relabel_labels(table=adata, instance_key="instance_key3")
+        assert expected_log_message in caplog.text
+    logger.propagate = False
+
+
+if __name__ == "__main__":
+    test_relabel_labels()