diff --git a/docs/api.md b/docs/api.md index d34ad726e..b27b07bd5 100644 --- a/docs/api.md +++ b/docs/api.md @@ -40,6 +40,7 @@ Operations on `SpatialData` objects. to_circles to_polygons aggregate + map_raster ``` ### Operations Utilities diff --git a/src/spatialdata/__init__.py b/src/spatialdata/__init__.py index 16e44d5a4..2d27f0a41 100644 --- a/src/spatialdata/__init__.py +++ b/src/spatialdata/__init__.py @@ -31,6 +31,7 @@ "save_transformations", "get_dask_backing_files", "are_extents_equal", + "map_raster", "deepcopy", ] @@ -40,6 +41,7 @@ from spatialdata._core.concatenate import concatenate from spatialdata._core.data_extent import are_extents_equal, get_extent from spatialdata._core.operations.aggregate import aggregate +from spatialdata._core.operations.map import map_raster from spatialdata._core.operations.rasterize import rasterize from spatialdata._core.operations.rasterize_bins import rasterize_bins from spatialdata._core.operations.transform import transform diff --git a/src/spatialdata/_core/centroids.py b/src/spatialdata/_core/centroids.py index 2736ef150..813e963c1 100644 --- a/src/spatialdata/_core/centroids.py +++ b/src/spatialdata/_core/centroids.py @@ -10,7 +10,6 @@ from datatree import DataTree from geopandas import GeoDataFrame from shapely import MultiPolygon, Point, Polygon -from spatial_image import SpatialImage from xarray import DataArray from spatialdata._core.operations.transform import transform @@ -105,7 +104,7 @@ def _( if isinstance(e, DataTree): assert len(e["scale0"]) == 1 - e = SpatialImage(next(iter(e["scale0"].values()))) + e = next(iter(e["scale0"].values())) dfs = [] for axis in get_axes_names(e): diff --git a/src/spatialdata/_core/operations/map.py b/src/spatialdata/_core/operations/map.py new file mode 100644 index 000000000..c064eaa10 --- /dev/null +++ b/src/spatialdata/_core/operations/map.py @@ -0,0 +1,132 @@ +from __future__ import annotations + +from collections.abc import Iterable, Mapping +from types import MappingProxyType +from typing import TYPE_CHECKING, Any, Callable + +import dask.array as da +from dask.array.overlap import coerce_depth +from datatree import DataTree +from xarray import DataArray + +from spatialdata.models._utils import get_axes_names, get_channels, get_raster_model_from_data_dims +from spatialdata.models.models import Labels2DModel, Labels3DModel, get_model +from spatialdata.transformations import get_transformation + +__all__ = ["map_raster"] + + +def map_raster( + data: DataArray | DataTree, + func: Callable[[da.Array], da.Array], + func_kwargs: Mapping[str, Any] = MappingProxyType({}), + blockwise: bool = True, + depth: int | tuple[int, ...] | dict[int, int] | None = None, + chunks: tuple[tuple[int, ...], ...] | None = None, + c_coords: Iterable[int] | Iterable[str] | None = None, + dims: tuple[str, ...] | None = None, + transformations: dict[str, Any] | None = None, + **kwargs: Any, +) -> DataArray: + """ + Apply a callable to raster data. + + Applies a `func` callable to raster data. If `blockwise` is set to `True`, + distributed processing will be achieved with: + + - :func:`dask.array.map_overlap` if `depth` is not `None` + - :func:`dask.array.map_blocks`, if `depth` is `None` + + otherwise `func` is applied to the full data. + + Parameters + ---------- + data + The data to process. It can be a :class:`xarray.DataArray` or :class:`datatree.DataTree`. + If it's a `DataTree`, the callable is applied to the first scale (`scale0`, the full-resolution data). + func + The callable that is applied to the data. + func_kwargs + Additional keyword arguments to pass to the callable `func`. + blockwise + If `True`, `func` will be distributed with :func:`dask.array.map_overlap` or :func:`dask.array.map_blocks`, + otherwise `func` is applied to the full data. If `False`, `depth`, `chunks` and `kwargs` are ignored. + depth + Specifies the overlap between chunks, i.e. the number of elements that each chunk + should share with its neighboring chunks. If not `None`, distributed processing will be achieved with + :func:`dask.array.map_overlap`, otherwise with :func:`dask.array.map_blocks`. + chunks + Chunk shape of resulting blocks if the callable does not preserve the data shape. + For example, if the input block has `shape: (3,100,100)` and the resulting block after the `map_raster` + call has `shape: (1, 100,100)`, the argument `chunks` should be passed accordingly. + Passed to :func:`dask.array.map_overlap` or :func:`dask.array.map_blocks`. Ignored if `blockwise` is `False`. + c_coords + The channel coordinates for the output data. If not provided, the channel coordinates of the input data are + used. If the callable `func` is expected to change the number of channel coordinates, + this argument should be provided, otherwise will default to `range(len(output_coords))`. + dims + The dimensions of the output data. If not provided, the dimensions of the input data are used. It must be + specified if the callable changes the data dimensions, e.g. `('c', 'y', 'x') -> ('y', 'x')`. + transformations + The transformations of the output data. If not provided, the transformations of the input data are copied to the + output data. It should be specified if the callable changes the data transformations. + kwargs + Additional keyword arguments to pass to :func:`dask.array.map_overlap` or :func:`dask.array.map_blocks`. + Ignored if `blockwise` is set to `False`. + + Returns + ------- + The processed data as a :class:`xarray.DataArray`. + """ + if isinstance(data, DataArray): + arr = data.data + elif isinstance(data, DataTree): + arr = data["scale0"].values().__iter__().__next__().data + else: + raise ValueError("Only 'DataArray' and 'DataTree' are supported.") + + model = get_model(data) + if model in (Labels2DModel, Labels3DModel) and c_coords is not None: + raise ValueError("Channel coordinates can not be provided for labels data.") + + kwargs = kwargs.copy() + kwargs["chunks"] = chunks + + if not blockwise: + arr = func(arr, **func_kwargs) + else: + if depth is not None: + kwargs.setdefault("boundary", "reflect") + + if not isinstance(depth, int) and len(depth) != arr.ndim: + raise ValueError( + f"Depth {depth} is provided for {len(depth)} dimensions. " + f"Please provide depth for {arr.ndim} dimensions." + ) + kwargs["depth"] = coerce_depth(arr.ndim, depth) + map_func = da.map_overlap + else: + map_func = da.map_blocks + + arr = map_func(func, arr, **func_kwargs, **kwargs, dtype=arr.dtype) + + dims = dims if dims is not None else get_axes_names(data) + if model not in (Labels2DModel, Labels3DModel): + if c_coords is None: + c_coords = range(arr.shape[0]) if arr.shape[0] != len(get_channels(data)) else get_channels(data) + else: + c_coords = None + if transformations is None: + d = get_transformation(data, get_all=True) + if TYPE_CHECKING: + assert isinstance(d, dict) + transformations = d + + model_kwargs = { + "chunks": arr.chunksize, + "c_coords": c_coords, + "dims": dims, + "transformations": transformations, + } + model = get_raster_model_from_data_dims(dims) + return model.parse(arr, **model_kwargs) diff --git a/src/spatialdata/_core/operations/vectorize.py b/src/spatialdata/_core/operations/vectorize.py index caf4bbb73..2b107bcd0 100644 --- a/src/spatialdata/_core/operations/vectorize.py +++ b/src/spatialdata/_core/operations/vectorize.py @@ -264,9 +264,14 @@ def _(gdf: GeoDataFrame, buffer_resolution: int = 16) -> GeoDataFrame: if isinstance(gdf.geometry.iloc[0], Point): ShapesModel.validate_shapes_not_mixed_types(gdf) if isinstance(gdf.geometry.iloc[0], Point): - buffered_df = gdf.set_geometry( - gdf.geometry.buffer(gdf[ShapesModel.RADIUS_KEY], resolution=buffer_resolution) + buffered_df = gdf.copy() + buffered_df["geometry"] = buffered_df.apply( + lambda row: row.geometry.buffer(row[ShapesModel.RADIUS_KEY], resolution=buffer_resolution), axis=1 ) + + # Ensure the GeoDataFrame recognizes the updated geometry column + buffered_df = buffered_df.set_geometry("geometry") + # 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 diff --git a/src/spatialdata/_io/io_raster.py b/src/spatialdata/_io/io_raster.py index ced112f95..136d14a6a 100644 --- a/src/spatialdata/_io/io_raster.py +++ b/src/spatialdata/_io/io_raster.py @@ -13,7 +13,6 @@ 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 spatial_image import SpatialImage from xarray import DataArray from spatialdata._io import SpatialDataFormatV01 @@ -92,13 +91,12 @@ def _read_multiscale( _set_transformations(msi, transformations) return compute_coordinates(msi) data = node.load(Multiscales).array(resolution=datasets[0], version=fmt.version) - si = SpatialImage( + si = DataArray( data, name="image", dims=axes, coords={"c": channels} if channels is not None else {}, ) - si.__class__ = DataArray _set_transformations(si, transformations) return compute_coordinates(si) diff --git a/src/spatialdata/models/_utils.py b/src/spatialdata/models/_utils.py index ebbacd3ad..8f354c9b2 100644 --- a/src/spatialdata/models/_utils.py +++ b/src/spatialdata/models/_utils.py @@ -1,7 +1,7 @@ from __future__ import annotations from functools import singledispatch -from typing import Any, Union +from typing import TYPE_CHECKING, Any, Union import dask.dataframe as dd import geopandas @@ -25,6 +25,9 @@ Y = "y" X = "x" +if TYPE_CHECKING: + from spatialdata.models.models import RasterSchema + # mypy says that we can't do isinstance(something, SpatialElement), # even if the code works fine in my machine. Since the solution described here don't work: @@ -345,3 +348,26 @@ def force_2d(gdf: GeoDataFrame) -> None: new_shapes.append(shape) if any_3d: gdf.geometry = new_shapes + + +def get_raster_model_from_data_dims(dims: tuple[str, ...]) -> type[RasterSchema]: + """ + Get the raster model from the dimensions of the data. + + Parameters + ---------- + dims + The dimensions of the data + + Returns + ------- + The raster model corresponding to the dimensions of the data. + """ + from spatialdata.models.models import Image2DModel, Image3DModel, Labels2DModel, Labels3DModel + + if not set(dims).issubset({C, Z, Y, X}): + raise ValueError(f"Invalid dimensions: {dims}") + + if C in dims: + return Image3DModel if Z in dims else Image2DModel + return Labels3DModel if Z in dims else Labels2DModel diff --git a/tests/conftest.py b/tests/conftest.py index 7ad080c97..0e608387f 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -11,14 +11,13 @@ import pytest from anndata import AnnData from dask.dataframe.core import DataFrame as DaskDataFrame +from datatree import DataTree from geopandas import GeoDataFrame -from multiscale_spatial_image import MultiscaleSpatialImage 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 spatial_image import SpatialImage from spatialdata._core._deepcopy import deepcopy as _deepcopy from spatialdata._core.spatialdata import SpatialData from spatialdata._types import ArrayLike @@ -133,7 +132,7 @@ def sdata(request) -> SpatialData: return request.getfixturevalue(request.param) -def _get_images() -> dict[str, SpatialImage | MultiscaleSpatialImage]: +def _get_images() -> dict[str, DataArray | DataTree]: out = {} dims_2d = ("c", "y", "x") dims_3d = ("z", "y", "x", "c") @@ -163,7 +162,7 @@ def _get_images() -> dict[str, SpatialImage | MultiscaleSpatialImage]: return out -def _get_labels() -> dict[str, SpatialImage | MultiscaleSpatialImage]: +def _get_labels() -> dict[str, DataArray | DataTree]: out = {} dims_2d = ("y", "x") dims_3d = ("z", "y", "x") diff --git a/tests/core/operations/test_map.py b/tests/core/operations/test_map.py new file mode 100644 index 000000000..01e7081d6 --- /dev/null +++ b/tests/core/operations/test_map.py @@ -0,0 +1,203 @@ +import re + +import numpy as np +import pytest +from spatialdata._core.operations.map import map_raster +from spatialdata.transformations import Translation, get_transformation, set_transformation +from xarray import DataArray + + +def _multiply(arr, parameter=10): + return arr * parameter + + +def _multiply_alter_c(arr, parameter=10): + arr = arr * parameter + arr = arr[0] + return arr[None, ...] + + +def _multiply_squeeze_z(arr, parameter=10): + arr = arr * parameter + return arr[:, 0, ...] + + +@pytest.mark.parametrize( + "depth", + [ + None, + (0, 60, 60), + ], +) +@pytest.mark.parametrize("element_name", ["blobs_image", "blobs_labels"]) +def test_map_raster(sdata_blobs, depth, element_name): + if element_name == "blobs_labels" and depth is not None: + depth = (60, 60) + + func_kwargs = {"parameter": 20} + se = map_raster( + sdata_blobs[element_name], + func=_multiply, + func_kwargs=func_kwargs, + c_coords=None, + depth=depth, + ) + + assert isinstance(se, DataArray) + data = sdata_blobs[element_name].data.compute() + res = se.data.compute() + assert np.array_equal(data * func_kwargs["parameter"], res) + + +@pytest.mark.parametrize( + "depth", + [ + None, + (0, 60, 60), + ], +) +def test_map_raster_multiscale(sdata_blobs, depth): + img_layer = "blobs_multiscale_image" + func_kwargs = {"parameter": 20} + se = map_raster( + sdata_blobs[img_layer], + func=_multiply, + func_kwargs=func_kwargs, + c_coords=None, + depth=depth, + ) + + data = sdata_blobs[img_layer]["scale0"]["image"].data.compute() + res = se.data.compute() + assert np.array_equal(data * func_kwargs["parameter"], res) + + +def test_map_raster_no_blockwise(sdata_blobs): + img_layer = "blobs_image" + func_kwargs = {"parameter": 20} + se = map_raster( + sdata_blobs[img_layer], + func=_multiply, + func_kwargs=func_kwargs, + blockwise=False, + c_coords=None, + depth=None, + ) + + assert isinstance(se, DataArray) + data = sdata_blobs[img_layer].data.compute() + res = se.data.compute() + assert np.array_equal(data * func_kwargs["parameter"], res) + + +def test_map_raster_output_chunks(sdata_blobs): + depth = 60 + func_kwargs = {"parameter": 20} + output_channels = ["test"] + se = map_raster( + sdata_blobs["blobs_image"].chunk((3, 100, 100)), + func=_multiply_alter_c, + func_kwargs=func_kwargs, + chunks=( + (1,), + (100 + 2 * depth, 96 + 2 * depth, 60 + 2 * depth), + (100 + 2 * depth, 96 + 2 * depth, 60 + 2 * depth), + ), # account for rechunking done by map_overlap to ensure minimum chunksize + c_coords=["test"], + depth=(0, depth, depth), + ) + + assert isinstance(se, DataArray) + assert np.array_equal(np.array(output_channels), se.c.data) + data = sdata_blobs["blobs_image"].data.compute() + res = se.data.compute() + assert np.array_equal(data[0] * func_kwargs["parameter"], res[0]) + + +@pytest.mark.parametrize("img_layer", ["blobs_image", "blobs_multiscale_image"]) +def test_map_transformation(sdata_blobs, img_layer): + func_kwargs = {"parameter": 20} + target_coordinate_system = "my_other_space0" + transformation = Translation(translation=[10, 12], axes=["y", "x"]) + + se = sdata_blobs[img_layer] + + set_transformation(se, transformation=transformation, to_coordinate_system=target_coordinate_system) + se = map_raster( + se, + func=_multiply, + func_kwargs=func_kwargs, + blockwise=False, + c_coords=None, + depth=None, + ) + assert transformation == get_transformation(se, to_coordinate_system=target_coordinate_system) + + +def test_map_squeeze_z(full_sdata): + img_layer = "image3d_numpy" + func_kwargs = {"parameter": 20} + + se = map_raster( + full_sdata[img_layer].chunk((3, 2, 64, 64)), + func=_multiply_squeeze_z, + func_kwargs=func_kwargs, + chunks=((3,), (64,), (64,)), + drop_axis=1, + c_coords=None, + dims=("c", "y", "x"), + depth=None, + ) + + assert isinstance(se, DataArray) + data = full_sdata[img_layer].data.compute() + res = se.data.compute() + assert np.array_equal(data[:, 0, ...] * func_kwargs["parameter"], res) + + +def test_map_squeeze_z_fails(full_sdata): + img_layer = "image3d_numpy" + func_kwargs = {"parameter": 20} + + with pytest.raises(IndexError): + map_raster( + full_sdata[img_layer].chunk((3, 2, 64, 64)), + func=_multiply_squeeze_z, + func_kwargs=func_kwargs, + chunks=((3,), (64,), (64,)), + drop_axis=1, + c_coords=None, + depth=None, + ) + + +def test_invalid_map_raster(sdata_blobs): + with pytest.raises(ValueError, match="Only 'DataArray' and 'DataTree' are supported."): + map_raster( + sdata_blobs["blobs_points"], + func=_multiply, + func_kwargs={"parameter": 20}, + c_coords=None, + depth=(0, 60), + ) + + with pytest.raises( + ValueError, + match=re.escape("Depth (0, 60) is provided for 2 dimensions. Please provide depth for 3 dimensions."), + ): + map_raster( + sdata_blobs["blobs_image"], + func=_multiply, + func_kwargs={"parameter": 20}, + c_coords=None, + depth=(0, 60), + ) + + with pytest.raises(ValueError, match="Channel coordinates can not be provided for labels data."): + map_raster( + sdata_blobs["blobs_labels"], + func=_multiply, + func_kwargs={"parameter": 20}, + c_coords=["c"], + depth=(0, 60, 60), + )