From b3e4487cbdfe34ccf38b5a311de1eb2a4b5427c4 Mon Sep 17 00:00:00 2001
From: ArneDefauw <arnedefauw@gmail.com>
Date: Mon, 17 Jun 2024 10:58:16 +0200
Subject: [PATCH 01/19] map raster

---
 src/spatialdata/_core/operations/map.py | 149 +++++++++++++++++++++
 tests/core/operations/test_map.py       | 165 ++++++++++++++++++++++++
 2 files changed, 314 insertions(+)
 create mode 100644 src/spatialdata/_core/operations/map.py
 create mode 100644 tests/core/operations/test_map.py

diff --git a/src/spatialdata/_core/operations/map.py b/src/spatialdata/_core/operations/map.py
new file mode 100644
index 000000000..445d28bce
--- /dev/null
+++ b/src/spatialdata/_core/operations/map.py
@@ -0,0 +1,149 @@
+from __future__ import annotations
+
+from collections.abc import Iterable, Mapping
+from types import MappingProxyType
+from typing import Any, Callable
+
+import dask.array as da
+from dask.array import Array
+from dask.array.overlap import coerce_depth
+from multiscale_spatial_image import MultiscaleSpatialImage
+from numpy.typing import NDArray
+from spatial_image import SpatialImage
+
+import spatialdata
+from spatialdata.models.models import ScaleFactors_t
+from spatialdata.transformations import get_transformation
+
+__all__ = ["map_raster"]
+
+
+def map_raster(
+    data: SpatialImage | MultiscaleSpatialImage,
+    func: Callable,
+    fn_kwargs: Mapping[str, Any] = MappingProxyType({}),
+    chunks: str | int | tuple[int, ...] | tuple[tuple[int, ...], ...] | None = None,
+    output_chunks: tuple[tuple[int, ...], ...] | None = None,
+    depth: str | int | tuple[int, ...] | dict[int:int] | None = None,
+    scale_factors: ScaleFactors_t | None = None,  # if specified will return multiscale
+    c_coords: int | str | Iterable[int | str] | None = None,
+    **kwargs,
+) -> SpatialImage | MultiscaleSpatialImage:
+    """
+    Apply a function to raster data.
+
+    Parameters
+    ----------
+    data
+        The data to process. Can be a `SpatialImage` or `MultiscaleSpatialImage`.
+    func
+        The function to apply to the data.
+    fn_kwargs
+        Additional keyword arguments to pass to the function `func`.
+    chunks
+        If specified, data will be rechunked and processed via `dask.array.map_blocks` or `dask.array.map_overlap`.
+        If `None`, `func` is applied to the data without use of `dask.array.map_blocks`/`dask.array.map_overlap`.
+    output_chunks
+        Chunk shape of resulting blocks if the function does not preserve
+        shape. If not provided, the resulting array is assumed to have the same
+        block structure as the first input array.
+        Passed to `dask.array.map_overlap`/`dask.array.map_blocks` as `chunks`.
+        Ignored when `chunks` is `None`.
+        E.g. ( (3,), (256,) , (256,)  ).
+    depth
+        If not `None` and `chunks` is not `None`, will use `dask.array.map_overlap` for distributed processing.
+        Specifies the number of elements that each block should share with its neighbors
+    scale_factors
+        If specified, the function returns a `MultiscaleSpatialImage`.
+    c_coords
+        Can be used to set the channel coordinates for the output data.
+        If the number of channels is altered, `c_coords` should match the output dimension.
+    kwargs
+        Additional keyword arguments to pass to `dask.array.map_overlap` or `dask.array.map_blocks`.
+
+    Returns
+    -------
+    The processed data. If `scale_factors` is provided, returns a `MultiscaleSpatialImage`, else `SpatialImage`.
+
+    Notes
+    -----
+    The transformations of the input data are preserved and applied to the output data.
+    """
+
+    def _map_func(
+        func: Callable[..., NDArray | Array],
+        arr: NDArray | Array,
+        fn_kwargs: Mapping[str, Any] = MappingProxyType({}),
+    ) -> Array:
+        if chunks is None:
+            # if dask array, we want to rechunk
+            if isinstance(arr, Array):
+                arr = arr.rechunk(arr.chunksize)
+            arr = func(arr, **fn_kwargs)
+            arr = da.asarray(arr)
+            # func could have cause irregular chunking
+            return arr.rechunk(arr.chunksize)
+        if output_chunks is not None:
+            kwargs["chunks"] = output_chunks
+        arr = da.asarray(arr).rechunk(chunks)
+        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 (only) provide depth for {arr.ndim} dimensions."
+                )
+
+            kwargs["depth"] = coerce_depth(arr.ndim, depth)
+
+            arr = da.map_overlap(func, arr, **fn_kwargs, **kwargs, dtype=arr.dtype)
+        else:
+            arr = da.map_blocks(func, arr, **fn_kwargs, **kwargs, dtype=arr.dtype)
+        # not sure if we want to rechunk here; it fixes irregular chunk sizes, necessary when wanting to save to zarr
+        return arr.rechunk(arr.chunksize)
+
+    # pass transformations as parameter to map_raster?
+    # If transformations is not None, then we can use this transformation when parsing dask array
+    # necessary if dimension is altered of spatialimage (via output_chunks parameter)
+    transformations = get_transformation(data, get_all=True)
+
+    if isinstance(data, SpatialImage):
+        arr = data.data
+    elif isinstance(data, MultiscaleSpatialImage):
+        scale_0 = data.__iter__().__next__()
+        name = data[scale_0].__iter__().__next__()
+        data = data[scale_0][name]
+        arr = data.data
+    else:
+        raise ValueError("Currently only supports 'SpatialImage' and 'MultiscaleSpatialImage'.")
+
+    arr = _map_func(func=func, arr=arr, fn_kwargs=fn_kwargs)
+
+    # should we add this line? if added, user needs to pass c_coords when nr of channels is altered,
+    # but doing this, allows users to not pass c_coords, and still c_coords are preserveed.
+    # probably remove, user can just copy coordinates from input image
+    # if c_coords is None:
+    #    c_coords = se.c.data
+
+    if "z" in data.dims:
+        data = spatialdata.models.Image3DModel.parse(
+            arr,
+            dims=data.dims,  # currently does not allow changing dims, we could allow passing dims to map_raster
+            scale_factors=scale_factors,
+            chunks=arr.chunksize,
+            c_coords=c_coords,  # Note that if c_coords is not None, it should match the output channels.
+            transformations=transformations,
+        )
+
+    else:
+        data = spatialdata.models.Image2DModel.parse(
+            arr,
+            dims=data.dims,
+            scale_factors=scale_factors,
+            chunks=arr.chunksize,
+            c_coords=c_coords,
+            transformations=transformations,
+        )
+
+    return data
diff --git a/tests/core/operations/test_map.py b/tests/core/operations/test_map.py
new file mode 100644
index 000000000..a862b66b5
--- /dev/null
+++ b/tests/core/operations/test_map.py
@@ -0,0 +1,165 @@
+import numpy as np
+import pytest
+from multiscale_spatial_image import MultiscaleSpatialImage
+from spatial_image import SpatialImage
+from spatialdata._core.operations.map import map_raster
+from spatialdata.transformations import Translation, get_transformation, set_transformation
+
+
+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),
+    ],
+)
+def test_map_raster(sdata_blobs, depth):
+    img_layer = "blobs_image"
+    fn_kwargs = {"parameter": 20}
+    se = map_raster(
+        sdata_blobs[img_layer],
+        func=_multiply,
+        fn_kwargs=fn_kwargs,
+        chunks=(3, 100, 100),
+        c_coords=None,
+        scale_factors=None,
+        depth=depth,
+    )
+
+    assert isinstance(se, SpatialImage)
+    data = sdata_blobs[img_layer].data.compute()
+    res = se.data.compute()
+    assert np.array_equal(data * fn_kwargs["parameter"], res)
+
+
+@pytest.mark.parametrize(
+    "depth",
+    [
+        None,
+        (0, 60, 60),
+    ],
+)
+def test_map_raster_multiscale(sdata_blobs, depth):
+    img_layer = "blobs_multiscale_image"
+    fn_kwargs = {"parameter": 20}
+    se = map_raster(
+        sdata_blobs[img_layer],
+        func=_multiply,
+        fn_kwargs=fn_kwargs,
+        chunks=(3, 100, 100),
+        c_coords=None,
+        scale_factors=[2, 2, 2, 2],
+        depth=depth,
+    )
+
+    assert isinstance(se, MultiscaleSpatialImage)
+    data = sdata_blobs[img_layer]["scale0"]["image"].data.compute()
+    res = se["scale0"]["image"].data.compute()
+    assert np.array_equal(data * fn_kwargs["parameter"], res)
+
+
+def test_map_raster_chunks_none(sdata_blobs):
+    img_layer = "blobs_image"
+    fn_kwargs = {"parameter": 20}
+    se = map_raster(
+        sdata_blobs[img_layer],
+        func=_multiply,
+        fn_kwargs=fn_kwargs,
+        chunks=None,
+        c_coords=None,
+        scale_factors=None,
+        depth=None,
+    )
+
+    assert isinstance(se, SpatialImage)
+    data = sdata_blobs[img_layer].data.compute()
+    res = se.data.compute()
+    assert np.array_equal(data * fn_kwargs["parameter"], res)
+
+
+def test_map_raster_output_chunks(sdata_blobs):
+    depth = 60
+    fn_kwargs = {"parameter": 20}
+    output_channels = ["test"]
+    se = map_raster(
+        sdata_blobs["blobs_image"],
+        func=_multiply_alter_c,
+        fn_kwargs=fn_kwargs,
+        chunks=(3, 100, 100),
+        output_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"],
+        scale_factors=None,
+        depth=(0, depth, depth),
+    )
+
+    assert isinstance(se, SpatialImage)
+    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] * fn_kwargs["parameter"], res[0])
+
+
+@pytest.mark.parametrize(
+    "img_layer, expected_type, scale_factors",
+    [
+        ("blobs_image", SpatialImage, None),
+        ("blobs_multiscale_image", MultiscaleSpatialImage, [2, 2, 2, 2]),
+    ],
+)
+def test_map_transformation(sdata_blobs, img_layer, expected_type, scale_factors):
+    fn_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(
+        sdata_blobs[img_layer],
+        func=_multiply,
+        fn_kwargs=fn_kwargs,
+        chunks=None,
+        c_coords=None,
+        scale_factors=scale_factors,
+        depth=None,
+    )
+    assert isinstance(se, expected_type)
+    assert transformation == get_transformation(se, to_coordinate_system=target_coordinate_system)
+
+
+def test_map_remove_z_fails(full_sdata):
+    fn_kwargs = {"parameter": 20}
+
+    # currently can not alter dims, e.g. ("c","z","y","x") -> ("c","y","x") fails
+    # could be supported by adding dims (and possibly transformations) to parameters of map_raster
+    with pytest.raises(IndexError):
+        map_raster(
+            full_sdata["image3d_numpy"],
+            func=_multiply_squeeze_z,
+            fn_kwargs=fn_kwargs,
+            chunks=100,
+            output_chunks=((3,), (64,), (64,)),
+            drop_axis=1,
+            c_coords=None,
+            scale_factors=None,
+            depth=None,
+        )

