fix: adding ome-zarr v0.4 back

Signed-off-by: Sricharan Reddy Varra <sricharan.varra@biohub.org>

Author: srivarra

src/iohub/core/config.py

@@ -36,6 +36,7 @@ class ZarrConfig(BaseModel):
 class TensorStoreConfig(BaseModel):
     """Config for the TensorStore implementation."""
 
+    compressor: CompressorConfig = Field(default_factory=CompressorConfig)
     data_copy_concurrency: int = Field(default=4, ge=1)
     context: dict | None = None
     file_io_concurrency: int | None = None

src/iohub/core/implementations/tensorstore.py

@@ -37,15 +37,20 @@ def __init__(self, group: _TsGroup):
         super().__init__(self._load())
 
     def _load(self) -> dict:
+        if self._group.zarr_driver == "zarr2":
+            zattrs = self._group.path / ".zattrs"
+            if zattrs.exists():
+                return json.loads(zattrs.read_text())
+            return {}
         zarr_json = self._group.path / "zarr.json"
         if zarr_json.exists():
             return json.loads(zarr_json.read_text()).get("attributes", {})
-        zattrs = self._group.path / ".zattrs"
-        if zattrs.exists():
-            return json.loads(zattrs.read_text())
         return {}
 
     def _save(self) -> None:
+        if self._group.zarr_driver == "zarr2":
+            (self._group.path / ".zattrs").write_text(json.dumps(dict(self)))
+            return
         zarr_json = self._group.path / "zarr.json"
         if zarr_json.exists():
             meta = json.loads(zarr_json.read_text())
@@ -63,13 +68,17 @@ def update(self, *args, **kwargs):
         self._save()
 
 
-def _detect_zarr_driver(path: Path) -> str:
+def _detect_zarr_driver(path: Path, zarr_format: int | None = None) -> str:
     """Detect zarr format for a store root. Called once at open_group time."""
+    if zarr_format == 2:
+        return "zarr2"
+    if zarr_format == 3:
+        return "zarr3"
     if (path / "zarr.json").exists():
         return "zarr3"
     if (path / ".zattrs").exists() or (path / ".zgroup").exists():
         return "zarr2"
-    return "zarr3"  # new stores  will be created as v3 always
+    return "zarr3"
 
 
 class _TsGroup:
