test and add support for channels last and its interaction with slicing operations

Author: pattonw

funlib/persistence/arrays/array.py

@@ -97,17 +97,22 @@ def chunk_shape(self) -> Coordinate:
 
     def uncollapsed_dims(self, physical: bool = False) -> list[bool]:
         if physical:
-            return self._uncollapsed_dims[-self._metadata.voxel_size.dims :]
+            return [
+                x
+                for x, c in zip(self._uncollapsed_dims, self._metadata.axis_names)
+                if not c.endswith("^")
+            ]
         else:
             return self._uncollapsed_dims
 
     @property
     def offset(self) -> Coordinate:
         """Get the offset of this array in world units."""
+        udims = self.uncollapsed_dims(physical=True)
         return Coordinate(
             [
                 self._metadata.offset[ii]
-                for ii, uncollapsed in enumerate(self.uncollapsed_dims(physical=True))
+                for ii, uncollapsed in enumerate(udims)
                 if uncollapsed
             ]
         )
@@ -119,10 +124,11 @@ def offset(self, offset: Iterable[int]) -> None:
     @property
     def voxel_size(self) -> Coordinate:
         """Get the size of a voxel in world units."""
+        udims = self.uncollapsed_dims(physical=True)
         return Coordinate(
             [
                 self._metadata.voxel_size[ii]
-                for ii, uncollapsed in enumerate(self.uncollapsed_dims(physical=True))
+                for ii, uncollapsed in enumerate(udims)
                 if uncollapsed
             ]
         )
@@ -133,9 +139,10 @@ def voxel_size(self, voxel_size: Iterable[int]) -> None:
 
     @property
     def units(self) -> list[str]:
+        udims = self.uncollapsed_dims(physical=True)
         return [
             self._metadata.units[ii]
-            for ii, uncollapsed in enumerate(self.uncollapsed_dims(physical=True))
+            for ii, uncollapsed in enumerate(udims)
             if uncollapsed
         ]
 
@@ -145,6 +152,7 @@ def units(self, units: list[str]) -> None:
 
     @property
     def axis_names(self) -> list[str]:
+        print(self._metadata.axis_names, self._uncollapsed_dims)
         return [
             self._metadata.axis_names[ii]
             for ii, uncollapsed in enumerate(self.uncollapsed_dims(physical=False))
@@ -205,7 +213,12 @@ def apply_adapter(self, adapter: Adapter):
                 adapter = (adapter,)
             for ii, a in enumerate(adapter):
                 if isinstance(a, int):
-                    self._uncollapsed_dims[ii] = False
+                    for i, uc in enumerate(self._uncollapsed_dims):
+                        if uc:
+                            if ii == 0:
+                                self._uncollapsed_dims[i] = False
+                                break
+                            ii -= 1
             self.data = self.data[adapter]
         elif callable(adapter):
             self.data = adapter(self.data)

tests/test_array.py

@@ -273,3 +273,80 @@ def test_slicing():
 
     with pytest.raises(RuntimeError):
         a[:, :] = np.array([42, 43, 44]).reshape(3, 1)
+
+
+def test_slicing_channel_dim_last():
+    a = Array(
+        np.arange(0, 4 * 4).reshape(2, 2, 4),
+        (0, 0),
+        (1, 1),
+        axis_names=["d0", "d1", "c0^"],
+    )
+
+    a.adapt(np.s_[1, :, 0:3])
+    assert a.shape == (2, 3)
+    assert a.axis_names == ["d1", "c0^"], a.axis_names
+    assert a.units == [""]
+
+    a.adapt(np.s_[:, 2])
+    assert a.shape == (2,)
+    assert a.axis_names == ["d1"]
+    assert a.units == [""]
+
+    a[:] = 42
+
+    assert all([x == 42 for x in a._source_data[1, :, 2]]), a._source_data[1, :, 2]
+
+    # test with list indexing
+    a = Array(
+        np.arange(0, 4 * 4).reshape(2, 2, 4),
+        (0, 0),
+        (1, 1),
+        axis_names=["d0", "d1", "c0^"],
+    )
+
+    a.adapt(np.s_[[0, 1], 1, :])
+    assert a.shape == (2, 4)
+    assert a.axis_names == ["d0", "c0^"]
+    assert a.units == [""]
+
+    a.adapt(np.s_[1, :])
+    assert a.shape == (4,)
+    assert a.axis_names == ["c0^"]
+    assert a.units == []
+
+    a[:] = 42
+
+    assert all([x == 42 for x in a._source_data[1, 1, :]]), a._source_data[1, 1, :]
+
+    # test weird case
+    a = Array(
+        np.arange(0, 4 * 4).reshape(4, 2, 2),
+        (0, 0),
+        (1, 1),
+        axis_names=["d0", "d1", "c0^"],
+    )
+
+    a.adapt(np.s_[[2, 2, 2], 1, :])
+    assert a.shape == (3, 2)
+    assert a.axis_names == ["d0", "c0^"]
+    assert a.units == [""]
+
+    a[:, :] = np.array([42, 43, 44]).reshape(3, 1)
+    assert all([x == 44 for x in a._source_data[2, 1, :]]), a._source_data[2, 1, :]
+
+    # test_bool_indexing
+    a = Array(
+        np.arange(0, 4 * 4).reshape(2, 2, 4),
+        (0, 0),
+        (1, 1),
+        axis_names=["d0", "d1", "c0^"],
+    )
+
+    a.adapt(np.s_[1, :, np.array([True, True, True, False])])
+    assert a.shape == (2, 3)
+    assert a.axis_names == ["d1", "c0^"]
+    assert a.units == [""]
+
+    with pytest.raises(RuntimeError):
+        a[:, :] = np.array([42, 43, 44]).reshape(3, 1)

tests/test_datasets.py

@@ -1,8 +1,7 @@
 from funlib.persistence.arrays.metadata import MetaDataFormat
-from funlib.persistence.arrays.datasets import open_ds, prepare_ds
+from funlib.persistence.arrays.datasets import open_ds, prepare_ds, ArrayNotFoundError
 from funlib.geometry import Coordinate, Roi
 
-from zarr.errors import ArrayNotFoundError
 import numpy as np
 
 import pytest
@@ -30,7 +29,7 @@ def test_helpers(tmpdir, store, dtype):
             "voxel_size": [1, 2, 3],
             "axis_names": ["sample^", "channel^", "z", "y", "x"],
             "units": ["nm", "nm", "nm"],
-        }
+        },
     )
 
     # test prepare_ds fails if array does not exist and mode is read