From 8722e8f965f2ed0711ec722c3b2a9880da04dd79 Mon Sep 17 00:00:00 2001
From: Luca Marconato <m.lucalmer@gmail.com>
Date: Tue, 18 Jun 2024 16:12:46 +0200
Subject: [PATCH 02/19] code review + refactoring

---
 src/spatialdata/_core/operations/map.py | 175 +++++++++++-------------
 src/spatialdata/models/_utils.py        |  28 +++-
 tests/core/operations/test_map.py       |  33 ++---
 3 files changed, 112 insertions(+), 124 deletions(-)

diff --git a/src/spatialdata/_core/operations/map.py b/src/spatialdata/_core/operations/map.py
index 445d28bce..c0ac18529 100644
--- a/src/spatialdata/_core/operations/map.py
+++ b/src/spatialdata/_core/operations/map.py
@@ -5,14 +5,11 @@
 from typing import Any, Callable
 
 import dask.array as da
-from dask.array import Array
 from dask.array.overlap import coerce_depth
 from multiscale_spatial_image import MultiscaleSpatialImage
-from numpy.typing import NDArray
 from spatial_image import SpatialImage
 
-import spatialdata
-from spatialdata.models.models import ScaleFactors_t
+from spatialdata.models._utils import get_axes_names, get_channels, get_raster_model_from_data_dims
 from spatialdata.transformations import get_transformation
 
 __all__ = ["map_raster"]
@@ -20,72 +17,85 @@
 
 def map_raster(
     data: SpatialImage | MultiscaleSpatialImage,
-    func: Callable,
+    func: Callable[[da.Array], da.Array],
     fn_kwargs: Mapping[str, Any] = MappingProxyType({}),
-    chunks: str | int | tuple[int, ...] | tuple[tuple[int, ...], ...] | None = None,
+    chunkwise: bool = True,
+    depth: str | int | tuple[int, ...] | dict[int, int] | None = None,
+    input_chunks: tuple[tuple[int, ...], ...] | None = None,
     output_chunks: tuple[tuple[int, ...], ...] | None = None,
-    depth: str | int | tuple[int, ...] | dict[int:int] | None = None,
-    scale_factors: ScaleFactors_t | None = None,  # if specified will return multiscale
     c_coords: int | str | Iterable[int | str] | None = None,
-    **kwargs,
-) -> SpatialImage | MultiscaleSpatialImage:
+    dims: tuple[str, ...] | None = None,
+    transformations: dict[str, Any] | None = None,
+    **kwargs: Any,
+) -> SpatialImage:
     """
-    Apply a function to raster data.
+    Apply a function to raster data, for each chunk and each scale.
 
     Parameters
     ----------
     data
-        The data to process. Can be a `SpatialImage` or `MultiscaleSpatialImage`.
+        The data to process. It can be a `SpatialImage` or `MultiscaleSpatialImage`. If it's a `MultiscaleSpatialImage`,
+        the function is applied to the first scale (full-resolution data).
     func
         The function to apply to the data.
     fn_kwargs
         Additional keyword arguments to pass to the function `func`.
-    chunks
-        If specified, data will be rechunked and processed via `dask.array.map_blocks` or `dask.array.map_overlap`.
-        If `None`, `func` is applied to the data without use of `dask.array.map_blocks`/`dask.array.map_overlap`.
+    chunkwise
+        If `True`, distributed processing will be achieved with `dask.array.map_overlap`/`dask.array.map_blocks`,
+        otherwise the function is applied to the full data. If `False`, `depth` and `input_chunks` are ignored.
+    depth
+        If not `None`, distributed processing will be achieved with `dask.array.map_overlap`, otherwise with
+        `dask.array.map_blocks`. Specifies the overlap between chunks, i.e. the number of elements that each chunk
+        should share with its neighbor chunks.
+        # TODO: Add examples for each data type in the signature.
+    input_chunks
+        If specified, rechunks the input data before applying the function using `dask.array.rechunk`.
     output_chunks
-        Chunk shape of resulting blocks if the function does not preserve
-        shape. If not provided, the resulting array is assumed to have the same
-        block structure as the first input array.
+        Chunk shape of resulting blocks if the function does not preserve the data shape. If not provided, the resulting
+        array is assumed to have the same chunk structure as the first input array.
         Passed to `dask.array.map_overlap`/`dask.array.map_blocks` as `chunks`.
-        Ignored when `chunks` is `None`.
-        E.g. ( (3,), (256,) , (256,)  ).
-    depth
-        If not `None` and `chunks` is not `None`, will use `dask.array.map_overlap` for distributed processing.
-        Specifies the number of elements that each block should share with its neighbors
-    scale_factors
-        If specified, the function returns a `MultiscaleSpatialImage`.
+        E.g. ( (3,), (256,), (256,) ).
+        # TODO: Add examples for each data type in the signature.
     c_coords
-        Can be used to set the channel coordinates for the output data.
-        If the number of channels is altered, `c_coords` should match the output dimension.
+        The channel coordinates for the output data. If not provided, the channel coordinates of the input data are
+        used. It should be specified if the function changes the number of channels.
+        # TODO: Add examples for each data type in the signature.
+    dims
+        The dimensions of the output data. If not provided, the dimensions of the input data are used. It must be
+        specified if the function changes the data dimensions.
+        E.g. ('c', '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 function changes the data transformations.
     kwargs
         Additional keyword arguments to pass to `dask.array.map_overlap` or `dask.array.map_blocks`.
 
     Returns
     -------
-    The processed data. If `scale_factors` is provided, returns a `MultiscaleSpatialImage`, else `SpatialImage`.
-
-    Notes
-    -----
-    The transformations of the input data are preserved and applied to the output data.
+    The processed data as a `SpatialImage`.
     """
