Feature/OmeZarr Chunk Dims (#519)

* chunk dim param

* chunk dims param

* lint

* Update aicsimageio/writers/ome_zarr_writer.py

Co-authored-by: Sean LeRoy <41307451+SeanLeRoy@users.noreply.github.com>

* Update aicsimageio/writers/ome_zarr_writer.py

Co-authored-by: Sean LeRoy <41307451+SeanLeRoy@users.noreply.github.com>

* Update aicsimageio/writers/ome_zarr_writer.py

Co-authored-by: Sean LeRoy <41307451+SeanLeRoy@users.noreply.github.com>

* DimensionNames

* pass dimension order to build_chunk_dims

* OmeZarrWriter chunk tests

* chunk comparison

* pytest tmpdir fixture

* DimensionNames

* chunk_dim scaling for levels

* nplanes inside chunk dims = None

* ome-types 0.4.0 breaks build-main

---------

Co-authored-by: Sean LeRoy <41307451+SeanLeRoy@users.noreply.github.com>

Author: BrianWhitneyAI

aicsimageio/tests/writers/test_ome_zarr_writer.py

@@ -1,7 +1,7 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
 
-import shutil
+import pathlib
 from typing import Callable, List, Optional, Tuple
 
 import numpy as np
@@ -12,7 +12,7 @@
 from aicsimageio import exceptions
 from aicsimageio.writers import OmeZarrWriter
 
-from ..conftest import LOCAL, array_constructor, get_resource_write_full_path
+from ..conftest import array_constructor
 
 
 @array_constructor
@@ -69,14 +69,12 @@ def test_ome_zarr_writer_dims(
     expected_read_shape: Tuple[int, ...],
     expected_read_dim_order: str,
     filename: str,
+    tmpdir: pathlib.Path,
 ) -> None:
     # Create array
     arr = array_constructor(write_shape, dtype=np.uint8)
 
-    # Construct save end point
-    save_uri = get_resource_write_full_path(filename, LOCAL)
-    # clear out anything left over
-    shutil.rmtree(save_uri, ignore_errors=True)
+    save_uri = str(tmpdir / filename)
 
     # Normal save
     writer = OmeZarrWriter(save_uri)
@@ -132,14 +130,11 @@ def test_ome_zarr_writer_scaling(
     expected_read_shapes: List[Tuple[int, ...]],
     expected_read_scales: List[List[int]],
     filename: str,
+    tmpdir: pathlib.Path,
 ) -> None:
     # Create array
     arr = array_constructor(write_shape, dtype=np.uint8)
-
-    # Construct save end point
-    save_uri = get_resource_write_full_path(filename, LOCAL)
-    # clear out anything left over
-    shutil.rmtree(save_uri, ignore_errors=True)
+    save_uri = str(tmpdir / filename)
 
     # Normal save
     writer = OmeZarrWriter(save_uri)
@@ -160,3 +155,63 @@ def test_ome_zarr_writer_scaling(
         assert len(xforms) == 1
         assert xforms[0]["type"] == "scale"
         assert xforms[0]["scale"] == expected_read_scales[i]
+
+
+@array_constructor
+@pytest.mark.parametrize(
+    "write_shape, chunk_dims, num_levels, expected_read_shapes",
+    [
+        (
+            (2, 4, 8, 16, 32),
+            (1, 1, 2, 16, 16),
+            2,
+            [(2, 4, 8, 16, 32), (2, 4, 8, 8, 16), (2, 4, 8, 4, 8)],
+        ),
+        (
+            (16, 32),
+            (2, 4),
+            2,
+            [(16, 32), (8, 16), (4, 8)],
+        ),
+    ],
+)
+@pytest.mark.parametrize("filename", ["e.zarr"])
+def test_ome_zarr_writer_chunks(
+    array_constructor: Callable,
+    write_shape: Tuple[int, ...],
+    chunk_dims: Tuple[int, ...],
+    num_levels: int,
+    filename: str,
+    expected_read_shapes: List[Tuple[int, ...]],
+    tmpdir: pathlib.Path,
+) -> None:
+    arr = array_constructor(write_shape, dtype=np.uint8)
+
+    # Construct save end point
+
+    baseline_save_uri = str(tmpdir / f"baseline_{filename}")
+    save_uri = str(tmpdir / filename)
+
+    # Normal save
+    writer = OmeZarrWriter(save_uri)
+    writer.write_image(
+        arr, "", None, None, None, chunk_dims=chunk_dims, scale_num_levels=num_levels
+    )
+    reader = Reader(parse_url(save_uri))
+    node = list(reader())[0]
+
+    # Check expected shapes
+    for level in range(num_levels):
+        shape = node.data[level].shape
+        assert shape == expected_read_shapes[level]
+
+    # Create baseline chunking to compare against manual.
+    writer = OmeZarrWriter(baseline_save_uri)
+    writer.write_image(arr, "", None, None, None, scale_num_levels=num_levels)
+    reader_baseline = Reader(parse_url(baseline_save_uri))
+    node_baseline = list(reader_baseline())[0]
+
+    data = node.data[0]
+    baseline_data = node_baseline.data[0]
+
+    assert np.all(np.equal(data, baseline_data))

aicsimageio/writers/ome_zarr_writer.py

@@ -8,6 +8,7 @@
 from zarr.storage import default_compressor
 
 from .. import exceptions, types
+from ..dimensions import DEFAULT_DIMENSION_ORDER, DimensionNames
 from ..metadata import utils
 from ..utils import io_utils
 
@@ -105,6 +106,13 @@ def build_ome(
         }
         return omero
 
+    @staticmethod
+    def _build_chunk_dims(
+        chunk_dim_map: Dict[str, int],
+        dimension_order: str = DEFAULT_DIMENSION_ORDER,
+    ) -> Tuple[int, ...]:
+        return tuple(chunk_dim_map[d] for d in dimension_order)
+
     def write_image(
         self,
         # TODO how to pass in precomputed multiscales?
@@ -113,6 +121,7 @@ def write_image(
         physical_pixel_sizes: Optional[types.PhysicalPixelSizes],
         channel_names: Optional[List[str]],
         channel_colors: Optional[List[int]],
+        chunk_dims: Optional[Tuple] = None,
         scale_num_levels: int = 1,
         scale_factor: float = 2.0,
         dimension_order: Optional[str] = None,
@@ -177,27 +186,34 @@ def write_image(
                 f"Received image data with shape: {image_data.shape}"
             )
         if dimension_order is None:
-            dimension_order = "TCZYX"[-ndims:]
+            dimension_order = DEFAULT_DIMENSION_ORDER[-ndims:]
         if len(dimension_order) != ndims:
             raise exceptions.InvalidDimensionOrderingError(
                 f"Dimension order {dimension_order} does not match data "
                 f"shape: {image_data.shape}"
             )
-        if (len(set(dimension_order) - set("TCZYX")) > 0) or len(
+        if (len(set(dimension_order) - set(DEFAULT_DIMENSION_ORDER)) > 0) or len(
             dimension_order
         ) != len(set(dimension_order)):
             raise exceptions.InvalidDimensionOrderingError(
-                f"Dimension order {dimension_order} is invalid or "
-                "contains unexpected dimensions. Only TCZYX currently supported."
+                f"Dimension order {dimension_order} is invalid or contains"
+                f"unexpected dimensions. Only {DEFAULT_DIMENSION_ORDER}"
+                f"currently supported."
             )
-        xdimindex = dimension_order.find("X")
-        ydimindex = dimension_order.find("Y")
-        zdimindex = dimension_order.find("Z")
-        cdimindex = dimension_order.find("C")
+        xdimindex = dimension_order.find(DimensionNames.SpatialX)
+        ydimindex = dimension_order.find(DimensionNames.SpatialY)
+        zdimindex = dimension_order.find(DimensionNames.SpatialZ)
+        cdimindex = dimension_order.find(DimensionNames.Channel)
         if cdimindex > min(i for i in [xdimindex, ydimindex, zdimindex] if i > -1):
             raise exceptions.InvalidDimensionOrderingError(
                 f"Dimension order {dimension_order} is invalid. Channel dimension "
-                "must be before X, Y, and Z."
+                f"must be before X, Y, and Z."
+            )
+
+        if chunk_dims is not None and len(chunk_dims) != ndims:
+            raise exceptions.UnexpectedShapeError(
+                f"Chunk dimensions:{chunk_dims} do not match data. "
+                f"Expected chunk dimension length:{ndims}"
             )
 
         if physical_pixel_sizes is None:
@@ -227,12 +243,13 @@ def write_image(
                 if cdimindex > -1
                 else [0]
             )
+        # Chunk spatial dimensions
         scale_dim_map = {
-            "T": 1.0,
-            "C": 1.0,
-            "Z": pixelsizes[0],
-            "Y": pixelsizes[1],
-            "X": pixelsizes[2],
+            DimensionNames.Time: 1.0,
+            DimensionNames.Channel: 1.0,
+            DimensionNames.SpatialZ: pixelsizes[0],
+            DimensionNames.SpatialY: pixelsizes[1],
+            DimensionNames.SpatialX: pixelsizes[2],
         }
         transforms = [
             [
@@ -249,26 +266,40 @@ def write_image(
             * image_data.shape[ydimindex]
             * image_data.itemsize
         )
+
         target_chunk_size = 16 * (1024 * 1024)  # 16 MB
         # this is making an assumption of chunking whole XY planes.
-        # TODO allow callers to configure chunk dims?
-        nplanes_per_chunk = int(math.ceil(target_chunk_size / plane_size))
-        nplanes_per_chunk = (
-            min(nplanes_per_chunk, image_data.shape[zdimindex]) if zdimindex > -1 else 1
-        )
-        chunk_dim_map = {
-            "T": 1,
-            "C": 1,
-            "Z": nplanes_per_chunk,
-            "Y": image_data.shape[ydimindex],
-            "X": image_data.shape[xdimindex],
-        }
-        chunk_dims = [
-            dict(
-                chunks=tuple(chunk_dim_map[d] for d in dimension_order),
-                compressor=default_compressor,
+
+        if chunk_dims is None:
+            nplanes_per_chunk = int(math.ceil(target_chunk_size / plane_size))
+            nplanes_per_chunk = (
+                min(nplanes_per_chunk, image_data.shape[zdimindex])
+                if zdimindex > -1
+                else 1
             )
-        ]
+            chunk_dim_map = {
+                DimensionNames.Time: 1,
+                DimensionNames.Channel: 1,
+                DimensionNames.SpatialZ: nplanes_per_chunk,
+                DimensionNames.SpatialY: image_data.shape[ydimindex],
+                DimensionNames.SpatialX: image_data.shape[xdimindex],
+            }
+            chunks = [
+                dict(
+                    chunks=OmeZarrWriter._build_chunk_dims(
+                        chunk_dim_map=chunk_dim_map, dimension_order=dimension_order
+                    ),
+                    compressor=default_compressor,
+                )
+            ]
+        else:
+            chunks = [
+                dict(
+                    chunks=chunk_dims,
+                    compressor=default_compressor,
+                )
+            ]
+
         lasty = image_data.shape[ydimindex]
         lastx = image_data.shape[xdimindex]
         # TODO scaler might want to use different method for segmentations than raw
@@ -281,9 +312,9 @@ def write_image(
             scaler.method = "nearest"
             scaler.max_layer = scale_num_levels - 1
             scaler.downscale = scale_factor if scale_factor is not None else 2
-            for i in range(scale_num_levels - 1):
-                scale_dim_map["Y"] *= scaler.downscale
-                scale_dim_map["X"] *= scaler.downscale
+            for _ in range(scale_num_levels - 1):
+                scale_dim_map[DimensionNames.SpatialY] *= scaler.downscale
+                scale_dim_map[DimensionNames.SpatialX] *= scaler.downscale
                 transforms.append(
                     [
                         {
@@ -292,24 +323,47 @@ def write_image(
                         }
                     ]
                 )
-                lasty = int(math.ceil(lasty / scaler.downscale))
-                lastx = int(math.ceil(lastx / scaler.downscale))
-                plane_size = lasty * lastx * image_data.itemsize
-                nplanes_per_chunk = int(math.ceil(target_chunk_size / plane_size))
-                nplanes_per_chunk = (
-                    min(nplanes_per_chunk, image_data.shape[zdimindex])
-                    if zdimindex > -1
-                    else 1
-                )
-                chunk_dim_map["Z"] = nplanes_per_chunk
-                chunk_dim_map["Y"] = lasty
-                chunk_dim_map["X"] = lastx
-                chunk_dims.append(
-                    dict(
-                        chunks=tuple(chunk_dim_map[d] for d in dimension_order),
-                        compressor=default_compressor,
+
+                if chunk_dims is None:
+                    lasty = int(math.ceil(lasty / scaler.downscale))
+                    lastx = int(math.ceil(lastx / scaler.downscale))
+                    chunk_dim_map = {
+                        DimensionNames.Time: 1,
+                        DimensionNames.Channel: 1,
+                    }
+                    plane_size = lasty * lastx * image_data.itemsize
+                    nplanes_per_chunk = int(math.ceil(target_chunk_size / plane_size))
+                    nplanes_per_chunk = (
+                        min(nplanes_per_chunk, image_data.shape[zdimindex])
+                        if zdimindex > -1
+                        else 1
                     )
-                )
+
+                    chunk_dim_map[DimensionNames.SpatialZ] = nplanes_per_chunk
+                    chunk_dim_map[DimensionNames.SpatialY] = lasty
+                    chunk_dim_map[DimensionNames.SpatialX] = lastx
+
+                    chunks.append(
+                        dict(
+                            chunks=OmeZarrWriter._build_chunk_dims(
+                                chunk_dim_map=chunk_dim_map,
+                                dimension_order=dimension_order,
+                            ),
+                            compressor=default_compressor,
+                        )
+                    )
+                else:
+                    rescaley = int(math.ceil(chunk_dims[ydimindex] / scaler.downscale))
+                    rescalex = int(math.ceil(chunk_dims[xdimindex] / scaler.downscale))
+                    chunk_dims = tuple(list(chunk_dims[:-2]) + [rescaley, rescalex])
+
+                    chunks.append(
+                        dict(
+                            chunks=chunk_dims,
+                            compressor=default_compressor,
+                        )
+                    )
+
         else:
             scaler = None
 
@@ -328,11 +382,23 @@ def write_image(
         )
         # TODO user supplies units?
         dim_to_axis = {
-            "T": {"name": "t", "type": "time", "unit": "millisecond"},
-            "C": {"name": "c", "type": "channel"},
-            "Z": {"name": "z", "type": "space", "unit": "micrometer"},
-            "Y": {"name": "y", "type": "space", "unit": "micrometer"},
-            "X": {"name": "x", "type": "space", "unit": "micrometer"},
+            DimensionNames.Time: {"name": "t", "type": "time", "unit": "millisecond"},
+            DimensionNames.Channel: {"name": "c", "type": "channel"},
+            DimensionNames.SpatialZ: {
+                "name": "z",
+                "type": "space",
+                "unit": "micrometer",
+            },
+            DimensionNames.SpatialY: {
+                "name": "y",
+                "type": "space",
+                "unit": "micrometer",
+            },
+            DimensionNames.SpatialX: {
+                "name": "x",
+                "type": "space",
+                "unit": "micrometer",
+            },
         }
 
         axes = [dim_to_axis[d] for d in dimension_order]
@@ -353,5 +419,5 @@ def write_image(
             # match the number of datasets in a multiresolution pyramid. One can
             # provide different chunk size for each level of a pyramid using this
             # option.
-            storage_options=chunk_dims,
+            storage_options=chunks,
         )

setup.py

@@ -105,7 +105,7 @@ def run(self):
     "imagecodecs>=2020.5.30",
     "lxml>=4.6,<5",
     "numpy>=1.16,<2",
-    "ome-types>=0.3.3",
+    "ome-types>=0.3.3,<0.4.0",
     "ome-zarr>=0.6.1",
     "PyYAML>=6.0",
     "wrapt>=1.12",