Merge pull request #124 from thewtex/dimension-names

BUG: Write out array dimension_names for OME-Zarr 0.5
thewtex · Dec 6, 2024 · 1823516 · 1823516
2 parents 9a3b09e + 50cc8f6
commit 1823516
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 1 deletion.
diff --git a/ngff_zarr/to_ngff_zarr.py b/ngff_zarr/to_ngff_zarr.py
@@ -105,7 +105,7 @@ def _numpy_to_zarr_dtype(dtype):
 
 
 def _write_with_tensorstore(
-    store_path: str, array, region, chunks, zarr_format
+    store_path: str, array, region, chunks, zarr_format, dimension_names=None
 ) -> None:
     """Write array using tensorstore backend"""
     import tensorstore as ts
@@ -131,6 +131,8 @@ def _write_with_tensorstore(
             "configuration": {"chunk_shape": chunks},
         }
         spec["metadata"]["data_type"] = _numpy_to_zarr_dtype(array.dtype)
+        if dimension_names:
+            spec["metadata"]["dimension_names"] = dimension_names
     else:
         raise ValueError(f"Unsupported zarr format: {zarr_format}")
     dataset = ts.open(spec, create=True, dtype=array.dtype).result()
@@ -199,6 +201,10 @@ def to_ngff_zarr(
             coordinateTransformations=metadata.coordinateTransformations,
             name=metadata.name,
         )
+    dimension_names = tuple([ax.name for ax in metadata.axes])
+    dimension_names_kwargs = (
+        {"dimension_names": dimension_names} if version != "0.4" else {}
+    )
 
     metadata_dict = asdict(metadata)
     metadata_dict = _pop_metadata_optionals(metadata_dict)
@@ -274,6 +280,7 @@ def to_ngff_zarr(
                 path=path,
                 mode="a",
                 **zarr_kwargs,
+                **dimension_names_kwargs,
             )
 
             shape = image.data.shape
@@ -407,6 +414,7 @@ def to_ngff_zarr(
                             region,
                             [c[0] for c in arr_region.chunks],
                             zarr_format=zarr_format,
+                            dimension_names=dimension_names,
                             **kwargs,
                         )
                     else:
@@ -419,6 +427,7 @@ def to_ngff_zarr(
                             compute=True,
                             return_stored=False,
                             **zarr_kwargs,
+                            **dimension_names_kwargs,
                             **kwargs,
                         )
         else:
@@ -435,6 +444,7 @@ def to_ngff_zarr(
                     region,
                     [c[0] for c in arr.chunks],
                     zarr_format=zarr_format,
+                    dimension_names=dimension_names,
                     **kwargs,
                 )
             else:
@@ -447,6 +457,7 @@ def to_ngff_zarr(
                     compute=True,
                     return_stored=False,
                     **zarr_kwargs,
+                    **dimension_names_kwargs,
                     **kwargs,
                 )
 

diff --git a/test/test_to_ngff_zarr_rfc2_zarr_v3.py b/test/test_to_ngff_zarr_rfc2_zarr_v3.py
@@ -31,6 +31,11 @@ def test_gaussian_isotropic_scale_factors(input_images):
     # store_new_multiscales(dataset_name, baseline_name, multiscales, version=version)
     verify_against_baseline(dataset_name, baseline_name, multiscales, version=version)
 
+    array0 = zarr.open_array(store=store, path="scale0/image", mode="r", zarr_format=3)
+    dimension_names = array0.metadata.dimension_names
+    for idx, ax in enumerate(multiscales.metadata.axes):
+        assert ax.name == dimension_names[idx]
+
 
 def test_gaussian_isotropic_scale_factors_tensorstore(input_images):
     pytest.importorskip("tensorstore")
@@ -48,3 +53,10 @@ def test_gaussian_isotropic_scale_factors_tensorstore(input_images):
         verify_against_baseline(
             dataset_name, baseline_name, multiscales, version=version
         )
+
+        array0 = zarr.open_array(
+            store=tmpdir, path="scale0/image", mode="r", zarr_format=3
+        )
+        dimension_names = array0.metadata.dimension_names
+        for idx, ax in enumerate(multiscales.metadata.axes):
+            assert ax.name == dimension_names[idx]