Zarr python v3 [DO NOT MERGE]

CHANGELOG.md

@@ -1,5 +1,16 @@
 # Changelog
 
+## [v0.5.0]
+
+### Features
+- Add support for OME-NGFF v0.5
+- Move to zarr-python v3
+
+### API Changes
+
+- The `compressor` argument has been renamed to `compressors` in all relevant functions and methods to reflect the support for multiple compressors in zarr v3.
+- The `version` argument has been renamed to `ngff_version` in all relevant functions and methods to specify the OME-NGFF version.
+
 ## [v0.4.4]
 
 ### Bug Fixes

pyproject.toml

@@ -35,8 +35,8 @@ classifiers = [
 dependencies = [
     "numpy",
     "filelock",
-    "zarr<3",
-    "anndata>=0.8.0,<0.11.4", # To be removed when we transition to zarr v3
+    "zarr>3",
+    "anndata",
     "pydantic",
     "pandas>=1.2.0",
     "requests",
@@ -155,7 +155,7 @@ minversion = "7.0"
 testpaths = ["tests"]
 filterwarnings = [
     "error",
-    "ignore::FutureWarning", # TODO remove after zarr-python v3 
+    "ignore::zarr.errors.ZarrUserWarning", # required for anndata
 ]
 addopts = [
     "-vv",

src/ngio/common/_pyramid.py

@@ -5,8 +5,9 @@
 import dask.array as da
 import numpy as np
 import zarr
-from zarr.types import DIMENSION_SEPARATOR
+from zarr.core.array import CompressorLike
 
+# from zarr.types import DIMENSION_SEPARATOR
 from ngio.common._zoom import (
     InterpolationOrder,
     _zoom_inputs_check,
@@ -26,7 +27,10 @@ def _on_disk_numpy_zoom(
     target: zarr.Array,
     order: InterpolationOrder,
 ) -> None:
-    target[...] = numpy_zoom(source[...], target_shape=target.shape, order=order)
+    source_array = source[...]
+    if not isinstance(source_array, np.ndarray):
+        raise NgioValueError("source zarr array could not be read as a numpy array")
+    target[...] = numpy_zoom(source_array, target_shape=target.shape, order=order)
 
 
 def _on_disk_dask_zoom(
@@ -37,7 +41,7 @@ def _on_disk_dask_zoom(
     source_array = da.from_zarr(source)
     target_array = dask_zoom(source_array, target_shape=target.shape, order=order)
 
-    target_array = target_array.rechunk(target.chunks)
+    target_array = target_array.rechunk(target.chunks)  # type: ignore
     target_array.compute_chunk_sizes()
     target_array.to_zarr(target)
 
@@ -199,20 +203,24 @@ def init_empty_pyramid(
     paths: list[str],
     ref_shape: Sequence[int],
     scaling_factors: Sequence[float],
-    chunks: Sequence[int] | None = None,
+    axes: Sequence[str],
+    chunks: Sequence[int] | Literal["auto"] = "auto",
     dtype: str = "uint16",
     mode: AccessModeLiteral = "a",
-    dimension_separator: DIMENSION_SEPARATOR = "/",
-    compressor="default",
+    dimension_separator: Literal[".", "/"] = "/",
+    compressors: CompressorLike = "auto",
+    zarr_format: Literal[2, 3] = 2,
 ) -> None:
     # Return the an Image object
-    if chunks is not None and len(chunks) != len(ref_shape):
+    if chunks != "auto" and len(chunks) != len(ref_shape):
         raise NgioValueError(
             "The shape and chunks must have the same number of dimensions."
         )
 
-    if chunks is not None:
-        chunks = [min(c, s) for c, s in zip(chunks, ref_shape, strict=True)]
+    if chunks != "auto":
+        chunks = tuple(min(c, s) for c, s in zip(chunks, ref_shape, strict=True))
+    else:
+        chunks = "auto"
 
     if len(ref_shape) != len(scaling_factors):
         raise NgioValueError(
@@ -223,31 +231,43 @@ def init_empty_pyramid(
     # To reduce the risk of floating point issues
     scaling_factors = [_maybe_int(s) for s in scaling_factors]
 
-    root_group = open_group_wrapper(store, mode=mode)
+    root_group = open_group_wrapper(store, mode=mode, zarr_format=zarr_format)
+
+    array_static_kwargs = {
+        "dtype": dtype,
+        "overwrite": True,
+        "compressors": compressors,
+    }
+
+    if zarr_format == 2:
+        array_static_kwargs["chunk_key_encoding"] = {
+            "name": "v2",
+            "separator": dimension_separator,
+        }
+    else:
+        array_static_kwargs["chunk_key_encoding"] = {
+            "name": "default",
+            "separator": dimension_separator,
+        }
+        array_static_kwargs["dimension_names"] = axes
 
     for path in paths:
         if any(s < 1 for s in ref_shape):
             raise NgioValueError(
                 "Level shape must be at least 1 on all dimensions. "
                 f"Calculated shape: {ref_shape} at level {path}."
             )
-        new_arr = root_group.zeros(
+        new_arr = root_group.create_array(
             name=path,
-            shape=ref_shape,
-            dtype=dtype,
+            shape=tuple(ref_shape),
             chunks=chunks,
-            dimension_separator=dimension_separator,
-            overwrite=True,
-            compressor=compressor,
+            **array_static_kwargs,
         )
 
-        _shape = [
+        ref_shape = [
             math.floor(s / sc) for s, sc in zip(ref_shape, scaling_factors, strict=True)
         ]
-        ref_shape = _shape
-
-        if chunks is None:
-            chunks = new_arr.chunks
-            assert chunks is not None
-        chunks = [min(c, s) for c, s in zip(chunks, ref_shape, strict=True)]
+        chunks = tuple(
+            min(c, s) for c, s in zip(new_arr.chunks, ref_shape, strict=True)
+        )
     return None

src/ngio/images/_create.py

@@ -1,9 +1,9 @@
 """Utility functions for working with OME-Zarr images."""
 
 from collections.abc import Sequence
-from typing import TypeVar
+from typing import Literal, TypeVar
 
-from zarr.types import DIMENSION_SEPARATOR
+from zarr.core.array import CompressorLike
 
 from ngio.common._pyramid import init_empty_pyramid
 from ngio.ome_zarr_meta import (
@@ -40,7 +40,7 @@ def _init_generic_meta(
     space_unit: SpaceUnits | str | None = DefaultSpaceUnit,
     time_unit: TimeUnits | str | None = DefaultTimeUnit,
     name: str | None = None,
-    version: NgffVersions = DefaultNgffVersion,
+    ngff_version: NgffVersions = DefaultNgffVersion,
 ) -> tuple[_image_or_label_meta, list[float]]:
     """Initialize the metadata for an image or label."""
     scaling_factors = []
@@ -75,7 +75,7 @@ def _init_generic_meta(
         axes_names=axes_names,
         pixel_size=pixel_sizes,
         scaling_factors=scaling_factors,
-        version=version,
+        version=ngff_version,
     )
     return meta, scaling_factors
 
@@ -93,12 +93,12 @@ def create_empty_label_container(
     time_unit: TimeUnits | str | None = DefaultTimeUnit,
     axes_names: Sequence[str] | None = None,
     name: str | None = None,
-    chunks: Sequence[int] | None = None,
+    chunks: Sequence[int] | Literal["auto"] = "auto",
     dtype: str = "uint32",
-    dimension_separator: DIMENSION_SEPARATOR = "/",
-    compressor="default",
+    dimension_separator: Literal[".", "/"] = "/",
+    compressors: CompressorLike = "auto",
     overwrite: bool = False,
-    version: NgffVersions = DefaultNgffVersion,
+    ngff_version: NgffVersions = DefaultNgffVersion,
 ) -> ZarrGroupHandler:
     """Create an empty label with the given shape and metadata.
 
@@ -122,15 +122,15 @@ def create_empty_label_container(
         axes_names (Sequence[str] | None, optional): The names of the axes.
             If None the canonical names are used. Defaults to None.
         name (str | None, optional): The name of the image. Defaults to None.
-        chunks (Sequence[int] | None, optional): The chunk shape. If None the shape
+        chunks (Sequence[int] | Literal["auto"]): The chunk shape. If None the shape
             is used. Defaults to None.
         dimension_separator (DIMENSION_SEPARATOR): The separator to use for
             dimensions. Defaults to "/".
-        compressor: The compressor to use. Defaults to "default".
+        compressors (CompressorLike): The compressors to use. Defaults to "auto".
         dtype (str, optional): The data type of the image. Defaults to "uint16".
         overwrite (bool, optional): Whether to overwrite an existing image.
             Defaults to True.
-        version (str, optional): The version of the OME-Zarr specification.
+        ngff_version (str, optional): The version of the OME-Zarr specification.
             Defaults to DefaultVersion.
 
     """
@@ -155,12 +155,16 @@ def create_empty_label_container(
         time_unit=time_unit,
         axes_names=axes_names,
         name=name,
-        version=version,
+        ngff_version=ngff_version,
     )
 
     mode = "w" if overwrite else "w-"
-    group_handler = ZarrGroupHandler(store=store, mode=mode, cache=False)
-    image_handler = get_label_meta_handler(version=version, group_handler=group_handler)
+    group_handler = ZarrGroupHandler(
+        store=store, mode=mode, cache=False, zarr_format=meta.zarr_format
+    )
+    image_handler = get_label_meta_handler(
+        version=ngff_version, group_handler=group_handler
+    )
     image_handler.write_meta(meta)
 
     init_empty_pyramid(
@@ -169,10 +173,11 @@ def create_empty_label_container(
         scaling_factors=scaling_factors,
         ref_shape=shape,
         chunks=chunks,
+        axes=axes_names,
         dtype=dtype,
         mode="a",
         dimension_separator=dimension_separator,
-        compressor=compressor,
+        compressors=compressors,
     )
     group_handler._mode = "r+"
     return group_handler
@@ -191,12 +196,12 @@ def create_empty_image_container(
     time_unit: TimeUnits | str | None = DefaultTimeUnit,
     axes_names: Sequence[str] | None = None,
     name: str | None = None,
-    chunks: Sequence[int] | None = None,
+    chunks: Sequence[int] | Literal["auto"] = "auto",
     dtype: str = "uint16",
-    dimension_separator: DIMENSION_SEPARATOR = "/",
-    compressor="default",
+    dimension_separator: Literal[".", "/"] = "/",
+    compressors: CompressorLike = "auto",
     overwrite: bool = False,
-    version: NgffVersions = DefaultNgffVersion,
+    ngff_version: NgffVersions = DefaultNgffVersion,
 ) -> ZarrGroupHandler:
     """Create an empty OME-Zarr image with the given shape and metadata.
 
@@ -220,15 +225,15 @@ def create_empty_image_container(
         axes_names (Sequence[str] | None, optional): The names of the axes.
             If None the canonical names are used. Defaults to None.
         name (str | None, optional): The name of the image. Defaults to None.
-        chunks (Sequence[int] | None, optional): The chunk shape. If None the shape
+        chunks (Sequence[int] | Literal["auto"]): The chunk shape. If None the shape
             is used. Defaults to None.
         dtype (str, optional): The data type of the image. Defaults to "uint16".
         dimension_separator (DIMENSION_SEPARATOR): The separator to use for
             dimensions. Defaults to "/".
-        compressor: The compressor to use. Defaults to "default".
+        compressors (CompressorLike): The compressors to use. Defaults to "auto".
         overwrite (bool, optional): Whether to overwrite an existing image.
             Defaults to True.
-        version (str, optional): The version of the OME-Zarr specification.
+        ngff_version (str, optional): The version of the OME-Zarr specification.
             Defaults to DefaultVersion.
 
     """
@@ -253,11 +258,15 @@ def create_empty_image_container(
         time_unit=time_unit,
         axes_names=axes_names,
         name=name,
-        version=version,
+        ngff_version=ngff_version,
     )
     mode = "w" if overwrite else "w-"
-    group_handler = ZarrGroupHandler(store=store, mode=mode, cache=False)
-    image_handler = get_image_meta_handler(version=version, group_handler=group_handler)
+    group_handler = ZarrGroupHandler(
+        store=store, mode=mode, cache=False, zarr_format=meta.zarr_format
+    )
+    image_handler = get_image_meta_handler(
+        version=ngff_version, group_handler=group_handler
+    )
     image_handler.write_meta(meta)
 
     init_empty_pyramid(
@@ -266,10 +275,12 @@ def create_empty_image_container(
         scaling_factors=scaling_factors,
         ref_shape=shape,
         chunks=chunks,
+        axes=axes_names,
         dtype=dtype,
         mode="a",
         dimension_separator=dimension_separator,
-        compressor=compressor,
+        compressors=compressors,
+        zarr_format=meta.zarr_format,
     )
 
     group_handler._mode = "r+"

src/ngio/images/_create_synt_container.py

@@ -1,10 +1,11 @@
 """Abstract class for handling OME-NGFF images."""
 
 from collections.abc import Sequence
+from typing import Literal
 
 import numpy as np
 import PIL.Image
-from zarr.types import DIMENSION_SEPARATOR
+from zarr.core.array import CompressorLike
 
 from ngio.common._synt_images_utils import fit_to_shape
 from ngio.images._ome_zarr_container import OmeZarrContainer, create_ome_zarr_from_array
@@ -30,16 +31,16 @@ def create_synthetic_ome_zarr(
     xy_scaling_factor: float = 2,
     z_scaling_factor: float = 1.0,
     axes_names: Sequence[str] | None = None,
-    chunks: Sequence[int] | None = None,
+    chunks: Sequence[int] | Literal["auto"] = "auto",
     channel_labels: list[str] | None = None,
     channel_wavelengths: list[str] | None = None,
     channel_colors: Sequence[str] | None = None,
     channel_active: Sequence[bool] | None = None,
     table_backend: TableBackend = DefaultTableBackend,
-    dimension_separator: DIMENSION_SEPARATOR = "/",
-    compressor="default",
+    dimension_separator: Literal[".", "/"] = "/",
+    compressors: CompressorLike = "auto",
     overwrite: bool = False,
-    version: NgffVersions = DefaultNgffVersion,
+    ngff_version: NgffVersions = DefaultNgffVersion,
 ) -> OmeZarrContainer:
     """Create an empty OME-Zarr image with the given shape and metadata.
 
@@ -55,8 +56,8 @@ def create_synthetic_ome_zarr(
             Defaults to 1.0.
         axes_names (Sequence[str] | None, optional): The names of the axes.
             If None the canonical names are used. Defaults to None.
-        chunks (Sequence[int] | None, optional): The chunk shape. If None the shape
-            is used. Defaults to None.
+        chunks (Sequence[int] | Literal["auto"]): The chunk shape. If None the shape
+            is used. Defaults to "auto".
         channel_labels (list[str] | None, optional): The labels of the channels.
             Defaults to None.
         channel_wavelengths (list[str] | None, optional): The wavelengths of the
@@ -68,10 +69,10 @@ def create_synthetic_ome_zarr(
         table_backend (TableBackend): Table backend to be used to store tables
         dimension_separator (DIMENSION_SEPARATOR): The separator to use for
             dimensions. Defaults to "/".
-        compressor: The compressor to use. Defaults to "default".
+        compressors (CompressorLike): The compressors to use. Defaults to "auto".
         overwrite (bool, optional): Whether to overwrite an existing image.
             Defaults to True.
-        version (NgffVersion, optional): The version of the OME-Zarr specification.
+        ngff_version (NgffVersion, optional): The version of the OME-Zarr specification.
             Defaults to DefaultNgffVersion.
     """
     if isinstance(reference_sample, str):
@@ -103,8 +104,8 @@ def create_synthetic_ome_zarr(
         chunks=chunks,
         overwrite=overwrite,
         dimension_separator=dimension_separator,
-        compressor=compressor,
-        version=version,
+        compressors=compressors,
+        ngff_version=ngff_version,
     )
 
     image = ome_zarr.get_image()