Added optional parameter open_array_kwargs to build_pyramid, to be passed as kwargs to zarr.open when creating the pyramid level arrays.

Author: m-albert

fractal_tasks_core/pyramids.py

@@ -14,7 +14,9 @@
 """
 import logging
 import pathlib
+import zarr
 from typing import Callable
+from typing import Mapping
 from typing import Optional
 from typing import Sequence
 from typing import Union
@@ -33,6 +35,7 @@ def build_pyramid(
     coarsening_xy: int = 2,
     chunksize: Optional[Sequence[int]] = None,
     aggregation_function: Optional[Callable] = None,
+    open_array_kwargs: Optional[Mapping] = None,
 ) -> None:
 
     """
@@ -48,6 +51,7 @@ def build_pyramid(
         coarsening_xy: Linear coarsening factor between subsequent levels.
         chunksize: Shape of a single chunk.
         aggregation_function: Function to be used when downsampling.
+        open_array_kwargs: Additional arguments for zarr.open.
     """
 
     # Clean up zarrurl
@@ -100,10 +104,21 @@ def build_pyramid(
             f"{str(newlevel_rechunked)}"
         )
 
+        if open_array_kwargs is None:
+            open_array_kwargs = {}
+
+        zarrarr = zarr.open(
+            f"{zarrurl}/{ind_level}",
+            shape=newlevel_rechunked.shape,
+            chunks=newlevel_rechunked.chunksize,
+            dtype=newlevel_rechunked.dtype,
+            mode="w",
+            **open_array_kwargs,
+        )
+
         # Write zarr and store output (useful to construct next level)
         previous_level = newlevel_rechunked.to_zarr(
-            zarrurl,
-            component=f"{ind_level}",
+            zarrarr,
             overwrite=overwrite,
             compute=True,
             return_stored=True,

tests/test_unit_pyramid_creation.py

@@ -64,3 +64,16 @@ def test_build_pyramid(tmp_path):
     assert level_3.chunksize == (9, 9)
     assert level_4.shape == (3, 3)
     assert level_5.shape == (1, 1)
+
+    # Succeed
+    zarrurl = str(tmp_path / "F.zarr")
+    da.zeros(shape=(8, 8)).to_zarr(f"{zarrurl}/0")
+    build_pyramid(zarrurl=zarrurl, coarsening_xy=2, num_levels=3,
+        open_array_kwargs={"write_empty_chunks": False, "fill_value": 0})
+    level_1 = da.from_zarr(f"{zarrurl}/1")
+    level_2 = da.from_zarr(f"{zarrurl}/2")
+    assert level_1.shape == (4, 4)
+    assert level_2.shape == (2, 2)
+    # check that the empty chunks are not written to disk
+    assert not (tmp_path / "F.zarr/1/0.0").exists()
+    assert not (tmp_path / "F.zarr/2/0.0").exists()