+    if isinstance(data, SpatialImage):
+        arr = data.data
+    elif isinstance(data, MultiscaleSpatialImage):
+        arr = data["scale0"].values().__iter__().__next__().data
+    else:
+        raise ValueError("Only 'SpatialImage' and 'MultiscaleSpatialImage' are supported.")
 
-    def _map_func(
-        func: Callable[..., NDArray | Array],
-        arr: NDArray | Array,
-        fn_kwargs: Mapping[str, Any] = MappingProxyType({}),
-    ) -> Array:
-        if chunks is None:
-            # if dask array, we want to rechunk
-            if isinstance(arr, Array):
-                arr = arr.rechunk(arr.chunksize)
-            arr = func(arr, **fn_kwargs)
-            arr = da.asarray(arr)
-            # func could have cause irregular chunking
-            return arr.rechunk(arr.chunksize)
+    if "depth" in kwargs or "chunks" in kwargs:
+        raise ValueError(
+            "Please provide 'depth' and 'chunks' as arguments to 'map_raster' (respectively 'depth' and 'output_chunks'"
+            ", and not as 'kwargs'."
+        )
+    kwargs = kwargs.copy()
+    kwargs["chunks"] = output_chunks
+
+    if not chunkwise:
+        arr = func(arr, **fn_kwargs)
         if output_chunks is not None:
-            kwargs["chunks"] = output_chunks
-        arr = da.asarray(arr).rechunk(chunks)
+            arr = arr.rechunk(output_chunks)
+    else:
+        if input_chunks is not None:
+            arr = arr.rechunk(input_chunks)
         if depth is not None:
             kwargs.setdefault("boundary", "reflect")
 
@@ -94,56 +104,25 @@ def _map_func(
                     f"Depth ({depth}) is provided for {len(depth)} dimensions. "
                     f"Please (only) provide depth for {arr.ndim} dimensions."
                 )
-
             kwargs["depth"] = coerce_depth(arr.ndim, depth)
-
-            arr = da.map_overlap(func, arr, **fn_kwargs, **kwargs, dtype=arr.dtype)
+            map_func = da.map_overlap
         else:
-            arr = da.map_blocks(func, arr, **fn_kwargs, **kwargs, dtype=arr.dtype)
-        # not sure if we want to rechunk here; it fixes irregular chunk sizes, necessary when wanting to save to zarr
-        return arr.rechunk(arr.chunksize)
-
-    # pass transformations as parameter to map_raster?
-    # If transformations is not None, then we can use this transformation when parsing dask array
-    # necessary if dimension is altered of spatialimage (via output_chunks parameter)
-    transformations = get_transformation(data, get_all=True)
-
-    if isinstance(data, SpatialImage):
-        arr = data.data
-    elif isinstance(data, MultiscaleSpatialImage):
-        scale_0 = data.__iter__().__next__()
-        name = data[scale_0].__iter__().__next__()
-        data = data[scale_0][name]
-        arr = data.data
-    else:
-        raise ValueError("Currently only supports 'SpatialImage' and 'MultiscaleSpatialImage'.")
-
-    arr = _map_func(func=func, arr=arr, fn_kwargs=fn_kwargs)
-
-    # should we add this line? if added, user needs to pass c_coords when nr of channels is altered,
-    # but doing this, allows users to not pass c_coords, and still c_coords are preserveed.
-    # probably remove, user can just copy coordinates from input image
-    # if c_coords is None:
-    #    c_coords = se.c.data
-
-    if "z" in data.dims:
-        data = spatialdata.models.Image3DModel.parse(
-            arr,
-            dims=data.dims,  # currently does not allow changing dims, we could allow passing dims to map_raster
-            scale_factors=scale_factors,
-            chunks=arr.chunksize,
-            c_coords=c_coords,  # Note that if c_coords is not None, it should match the output channels.
-            transformations=transformations,
-        )
-
-    else:
-        data = spatialdata.models.Image2DModel.parse(
-            arr,
-            dims=data.dims,
-            scale_factors=scale_factors,
-            chunks=arr.chunksize,
-            c_coords=c_coords,
-            transformations=transformations,
-        )
-
-    return data
+            map_func = da.map_blocks
+
+        arr = map_func(func, arr, **fn_kwargs, **kwargs, dtype=arr.dtype)
+
+    dims = dims if dims is not None else get_axes_names(data)
+    c_coords = c_coords if c_coords is not None else get_channels(data)
+    if transformations is None:
+        d = get_transformation(data, get_all=True)
+        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/models/_utils.py b/src/spatialdata/models/_utils.py
index 521e8233c..34bce8b8b 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/core/operations/test_map.py b/tests/core/operations/test_map.py
index a862b66b5..b33e181f7 100644
--- a/tests/core/operations/test_map.py
+++ b/tests/core/operations/test_map.py
@@ -1,6 +1,5 @@
 import numpy as np
 import pytest
-from multiscale_spatial_image import MultiscaleSpatialImage
 from spatial_image import SpatialImage
 from spatialdata._core.operations.map import map_raster
 from spatialdata.transformations import Translation, get_transformation, set_transformation
@@ -35,9 +34,7 @@ def test_map_raster(sdata_blobs, depth):
         sdata_blobs[img_layer],
         func=_multiply,
         fn_kwargs=fn_kwargs,
-        chunks=(3, 100, 100),
         c_coords=None,
-        scale_factors=None,
         depth=depth,
     )
 
@@ -61,28 +58,24 @@ def test_map_raster_multiscale(sdata_blobs, depth):
         sdata_blobs[img_layer],
         func=_multiply,
         fn_kwargs=fn_kwargs,
-        chunks=(3, 100, 100),
         c_coords=None,
-        scale_factors=[2, 2, 2, 2],
         depth=depth,
     )
 
-    assert isinstance(se, MultiscaleSpatialImage)
     data = sdata_blobs[img_layer]["scale0"]["image"].data.compute()
-    res = se["scale0"]["image"].data.compute()
+    res = se.data.compute()
     assert np.array_equal(data * fn_kwargs["parameter"], res)
 
 
-def test_map_raster_chunks_none(sdata_blobs):
+def test_map_raster_no_chunkwise(sdata_blobs):
     img_layer = "blobs_image"
     fn_kwargs = {"parameter": 20}
     se = map_raster(
         sdata_blobs[img_layer],
         func=_multiply,
         fn_kwargs=fn_kwargs,
-        chunks=None,
+        chunkwise=False,
         c_coords=None,
-        scale_factors=None,
         depth=None,
     )
 
@@ -100,14 +93,13 @@ def test_map_raster_output_chunks(sdata_blobs):
         sdata_blobs["blobs_image"],
         func=_multiply_alter_c,
         fn_kwargs=fn_kwargs,
-        chunks=(3, 100, 100),
+        input_chunks=(3, 100, 100),
         output_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"],
-        scale_factors=None,
         depth=(0, depth, depth),
     )
 
@@ -118,14 +110,8 @@ def test_map_raster_output_chunks(sdata_blobs):
     assert np.array_equal(data[0] * fn_kwargs["parameter"], res[0])
 
 
-@pytest.mark.parametrize(
-    "img_layer, expected_type, scale_factors",
-    [
-        ("blobs_image", SpatialImage, None),
-        ("blobs_multiscale_image", MultiscaleSpatialImage, [2, 2, 2, 2]),
-    ],
-)
-def test_map_transformation(sdata_blobs, img_layer, expected_type, scale_factors):
+@pytest.mark.parametrize("img_layer", ["blobs_image", "blobs_multiscale_image"])
+def test_map_transformation(sdata_blobs, img_layer):
     fn_kwargs = {"parameter": 20}
     target_coordinate_system = "my_other_space0"
     transformation = Translation(translation=[10, 12], axes=["y", "x"])
@@ -137,12 +123,10 @@ def test_map_transformation(sdata_blobs, img_layer, expected_type, scale_factors
         sdata_blobs[img_layer],
         func=_multiply,
         fn_kwargs=fn_kwargs,
-        chunks=None,
+        chunkwise=False,
         c_coords=None,
-        scale_factors=scale_factors,
         depth=None,
     )
-    assert isinstance(se, expected_type)
     assert transformation == get_transformation(se, to_coordinate_system=target_coordinate_system)
 
 
@@ -156,10 +140,9 @@ def test_map_remove_z_fails(full_sdata):
             full_sdata["image3d_numpy"],
             func=_multiply_squeeze_z,
             fn_kwargs=fn_kwargs,
-            chunks=100,
+            input_chunks=100,
             output_chunks=((3,), (64,), (64,)),
             drop_axis=1,
             c_coords=None,
