add and test operations

multiscale_spatial_image/multiscale_spatial_image.py

@@ -5,7 +5,11 @@
 from collections.abc import MutableMapping, Hashable
 from pathlib import Path
 from zarr.storage import BaseStore
-from multiscale_spatial_image.operations import transpose
+from multiscale_spatial_image.operations import (
+    transpose,
+    reindex_data_arrays,
+    assign_coords,
+)
 
 
 @register_datatree_accessor("msi")
@@ -137,3 +141,25 @@ def transpose(self, *dims: Hashable) -> DataTree:
             reorder the dimensions to the order that the `dims` are specified..
         """
         return self._dt.map_over_datasets(transpose, *dims)
+
+    def reindex_data_arrays(
+        self,
+        indexers,
+        method=None,
+        tolerance=None,
+        copy=False,
+        fill_value=None,
+        **indexer_kwargs,
+    ):
+        return self._dt.map_over_datasets(
+            reindex_data_arrays,
+            indexers,
+            method,
+            tolerance,
+            copy,
+            fill_value,
+            *indexer_kwargs,
+        )
+
+    def assign_coords(self, coords, **coords_kwargs):
+        return self._dt.map_over_datasets(assign_coords, coords, *coords_kwargs)

multiscale_spatial_image/operations/__init__.py

@@ -1,3 +1,3 @@
-from .operations import assign_coords, transpose, reindex
+from .operations import assign_coords, transpose, reindex_data_arrays
 
-__all__ = ["assign_coords", "transpose", "reindex"]
+__all__ = ["assign_coords", "transpose", "reindex_data_arrays"]

multiscale_spatial_image/operations/operations.py

@@ -14,9 +14,5 @@ def transpose(ds: Dataset, *args: Any, **kwargs: Any) -> Dataset:
 
 
 @skip_non_dimension_nodes
-def reindex(ds: Dataset, *args: Any, **kwargs: Any) -> Dataset:
-    # A copy is required as a dataset view as used in map_over_datasets is not mutable
-    # TODO: Check whether setting item on wrapping datatree node would be better than copy or this can be dropped.
-    ds_copy = ds.copy()
-    ds_copy["image"] = ds_copy["image"].reindex(*args, **kwargs)
-    return ds_copy
+def reindex_data_arrays(ds: Dataset, *args: Any, **kwargs: Any) -> Dataset:
+    return ds["image"].reindex(*args, **kwargs).to_dataset()

test/conftest.py

@@ -0,0 +1,13 @@
+import pytest
+import numpy as np
+from spatial_image import to_spatial_image
+from multiscale_spatial_image import to_multiscale
+
+
+@pytest.fixture()
+def multiscale_data():
+    data = np.zeros((3, 200, 200))
+    dims = ("c", "y", "x")
+    scale_factors = [2, 2]
+    image = to_spatial_image(array_like=data, dims=dims)
+    return to_multiscale(image, scale_factors=scale_factors)

test/test_operations.py

@@ -0,0 +1,17 @@
+def test_transpose(multiscale_data):
+    multiscale_data = multiscale_data.msi.transpose("y", "x", "c")
+
+    for scale in list(multiscale_data.keys()):
+        assert multiscale_data[scale]["image"].dims == ("y", "x", "c")
+
+
+def test_reindex_arrays(multiscale_data):
+    multiscale_data = multiscale_data.msi.reindex_data_arrays({"c": ["r", "g", "b"]})
+    for scale in list(multiscale_data.keys()):
+        assert multiscale_data[scale].c.data.tolist() == ["r", "g", "b"]
+
+
+def test_assign_coords(multiscale_data):
+    multiscale_data = multiscale_data.msi.assign_coords({"c": ["r", "g", "b"]})
+    for scale in list(multiscale_data.keys()):
+        assert multiscale_data[scale].c.data.tolist() == ["r", "g", "b"]

test/test_utils.py

@@ -1,23 +1,15 @@
-import numpy as np
-from spatial_image import to_spatial_image
-from multiscale_spatial_image import skip_non_dimension_nodes, to_multiscale
+from multiscale_spatial_image import skip_non_dimension_nodes
 
 
-def test_skip_nodes():
-    data = np.zeros((2, 200, 200))
-    dims = ("c", "y", "x")
-    scale_factors = [2, 2]
-    image = to_spatial_image(array_like=data, dims=dims)
-    multiscale_img = to_multiscale(image, scale_factors=scale_factors)
-
+def test_skip_nodes(multiscale_data):
     @skip_non_dimension_nodes
     def transpose(ds, *args, **kwargs):
         return ds.transpose(*args, **kwargs)
 
-    for scale in list(multiscale_img.keys()):
-        assert multiscale_img[scale]["image"].dims == ("c", "y", "x")
-    multiscale_img.msi.transpose()
+    for scale in list(multiscale_data.keys()):
+        assert multiscale_data[scale]["image"].dims == ("c", "y", "x")
+
     # applying this function without skipping the root node would fail as the root node does not have dimensions.
-    result = multiscale_img.map_over_datasets(transpose, "y", "x", "c")
+    result = multiscale_data.map_over_datasets(transpose, "y", "x", "c")
     for scale in list(result.keys()):
         assert result[scale]["image"].dims == ("y", "x", "c")