Merge pull request #301 from czbiohub-sf/zarr3-dev

OME-Zarr v0.5 (Zarr v3) support

Author: ziw-liu

.github/workflows/pr.yml

@@ -6,7 +6,7 @@ name: lint, style, and tests
 on:
   pull_request:
     branches:
-      - main
+      - "*"
 
 jobs:
   style:
@@ -94,7 +94,7 @@ jobs:
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
-          pip install ".[dev]"
+          pip install ".[dev,acquire-zarr]"
 
       - name: Test with pytest
         env:

.github/workflows/test.yml

@@ -24,7 +24,7 @@ jobs:
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
-          pip install ".[dev]"
+          pip install ".[dev,acquire-zarr]"
 
       - name: Test with pytest
         env:

CONTRIBUTING.md

@@ -91,11 +91,13 @@ git clone https://github.com/czbiohub-sf/iohub.git
 Otherwise, you can follow [these instructions](https://docs.github.com/en/get-started/quickstart/fork-a-repo)
 to [fork](https://github.com/czbiohub-sf/iohub/fork) the repository.
 
-Then install the package in editable mode with the development dependencies:
+Then install the package in editable mode with the development dependencies.
+Remove acquire-zarr if you do not have glibc version 2.35 or later,
+for example on the Bruno cluster (Rocky Linux 8).
 
 ```sh
 cd iohub/ # or the renamed project root directory
-pip install -e ".[dev]"
+pip install -e ".[dev,acquire-zarr]"
 ```
 
 Then make the changes and [track them with Git](https://docs.github.com/en/get-started/using-git/about-git#example-contribute-to-an-existing-repository).

docs/examples/run_coordinate_transform.py

@@ -58,9 +58,9 @@
 ) as dataset:
     # Create and write to positions
     # This affects the tile arrangement in visualization
-    position = dataset.create_position(0, 0, 0)
+    position = dataset.create_position("0", "0", "0")
     position.create_image("0", tczyx_1, transform=[translation[0]])
-    position = dataset.create_position(0, 0, 1)
+    position = dataset.create_position("0", "0", "1")
     position.create_image("0", tczyx_2, transform=[translation[1], scaling[0]])
     # Print dataset summary
     dataset.print_tree()
@@ -72,4 +72,4 @@
 
 # %%
 # Clean up
-tmp_dir.cleanup()
+tmp_dir.cleanup()

docs/examples/run_multi_fov_hcs_ome_zarr.py

@@ -32,9 +32,9 @@
 
 position_list = (
     ("A", "1", "0"),
-    ("H", 1, "0"),
+    ("H", "1", "0"),
     ("H", "12", "CannotVisualize"),
-    ("Control", "Blank", 0),
+    ("Control", "Blank", "0"),
 )
 
 with open_ome_zarr(

docs/examples/run_update_ome_zarr.py

@@ -0,0 +1,81 @@
+"""
+Update OME-Zarr Version
+=======================
+
+This script shows how to write the same OME-Zarr image
+using a new version.
+"""
+
+# %%
+from pathlib import Path
+from tempfile import TemporaryDirectory
+
+import numpy as np
+
+from iohub.ngff import TransformationMeta, open_ome_zarr
+
+# %%
+# Set storage path
+tmp_dir = TemporaryDirectory()
+old_store_path = Path(tmp_dir.name) / "old.zarr"
+new_store_path = Path(tmp_dir.name) / "new.zarr"
+
+# %%
+# Create a version 0.4 OME-Zarr dataset
+random_image = np.random.randint(
+    0, np.iinfo(np.uint16).max, size=(10, 2, 32, 128, 128), dtype=np.uint16
+)
+scale = [2.0, 3.0, 4.0, 5.0, 6.0]
+
+
+with open_ome_zarr(
+    old_store_path,
+    layout="hcs",
+    mode="w-",
+    channel_names=["DAPI", "GFP"],
+    version="0.4",
+) as old_dataset:
+    position = old_dataset.create_position("A", "1", "0")
+    image = position.create_image(
+        "0",
+        random_image,
+        chunks=(1, 1, 4, 32, 32),
+        transform=[TransformationMeta(type="scale", scale=scale)],
+    )
+
+# %%
+# Write the same image with version 0.5 and sharding
+
+with open_ome_zarr(old_store_path, mode="r", layout="hcs") as old_dataset:
+    with open_ome_zarr(
+        new_store_path,
+        layout="hcs",
+        mode="w",
+        channel_names=old_dataset.channel_names,
+        version="0.5",
+    ) as new_dataset:
+        for name, old_position in old_dataset.positions():
+            row, col, fov = name.split("/")
+            new_position = new_dataset.create_position(row, col, fov)
+            old_image = old_position["0"]
+            new_image = new_position.create_image(
+                "0",
+                data=old_image.numpy(),
+                chunks=(1, 1, 4, 32, 32),
+                shards_ratio=(2, 1, 8, 4, 4),
+                transform=old_position.metadata.multiscales[0]
+                .datasets[0]
+                .coordinate_transformations,
+            )
+
+# %%
+# Read the new FOV to verify it was written correctly
+with open_ome_zarr(new_store_path / "A/1/0", mode="r") as dataset:
+    assert dataset.scale == scale
+    image = dataset["0"]
+    assert image.shards == (2, 1, 32, 128, 128)
+    assert np.array_equal(image.numpy(), random_image)
+
+# %%
+# Clean up
+tmp_dir.cleanup()

iohub/_deprecated/singlepagetiff.py

@@ -3,17 +3,19 @@
 import json
 import logging
 import os
+from pathlib import Path
 
 import natsort
 import numpy as np
 import tifffile as tiff
 import zarr
 
 from iohub._deprecated.reader_base import ReaderBase
+from iohub.mmstack import _tiff_to_fsspec_store
 
 
 class MicromanagerSequenceReader(ReaderBase):
-    def __init__(self, folder, extract_data=False):
+    def __init__(self, folder, extract_data=False, strict=False):
         super().__init__()
 
         """
@@ -32,13 +34,16 @@ def __init__(self, folder, extract_data=False):
             which contain singlepage tiff sequences
         extract_data    (bool)
             True if zarr arrays should be extracted immediately
+        strict:         (bool)
+            True if failures in getting images should raise exceptions
         """
 
         if not os.path.isdir(folder):
             raise NotImplementedError(
                 "supplied path for singlepage tiff sequence reader "
                 "is not a folder"
             )
+        self._strict = strict
 
         self.log = logging.getLogger(__name__)
         self.positions = {}
@@ -206,20 +211,32 @@ def _create_stores(self, p):
             if c[0] == p:
                 self.log.info(f"reading coord = {c} from filename = {fn}")
                 with tiff.imread(fn, aszarr=True) as store:
-                    z[c[1], c[2], c[3]] = zarr.open(store)
+                    try:
+                        array = zarr.open(
+                            _tiff_to_fsspec_store(
+                                store, root_uri=Path(fn).parent.as_uri()
+                            ),
+                            mode="r",
+                        )[:]
+                        z[c[1], c[2], c[3]] = array
+                    except Exception:
+                        self.log.error(
+                            f"error reading file {fn} for coordinate {c}"
+                        )
 
         # check that the array was assigned
-        if z == zarr.zeros(
-            shape=(
-                self.frames,
-                self.channels,
-                self.slices,
-                self.height,
-                self.width,
-            ),
-            chunks=(1, 1, 1, self.height, self.width),
-        ):
-            raise IOError(f"array at position {p} can not be found")
+        if self._strict:
+            if z == zarr.zeros(
+                shape=(
+                    self.frames,
+                    self.channels,
+                    self.slices,
+                    self.height,
+                    self.width,
+                ),
+                chunks=(1, 1, 1, self.height, self.width),
+            ):
+                raise IOError(f"array at position {p} can not be found")
 
         self.positions[p] = z
 

iohub/_deprecated/zarrfile.py

@@ -8,6 +8,7 @@
 
 import numpy as np
 import zarr
+import zarr.storage
 
 from iohub._deprecated.reader_base import ReaderBase
 
@@ -27,7 +28,7 @@ class ZarrReader(ReaderBase):
     """
 
     def __init__(
-        self, store_path: str, version: Literal["0.1", "0.4"] = "0.1"
+        self, store_path: str, version: Literal["0.1", "0.4", "0.5"] = "0.1"
     ):
         super().__init__()
 
@@ -43,17 +44,11 @@ def __init__(
         # zarr files (.zarr) are directories
         if not os.path.isdir(store_path):
             raise ValueError("file does not exist")
-        if version == "0.4":
-            dimension_separator = "/"
-        elif version == "0.1":
-            dimension_separator = "."
-        else:
+        if version not in ("0.1", "0.4", "0.5"):
             raise ValueError(f"Invalid NGFF version: {version}")
         try:
-            self.store = zarr.DirectoryStore(
-                store_path, dimension_separator=dimension_separator
-            )
-            self.root = zarr.open(self.store, "r")
+            self.store = zarr.storage.LocalStore(store_path)
+            self.root = zarr.open(self.store, mode="r")
         except Exception:
             raise FileNotFoundError("Supplies path is not a valid zarr root")
         try:

iohub/convert.py

@@ -309,7 +309,9 @@ def _init_hcs_arrays(self, arr_kwargs):
     def _init_grid_arrays(self, arr_kwargs):
         for row, columns in enumerate(self.position_grid):
             for column in columns:
-                self._create_zeros_array(row, column, "0", arr_kwargs)
+                self._create_zeros_array(
+                    str(row), str(column), "0", arr_kwargs
+                )
 
     def _create_zeros_array(
         self, row_name: str, col_name: str, pos_name: str, arr_kwargs: dict

iohub/mmstack.py

@@ -1,17 +1,21 @@
 from __future__ import annotations
 
+import io
+import json
 import logging
 from copy import copy
 from pathlib import Path
 from typing import TYPE_CHECKING, Iterable
 from warnings import catch_warnings, filterwarnings
 
 import dask.array as da
+import fsspec
 import numpy as np
 import zarr
+import zarr.storage
 from natsort import natsorted
 from numpy.typing import ArrayLike
-from tifffile import TiffFile
+from tifffile import TiffFile, ZarrTiffStore
 from xarray import DataArray
 
 from iohub.mm_fov import MicroManagerFOV, MicroManagerFOVMapping
@@ -31,6 +35,34 @@ def _normalize_mm_pos_key(key: str | int) -> int:
         raise TypeError("Micro-Manager position keys must be integers.")
 
 
+def _tiff_to_fsspec_store(
+    zarr_tiff_store: ZarrTiffStore, root_uri: str
+) -> zarr.storage.FsspecStore:
+    """Bridge tifffile (zarr-python v2 interface) with zarr-python v3.
+
+    Parameters
+    ----------
+    zarr_tiff_store : ZarrTiffStore
+        Zarr (v2) wrapper for a TIFF series
+    root_uri : str
+        `file://` URI to the directory containing the TIFF files
+
+    Returns
+    -------
+    zarr.storage.FsspecStore
+        Zarr (v3) wrapper for a TIFF series
+    """
+    spec_container = io.StringIO()
+    zarr_tiff_store.write_fsspec(spec_container, url=root_uri)
+    fs, _ = fsspec.url_to_fs(
+        "reference://",
+        fo=json.loads(spec_container.getvalue()),
+        target_protocol="file",
+        asynchronous=True,
+    )
+    return zarr.storage.FsspecStore(fs=fs)
+
+
 def find_first_ome_tiff_in_mmstack(data_path: Path) -> Path:
     if data_path.is_file():
         if "ome.tif" in data_path.name:
@@ -120,9 +152,12 @@ def _parse_data(self):
             self.width,
         ) = dims.values()
         self._set_mm_meta(self._first_tif.micromanager_metadata)
-        self._store = series.aszarr()
+        zarr_tiff_store = series.aszarr(multiscales=True)
+        self._store = _tiff_to_fsspec_store(
+            zarr_tiff_store, root_uri=self._root.as_uri()
+        )
         _logger.debug(f"Opened {self._store}.")
-        data = da.from_zarr(zarr.open(self._store))
+        data = da.from_zarr(zarr.open(self._store, mode="r")["0"])
         self.dtype = data.dtype
         img = DataArray(data, dims=raw_dims, name=self.dirname)
         xarr = img.expand_dims(