-            scale_factors=None,
             depth=None,
         )

From 3215e349f4459f32c002e1e0749c3939df4d3e86 Mon Sep 17 00:00:00 2001
From: Luca Marconato <m.lucalmer@gmail.com>
Date: Tue, 18 Jun 2024 23:45:40 +0200
Subject: [PATCH 03/19] fix docstrings and types of dask arguments

---
 src/spatialdata/_core/operations/map.py | 10 ++++------
 1 file changed, 4 insertions(+), 6 deletions(-)

diff --git a/src/spatialdata/_core/operations/map.py b/src/spatialdata/_core/operations/map.py
index c0ac18529..ab9a67a76 100644
--- a/src/spatialdata/_core/operations/map.py
+++ b/src/spatialdata/_core/operations/map.py
@@ -20,10 +20,10 @@ def map_raster(
     func: Callable[[da.Array], da.Array],
     fn_kwargs: Mapping[str, Any] = MappingProxyType({}),
     chunkwise: bool = True,
-    depth: str | int | tuple[int, ...] | dict[int, int] | None = None,
+    depth: int | tuple[int, ...] | dict[int, int] | None = None,
     input_chunks: tuple[tuple[int, ...], ...] | None = None,
     output_chunks: tuple[tuple[int, ...], ...] | None = None,
-    c_coords: int | str | Iterable[int | str] | None = None,
+    c_coords: Iterable[int] | Iterable[str] | None = None,
     dims: tuple[str, ...] | None = None,
     transformations: dict[str, Any] | None = None,
     **kwargs: Any,
@@ -46,8 +46,8 @@ def map_raster(
     depth
         If not `None`, distributed processing will be achieved with `dask.array.map_overlap`, otherwise with
         `dask.array.map_blocks`. Specifies the overlap between chunks, i.e. the number of elements that each chunk
-        should share with its neighbor chunks.
-        # TODO: Add examples for each data type in the signature.
+        should share with its neighbor chunks. Please see `dask.array.map_overlap` for more information on the accepted
+        values.
     input_chunks
         If specified, rechunks the input data before applying the function using `dask.array.rechunk`.
     output_chunks
@@ -55,11 +55,9 @@ def map_raster(
         array is assumed to have the same chunk structure as the first input array.
         Passed to `dask.array.map_overlap`/`dask.array.map_blocks` as `chunks`.
         E.g. ( (3,), (256,), (256,) ).
-        # TODO: Add examples for each data type in the signature.
     c_coords
         The channel coordinates for the output data. If not provided, the channel coordinates of the input data are
         used. It should be specified if the function changes the number of channels.
-        # TODO: Add examples for each data type in the signature.
     dims
         The dimensions of the output data. If not provided, the dimensions of the input data are used. It must be
         specified if the function changes the data dimensions.

From 3e9a6cb498eefde83341d06254853b1084301dd9 Mon Sep 17 00:00:00 2001
From: Luca Marconato <m.lucalmer@gmail.com>
Date: Tue, 18 Jun 2024 23:52:46 +0200
Subject: [PATCH 04/19] add tests for exceptions

---
 src/spatialdata/_core/operations/map.py |  7 +------
 tests/core/operations/test_map.py       | 25 +++++++++++++++++++++++++
 2 files changed, 26 insertions(+), 6 deletions(-)

diff --git a/src/spatialdata/_core/operations/map.py b/src/spatialdata/_core/operations/map.py
index ab9a67a76..79a1a816f 100644
--- a/src/spatialdata/_core/operations/map.py
+++ b/src/spatialdata/_core/operations/map.py
@@ -79,11 +79,6 @@ def map_raster(
     else:
         raise ValueError("Only 'SpatialImage' and 'MultiscaleSpatialImage' are supported.")
 
-    if "depth" in kwargs or "chunks" in kwargs:
-        raise ValueError(
-            "Please provide 'depth' and 'chunks' as arguments to 'map_raster' (respectively 'depth' and 'output_chunks'"
-            ", and not as 'kwargs'."
-        )
     kwargs = kwargs.copy()
     kwargs["chunks"] = output_chunks
 
@@ -99,7 +94,7 @@ def map_raster(
 
             if not isinstance(depth, int) and len(depth) != arr.ndim:
                 raise ValueError(
-                    f"Depth ({depth}) is provided for {len(depth)} dimensions. "
+                    f"Depth {depth} is provided for {len(depth)} dimensions. "
                     f"Please (only) provide depth for {arr.ndim} dimensions."
                 )
             kwargs["depth"] = coerce_depth(arr.ndim, depth)
diff --git a/tests/core/operations/test_map.py b/tests/core/operations/test_map.py
index b33e181f7..09fe188da 100644
--- a/tests/core/operations/test_map.py
+++ b/tests/core/operations/test_map.py
@@ -1,3 +1,5 @@
+import re
+
 import numpy as np
 import pytest
 from spatial_image import SpatialImage
@@ -146,3 +148,26 @@ def test_map_remove_z_fails(full_sdata):
             c_coords=None,
             depth=None,
         )
+
+
+def test_invalid_map_raster(sdata_blobs):
+    with pytest.raises(ValueError, match="Only 'SpatialImage' and 'MultiscaleSpatialImage' are supported."):
+        map_raster(
+            sdata_blobs["blobs_points"],
+            func=_multiply,
+            fn_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 (only) provide depth for 3 dimensions."),
+    ):
+        map_raster(
+            sdata_blobs["blobs_image"],
+            func=_multiply,
+            fn_kwargs={"parameter": 20},
+            c_coords=None,
+            depth=(0, 60),
+        )

From 0b21ff0bb456fe8793c92e4152ae56454f987365 Mon Sep 17 00:00:00 2001
From: Luca Marconato <m.lucalmer@gmail.com>
Date: Tue, 18 Jun 2024 23:59:10 +0200
Subject: [PATCH 05/19] map raster support and tests for labels

---
 src/spatialdata/_core/operations/map.py | 10 +++++++++-
 tests/core/operations/test_map.py       | 20 ++++++++++++++++----
 2 files changed, 25 insertions(+), 5 deletions(-)

diff --git a/src/spatialdata/_core/operations/map.py b/src/spatialdata/_core/operations/map.py
index 79a1a816f..e43d7086b 100644
--- a/src/spatialdata/_core/operations/map.py
+++ b/src/spatialdata/_core/operations/map.py
@@ -10,6 +10,7 @@
 from spatial_image import SpatialImage
 
 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"]
@@ -79,6 +80,10 @@ def map_raster(
     else:
         raise ValueError("Only 'SpatialImage' and 'MultiscaleSpatialImage' 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"] = output_chunks
 
@@ -105,7 +110,10 @@ def map_raster(
         arr = map_func(func, arr, **fn_kwargs, **kwargs, dtype=arr.dtype)
 
     dims = dims if dims is not None else get_axes_names(data)
-    c_coords = c_coords if c_coords is not None else get_channels(data)
+    if model not in (Labels2DModel, Labels3DModel):
+        c_coords = c_coords if c_coords is not None else get_channels(data)
+    else:
+        c_coords = None
     if transformations is None:
         d = get_transformation(data, get_all=True)
         assert isinstance(d, dict)
diff --git a/tests/core/operations/test_map.py b/tests/core/operations/test_map.py
index 09fe188da..f5c12944a 100644
--- a/tests/core/operations/test_map.py
+++ b/tests/core/operations/test_map.py
@@ -29,11 +29,14 @@ def _multiply_squeeze_z(arr, parameter=10):
         (0, 60, 60),
     ],
 )
-def test_map_raster(sdata_blobs, depth):
-    img_layer = "blobs_image"
+@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)
+
     fn_kwargs = {"parameter": 20}
     se = map_raster(
-        sdata_blobs[img_layer],
+        sdata_blobs[element_name],
         func=_multiply,
         fn_kwargs=fn_kwargs,
         c_coords=None,
@@ -41,7 +44,7 @@ def test_map_raster(sdata_blobs, depth):
     )
 
     assert isinstance(se, SpatialImage)
-    data = sdata_blobs[img_layer].data.compute()
+    data = sdata_blobs[element_name].data.compute()
     res = se.data.compute()
     assert np.array_equal(data * fn_kwargs["parameter"], res)
 
@@ -171,3 +174,12 @@ def test_invalid_map_raster(sdata_blobs):
             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,
+            fn_kwargs={"parameter": 20},
+            c_coords=["c"],
+            depth=(0, 60, 60),
+        )

From 069d5db0342a20906b45dd37dc8153f5eab7f6cb Mon Sep 17 00:00:00 2001
From: Wouter-Michiel Vierdag <michiel.vierdag@embl.de>
Date: Wed, 19 Jun 2024 00:29:31 +0200
Subject: [PATCH 06/19] adjust for geopandas 1.0.0

---
 src/spatialdata/_core/operations/vectorize.py | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

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

From a07b2863cc0c78d508cf5a4dc6b8e481ebd472f3 Mon Sep 17 00:00:00 2001
From: Luca Marconato <m.lucalmer@gmail.com>
Date: Wed, 19 Jun 2024 00:29:42 +0200
Subject: [PATCH 07/19] fix tests; fix bugs around DataArray/DataTree

---
 src/spatialdata/_core/centroids.py      |  3 +--
 src/spatialdata/_core/operations/map.py | 18 +++++++++---------
 src/spatialdata/_io/io_raster.py        |  4 +---
 src/spatialdata/models/models.py        |  2 +-
 tests/conftest.py                       |  7 +++----
 tests/core/operations/test_map.py       | 10 +++++-----
 6 files changed, 20 insertions(+), 24 deletions(-)

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
index e43d7086b..6806bee8f 100644
--- a/src/spatialdata/_core/operations/map.py
+++ b/src/spatialdata/_core/operations/map.py
@@ -6,8 +6,8 @@
 
 import dask.array as da
 from dask.array.overlap import coerce_depth
-from multiscale_spatial_image import MultiscaleSpatialImage
-from spatial_image import SpatialImage
+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
@@ -17,7 +17,7 @@
 
 
 def map_raster(
-    data: SpatialImage | MultiscaleSpatialImage,
+    data: DataArray | DataTree,
     func: Callable[[da.Array], da.Array],
     fn_kwargs: Mapping[str, Any] = MappingProxyType({}),
     chunkwise: bool = True,
@@ -28,14 +28,14 @@ def map_raster(
     dims: tuple[str, ...] | None = None,
     transformations: dict[str, Any] | None = None,
     **kwargs: Any,
-) -> SpatialImage:
+) -> DataArray:
     """
     Apply a function to raster data, for each chunk and each scale.
 
     Parameters
     ----------
     data
-        The data to process. It can be a `SpatialImage` or `MultiscaleSpatialImage`. If it's a `MultiscaleSpatialImage`,
+        The data to process. It can be a `DataArray` or `DataTree`. If it's a `DataTree`,
         the function is applied to the first scale (full-resolution data).
     func
         The function to apply to the data.
@@ -71,14 +71,14 @@ def map_raster(
 
     Returns
     -------
-    The processed data as a `SpatialImage`.
+    The processed data as a `DataArray`.
     """
-    if isinstance(data, SpatialImage):
+    if isinstance(data, DataArray):
         arr = data.data
-    elif isinstance(data, MultiscaleSpatialImage):
+    elif isinstance(data, DataTree):
         arr = data["scale0"].values().__iter__().__next__().data
     else:
-        raise ValueError("Only 'SpatialImage' and 'MultiscaleSpatialImage' are supported.")
+        raise ValueError("Only 'DataArray' and 'DataTree' are supported.")
 
     model = get_model(data)
     if model in (Labels2DModel, Labels3DModel) and c_coords is not None:
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/models.py b/src/spatialdata/models/models.py
index 27dc13a31..91f939343 100644
--- a/src/spatialdata/models/models.py
+++ b/src/spatialdata/models/models.py
@@ -182,7 +182,7 @@ def parse(
                 ) from e
 
         # finally convert to spatial image
-        data = to_spatial_image(array_like=data, dims=cls.dims.dims, **kwargs)
+        data = DataArray(to_spatial_image(array_like=data, dims=cls.dims.dims, **kwargs))
         # parse transformations
         _parse_transformations(data, transformations)
         # convert to multiscale if needed
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
index f5c12944a..bdc743b67 100644
--- a/tests/core/operations/test_map.py
+++ b/tests/core/operations/test_map.py
@@ -2,9 +2,9 @@
 
 import numpy as np
 import pytest
-from spatial_image import SpatialImage
 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):
@@ -43,7 +43,7 @@ def test_map_raster(sdata_blobs, depth, element_name):
         depth=depth,
     )
 
