adjust to xarray datatree (scverse#752)

* adjust to xarray datatree

* add decorator skipping nodes without dimensions

* update docs

* move skip_non_dimension_nodes -> multiscale-spatial-image v2.0.1

* update multiscale-spatial-image dependency

* update python version

* adjust test CI

* update black python version

* adjust mypy python 3.10

* fix tests

* further pre-commit fixes

* more fixes

* fix callable import

* fixes

* last fixes

.github/workflows/test.yaml

@@ -18,11 +18,11 @@ jobs:
         strategy:
             fail-fast: false
             matrix:
-                python: ["3.9", "3.12"]
+                python: ["3.10", "3.12"]
                 os: [ubuntu-latest]
                 include:
                     - os: macos-latest
-                      python: "3.9"
+                      python: "3.10"
                     - os: macos-latest
                       python: "3.12"
                       pip-flags: "--pre"

.mypy.ini

@@ -1,5 +1,5 @@
 [mypy]
-python_version = 3.9
+python_version = 3.10
 plugins = numpy.typing.mypy_plugin
 
 ignore_errors = False

pyproject.toml

@@ -15,7 +15,7 @@ maintainers = [
 urls.Documentation = "https://spatialdata.scverse.org/en/latest"
 urls.Source = "https://github.com/scverse/spatialdata.git"
 urls.Home-page = "https://github.com/scverse/spatialdata.git"
-requires-python = ">=3.9"
+requires-python = ">=3.10, <3.13" # include 3.13 once multiscale-spatial-image conflicts are resolved
 dynamic= [
   "version" # allow version to be set by git tags
 ]
@@ -28,7 +28,7 @@ dependencies = [
     "dask>=2024.4.1",
     "fsspec<=2023.6",
     "geopandas>=0.14",
-    "multiscale_spatial_image>=1.0.0,<2.0.0",
+    "multiscale_spatial_image>=2.0.1",
     "networkx",
     "numba",
     "numpy",
@@ -43,8 +43,7 @@ dependencies = [
     "scikit-image",
     "scipy",
     "typing_extensions>=4.8.0",
-    "xarray<2024.10.0",
-    "xarray-datatree",
+    "xarray>=2024.10.0",
     "xarray-schema",
     "xarray-spatial>=0.3.5",
     "zarr",
@@ -102,7 +101,7 @@ filterwarnings = [
 
 [tool.black]
 line-length = 120
-target-version = ['py39']
+target-version = ['py310']
 include = '\.pyi?$'
 exclude = '''
 (
@@ -147,7 +146,7 @@ exclude = [
     "setup.py",
 ]
 line-length = 120
-target-version = "py39"
+target-version = "py310"
 
 [tool.ruff.lint]
 ignore = [

src/spatialdata/_core/_deepcopy.py

@@ -7,9 +7,8 @@
 from dask.array.core import Array as DaskArray
 from dask.array.core import from_array
 from dask.dataframe import DataFrame as DaskDataFrame
-from datatree import DataTree
 from geopandas import GeoDataFrame
-from xarray import DataArray
+from xarray import DataArray, DataTree
 
 from spatialdata._core.spatialdata import SpatialData
 from spatialdata.models._utils import SpatialElement

src/spatialdata/_core/centroids.py

@@ -7,10 +7,9 @@
 import pandas as pd
 import xarray as xr
 from dask.dataframe import DataFrame as DaskDataFrame
-from datatree import DataTree
 from geopandas import GeoDataFrame
 from shapely import MultiPolygon, Point, Polygon
-from xarray import DataArray
+from xarray import DataArray, DataTree
 
 from spatialdata._core.operations.transform import transform
 from spatialdata.models import get_axes_names
@@ -86,7 +85,7 @@ def _get_centroids_for_axis(xdata: xr.DataArray, axis: str) -> pd.DataFrame:
             centroids[label_value] += count * i.values.item()
 
     all_labels_values, all_labels_counts = da.unique(xdata.data, return_counts=True)
-    all_labels = dict(zip(all_labels_values.compute(), all_labels_counts.compute()))
+    all_labels = dict(zip(all_labels_values.compute(), all_labels_counts.compute(), strict=True))
     for label_value in centroids:
         centroids[label_value] /= all_labels[label_value]
     centroids = dict(sorted(centroids.items(), key=lambda x: x[0]))
@@ -132,7 +131,7 @@ def _(e: GeoDataFrame, coordinate_system: str = "global") -> DaskDataFrame:
     if isinstance(first_geometry, Point):
         xy = e.geometry.get_coordinates().values
     else:
-        assert isinstance(first_geometry, (Polygon, MultiPolygon)), (
+        assert isinstance(first_geometry, Polygon | MultiPolygon), (
             f"Expected a GeoDataFrame either composed entirely of circles (Points with the `radius` column) or"
             f" Polygons/MultiPolygons. Found {type(first_geometry)} instead."
         )

src/spatialdata/_core/concatenate.py

@@ -51,7 +51,7 @@ def _concatenate_tables(
             raise ValueError("`instance_key` must be specified if tables have different instance keys")
 
     tables_l = []
-    for table_region_key, table_instance_key, table in zip(region_keys, instance_keys, tables):
+    for table_region_key, table_instance_key, table in zip(region_keys, instance_keys, tables, strict=True):
         rename_dict = {}
         if table_region_key != region_key:
             rename_dict[table_region_key] = region_key
@@ -239,7 +239,7 @@ def _fix_ensure_unique_element_names(
             tables[new_name] = table
         tables_by_sdata.append(tables)
     sdatas_fixed = []
-    for elements, tables in zip(elements_by_sdata, tables_by_sdata):
+    for elements, tables in zip(elements_by_sdata, tables_by_sdata, strict=True):
         sdata = SpatialData.init_from_elements(elements, tables=tables)
         sdatas_fixed.append(sdata)
     return sdatas_fixed

src/spatialdata/_core/data_extent.py

@@ -8,10 +8,9 @@
 import numpy as np
 import pandas as pd
 from dask.dataframe import DataFrame as DaskDataFrame
-from datatree import DataTree
 from geopandas import GeoDataFrame
 from shapely import MultiPolygon, Point, Polygon
-from xarray import DataArray
+from xarray import DataArray, DataTree
 
 from spatialdata._core.operations.transform import transform
 from spatialdata._core.spatialdata import SpatialData
@@ -71,7 +70,7 @@ def _get_extent_of_polygons_multipolygons(
     -------
     The bounding box description.
     """
-    assert isinstance(shapes.geometry.iloc[0], (Polygon, MultiPolygon))
+    assert isinstance(shapes.geometry.iloc[0], Polygon | MultiPolygon)
     axes = get_axes_names(shapes)
     bounds = shapes["geometry"].bounds
     return {ax: (bounds[f"min{ax}"].min(), bounds[f"max{ax}"].max()) for ax in axes}
@@ -201,7 +200,7 @@ def _(
     new_max_coordinates_dict: dict[str, list[float]] = defaultdict(list)
     mask = [has_images, has_labels, has_points, has_shapes]
     include_spatial_elements = ["images", "labels", "points", "shapes"]
-    include_spatial_elements = [i for (i, v) in zip(include_spatial_elements, mask) if v]
+    include_spatial_elements = [i for (i, v) in zip(include_spatial_elements, mask, strict=True) if v]
 
     if elements is None:  # to shut up ruff
         elements = []
@@ -217,7 +216,7 @@ def _(
             assert isinstance(transformations, dict)
             coordinate_systems = list(transformations.keys())
             if coordinate_system in coordinate_systems:
-                if isinstance(element_obj, (DaskDataFrame, GeoDataFrame)):
+                if isinstance(element_obj, DaskDataFrame | GeoDataFrame):
                     extent = get_extent(element_obj, coordinate_system=coordinate_system, exact=exact)
                 else:
                     extent = get_extent(element_obj, coordinate_system=coordinate_system)
@@ -254,7 +253,7 @@ def _get_extent_of_shapes(e: GeoDataFrame) -> BoundingBoxDescription:
     first_geometry = e_temp["geometry"].iloc[0]
     if isinstance(first_geometry, Point):
         return _get_extent_of_circles(e)
-    assert isinstance(first_geometry, (Polygon, MultiPolygon))
+    assert isinstance(first_geometry, Polygon | MultiPolygon)
     return _get_extent_of_polygons_multipolygons(e)
 
 

src/spatialdata/_core/operations/_utils.py

@@ -2,8 +2,7 @@
 
 from typing import TYPE_CHECKING
 
-from datatree import DataTree
-from xarray import DataArray
+from xarray import DataArray, DataTree
 
 from spatialdata.models import SpatialElement
 
@@ -115,7 +114,7 @@ def transform_to_data_extent(
     }
 
     for _, element_name, element in sdata_raster.gen_spatial_elements():
-        if isinstance(element, (DataArray, DataTree)):
+        if isinstance(element, DataArray | DataTree):
             rasterized = rasterize(
                 element,
                 axes=data_extent_axes,

src/spatialdata/_core/operations/aggregate.py

@@ -10,11 +10,10 @@
 import numpy as np
 import pandas as pd
 from dask.dataframe import DataFrame as DaskDataFrame
-from datatree import DataTree
 from geopandas import GeoDataFrame
 from scipy import sparse
 from shapely import Point
-from xarray import DataArray
+from xarray import DataArray, DataTree
 from xrspatial import zonal_stats
 
 from spatialdata._core.operations._utils import _parse_element
@@ -246,7 +245,7 @@ def _create_sdata_from_table_and_shapes(
     table = TableModel.parse(table, region=shapes_name, region_key=region_key, instance_key=instance_key)
 
     # labels case, needs conversion from str to int
-    if isinstance(shapes, (DataArray, DataTree)):
+    if isinstance(shapes, DataArray | DataTree):
         table.obs[instance_key] = table.obs[instance_key].astype(int)
 
     if deepcopy:

src/spatialdata/_core/operations/map.py

@@ -1,13 +1,12 @@
 from __future__ import annotations
 
-from collections.abc import Iterable, Mapping
+from collections.abc import Callable, Iterable, Mapping
 from types import MappingProxyType
-from typing import TYPE_CHECKING, Any, Callable
+from typing import TYPE_CHECKING, Any
 
 import dask.array as da
 from dask.array.overlap import coerce_depth
-from datatree import DataTree
-from xarray import DataArray
+from xarray import DataArray, DataTree
 
 from spatialdata.models._utils import get_axes_names, get_channels, get_raster_model_from_data_dims
 from spatialdata.transformations import get_transformation

src/spatialdata/_core/operations/rasterize.py

@@ -5,10 +5,9 @@
 import numpy as np
 from dask.array import Array as DaskArray
 from dask.dataframe import DataFrame as DaskDataFrame
-from datatree import DataTree
 from geopandas import GeoDataFrame
 from shapely import Point
-from xarray import DataArray
+from xarray import DataArray, DataTree
 
 from spatialdata._core.operations._utils import _parse_element
 from spatialdata._core.operations.transform import transform

src/spatialdata/_core/operations/rasterize_bins.py

@@ -73,7 +73,7 @@ def rasterize_bins(
     """
     element = sdata[bins]
     table = sdata.tables[table_name]
-    if not isinstance(element, (GeoDataFrame, DaskDataFrame)):
+    if not isinstance(element, GeoDataFrame | DaskDataFrame):
         raise ValueError("The bins should be a GeoDataFrame or a DaskDataFrame.")
 
     _, region_key, instance_key = get_table_keys(table)
@@ -94,7 +94,7 @@ def rasterize_bins(
     keys = ([value_key] if isinstance(value_key, str) else value_key) if value_key is not None else table.var_names
 
     if (value_key is None or any(key in table.var_names for key in keys)) and not isinstance(
-        table.X, (csc_matrix, np.ndarray)
+        table.X, csc_matrix | np.ndarray
     ):
         raise ValueError(
             "To speed up bins rasterization, the X matrix in the table, when sparse, should be a csc_matrix matrix. "
@@ -162,7 +162,7 @@ def channel_rasterization(block_id: tuple[int, int, int] | None) -> ArrayLike:
             sub_x = sub_df.geometry.x.values
             sub_y = sub_df.geometry.y.values
         else:
-            assert isinstance(sub_df.iloc[0].geometry, (Polygon, MultiPolygon))
+            assert isinstance(sub_df.iloc[0].geometry, Polygon | MultiPolygon)
             sub_x = sub_df.centroid.x
             sub_y = sub_df.centroid.y
     else:

src/spatialdata/_core/operations/transform.py

@@ -10,10 +10,9 @@
 import numpy as np
 from dask.array.core import Array as DaskArray
 from dask.dataframe import DataFrame as DaskDataFrame
-from datatree import DataTree
 from geopandas import GeoDataFrame
 from shapely import Point
-from xarray import DataArray
+from xarray import DataArray, Dataset, DataTree
 
 from spatialdata._core.spatialdata import SpatialData
 from spatialdata._types import ArrayLike
@@ -161,13 +160,13 @@ def _set_transformation_for_transformed_elements(
         assert to_coordinate_system is None
 
     to_prepend: BaseTransformation | None
-    if isinstance(element, (DataArray, DataTree)):
+    if isinstance(element, DataArray | DataTree):
         if maintain_positioning:
             assert raster_translation is not None
             to_prepend = Sequence([raster_translation, transformation.inverse()])
         else:
             to_prepend = raster_translation
-    elif isinstance(element, (GeoDataFrame, DaskDataFrame)):
+    elif isinstance(element, GeoDataFrame | DaskDataFrame):
         assert raster_translation is None
         to_prepend = transformation.inverse() if maintain_positioning else Identity()
     else:
@@ -393,8 +392,11 @@ def _(
             raster_translation = raster_translation_single_scale
         # we set a dummy empty dict for the transformation that will be replaced with the correct transformation for
         # each scale later in this function, when calling set_transformation()
-        transformed_dict[k] = DataArray(transformed_dask, dims=xdata.dims, name=xdata.name, attrs={TRANSFORM_KEY: {}})
+        transformed_dict[k] = Dataset(
+            {"image": DataArray(transformed_dask, dims=xdata.dims, name=xdata.name, attrs={TRANSFORM_KEY: {}})}
+        )
         if channel_names is not None:
+            # This expression returns a dataset now.
             transformed_dict[k] = transformed_dict[k].assign_coords(c=channel_names)
 
     # mypy thinks that schema could be ShapesModel, PointsModel, ...

src/spatialdata/_core/operations/vectorize.py

@@ -9,11 +9,10 @@
 import shapely
 import skimage.measure
 from dask.dataframe import DataFrame as DaskDataFrame
-from datatree import DataTree
 from geopandas import GeoDataFrame
 from shapely import MultiPolygon, Point, Polygon
 from skimage.measure._regionprops import RegionProperties
-from xarray import DataArray
+from xarray import DataArray, DataTree
 
 from spatialdata._core.centroids import get_centroids
 from spatialdata._core.operations.aggregate import aggregate
@@ -106,7 +105,7 @@ def _(element: DataArray | DataTree, **kwargs: Any) -> GeoDataFrame:
 @to_circles.register(GeoDataFrame)
 def _(element: GeoDataFrame, **kwargs: Any) -> GeoDataFrame:
     assert len(kwargs) == 0
-    if isinstance(element.geometry.iloc[0], (Polygon, MultiPolygon)):
+    if isinstance(element.geometry.iloc[0], Polygon | MultiPolygon):
         radius = np.sqrt(element.geometry.area / np.pi)
         centroids = _get_centroids(element)
         obs = pd.DataFrame({"radius": radius})
@@ -258,7 +257,7 @@ def _dissolve_on_overlaps(label: int, group: GeoDataFrame) -> GeoDataFrame:
 
 @to_polygons.register(GeoDataFrame)
 def _(gdf: GeoDataFrame, buffer_resolution: int = 16) -> GeoDataFrame:
-    if isinstance(gdf.geometry.iloc[0], (Polygon, MultiPolygon)):
+    if isinstance(gdf.geometry.iloc[0], Polygon | MultiPolygon):
         return gdf
     if isinstance(gdf.geometry.iloc[0], Point):
         ShapesModel.validate_shapes_not_mixed_types(gdf)
@@ -274,7 +273,7 @@ def _(gdf: GeoDataFrame, buffer_resolution: int = 16) -> GeoDataFrame:
             # TODO replace with a function to copy the metadata (the parser could also do this): https://github.com/scverse/spatialdata/issues/258
             buffered_df.attrs[ShapesModel.TRANSFORM_KEY] = gdf.attrs[ShapesModel.TRANSFORM_KEY]
             return buffered_df
-        assert isinstance(gdf.geometry.iloc[0], (Polygon, MultiPolygon))
+        assert isinstance(gdf.geometry.iloc[0], Polygon | MultiPolygon)
         return gdf
     raise RuntimeError("Unsupported geometry type: " f"{type(gdf.geometry.iloc[0])}")
 

src/spatialdata/_core/query/_utils.py

@@ -5,8 +5,7 @@
 import numba as nb
 import numpy as np
 from anndata import AnnData
-from datatree import DataTree
-from xarray import DataArray
+from xarray import DataArray, Dataset, DataTree
 
 from spatialdata._core._elements import Tables
 from spatialdata._core.spatialdata import SpatialData
@@ -138,7 +137,7 @@ def _process_data_tree_query_result(query_result: DataTree) -> DataTree | None:
     if len(scales_to_keep) == 0:
         return None
 
-    d = {k: d[k] for k in scales_to_keep}
+    d = {k: Dataset({"image": d[k]}) for k in scales_to_keep}
     result = DataTree.from_dict(d)
 
     # Rechunk the data to avoid irregular chunks

src/spatialdata/_core/query/relational_query.py

@@ -13,9 +13,8 @@
 import pandas as pd
 from anndata import AnnData
 from dask.dataframe import DataFrame as DaskDataFrame
-from datatree import DataTree
 from geopandas import GeoDataFrame
-from xarray import DataArray
+from xarray import DataArray, DataTree
 
 from spatialdata._core.spatialdata import SpatialData
 from spatialdata._types import ArrayLike
@@ -653,7 +652,7 @@ def join_spatialelement_table(
     if sdata is not None:
         elements_dict = _create_sdata_elements_dict_for_join(sdata, spatial_element_names)
     else:
-        derived_sdata = SpatialData.from_elements_dict(dict(zip(spatial_element_names, spatial_elements)))
+        derived_sdata = SpatialData.from_elements_dict(dict(zip(spatial_element_names, spatial_elements, strict=True)))
         element_types = ["labels", "shapes", "points"]
         elements_dict = defaultdict(lambda: defaultdict(dict))
         for element_type in element_types:
@@ -919,7 +918,7 @@ def get_values(
             x = matched_table[:, value_key_values].X
             import scipy
 
-            if isinstance(x, (scipy.sparse.csr_matrix, scipy.sparse.csc_matrix, scipy.sparse.coo_matrix)):
+            if isinstance(x, scipy.sparse.csr_matrix | scipy.sparse.csc_matrix | scipy.sparse.coo_matrix):
                 x = x.todense()
             df = pd.DataFrame(x, columns=value_key_values)
         if origin == "obsm":