@@ -184,3 +183,76 @@ def test_helpers(tmpdir, store, dtype):
     assert array.offset == metadata.offset
     assert array.axis_names == metadata.axis_names
     assert array.units == metadata.units
+
+
+@pytest.mark.parametrize("store", stores.keys())
+@pytest.mark.parametrize("dtype", [np.float32, np.uint8, np.uint64])
+def test_open_ds(tmpdir, store, dtype):
+    shape = Coordinate(1, 1, 10, 20, 30)
+    store = tmpdir / store
+    metadata = MetaDataFormat().parse(
+        shape,
+        {
+            "offset": [100, 200, 400],
+            "voxel_size": [1, 2, 3],
+            "axis_names": ("sample^", "channel^", "z", "y", "x"),
+            "units": ("nm", "nm", "nm"),
+        },
+    )
+
+    # test open_ds fails if array does not exist and mode is read
+    with pytest.raises(ArrayNotFoundError):
+        open_ds(
+            store,
+            offset=metadata.offset,
+            voxel_size=metadata.voxel_size,
+            axis_names=metadata.axis_names,
+            units=metadata.units,
+            mode="r",
+        )
+
+    # test open_ds creates array if it does not exist and mode is write
+    array = prepare_ds(
+        store,
+        shape,
+        offset=metadata.offset,
+        voxel_size=metadata.voxel_size,
+        axis_names=metadata.axis_names,
+        units=metadata.units,
+        dtype=dtype,
+        mode="w",
+    )
+    assert array.roi == Roi(
+        metadata.offset, metadata.voxel_size * Coordinate(*shape[-3:])
+    )
+    assert array.voxel_size == metadata.voxel_size
+    assert array.offset == metadata.offset
+    assert array.axis_names == metadata.axis_names
+    assert array.units == metadata.units
+
+    # test open_ds opens array if it exists and mode is read
+    array = open_ds(
+        store,
+        offset=metadata.offset,
+        voxel_size=metadata.voxel_size,
+        axis_names=metadata.axis_names,
+        units=metadata.units,
+        mode="r",
+    )
+    assert array.roi == Roi(
+        metadata.offset, metadata.voxel_size * Coordinate(*shape[-3:])
+    )
+    assert array.voxel_size == metadata.voxel_size
+    assert array.offset == metadata.offset
+    assert array.axis_names == metadata.axis_names
+    assert array.units == metadata.units
+
+    # test open_ds fails if array exists and is opened in read mode
+    # with incompatible arguments
+    array = open_ds(
+        store,
+        offset=(1, 2, 3),
+        voxel_size=(1, 2, 3),
+        axis_names=None,
+        units=None,
+    )