-    assert isinstance(se, SpatialImage)
+    assert isinstance(se, DataArray)
     data = sdata_blobs[element_name].data.compute()
     res = se.data.compute()
     assert np.array_equal(data * fn_kwargs["parameter"], res)
@@ -84,7 +84,7 @@ def test_map_raster_no_chunkwise(sdata_blobs):
         depth=None,
     )
 
-    assert isinstance(se, SpatialImage)
+    assert isinstance(se, DataArray)
     data = sdata_blobs[img_layer].data.compute()
     res = se.data.compute()
     assert np.array_equal(data * fn_kwargs["parameter"], res)
@@ -108,7 +108,7 @@ def test_map_raster_output_chunks(sdata_blobs):
         depth=(0, depth, depth),
     )
 
-    assert isinstance(se, SpatialImage)
+    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()
@@ -154,7 +154,7 @@ def test_map_remove_z_fails(full_sdata):
 
 
 def test_invalid_map_raster(sdata_blobs):
-    with pytest.raises(ValueError, match="Only 'SpatialImage' and 'MultiscaleSpatialImage' are supported."):
+    with pytest.raises(ValueError, match="Only 'DataArray' and 'DataTree' are supported."):
         map_raster(
             sdata_blobs["blobs_points"],
             func=_multiply,

From 71c46654b5d85742d7ee71e3ee630533f276e339 Mon Sep 17 00:00:00 2001
From: Wouter-Michiel Vierdag <w-mv@hotmail.com>
Date: Wed, 19 Jun 2024 00:59:27 +0200
Subject: [PATCH 08/19] Update src/spatialdata/models/models.py

---
 src/spatialdata/models/models.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/spatialdata/models/models.py b/src/spatialdata/models/models.py
index 91f939343..27dc13a31 100644
--- a/src/spatialdata/models/models.py
+++ b/src/spatialdata/models/models.py
@@ -182,7 +182,7 @@ def parse(
                 ) from e
 
         # finally convert to spatial image
-        data = DataArray(to_spatial_image(array_like=data, dims=cls.dims.dims, **kwargs))
+        data = to_spatial_image(array_like=data, dims=cls.dims.dims, **kwargs)
         # parse transformations
         _parse_transformations(data, transformations)
         # convert to multiscale if needed

From 3f85ce14d887f4e16a7a37e9221ac622e4f6bfb7 Mon Sep 17 00:00:00 2001
From: ArneDefauw <arnedefauw@gmail.com>
Date: Wed, 19 Jun 2024 08:29:38 +0200
Subject: [PATCH 09/19] apply_raster squeeze z

---
 tests/core/operations/test_map.py | 26 +++++++++++++++++++++++---
 1 file changed, 23 insertions(+), 3 deletions(-)

diff --git a/tests/core/operations/test_map.py b/tests/core/operations/test_map.py
index bdc743b67..8606ac7bf 100644
--- a/tests/core/operations/test_map.py
+++ b/tests/core/operations/test_map.py
@@ -135,11 +135,31 @@ def test_map_transformation(sdata_blobs, img_layer):
     assert transformation == get_transformation(se, to_coordinate_system=target_coordinate_system)
 
 
-def test_map_remove_z_fails(full_sdata):
+def test_map_squeeze_z(full_sdata):
+    img_layer = "image3d_numpy"
+    fn_kwargs = {"parameter": 20}
+
+    se = map_raster(
+        full_sdata[img_layer],
+        func=_multiply_squeeze_z,
+        fn_kwargs=fn_kwargs,
+        input_chunks=100,
+        output_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, ...] * fn_kwargs["parameter"], res)
+
+
+def test_map_squeeze_z_fails(full_sdata):
     fn_kwargs = {"parameter": 20}
 
