Validate axes for writer

ome_zarr/writer.py

@@ -14,6 +14,47 @@
 LOGGER = logging.getLogger("ome_zarr.writer")
 
 
+def validate_axes_names(
+    ndim: int, axes: Union[str, List[str]] = None, fmt: Format = CurrentFormat()
+) -> Union[None, List[str]]:
+
+    if fmt.version not in ("0.1", "0.2"):
+        if axes is None:
+            if ndim == 2:
+                axes = ["y", "x"]
+            elif ndim == 5:
+                axes = ["t", "c", "z", "y", "x"]
+            else:
+                raise ValueError(
+                    "axes must be provided. Can't be guessed for 3D or 4D data"
+                )
+
+    if isinstance(axes, str):
+        axes = list(axes)
+
+    if axes is not None:
+        if len(axes) != ndim:
+            raise ValueError("axes length must match number of dimensions")
+        # from https://github.com/constantinpape/ome-ngff-implementations/
+        val_axes = tuple(axes)
+        if ndim == 2:
+            assert val_axes == ("y", "x"), str(val_axes)
+        elif ndim == 3:
+            assert val_axes in [("z", "y", "x"), ("c", "y", "x"), ("t", "y", "x")], str(
+                val_axes
+            )
+        elif ndim == 4:
+            assert val_axes in [
+                ("t", "z", "y", "x"),
+                ("c", "z", "y", "x"),
+                ("t", "c", "y", "x"),
+            ], str(val_axes)
+        else:
+            assert val_axes == ("t", "c", "z", "y", "x"), str(val_axes)
+
+    return axes
+
+
 def write_multiscale(
     pyramid: List,
     group: zarr.Group,
@@ -28,6 +69,8 @@ def write_multiscale(
     ----------
     pyramid: List of np.ndarray
       the image data to save. Largest level first
+      All image arrays MUST be up to 5-dimensional with dimensions
+      ordered (t, c, z, y, x)
     group: zarr.Group
       the group within the zarr store to store the data in
     chunks: int or tuple of ints,
@@ -41,25 +84,7 @@ def write_multiscale(
     """
 
     dims = len(pyramid[0].shape)
-    if fmt.version not in ("0.1", "0.2"):
-        if axes is None:
-            if dims == 2:
-                axes = ["y", "x"]
-            elif dims == 5:
-                axes = ["t", "c", "z", "y", "x"]
-            else:
-                raise ValueError(
-                    "axes must be provided. Can't be guessed for 3D or 4D data"
-                )
-        if len(axes) != dims:
-            raise ValueError("axes length must match number of dimensions")
-
-        if isinstance(axes, str):
-            axes = list(axes)
-
-        for dim in axes:
-            if dim not in ("t", "c", "z", "y", "x"):
-                raise ValueError("axes must each be one of 'x', 'y', 'z', 'c' or 't'")
+    axes = validate_axes_names(dims, axes, fmt)
 
     paths = []
     for path, dataset in enumerate(pyramid):
@@ -90,6 +115,8 @@ def write_image(
     image: np.ndarray
       the image data to save. A downsampling of the data will be computed
       if the scaler argument is non-None.
+      Image array MUST be up to 5-dimensional with dimensions
+      ordered (t, c, z, y, x)
     group: zarr.Group
       the group within the zarr store to store the data in
     chunks: int or tuple of ints,
@@ -115,6 +142,8 @@ def write_image(
         # v0.1 and v0.2 are strictly 5D
         shape_5d: Tuple[Any, ...] = (*(1,) * (5 - image.ndim), *image.shape)
         image = image.reshape(shape_5d)
+        # and we don't need axes
+        axes = None
 
     if chunks is not None:
         chunks = _retuple(chunks, image.shape)

tests/test_writer.py

@@ -8,7 +8,7 @@
 from ome_zarr.io import parse_url
 from ome_zarr.reader import Multiscales, Reader
 from ome_zarr.scale import Scaler
-from ome_zarr.writer import write_image
+from ome_zarr.writer import validate_axes_names, write_image
 
 
 class TestWriter:
@@ -77,3 +77,34 @@ def test_writer(self, shape, scaler, format_version):
         else:
             assert node.data[0].ndim == 5
         assert np.allclose(data, node.data[0][...].compute())
+
+    def test_dim_names(self):
+
+        v03 = FormatV03()
+
+        # v0.3 MUST specify axes for 3D or 4D data
+        with pytest.raises(ValueError):
+            validate_axes_names(3, axes=None, fmt=v03)
+
+        # ndims must match axes length
+        with pytest.raises(ValueError):
+            validate_axes_names(3, axes="yx", fmt=v03)
+
+        # axes must be ordered tczyx
+        with pytest.raises(AssertionError):
+            validate_axes_names(3, axes="yxt", fmt=v03)
+        with pytest.raises(AssertionError):
+            validate_axes_names(2, axes=["x", "y"], fmt=v03)
+
+        validate_axes_names(2, axes=["y", "x"], fmt=v03)
+        validate_axes_names(5, axes="tczyx", fmt=v03)
+
+        # check that write_image is checking axes
+        data = self.create_data((125, 125))
+        with pytest.raises(ValueError):
+            write_image(
+                image=data,
+                group=self.group,
+                fmt=v03,
+                axes="xyz",
+            )