adjust to xarray datatree

melonora · melonora · commit abd3a2fb65f6 · 2024-11-01T18:16:57.000+01:00
diff --git a/docs/tutorials/notebooks b/docs/tutorials/notebooks
@@ -1 +1 @@
-Subproject commit 9281cb127ccd849a14daf83fe79c7cb21613c662
+Subproject commit 00a9c79a87d012b8666c762946c553c659d80ac1
diff --git a/pyproject.toml b/pyproject.toml
@@ -28,7 +28,7 @@ dependencies = [
     "dask>=2024.4.1",
     "fsspec<=2023.6",
     "geopandas>=0.14",
-    "multiscale_spatial_image>=1.0.0",
+    #"multiscale_spatial_image>=1.0.0", Uncomment when new release using xr.DataTree is out
     "networkx",
     "numba",
     "numpy",
@@ -43,8 +43,7 @@ dependencies = [
     "scikit-image",
     "scipy",
     "typing_extensions>=4.8.0",
-    "xarray",
-    "xarray-datatree",
+    "xarray>=2024.10.0",
     "xarray-schema",
     "xarray-spatial>=0.3.5",
     "zarr",
diff --git a/src/spatialdata/_core/_deepcopy.py b/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
diff --git a/src/spatialdata/_core/centroids.py b/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
diff --git a/src/spatialdata/_core/data_extent.py b/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
diff --git a/src/spatialdata/_core/operations/_utils.py b/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
 
diff --git a/src/spatialdata/_core/operations/aggregate.py b/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
diff --git a/src/spatialdata/_core/operations/map.py b/src/spatialdata/_core/operations/map.py
@@ -6,8 +6,7 @@
 
 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
diff --git a/src/spatialdata/_core/operations/rasterize.py b/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
diff --git a/src/spatialdata/_core/operations/transform.py b/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
@@ -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, ...
diff --git a/src/spatialdata/_core/operations/vectorize.py b/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
diff --git a/src/spatialdata/_core/query/_utils.py b/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
diff --git a/src/spatialdata/_core/query/relational_query.py b/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
diff --git a/src/spatialdata/_core/query/spatial_query.py b/src/spatialdata/_core/query/spatial_query.py
@@ -11,10 +11,9 @@
 import dask.dataframe as dd
 import numpy as np
 from dask.dataframe import DataFrame as DaskDataFrame
-from datatree import DataTree
 from geopandas import GeoDataFrame
 from shapely.geometry import MultiPolygon, Point, Polygon
-from xarray import DataArray
+from xarray import DataArray, DataTree
 
 from spatialdata import to_polygons
 from spatialdata._core.query._utils import (
diff --git a/src/spatialdata/_core/spatialdata.py b/src/spatialdata/_core/spatialdata.py
@@ -14,12 +14,11 @@
 from dask.dataframe import DataFrame as DaskDataFrame
 from dask.dataframe import read_parquet
 from dask.delayed import Delayed
-from datatree import DataTree
 from geopandas import GeoDataFrame
 from ome_zarr.io import parse_url
 from ome_zarr.types import JSONDict
 from shapely import MultiPolygon, Polygon
-from xarray import DataArray
+from xarray import DataArray, DataTree
 
 from spatialdata._core._elements import Images, Labels, Points, Shapes, Tables
 from spatialdata._logging import logger
diff --git a/src/spatialdata/_io/_utils.py b/src/spatialdata/_io/_utils.py
@@ -16,9 +16,8 @@
 from anndata import AnnData
 from dask.array import Array as DaskArray
 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._utils import get_pyramid_levels
diff --git a/src/spatialdata/_io/io_raster.py b/src/spatialdata/_io/io_raster.py
@@ -4,7 +4,6 @@
 import dask.array as da
 import numpy as np
 import zarr
-from datatree import DataTree
 from ome_zarr.format import Format
 from ome_zarr.io import ZarrLocation
 from ome_zarr.reader import Label, Multiscales, Node, Reader
@@ -14,7 +13,7 @@
 from ome_zarr.writer import write_labels as write_labels_ngff
 from ome_zarr.writer import write_multiscale as write_multiscale_ngff
 from ome_zarr.writer import write_multiscale_labels as write_multiscale_labels_ngff
-from xarray import DataArray
+from xarray import DataArray, Dataset, DataTree
 
 from spatialdata._io._utils import (
     _get_transformations_from_ngff_dict,
@@ -93,11 +92,15 @@ def _read_multiscale(store: Union[str, Path], raster_type: Literal["image", "lab
         multiscale_image = {}
         for i, d in enumerate(datasets):
             data = node.load(Multiscales).array(resolution=d, version=format.version)
-            multiscale_image[f"scale{i}"] = DataArray(
-                data,
-                name="image",
-                dims=axes,
-                coords={"c": channels} if channels is not None else {},
+            multiscale_image[f"scale{i}"] = Dataset(
+                {
+                    "image": DataArray(
+                        data,
+                        name="image",
+                        dims=axes,
+                        coords={"c": channels} if channels is not None else {},
+                    )
+                }
             )
         msi = DataTree.from_dict(multiscale_image)
         _set_transformations(msi, transformations)
diff --git a/src/spatialdata/_types.py b/src/spatialdata/_types.py
@@ -3,8 +3,7 @@
 from typing import Union
 
 import numpy as np
-from datatree import DataTree
-from xarray import DataArray
+from xarray import DataArray, DataTree
 
 __all__ = ["ArrayLike", "ColorLike", "DTypeLike", "Raster_T"]
 
diff --git a/src/spatialdata/_utils.py b/src/spatialdata/_utils.py
@@ -11,8 +11,7 @@
 import pandas as pd
 from anndata import AnnData
 from dask import array as da
-from datatree import DataTree
-from xarray import DataArray
+from xarray import DataArray, Dataset, DataTree
 
 from spatialdata._types import ArrayLike
 from spatialdata.transformations import (
@@ -136,7 +135,7 @@ def _compute_paddings(data: DataArray, axis: str) -> tuple[int, int]:
             assert len(v.values()) == 1
             xdata = v.values().__iter__().__next__()
             if 0 not in xdata.shape:
-                d[k] = xdata
+                d[k] = Dataset({"image": xdata})
         unpadded = DataTree.from_dict(d)
     else:
         raise TypeError(f"Unsupported type: {type(raster)}")
diff --git a/src/spatialdata/dataloader/datasets.py b/src/spatialdata/dataloader/datasets.py
@@ -11,12 +11,11 @@
 import numpy as np
 import pandas as pd
 from anndata import AnnData
-from datatree import DataTree
 from geopandas import GeoDataFrame
 from pandas import CategoricalDtype
 from scipy.sparse import issparse
 from torch.utils.data import Dataset
-from xarray import DataArray
+from xarray import DataArray, DataTree
 
 from spatialdata._core.centroids import get_centroids
 from spatialdata._core.operations.transform import transform
diff --git a/src/spatialdata/datasets.py b/src/spatialdata/datasets.py
@@ -10,13 +10,12 @@
 import scipy
 from anndata import AnnData
 from dask.dataframe import DataFrame as DaskDataFrame
-from datatree import DataTree
 from geopandas import GeoDataFrame
 from numpy.random import default_rng
 from shapely.affinity import translate
 from shapely.geometry import MultiPolygon, Point, Polygon
 from skimage.segmentation import slic
-from xarray import DataArray
+from xarray import DataArray, DataTree
 
 from spatialdata._core.operations.aggregate import aggregate
 from spatialdata._core.query.relational_query import get_element_instances
diff --git a/src/spatialdata/models/_utils.py b/src/spatialdata/models/_utils.py
@@ -10,10 +10,9 @@
 import pandas as pd
 from anndata import AnnData
 from dask.dataframe import DataFrame as DaskDataFrame
-from datatree import DataTree
 from geopandas import GeoDataFrame
 from shapely.geometry import MultiPolygon, Point, Polygon
-from xarray import DataArray
+from xarray import DataArray, DataTree
 
 from spatialdata._logging import logger
 from spatialdata.transformations.transformations import BaseTransformation
diff --git a/src/spatialdata/models/models.py b/src/spatialdata/models/models.py
@@ -15,7 +15,6 @@
 from dask.array 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, GeoSeries
 from multiscale_spatial_image import to_multiscale
 from multiscale_spatial_image.to_multiscale.to_multiscale import Methods
@@ -25,7 +24,7 @@
 from shapely.geometry.collection import GeometryCollection
 from shapely.io import from_geojson, from_ragged_array
 from spatial_image import to_spatial_image
-from xarray import DataArray
+from xarray import DataArray, DataTree
 from xarray_schema.components import (
     ArrayTypeSchema,
     AttrSchema,
diff --git a/src/spatialdata/testing.py b/src/spatialdata/testing.py
@@ -4,12 +4,10 @@
 from anndata.tests.helpers import assert_equal as assert_anndata_equal
 from dask.dataframe import DataFrame as DaskDataFrame
 from dask.dataframe.tests.test_dataframe import assert_eq as assert_dask_dataframe_equal
-from datatree import DataTree
-from datatree.testing import assert_equal as assert_datatree_equal
 from geopandas import GeoDataFrame
 from geopandas.testing import assert_geodataframe_equal
-from xarray import DataArray
-from xarray.testing import assert_equal as assert_xarray_equal
+from xarray import DataArray, DataTree
+from xarray.testing import assert_equal
 
 from spatialdata import SpatialData
 from spatialdata._core._elements import Elements
@@ -115,10 +113,8 @@ def assert_elements_are_identical(
     # compare the elements
     if isinstance(element0, AnnData):
         assert_anndata_equal(element0, element1)
-    elif isinstance(element0, DataArray):
-        assert_xarray_equal(element0, element1)
-    elif isinstance(element0, DataTree):
-        assert_datatree_equal(element0, element1)
+    elif isinstance(element0, (DataArray, DataTree)):
+        assert_equal(element0, element1)
     elif isinstance(element0, GeoDataFrame):
         assert_geodataframe_equal(element0, element1, check_less_precise=True)
     else:
diff --git a/src/spatialdata/transformations/_utils.py b/src/spatialdata/transformations/_utils.py
@@ -5,9 +5,8 @@
 
 import numpy as np
 from dask.dataframe import DataFrame as DaskDataFrame
-from datatree import DataTree
 from geopandas import GeoDataFrame
-from xarray import DataArray
+from xarray import DataArray, Dataset, DataTree
 
 from spatialdata._logging import logger
 from spatialdata._types import ArrayLike
@@ -223,8 +222,10 @@ def _(data: DataTree) -> DataTree:
             offset = max_dim / n / 2
             coords = np.linspace(0, max_dim, n + 1)[:-1] + offset
             new_coords[ax] = coords
-        out[name] = dt[img_name].assign_coords(new_coords)
-    datatree = DataTree.from_dict(d=out)
+
+        # Xarray now only accepts Dataset as dictionary values for DataTree.from_dict.
+        out[name] = Dataset({img_name: dt[img_name].assign_coords(new_coords)})
+    datatree = DataTree.from_dict(out)
     # this is to trigger the validation of the dims
     _ = get_axes_names(datatree)
     return datatree
diff --git a/tests/conftest.py b/tests/conftest.py
@@ -14,14 +14,13 @@
 import pytest
 from anndata import AnnData
 from dask.dataframe import DataFrame as DaskDataFrame
-from datatree import DataTree
 from geopandas import GeoDataFrame
 from numpy.random import default_rng
 from scipy import ndimage as ndi
 from shapely import linearrings, polygons
 from shapely.geometry import MultiPolygon, Point, Polygon
 from skimage import data
-from xarray import DataArray
+from xarray import DataArray, DataTree
 
 from spatialdata._core._deepcopy import deepcopy
 from spatialdata._core.spatialdata import SpatialData
diff --git a/tests/core/operations/test_rasterize.py b/tests/core/operations/test_rasterize.py
diff --git a/tests/core/operations/test_transform.py b/tests/core/operations/test_transform.py
diff --git a/tests/core/query/test_spatial_query.py b/tests/core/query/test_spatial_query.py
diff --git a/tests/io/test_pyramids_performance.py b/tests/io/test_pyramids_performance.py
diff --git a/tests/models/test_models.py b/tests/models/test_models.py
diff --git a/tests/utils/test_element_utils.py b/tests/utils/test_element_utils.py
diff --git a/tests/utils/test_testing.py b/tests/utils/test_testing.py
