Adjusts spatialdata-plot to use DataArray and DataTree (#277)

* initial adoptation + minor fixes

* adjust to new locate values

* fix test

* fix test

Author: melonora

src/spatialdata_plot/pl/basic.py

@@ -14,14 +14,14 @@
 import spatialdata as sd
 from anndata import AnnData
 from dask.dataframe.core import DataFrame as DaskDataFrame
+from datatree import DataTree
 from geopandas import GeoDataFrame
 from matplotlib.axes import Axes
 from matplotlib.colors import Colormap, Normalize
 from matplotlib.figure import Figure
-from multiscale_spatial_image.multiscale_spatial_image import MultiscaleSpatialImage
-from spatial_image import SpatialImage
 from spatialdata._core.data_extent import get_extent
-from spatialdata._utils import deprecation_alias
+from spatialdata._utils import _deprecation_alias
+from xarray import DataArray
 
 from spatialdata_plot._accessor import register_spatial_data_accessor
 from spatialdata_plot.pl.render import (
@@ -95,8 +95,8 @@ def __init__(self, sdata: sd.SpatialData) -> None:
 
     def _copy(
         self,
-        images: dict[str, SpatialImage | MultiscaleSpatialImage] | None = None,
-        labels: dict[str, SpatialImage | MultiscaleSpatialImage] | None = None,
+        images: dict[str, DataArray | DataTree] | None = None,
+        labels: dict[str, DataArray | DataTree] | None = None,
         points: dict[str, DaskDataFrame] | None = None,
         shapes: dict[str, GeoDataFrame] | None = None,
         tables: dict[str, AnnData] | None = None,
@@ -105,11 +105,11 @@ def _copy(
 
         Parameters
         ----------
-        images : dict[str, SpatialImage | MultiscaleSpatialImage] | None, optional
+        images : dict[str, DataArray | DataTree] | None, optional
             A dictionary containing image data to replace the images in the
             original `SpatialData` object, or `None` to keep the original
             images. Defaults to `None`.
-        labels : dict[str, SpatialImage | MultiscaleSpatialImage] | None, optional
+        labels : dict[str, DataArray | DataTree] | None, optional
             A dictionary containing label data to replace the labels in the
             original `SpatialData` object, or `None` to keep the original
             labels. Defaults to `None`.
@@ -150,7 +150,7 @@ def _copy(
 
         return sdata
 
-    @deprecation_alias(elements="element", version="0.3.0")
+    @_deprecation_alias(elements="element", version="0.3.0")
     def render_shapes(
         self,
         element: str | None = None,
@@ -286,7 +286,7 @@ def render_shapes(
 
         return sdata
 
-    @deprecation_alias(elements="element", version="0.3.0")
+    @_deprecation_alias(elements="element", version="0.3.0")
     def render_points(
         self,
         element: str | None = None,
@@ -396,7 +396,7 @@ def render_points(
 
         return sdata
 
-    @deprecation_alias(elements="element", quantiles_for_norm="percentiles_for_norm", version="version 0.3.0")
+    @_deprecation_alias(elements="element", quantiles_for_norm="percentiles_for_norm", version="version 0.3.0")
     def render_images(
         self,
         element: str | None = None,
@@ -509,7 +509,7 @@ def render_images(
 
         return sdata
 
-    @deprecation_alias(elements="element", version="0.3.0")
+    @_deprecation_alias(elements="element", version="0.3.0")
     def render_labels(
         self,
         element: str | None = None,

src/spatialdata_plot/pl/render.py

@@ -14,8 +14,8 @@
 import scanpy as sc
 import spatialdata as sd
 from anndata import AnnData
+from datatree import DataTree
 from matplotlib.colors import ListedColormap, Normalize
-from multiscale_spatial_image.multiscale_spatial_image import MultiscaleSpatialImage
 from scanpy._settings import settings as sc_settings
 from spatialdata._core.data_extent import get_extent
 from spatialdata.models import PointsModel, get_table_keys
@@ -244,7 +244,10 @@ def _render_points(
         )
     else:
         adata = AnnData(
-            X=points[["x", "y"]].values, obs=sdata_filt[table_name].obs, dtype=points[["x", "y"]].values.dtype
+            X=points[["x", "y"]].values,
+            obs=sdata_filt[table_name].obs,
+            dtype=points[["x", "y"]].values.dtype,
+            uns=sdata_filt[table_name].uns,
         )
         sdata_filt[table_name] = adata
 
@@ -351,7 +354,7 @@ def _render_images(
     scale = render_params.scale
 
     # get best scale out of multiscale image
-    if isinstance(img, MultiscaleSpatialImage):
+    if isinstance(img, DataTree):
         img = _multiscale_to_spatial_image(
             multiscale_image=img,
             dpi=fig_params.fig.dpi,
@@ -541,7 +544,7 @@ def _render_labels(
     extent = get_extent(label, coordinate_system=coordinate_system)
 
     # get best scale out of multiscale label
-    if isinstance(label, MultiscaleSpatialImage):
+    if isinstance(label, DataTree):
         label = _multiscale_to_spatial_image(
             multiscale_image=label,
             dpi=fig_params.fig.dpi,

src/spatialdata_plot/pl/utils.py

@@ -18,11 +18,11 @@
 import numpy as np
 import pandas as pd
 import shapely
-import spatial_image
 import spatialdata as sd
 import xarray as xr
 from anndata import AnnData
 from cycler import Cycler, cycler
+from datatree import DataTree
 from geopandas import GeoDataFrame
 from matplotlib import colors, patheffects, rcParams
 from matplotlib.axes import Axes
@@ -40,7 +40,6 @@
 from matplotlib.gridspec import GridSpec
 from matplotlib.transforms import CompositeGenericTransform
 from matplotlib_scalebar.scalebar import ScaleBar
-from multiscale_spatial_image.multiscale_spatial_image import MultiscaleSpatialImage
 from numpy.ma.core import MaskedArray
 from numpy.random import default_rng
 from pandas.api.types import CategoricalDtype
@@ -52,13 +51,13 @@
 from skimage.morphology import erosion, square
 from skimage.segmentation import find_boundaries
 from skimage.util import map_array
-from spatial_image import SpatialImage
 from spatialdata import SpatialData
 from spatialdata._core.operations.rasterize import rasterize
 from spatialdata._core.query.relational_query import _get_element_annotators, _locate_value, _ValueOrigin, get_values
 from spatialdata._types import ArrayLike
 from spatialdata.models import Image2DModel, Labels2DModel, PointsModel, SpatialElement, get_model
 from spatialdata.transformations.operations import get_transformation
+from xarray import DataArray
 
 from spatialdata_plot._logging import logger
 from spatialdata_plot.pl.render_params import (
@@ -609,18 +608,6 @@ def _get_colors_for_categorical_obs(
     return palette[:len_cat]  # type: ignore[return-value]
 
 
-def _locate_points_value_in_table(value_key: str, sdata: SpatialData, table_name: str) -> _ValueOrigin:
-    table = sdata[table_name]
-
-    if value_key in table.obs.columns:
-        value = table.obs[value_key]
-        is_categorical = isinstance(value.dtype, CategoricalDtype)
-        return _ValueOrigin(origin="obs", is_categorical=is_categorical, value_key=value_key)
-
-    is_categorical = False
-    return _ValueOrigin(origin="var", is_categorical=is_categorical, value_key=value_key)
-
-
 # TODO consider move to relational query in spatialdata
 def get_values_point_table(sdata: SpatialData, origin: _ValueOrigin, table_name: str) -> pd.Series:
     """Get a particular column stored in _ValueOrigin from the table in the spatialdata object."""
@@ -651,10 +638,6 @@ def _set_color_source_vec(
 
     # Figure out where to get the color from
     origins = _locate_value(value_key=value_to_plot, sdata=sdata, element_name=element_name, table_name=table_name)
-    if model == PointsModel and table_name is not None:
-        origin = _locate_points_value_in_table(value_key=value_to_plot, sdata=sdata, table_name=table_name)
-        if origin is not None:
-            origins.append(origin)
 
     if len(origins) > 1:
         raise ValueError(
@@ -663,7 +646,7 @@ def _set_color_source_vec(
 
     if len(origins) == 1:
         if model == PointsModel and table_name is not None:
-            color_source_vector = get_values_point_table(sdata=sdata, origin=origin, table_name=table_name)
+            color_source_vector = get_values_point_table(sdata=sdata, origin=origins[0], table_name=table_name)
         else:
             vals = get_values(value_key=value_to_plot, sdata=sdata, element_name=element_name, table_name=table_name)
             color_source_vector = vals[value_to_plot]
@@ -765,8 +748,10 @@ def _map_color_seg(
         cols = colors.to_rgba_array(color_vector.categories)
 
     else:
-        val_im = map_array(seg, cell_id, cell_id)  # replace with same seg id to remove missing segs
+        val_im = map_array(seg.copy(), cell_id, cell_id)  # replace with same seg id to remove missing segs
 
+        if val_im.shape[0] == 1:
+            val_im = np.squeeze(val_im, axis=0)
         try:
             cols = cmap_params.cmap(cmap_params.norm(color_vector))
         except TypeError:
@@ -793,7 +778,9 @@ def _map_color_seg(
         seg_bound: ArrayLike = np.clip(seg_im - find_boundaries(seg)[:, :, None], 0, 1)
         return np.dstack((seg_bound, np.where(val_im > 0, 1, 0)))  # add transparency here
 
-    return np.dstack((seg_im, np.where(val_im > 0, 1, 0)))
+    if len(val_im.shape) != len(seg_im.shape):
+        val_im = np.expand_dims((val_im > 0).astype(int), axis=-1)
+    return np.dstack((seg_im, val_im))
 
 
 def _get_palette(
@@ -1029,7 +1016,7 @@ def _multiscale_to_image(sdata: sd.SpatialData) -> sd.SpatialData:
         raise ValueError("No images found in the SpatialData object.")
 
     for k, v in sdata.images.items():
-        if isinstance(v, msi.multiscale_spatial_image.MultiscaleSpatialImage):
+        if isinstance(v, msi.multiscale_spatial_image.DataTree):
             sdata.images[k] = Image2DModel.parse(v["scale0"].ds.to_array().squeeze(axis=0))
 
     return sdata
@@ -1047,9 +1034,9 @@ def _get_listed_colormap(color_dict: dict[str, str]) -> ListedColormap:
 
 
 def _translate_image(
-    image: spatial_image.SpatialImage,
+    image: DataArray,
     translation: sd.transformations.transformations.Translation,
-) -> spatial_image.SpatialImage:
+) -> DataArray:
     shifts: dict[str, int] = {axis: int(translation.translation[idx]) for idx, axis in enumerate(translation.axes)}
     img = image.values.copy()
     # for yx images (important for rasterized MultiscaleImages as labels)
@@ -1201,16 +1188,16 @@ def _get_valid_cs(
 
 
 def _rasterize_if_necessary(
-    image: SpatialImage,
+    image: DataArray,
     dpi: float,
     width: float,
     height: float,
     coordinate_system: str,
     extent: dict[str, tuple[float, float]],
-) -> SpatialImage:
+) -> DataArray:
     """Ensure fast rendering by adapting the resolution if necessary.
 
-    A SpatialImage is prepared for plotting. To improve performance, large images are rasterized.
+    A DataArray is prepared for plotting. To improve performance, large images are rasterized.
 
     Parameters
     ----------
@@ -1230,7 +1217,7 @@ def _rasterize_if_necessary(
 
     Returns
     -------
-    SpatialImage
+    DataArray
         Spatial image ready for rendering
     """
     has_c_dim = len(image.shape) == 3
@@ -1265,22 +1252,22 @@ def _rasterize_if_necessary(
 
 
 def _multiscale_to_spatial_image(
-    multiscale_image: MultiscaleSpatialImage,
+    multiscale_image: DataTree,
     dpi: float,
     width: float,
     height: float,
     scale: str | None = None,
     is_label: bool = False,
-) -> SpatialImage:
-    """Extract the SpatialImage to be rendered from a multiscale image.
+) -> DataArray:
+    """Extract the DataArray to be rendered from a multiscale image.
 
-    From the `MultiscaleSpatialImage`, the scale that fits the given image size and dpi most is selected
+    From the `DataTree`, the scale that fits the given image size and dpi most is selected
     and returned. In case the lowest resolution is still too high, a rasterization step is added.
 
     Parameters
     ----------
     multiscale_image
-        `MultiscaleSpatialImage` that should be rendered
+        `DataTree` that should be rendered
     dpi
         dpi of the target image
     width
@@ -1294,8 +1281,8 @@ def _multiscale_to_spatial_image(
 
     Returns
     -------
-    SpatialImage
-        To be rendered, extracted from the MultiscaleSpatialImage respecting the dpi and size of the target image.
+    DataArray
+        To be rendered, extracted from the DataTree respecting the dpi and size of the target image.
     """
     scales = [leaf.name for leaf in multiscale_image.leaves]
     x_dims = [multiscale_image[scale].dims["x"] for scale in scales]
@@ -1879,7 +1866,7 @@ def _validate_image_render_params(
         spatial_element = param_dict["sdata"][el]
 
         spatial_element_ch = (
-            spatial_element.c if isinstance(spatial_element, SpatialImage) else spatial_element["scale0"].c
+            spatial_element.c if isinstance(spatial_element, DataArray) else spatial_element["scale0"].c
         )
         if (channel := param_dict["channel"]) is not None and (
             (isinstance(channel[0], int) and max([abs(ch) for ch in channel]) <= len(spatial_element_ch))
@@ -1908,7 +1895,7 @@ def _validate_image_render_params(
                 cmap = None
         element_params[el]["cmap"] = cmap
         element_params[el]["norm"] = param_dict["norm"]
-        if (scale := param_dict["scale"]) and isinstance(sdata[el], MultiscaleSpatialImage):
+        if (scale := param_dict["scale"]) and isinstance(sdata[el], DataTree):
             if scale not in list(sdata[el].keys()) and scale != "full":
                 element_params[el]["scale"] = None
             else:

src/spatialdata_plot/pp/basic.py

@@ -4,10 +4,10 @@
 import spatialdata as sd
 from anndata import AnnData
 from dask.dataframe.core import DataFrame as DaskDataFrame
+from datatree import DataTree
 from geopandas import GeoDataFrame
-from multiscale_spatial_image.multiscale_spatial_image import MultiscaleSpatialImage
-from spatial_image import SpatialImage
 from spatialdata.models import get_table_keys
+from xarray import DataArray
 
 from spatialdata_plot._accessor import register_spatial_data_accessor
 from spatialdata_plot.pp.utils import (
@@ -43,8 +43,8 @@ def __init__(self, sdata: sd.SpatialData) -> None:
 
     def _copy(
         self,
-        images: Union[None, dict[str, Union[SpatialImage, MultiscaleSpatialImage]]] = None,
-        labels: Union[None, dict[str, Union[SpatialImage, MultiscaleSpatialImage]]] = None,
+        images: Union[None, dict[str, Union[DataArray, DataTree]]] = None,
+        labels: Union[None, dict[str, Union[DataArray, DataTree]]] = None,
         points: Union[None, dict[str, DaskDataFrame]] = None,
         shapes: Union[None, dict[str, GeoDataFrame]] = None,
         tables: Union[None, dict[str, AnnData]] = None,

tests/conftest.py

@@ -11,11 +11,10 @@
 import spatialdata as sd
 import spatialdata_plot  # noqa: F401
 from anndata import AnnData
+from datatree import DataTree
 from geopandas import GeoDataFrame
 from matplotlib.testing.compare import compare_images
-from multiscale_spatial_image import MultiscaleSpatialImage
 from shapely.geometry import MultiPolygon, Polygon
-from spatial_image import SpatialImage
 from spatialdata import SpatialData
 from spatialdata.datasets import blobs, raccoon
 from spatialdata.models import (
@@ -217,7 +216,7 @@ def sdata(request) -> SpatialData:
     return s
 
 
-def _get_images() -> dict[str, Union[SpatialImage, MultiscaleSpatialImage]]:
+def _get_images() -> dict[str, Union[DataArray, DataTree]]:
     out = {}
     dims_2d = ("c", "y", "x")
     dims_3d = ("z", "y", "x", "c")
@@ -244,7 +243,7 @@ def _get_images() -> dict[str, Union[SpatialImage, MultiscaleSpatialImage]]:
     return out
 
 
-def _get_labels() -> dict[str, Union[SpatialImage, MultiscaleSpatialImage]]:
+def _get_labels() -> dict[str, Union[DataArray, DataTree]]:
     out = {}
     dims_2d = ("y", "x")
     dims_3d = ("z", "y", "x")