-    # currently can not alter dims, e.g. ("c","z","y","x") -> ("c","y","x") fails
-    # could be supported by adding dims (and possibly transformations) to parameters of map_raster
     with pytest.raises(IndexError):
         map_raster(
             full_sdata["image3d_numpy"],

From 4bf6ff15dfd23994f27282d94e646d9fd95422b1 Mon Sep 17 00:00:00 2001
From: ArneDefauw <arnedefauw@gmail.com>
Date: Wed, 19 Jun 2024 13:51:05 +0200
Subject: [PATCH 10/19] replace input_chunks and output_chunks by chunks

---
 src/spatialdata/_core/operations/map.py | 27 +++++++++----------------
 tests/core/operations/test_map.py       | 24 ++++++++++------------
 2 files changed, 21 insertions(+), 30 deletions(-)

diff --git a/src/spatialdata/_core/operations/map.py b/src/spatialdata/_core/operations/map.py
index 6806bee8f..df264e7f5 100644
--- a/src/spatialdata/_core/operations/map.py
+++ b/src/spatialdata/_core/operations/map.py
@@ -20,17 +20,16 @@ def map_raster(
     data: DataArray | DataTree,
     func: Callable[[da.Array], da.Array],
     fn_kwargs: Mapping[str, Any] = MappingProxyType({}),
-    chunkwise: bool = True,
+    blockwise: bool = True,
     depth: int | tuple[int, ...] | dict[int, int] | None = None,
-    input_chunks: tuple[tuple[int, ...], ...] | None = None,
-    output_chunks: tuple[tuple[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 function to raster data, for each chunk and each scale.
+    Apply a function to raster data.
 
     Parameters
     ----------
@@ -41,21 +40,19 @@ def map_raster(
         The function to apply to the data.
     fn_kwargs
         Additional keyword arguments to pass to the function `func`.
-    chunkwise
+    blockwise
         If `True`, distributed processing will be achieved with `dask.array.map_overlap`/`dask.array.map_blocks`,
-        otherwise the function is applied to the full data. If `False`, `depth` and `input_chunks` are ignored.
+        otherwise the function is applied to the full data. If `False`, `depth` and `chunks` are ignored.
     depth
         If not `None`, distributed processing will be achieved with `dask.array.map_overlap`, otherwise with
         `dask.array.map_blocks`. Specifies the overlap between chunks, i.e. the number of elements that each chunk
         should share with its neighbor chunks. Please see `dask.array.map_overlap` for more information on the accepted
         values.
-    input_chunks
-        If specified, rechunks the input data before applying the function using `dask.array.rechunk`.
-    output_chunks
+    chunks
+        Passed to `dask.array.map_overlap`/`dask.array.map_blocks` as `chunks`. Ignored if `blockwise` is `False`.
         Chunk shape of resulting blocks if the function does not preserve the data shape. If not provided, the resulting
         array is assumed to have the same chunk structure as the first input array.
-        Passed to `dask.array.map_overlap`/`dask.array.map_blocks` as `chunks`.
-        E.g. ( (3,), (256,), (256,) ).
+        E.g. ( (3,), (100,100), (100,100) ).
     c_coords
         The channel coordinates for the output data. If not provided, the channel coordinates of the input data are
         used. It should be specified if the function changes the number of channels.
@@ -85,15 +82,11 @@ def map_raster(
         raise ValueError("Channel coordinates can not be provided for labels data.")
 
     kwargs = kwargs.copy()
-    kwargs["chunks"] = output_chunks
+    kwargs["chunks"] = chunks
 
-    if not chunkwise:
+    if not blockwise:
         arr = func(arr, **fn_kwargs)
-        if output_chunks is not None:
-            arr = arr.rechunk(output_chunks)
     else:
-        if input_chunks is not None:
-            arr = arr.rechunk(input_chunks)
         if depth is not None:
             kwargs.setdefault("boundary", "reflect")
 
diff --git a/tests/core/operations/test_map.py b/tests/core/operations/test_map.py
index 8606ac7bf..1629c6850 100644
--- a/tests/core/operations/test_map.py
+++ b/tests/core/operations/test_map.py
@@ -72,14 +72,14 @@ def test_map_raster_multiscale(sdata_blobs, depth):
     assert np.array_equal(data * fn_kwargs["parameter"], res)
 
 
-def test_map_raster_no_chunkwise(sdata_blobs):
+def test_map_raster_no_blockwise(sdata_blobs):
     img_layer = "blobs_image"
     fn_kwargs = {"parameter": 20}
     se = map_raster(
         sdata_blobs[img_layer],
         func=_multiply,
         fn_kwargs=fn_kwargs,
-        chunkwise=False,
+        blockwise=False,
         c_coords=None,
         depth=None,
     )
@@ -95,11 +95,10 @@ def test_map_raster_output_chunks(sdata_blobs):
     fn_kwargs = {"parameter": 20}
     output_channels = ["test"]
     se = map_raster(
-        sdata_blobs["blobs_image"],
+        sdata_blobs["blobs_image"].chunk((3, 100, 100)),
         func=_multiply_alter_c,
         fn_kwargs=fn_kwargs,
-        input_chunks=(3, 100, 100),
-        output_chunks=(
+        chunks=(
             (1,),
             (100 + 2 * depth, 96 + 2 * depth, 60 + 2 * depth),
             (100 + 2 * depth, 96 + 2 * depth, 60 + 2 * depth),
@@ -125,10 +124,10 @@ def test_map_transformation(sdata_blobs, img_layer):
 
     set_transformation(se, transformation=transformation, to_coordinate_system=target_coordinate_system)
     se = map_raster(
-        sdata_blobs[img_layer],
+        se,
         func=_multiply,
         fn_kwargs=fn_kwargs,
-        chunkwise=False,
+        blockwise=False,
         c_coords=None,
         depth=None,
     )
@@ -140,11 +139,10 @@ def test_map_squeeze_z(full_sdata):
     fn_kwargs = {"parameter": 20}
 
     se = map_raster(
-        full_sdata[img_layer],
+        full_sdata[img_layer].chunk((3, 2, 64, 64)),
         func=_multiply_squeeze_z,
         fn_kwargs=fn_kwargs,
-        input_chunks=100,
-        output_chunks=((3,), (64,), (64,)),
+        chunks=((3,), (64,), (64,)),
         drop_axis=1,
         c_coords=None,
         dims=("c", "y", "x"),
@@ -158,15 +156,15 @@ def test_map_squeeze_z(full_sdata):
 
 
 def test_map_squeeze_z_fails(full_sdata):
+    img_layer = "image3d_numpy"
     fn_kwargs = {"parameter": 20}
 
     with pytest.raises(IndexError):
         map_raster(
-            full_sdata["image3d_numpy"],
+            full_sdata[img_layer].chunk((3, 2, 64, 64)),
             func=_multiply_squeeze_z,
             fn_kwargs=fn_kwargs,
-            input_chunks=100,
-            output_chunks=((3,), (64,), (64,)),
+            chunks=((3,), (64,), (64,)),
             drop_axis=1,
             c_coords=None,
             depth=None,

From a9698ee633ae3f95b2c1f6ac5c375afc6a393444 Mon Sep 17 00:00:00 2001
From: ArneD <arnedefauw@gmail.com>
Date: Thu, 20 Jun 2024 08:51:20 +0200
Subject: [PATCH 11/19] Update src/spatialdata/_core/operations/map.py

Co-authored-by: Giovanni Palla <25887487+giovp@users.noreply.github.com>
---
 src/spatialdata/_core/operations/map.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/spatialdata/_core/operations/map.py b/src/spatialdata/_core/operations/map.py
index df264e7f5..afbe68a46 100644
--- a/src/spatialdata/_core/operations/map.py
+++ b/src/spatialdata/_core/operations/map.py
@@ -37,7 +37,7 @@ def map_raster(
         The data to process. It can be a `DataArray` or `DataTree`. If it's a `DataTree`,
         the function is applied to the first scale (full-resolution data).
     func
-        The function to apply to the data.
+        The callable that is applied to the data.
     fn_kwargs
         Additional keyword arguments to pass to the function `func`.
     blockwise

From 3d1860bca000465ebfadd72ae23f1880f60b28dc Mon Sep 17 00:00:00 2001
From: ArneD <arnedefauw@gmail.com>
Date: Thu, 20 Jun 2024 08:53:29 +0200
Subject: [PATCH 12/19] Update src/spatialdata/_core/operations/map.py

Co-authored-by: Giovanni Palla <25887487+giovp@users.noreply.github.com>
---
 src/spatialdata/_core/operations/map.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/spatialdata/_core/operations/map.py b/src/spatialdata/_core/operations/map.py
index afbe68a46..d581d6035 100644
--- a/src/spatialdata/_core/operations/map.py
+++ b/src/spatialdata/_core/operations/map.py
@@ -44,10 +44,10 @@ def map_raster(
         If `True`, distributed processing will be achieved with `dask.array.map_overlap`/`dask.array.map_blocks`,
         otherwise the function is applied to the full data. If `False`, `depth` and `chunks` are ignored.
     depth
-        If not `None`, distributed processing will be achieved with `dask.array.map_overlap`, otherwise with
-        `dask.array.map_blocks`. Specifies the overlap between chunks, i.e. the number of elements that each chunk
-        should share with its neighbor chunks. Please see `dask.array.map_overlap` for more information on the accepted
-        values.
+        Specifies the overlap between chunks, i.e. the number of elements that each chunk
+        should share with its neighbor chunks. If not `None`, distributed processing will be achieved with 
+        `dask.array.map_overlap`, otherwise with `dask.array.map_blocks`.  Please see 
+        :func:`dask.array.map_overlap` for more information on the accepted values.
     chunks
         Passed to `dask.array.map_overlap`/`dask.array.map_blocks` as `chunks`. Ignored if `blockwise` is `False`.
         Chunk shape of resulting blocks if the function does not preserve the data shape. If not provided, the resulting

From 27de8aafbea5a06952f2d62398386fa5e58d997b Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Thu, 20 Jun 2024 06:53:55 +0000
Subject: [PATCH 13/19] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 src/spatialdata/_core/operations/map.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/spatialdata/_core/operations/map.py b/src/spatialdata/_core/operations/map.py
index d581d6035..b2a2ae018 100644
--- a/src/spatialdata/_core/operations/map.py
+++ b/src/spatialdata/_core/operations/map.py
@@ -45,8 +45,8 @@ def map_raster(
         otherwise the function is applied to the full data. If `False`, `depth` and `chunks` are ignored.
     depth
         Specifies the overlap between chunks, i.e. the number of elements that each chunk
-        should share with its neighbor chunks. If not `None`, distributed processing will be achieved with 
-        `dask.array.map_overlap`, otherwise with `dask.array.map_blocks`.  Please see 
+        should share with its neighbor chunks. If not `None`, distributed processing will be achieved with
+        `dask.array.map_overlap`, otherwise with `dask.array.map_blocks`.  Please see
         :func:`dask.array.map_overlap` for more information on the accepted values.
     chunks
         Passed to `dask.array.map_overlap`/`dask.array.map_blocks` as `chunks`. Ignored if `blockwise` is `False`.

From 5f8db56740a862d9ab7b3894cfc4f7283930c384 Mon Sep 17 00:00:00 2001
From: ArneDefauw <arnedefauw@gmail.com>
Date: Thu, 20 Jun 2024 13:28:43 +0200
Subject: [PATCH 14/19] apply raster c_coords + docs

---
 src/spatialdata/_core/operations/map.py | 40 ++++++++++++---------
 tests/core/operations/test_map.py       | 46 ++++++++++++-------------
 2 files changed, 47 insertions(+), 39 deletions(-)

diff --git a/src/spatialdata/_core/operations/map.py b/src/spatialdata/_core/operations/map.py
index b2a2ae018..697f6940a 100644
--- a/src/spatialdata/_core/operations/map.py
+++ b/src/spatialdata/_core/operations/map.py
@@ -19,7 +19,7 @@
 def map_raster(
     data: DataArray | DataTree,
     func: Callable[[da.Array], da.Array],
-    fn_kwargs: Mapping[str, Any] = MappingProxyType({}),
+    func_kwargs: Mapping[str, Any] = MappingProxyType({}),
     blockwise: bool = True,
     depth: int | tuple[int, ...] | dict[int, int] | None = None,
     chunks: tuple[tuple[int, ...], ...] | None = None,
@@ -29,42 +29,49 @@ def map_raster(
     **kwargs: Any,
 ) -> DataArray:
     """
-    Apply a function to raster data.
+    Apply a callable to raster data.
+
+    Applies a callable (`func`) to raster data. If `blockwise` is set to True,
+    distributed processing will be achieved with `dask.array.map_overlap`/`dask.array.map_blocks`,
+    otherwise `func` is appplied to the full data.
 
     Parameters
     ----------
     data
         The data to process. It can be a `DataArray` or `DataTree`. If it's a `DataTree`,
-        the function is applied to the first scale (full-resolution data).
+        the callable is applied to the first scale (full-resolution data).
     func
         The callable that is applied to the data.
-    fn_kwargs
-        Additional keyword arguments to pass to the function `func`.
+    func_kwargs
+        Additional keyword arguments to pass to the callable `func`.
     blockwise
         If `True`, distributed processing will be achieved with `dask.array.map_overlap`/`dask.array.map_blocks`,
-        otherwise the function is applied to the full data. If `False`, `depth` and `chunks` are ignored.
+        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 neighbor chunks. If not `None`, distributed processing will be achieved with
         `dask.array.map_overlap`, otherwise with `dask.array.map_blocks`.  Please see
         :func:`dask.array.map_overlap` for more information on the accepted values.
     chunks
-        Passed to `dask.array.map_overlap`/`dask.array.map_blocks` as `chunks`. Ignored if `blockwise` is `False`.
-        Chunk shape of resulting blocks if the function does not preserve the data shape. If not provided, the resulting
+        Chunk shape of resulting blocks if the callable does not preserve the data shape. If not provided, the resulting
         array is assumed to have the same chunk structure as the first input array.
         E.g. ( (3,), (100,100), (100,100) ).
+        Passed to `dask.array.map_overlap`/`dask.array.map_blocks` as `chunks`. Ignored if `blockwise` is `False`.
+        Please see :func:`dask.array.map_blocks` for more information on the accepted values.
     c_coords
         The channel coordinates for the output data. If not provided, the channel coordinates of the input data are
-        used. It should be specified if the function changes the number of channels.
+        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 function changes the data dimensions.
+        specified if the callable changes the data dimensions.
         E.g. ('c', '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 function changes the data transformations.
+        output data. It should be specified if the callable changes the data transformations.
     kwargs
-        Additional keyword arguments to pass to `dask.array.map_overlap` or `dask.array.map_blocks`.
+        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
     -------
@@ -85,7 +92,7 @@ def map_raster(
     kwargs["chunks"] = chunks
 
     if not blockwise:
-        arr = func(arr, **fn_kwargs)
+        arr = func(arr, **func_kwargs)
     else:
         if depth is not None:
             kwargs.setdefault("boundary", "reflect")
@@ -93,18 +100,19 @@ def map_raster(
             if not isinstance(depth, int) and len(depth) != arr.ndim:
                 raise ValueError(
                     f"Depth {depth} is provided for {len(depth)} dimensions. "
-                    f"Please (only) provide depth for {arr.ndim} 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, **fn_kwargs, **kwargs, dtype=arr.dtype)
+        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):
-        c_coords = c_coords if c_coords is not None else get_channels(data)
+        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:
diff --git a/tests/core/operations/test_map.py b/tests/core/operations/test_map.py
index 1629c6850..01e7081d6 100644
--- a/tests/core/operations/test_map.py
+++ b/tests/core/operations/test_map.py
@@ -34,11 +34,11 @@ def test_map_raster(sdata_blobs, depth, element_name):
     if element_name == "blobs_labels" and depth is not None:
         depth = (60, 60)
 
-    fn_kwargs = {"parameter": 20}
+    func_kwargs = {"parameter": 20}
     se = map_raster(
         sdata_blobs[element_name],
         func=_multiply,
-        fn_kwargs=fn_kwargs,
+        func_kwargs=func_kwargs,
         c_coords=None,
         depth=depth,
     )
@@ -46,7 +46,7 @@ def test_map_raster(sdata_blobs, depth, element_name):
     assert isinstance(se, DataArray)
     data = sdata_blobs[element_name].data.compute()
     res = se.data.compute()
-    assert np.array_equal(data * fn_kwargs["parameter"], res)
+    assert np.array_equal(data * func_kwargs["parameter"], res)
 
 
 @pytest.mark.parametrize(
@@ -58,27 +58,27 @@ def test_map_raster(sdata_blobs, depth, element_name):
 )
 def test_map_raster_multiscale(sdata_blobs, depth):
     img_layer = "blobs_multiscale_image"
-    fn_kwargs = {"parameter": 20}
+    func_kwargs = {"parameter": 20}
     se = map_raster(
         sdata_blobs[img_layer],
         func=_multiply,
-        fn_kwargs=fn_kwargs,
+        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 * fn_kwargs["parameter"], res)
+    assert np.array_equal(data * func_kwargs["parameter"], res)
 
 
 def test_map_raster_no_blockwise(sdata_blobs):
     img_layer = "blobs_image"
-    fn_kwargs = {"parameter": 20}
+    func_kwargs = {"parameter": 20}
     se = map_raster(
         sdata_blobs[img_layer],
         func=_multiply,
-        fn_kwargs=fn_kwargs,
+        func_kwargs=func_kwargs,
         blockwise=False,
         c_coords=None,
         depth=None,
@@ -87,17 +87,17 @@ def test_map_raster_no_blockwise(sdata_blobs):
     assert isinstance(se, DataArray)
     data = sdata_blobs[img_layer].data.compute()
     res = se.data.compute()
-    assert np.array_equal(data * fn_kwargs["parameter"], res)
+    assert np.array_equal(data * func_kwargs["parameter"], res)
 
 
 def test_map_raster_output_chunks(sdata_blobs):
     depth = 60
-    fn_kwargs = {"parameter": 20}
+    func_kwargs = {"parameter": 20}
     output_channels = ["test"]
     se = map_raster(
         sdata_blobs["blobs_image"].chunk((3, 100, 100)),
         func=_multiply_alter_c,
-        fn_kwargs=fn_kwargs,
+        func_kwargs=func_kwargs,
         chunks=(
             (1,),
             (100 + 2 * depth, 96 + 2 * depth, 60 + 2 * depth),
@@ -111,12 +111,12 @@ def test_map_raster_output_chunks(sdata_blobs):
     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] * fn_kwargs["parameter"], res[0])
+    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):
-    fn_kwargs = {"parameter": 20}
+    func_kwargs = {"parameter": 20}
     target_coordinate_system = "my_other_space0"
     transformation = Translation(translation=[10, 12], axes=["y", "x"])
 
@@ -126,7 +126,7 @@ def test_map_transformation(sdata_blobs, img_layer):
     se = map_raster(
         se,
         func=_multiply,
-        fn_kwargs=fn_kwargs,
+        func_kwargs=func_kwargs,
         blockwise=False,
         c_coords=None,
         depth=None,
@@ -136,12 +136,12 @@ def test_map_transformation(sdata_blobs, img_layer):
 
 def test_map_squeeze_z(full_sdata):
     img_layer = "image3d_numpy"
-    fn_kwargs = {"parameter": 20}
+    func_kwargs = {"parameter": 20}
 
     se = map_raster(
         full_sdata[img_layer].chunk((3, 2, 64, 64)),
         func=_multiply_squeeze_z,
-        fn_kwargs=fn_kwargs,
+        func_kwargs=func_kwargs,
         chunks=((3,), (64,), (64,)),
         drop_axis=1,
         c_coords=None,
@@ -152,18 +152,18 @@ def test_map_squeeze_z(full_sdata):
     assert isinstance(se, DataArray)
     data = full_sdata[img_layer].data.compute()
     res = se.data.compute()
-    assert np.array_equal(data[:, 0, ...] * fn_kwargs["parameter"], res)
+    assert np.array_equal(data[:, 0, ...] * func_kwargs["parameter"], res)
 
 
 def test_map_squeeze_z_fails(full_sdata):
     img_layer = "image3d_numpy"
-    fn_kwargs = {"parameter": 20}
+    func_kwargs = {"parameter": 20}
 
     with pytest.raises(IndexError):
         map_raster(
             full_sdata[img_layer].chunk((3, 2, 64, 64)),
             func=_multiply_squeeze_z,
-            fn_kwargs=fn_kwargs,
+            func_kwargs=func_kwargs,
             chunks=((3,), (64,), (64,)),
             drop_axis=1,
             c_coords=None,
@@ -176,19 +176,19 @@ def test_invalid_map_raster(sdata_blobs):
         map_raster(
             sdata_blobs["blobs_points"],
             func=_multiply,
-            fn_kwargs={"parameter": 20},
+            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 (only) provide depth for 3 dimensions."),
+        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,
-            fn_kwargs={"parameter": 20},
+            func_kwargs={"parameter": 20},
             c_coords=None,
             depth=(0, 60),
         )
@@ -197,7 +197,7 @@ def test_invalid_map_raster(sdata_blobs):
         map_raster(
             sdata_blobs["blobs_labels"],
             func=_multiply,
-            fn_kwargs={"parameter": 20},
+            func_kwargs={"parameter": 20},
             c_coords=["c"],
             depth=(0, 60, 60),
         )

From f2cd4db0fa8e936f72f48d700af127bc768e5b29 Mon Sep 17 00:00:00 2001
From: giovp <giov.pll@gmail.com>
Date: Fri, 21 Jun 2024 11:47:56 +0200
Subject: [PATCH 15/19] improve docstrings

---
 src/spatialdata/_core/operations/map.py | 37 +++++++++++++------------
 1 file changed, 19 insertions(+), 18 deletions(-)

diff --git a/src/spatialdata/_core/operations/map.py b/src/spatialdata/_core/operations/map.py
index 697f6940a..db85c4771 100644
--- a/src/spatialdata/_core/operations/map.py
+++ b/src/spatialdata/_core/operations/map.py
@@ -31,41 +31,42 @@ def map_raster(
     """
     Apply a callable to raster data.
 
-    Applies a callable (`func`) to raster data. If `blockwise` is set to True,
-    distributed processing will be achieved with `dask.array.map_overlap`/`dask.array.map_blocks`,
-    otherwise `func` is appplied to the full 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 `DataArray` or `DataTree`. If it's a `DataTree`,
-        the callable is applied to the first scale (full-resolution 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`, distributed processing will be achieved with `dask.array.map_overlap`/`dask.array.map_blocks`,
+        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 neighbor chunks. If not `None`, distributed processing will be achieved with
-        `dask.array.map_overlap`, otherwise with `dask.array.map_blocks`.  Please see
-        :func:`dask.array.map_overlap` for more information on the accepted values.
+        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. If not provided, the resulting
-        array is assumed to have the same chunk structure as the first input array.
-        E.g. ( (3,), (100,100), (100,100) ).
-        Passed to `dask.array.map_overlap`/`dask.array.map_blocks` as `chunks`. Ignored if `blockwise` is `False`.
-        Please see :func:`dask.array.map_blocks` for more information on the accepted values.
+        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)).
+        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').
+        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.
@@ -75,7 +76,7 @@ def map_raster(
 
     Returns
     -------
-    The processed data as a `DataArray`.
+    The processed data as a :class:`xarray.DataArray`.
     """
     if isinstance(data, DataArray):
         arr = data.data

From ad7550cd4885143b97a4c331afcd7c3f330c64ef Mon Sep 17 00:00:00 2001
From: giovp <giov.pll@gmail.com>
Date: Fri, 21 Jun 2024 11:49:11 +0200
Subject: [PATCH 16/19] add  to api

---
 docs/api.md | 1 +
 1 file changed, 1 insertion(+)

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

From 9ac6d4090e901c348ddbe0e5887e8525b87bca02 Mon Sep 17 00:00:00 2001
From: giovp <giov.pll@gmail.com>
Date: Fri, 21 Jun 2024 11:59:37 +0200
Subject: [PATCH 17/19] add import

---
 src/spatialdata/__init__.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/spatialdata/__init__.py b/src/spatialdata/__init__.py
index 16e44d5a4..820dc95ee 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",
 ]
 

From 0594c42e8fbe34a3cc66e89c738d93cdfeb3e559 Mon Sep 17 00:00:00 2001
From: giovp <giov.pll@gmail.com>
Date: Fri, 21 Jun 2024 12:03:32 +0200
Subject: [PATCH 18/19] finish import

---
 src/spatialdata/__init__.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/spatialdata/__init__.py b/src/spatialdata/__init__.py
index 820dc95ee..2d27f0a41 100644
--- a/src/spatialdata/__init__.py
+++ b/src/spatialdata/__init__.py
@@ -41,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

From fbf25d5eca80f5d0852ddc7d46b5aec877394486 Mon Sep 17 00:00:00 2001
From: giovp <giov.pll@gmail.com>
Date: Fri, 21 Jun 2024 12:07:31 +0200
Subject: [PATCH 19/19] make mypy happy

---
 src/spatialdata/_core/operations/map.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/src/spatialdata/_core/operations/map.py b/src/spatialdata/_core/operations/map.py
index db85c4771..c064eaa10 100644
--- a/src/spatialdata/_core/operations/map.py
+++ b/src/spatialdata/_core/operations/map.py
@@ -2,7 +2,7 @@
 
 from collections.abc import Iterable, Mapping
 from types import MappingProxyType
-from typing import Any, Callable
+from typing import TYPE_CHECKING, Any, Callable
 
 import dask.array as da
 from dask.array.overlap import coerce_depth
@@ -118,7 +118,8 @@ def map_raster(
         c_coords = None
     if transformations is None:
         d = get_transformation(data, get_all=True)
-        assert isinstance(d, dict)
+        if TYPE_CHECKING:
+            assert isinstance(d, dict)
         transformations = d
 
     model_kwargs = {