@@ -87,7 +96,10 @@ def __init__(
             raise FileExistsError(f"Store already exists: {path}")
         if mode in ("w", "w-", "a") and not path.exists():
             path.mkdir(parents=True, exist_ok=True)
-            (path / "zarr.json").write_text('{"zarr_format": 3, "node_type": "group", "attributes": {}}')
+            if zarr_driver == "zarr2":
+                (path / ".zgroup").write_text('{"zarr_format": 2}')
+            else:
+                (path / "zarr.json").write_text('{"zarr_format": 3, "node_type": "group", "attributes": {}}')
         self.path = path
         self.mode = mode
         self._impl = impl
@@ -99,9 +111,14 @@ def create_group(self, name: str, overwrite: bool = False) -> _TsGroup:
         if sub.exists() and not overwrite:
             return _TsGroup(path=sub, mode="a", impl=self._impl, zarr_driver=self.zarr_driver)
         sub.mkdir(parents=True, exist_ok=True)
-        zarr_json = sub / "zarr.json"
-        if not zarr_json.exists() or overwrite:
-            zarr_json.write_text(json.dumps({"zarr_format": 3, "node_type": "group", "attributes": {}}))
+        if self.zarr_driver == "zarr2":
+            zgroup = sub / ".zgroup"
+            if not zgroup.exists() or overwrite:
+                zgroup.write_text('{"zarr_format": 2}')
+        else:
+            zarr_json = sub / "zarr.json"
+            if not zarr_json.exists() or overwrite:
+                zarr_json.write_text(json.dumps({"zarr_format": 3, "node_type": "group", "attributes": {}}))
         return _TsGroup(path=sub, mode="a", impl=self._impl, zarr_driver=self.zarr_driver)
 
     def __contains__(self, name: str) -> bool:
@@ -271,7 +288,8 @@ def _context(self) -> ts.Context:
 
     def open_group(self, path: StorePath, mode: str, zarr_format: int | None = None) -> _TsGroup:
         p = Path(path)
-        return _TsGroup(path=p, mode=mode, impl=self, zarr_driver=_detect_zarr_driver(p), root=p)
+        driver = _detect_zarr_driver(p, zarr_format)
+        return _TsGroup(path=p, mode=mode, impl=self, zarr_driver=driver, root=p)
 
     def _iter_children(self, group: _TsGroup, node_type: str) -> list[str]:
         """Return sorted child names matching node_type ('group' or 'array')."""
@@ -306,7 +324,7 @@ def close(self, group: _TsGroup) -> None:
         pass  # TensorStore handles are not persistent connections
 
     def get_zarr_format(self, group: _TsGroup) -> int:
-        return 3  # TensorStore only supports zarr v3
+        return 2 if group.zarr_driver == "zarr2" else 3
 
     # -- Array lifecycle ---------------------------------------------------
 
@@ -325,14 +343,21 @@ def create_array_v2(
         fill_value: int = 0,
         overwrite: bool = False,
     ) -> ts.TensorStore:
+        shuffle_map = {"noshuffle": 0, "shuffle": 1, "bitshuffle": 2}
+        comp = self.config.compressor
         spec = {
             "driver": "zarr2",
             "kvstore": {"driver": "file", "path": str(Path(group.path) / name)},
             "metadata": {
                 "shape": list(shape),
                 "chunks": list(chunks),
                 "dtype": np.dtype(dtype).str,  # zarr2 uses NumPy dtype strings e.g. "<u2"
-                "compressor": {"id": "blosc", "cname": "lz4", "clevel": 5, "shuffle": 1},
+                "compressor": {
+                    "id": "blosc",
+                    "cname": comp.cname,
+                    "clevel": comp.clevel,
+                    "shuffle": shuffle_map.get(comp.shuffle, 2),
+                },
                 "fill_value": fill_value,
                 "order": "C",
                 "filters": None,

src/iohub/ngff/nodes.py

@@ -92,7 +92,7 @@ def _open_store(
     try:
         zarr_format = None
         if mode in ("w", "w-") or (is_fs and mode == "a" and not store_path.exists()):
-            zarr_format = 3
+            zarr_format = 2 if version == "0.4" else 3
         root = impl.open_group(store_path, mode=mode, zarr_format=zarr_format)
     except (FileNotFoundError, FileExistsError, PermissionError):
         raise
@@ -382,6 +382,10 @@ def _create_zarr_array(
             shards = tuple(c * s for c, s in zip(chunks, shards_ratio, strict=False))
         else:
             shards = None
+        if shards is not None and self._zarr_format == 2:
+            raise ValueError(
+                "Sharding is not supported in Zarr v2 (OME-Zarr v0.4). Remove shards_ratio or use version='0.5'."
+            )
         if self._zarr_format == 3:
             spec = ArraySpec.create(
                 shape=shape,
@@ -3066,8 +3070,6 @@ def open_ome_zarr(
     if _is_fslike(store_path):
         store_path = Path(store_path)
     _is_new_store = mode in ("w", "w-") or (mode == "a" and _is_fslike(store_path) and not store_path.exists())
-    if version == "0.4" and _is_new_store:
-        raise ValueError("Creating new OME-Zarr v0.4 stores is not supported. Use version='0.5' instead.")
     parse_meta = _check_file_mode(store_path, mode, disable_path_checking=disable_path_checking)
     root, impl = _open_store(
         store_path,

src/iohub/ngff/utils.py

@@ -418,10 +418,12 @@ def process_single_position(
             partial_apply_transform_to_czyx_and_save(*args)
     else:
         with ThreadPoolExecutor(max_workers=num_workers) as executor:
-            list(executor.map(
-                lambda args: partial_apply_transform_to_czyx_and_save(*args),
-                flat_iterable,
-            ))
+            list(
+                executor.map(
+                    lambda args: partial_apply_transform_to_czyx_and_save(*args),
+                    flat_iterable,
+                )
+            )
     click.echo("Shut down thread pool")
 
 

tests/ngff/test_ngff.py

@@ -45,7 +45,7 @@
 y_dim_st = st.integers(1, 32)
 x_dim_st = st.integers(1, 32)
 channel_names_st = c_dim_st.flatmap(lambda c_dim: st.lists(short_text_st, min_size=c_dim, max_size=c_dim, unique=True))
-ngff_versions_st = st.just("0.5")
+ngff_versions_st = st.sampled_from(["0.4", "0.5"])
 short_alpha_numeric = st.text(
     alphabet=list(string.ascii_lowercase + string.ascii_uppercase + string.digits),
     min_size=1,
@@ -238,28 +238,33 @@ def test_init_ome_zarr(channel_names, version):
 
 
 @pytest.mark.parametrize("mode", ["w", "w-"])
-def test_open_ome_zarr_v04_write_raises(tmp_path, mode):
-    """Creating new v0.4 stores must raise ValueError for all write modes."""
-    with pytest.raises(ValueError, match=r"v0.4"):
-        open_ome_zarr(
-            tmp_path / "out.zarr",
-            layout="fov",
-            mode=mode,
-            channel_names=["DAPI"],
-            version="0.4",
-        )
-
-
-def test_open_ome_zarr_v04_append_new_path_raises(tmp_path):
-    """mode='a' on a nonexistent path is a new store and must also raise."""
-    with pytest.raises(ValueError, match=r"v0.4"):
-        open_ome_zarr(
-            tmp_path / "nonexistent.zarr",
-            layout="fov",
-            mode="a",
-            channel_names=["DAPI"],
-            version="0.4",
-        )
+def test_open_ome_zarr_v04_write_succeeds(tmp_path, mode):
+    """Creating new v0.4 stores must succeed."""
+    store_path = tmp_path / "out.zarr"
+    with open_ome_zarr(
+        store_path,
+        layout="fov",
+        mode=mode,
+        channel_names=["DAPI"],
+        version="0.4",
+    ) as ds:
+        assert ds.version == "0.4"
+        assert (store_path / ".zgroup").exists()
+        assert not (store_path / "zarr.json").exists()
+
+
+def test_open_ome_zarr_v04_append_new_path_succeeds(tmp_path):
+    """mode='a' on a nonexistent path should create a v0.4 store."""
+    store_path = tmp_path / "nonexistent.zarr"
+    with open_ome_zarr(
+        store_path,
+        layout="fov",
+        mode="a",
+        channel_names=["DAPI"],
+        version="0.4",
+    ) as ds:
+        assert ds.version == "0.4"
+        assert (store_path / ".zgroup").exists()
 
 
 @pytest.mark.parametrize("version", ["0.5"])
@@ -391,12 +396,14 @@ def test_write_ome_zarr(channels_and_random_5d, arr_name, version):
     channel_names, random_5d = channels_and_random_5d
     with _temp_ome_zarr(random_5d, channel_names, arr_name, version=version) as dataset:
         assert_allclose(dataset[arr_name][:], random_5d)
-        # round-trip test with the offical reader implementation
-        ext_reader = ome_zarr.reader.Reader(ome_zarr.io.parse_url(dataset.zgroup.store.root))
-        node = next(iter(ext_reader()))
-        assert node.metadata["channel_names"] == channel_names
-        assert node.specs[0].datasets == [arr_name]
-        assert_allclose(node.data[0], random_5d)
+        if version == "0.5":
+            # round-trip test with the official reader implementation
+            # ome-zarr-py reader requires zarr-python Group with .store.root
+            ext_reader = ome_zarr.reader.Reader(ome_zarr.io.parse_url(dataset.zgroup.store.root))
+            node = next(iter(ext_reader()))
+            assert node.metadata["channel_names"] == channel_names
+            assert node.specs[0].datasets == [arr_name]
+            assert_allclose(node.data[0], random_5d)
 
 
 @given(
@@ -421,9 +428,10 @@ def test_create_zeros(ch_shape_dtype, arr_name, version):
             version=version,
         )
         dataset.create_zeros(name=arr_name, shape=shape, dtype=dtype)
-        assert {e.name for e in (Path(store_path) / arr_name).iterdir()} == {
-            "zarr.json",
-        }
+        if version == "0.5":
+            assert (Path(store_path) / arr_name / "zarr.json").exists()
+        else:
+            assert (Path(store_path) / arr_name / ".zarray").exists()
         if version == "0.5":
             assert dataset[arr_name].metadata.dimension_names == (
                 "T",
@@ -1198,7 +1206,7 @@ def test_create_position(row, col, pos, version):
             version=version,
         )
         _ = dataset.create_position(row_name=row, col_name=col, pos_name=pos)
-        ome = dataset.zgroup.attrs["ome"]
+        ome = dict(dataset.zgroup.attrs) if version == "0.4" else dataset.zgroup.attrs["ome"]
         assert [c["name"] for c in ome["plate"]["columns"]] == [col]
         assert [r["name"] for r in ome["plate"]["rows"]] == [row]
         assert (store_path / row / col / pos).is_dir()
@@ -1892,3 +1900,69 @@ def test_initialize_pyramid_invalid_dims(implementation, tmp_path):
         pos.create_zeros("0", shape=(1, 1, 2, 8, 8), dtype=np.float32)
         with pytest.raises(ValueError, match="not in dataset axes"):
             pos.initialize_pyramid(levels=2, dims={"w"})
+
+
+# ---------- v0.4 dedicated tests ----------
+
+
+def test_write_ome_zarr_v04_fov_roundtrip(tmp_path):
+    """Full round-trip: create v0.4 FOV store, write image, read back."""
+    store_path = tmp_path / "v04.ome.zarr"
+    data = np.random.default_rng(42).random((1, 2, 3, 64, 64)).astype(np.float32)
+    with open_ome_zarr(
+        store_path,
+        layout="fov",
+        mode="w-",
+        channel_names=["A", "B"],
+        version="0.4",
+    ) as ds:
+        ds.create_image("0", data)
+        assert ds.version == "0.4"
+        # Verify v2 file structure
+        assert (store_path / ".zgroup").exists()
+        assert (store_path / ".zattrs").exists()
+        assert not (store_path / "zarr.json").exists()
+    # Re-open read-only
+    with open_ome_zarr(store_path, layout="fov", mode="r") as ds:
+        assert ds.version == "0.4"
+        assert_array_equal(ds["0"][:], data)
+        assert ds.channel_names == ["A", "B"]
+
+
+def test_write_ome_zarr_v04_hcs_roundtrip(tmp_path):
+    """HCS plate creation with v0.4."""
+    store_path = tmp_path / "v04_hcs.ome.zarr"
+    data = np.zeros((1, 2, 3, 32, 32), dtype=np.uint16)
+    with open_ome_zarr(
+        store_path,
+        layout="hcs",
+        mode="w-",
+        channel_names=["A", "B"],
+        version="0.4",
+    ) as plate:
+        pos = plate.create_position("A", "1", "0")
+        pos.create_image("0", data)
+        # Flat metadata, no "ome" wrapper
+        assert "plate" in plate.zattrs
+        assert "ome" not in plate.zattrs
+    with open_ome_zarr(store_path, layout="hcs", mode="r") as plate:
+        assert plate.version == "0.4"
+        assert_array_equal(plate["A/1/0"]["0"][:], data)
+
+
+def test_sharding_raises_on_v04(tmp_path):
+    """Sharding must raise ValueError for v0.4."""
+    store_path = tmp_path / "v04_shard.zarr"
+    with open_ome_zarr(
+        store_path,
+        layout="fov",
+        mode="w-",
+        channel_names=["A"],
+        version="0.4",
+    ) as ds:
+        with pytest.raises(ValueError, match="Sharding is not supported"):
+            ds.create_image(
+                "0",
+                np.zeros((1, 1, 1, 64, 64)),
+                shards_ratio=(1, 1, 1, 2, 2),
+            )