Added Position.compute_pyramids (#348)

* added Position.compute_pyramid

* moved import to compute_pyramid to avoid circular import issues

* added tensorstore optional dep to docs GHA workflow

* added tensorstore in pr.yml

* tensorstore not optional?

* implemented requested changes

* fixed example

* adjusted TIFFConverter._gen_chunks, had an Dask Chunk alignment problem

* added tensorstore optional dep to docs and pr GHA workflows

* added comma

* added tensorstore in pip_args in docs/poly.py

* undid changes

* Fix chunk validation order in `_gen_chunks`, updated tests, consolidated checking chunks for divisibility

* added delete and adjusted tests

* move _adjust_chunks_for_divisibility into ngff.utils

* reuse chunk limiter in utils

* updated delete pyramid to make sure to delete the metadata

* replace [:1] with [0]

* update deletion

* updated _limit_zyx_chunk_size to include chunks as an optional param

* keep shape and chunks as tuples at input to _adjust_chunks_for_divisibility

* make initialize_pyramid public, cleaned up a test for initializing the pyramid

* Add pyramid sharding (#362)

* wip: add pyramid sharding

* check if array is actually sharded or not

Signed-off-by: Sricharan Reddy Varra <[email protected]>

---------

Signed-off-by: Sricharan Reddy Varra <[email protected]>
Co-authored-by: Sricharan Reddy Varra <[email protected]>

* style: running formatting

Signed-off-by: Sricharan Reddy Varra <[email protected]>

---------

Signed-off-by: Sricharan Reddy Varra <[email protected]>
Co-authored-by: Ivan Ivanov <[email protected]>

Author: srivarra

docs/examples/run_create_pyramid.py

@@ -0,0 +1,37 @@
+"""
+Multiscale Pyramid Creation
+============================
+
+Create and compute a multiscale pyramid for efficient visualization.
+"""
+
+import os
+from tempfile import TemporaryDirectory
+
+import numpy as np
+
+from iohub import open_ome_zarr
+
+# Create temporary directory and sample data
+tmp_dir = TemporaryDirectory()
+store_path = os.path.join(tmp_dir.name, "pyramid.zarr")
+
+data = np.random.randint(0, 255, size=(1, 2, 32, 256, 256), dtype=np.uint16)
+print(f"Original data shape: {data.shape}\n")
+
+# Create position with data and compute pyramid
+with open_ome_zarr(
+    store_path, layout="fov", mode="a", channel_names=["DAPI", "GFP"]
+) as position:
+    # Create full resolution image
+    position.create_image("0", data)
+
+    # Fill pyramid with downsampled data
+    position.compute_pyramid(levels=3, method="mean")
+
+    # Print results
+    print("Pyramid levels:")
+    for level in range(3):
+        level_array = position[str(level)]
+        scale = position.get_effective_scale(str(level))
+        print(f"  Level {level}: {level_array.shape}, scale={scale[-3:]}")

docs/poly.py

@@ -33,7 +33,7 @@
 SOURCE_DIR = "docs/"
 
 #: Arguments to pass to `pip install`
-PIP_ARGS = [".[doc]"]
+PIP_ARGS = [".[doc,tensorstore]"]
 
 #: Mock data used for building local version
 MOCK_DATA = {

iohub/_deprecated/upti.py

@@ -28,7 +28,7 @@ def __init__(self, folder: str, extract_data: bool = False):
             self.files = glob.glob(os.path.join(folder, "*.tif"))
         info_img = tiff.imread(self.files[0])
         self.dtype = info_img.dtype
-        (self.height, self.width) = info_img.shape
+        self.height, self.width = info_img.shape
         self.positions = 1
         self.frames = 1
         self.patterns = 0

iohub/convert.py

@@ -12,6 +12,10 @@
 
 from iohub.ngff.models import TransformationMeta
 from iohub.ngff.nodes import Position, open_ome_zarr
+from iohub.ngff.utils import (
+    _adjust_chunks_for_divisibility,
+    _limit_zyx_chunk_size,
+)
 from iohub.reader import MMStack, NDTiffDataset, read_images
 
 __all__ = ["TIFFConverter"]
@@ -20,23 +24,6 @@
 MAX_CHUNK_SIZE = 500e6  # in bytes
 
 
-def _adjust_chunks_for_divisibility(
-    chunks: list[int], dims: list[int]
-) -> list[int]:
-    """Adjust chunks to divide evenly into dimensions for Dask."""
-    adjusted = []
-    for chunk, dim in zip(chunks, dims):
-        if chunk > dim:
-            adjusted.append(dim)
-        elif dim % chunk != 0:
-            while chunk > 1 and dim % chunk != 0:
-                chunk -= 1
-            adjusted.append(chunk)
-        else:
-            adjusted.append(chunk)
-    return adjusted
-
-
 def _create_grid_from_coordinates(
     xy_coords: list[tuple[float, float]], rows: int, columns: int
 ):
@@ -250,15 +237,19 @@ def _get_pos_names(self):
     def _gen_chunks(self, input_chunks):
         """Generate valid chunk sizes for the output Zarr array.
 
-        First limits chunk size to MAX_CHUNK_SIZE by halving Z, then adjusts
-        chunks to divide evenly into dimensions. Order matters because halving
-        Z (e.g. 10 -> 5 -> 2) may no longer divide evenly, so we validate
-        divisibility last.
+        input_chunks may be a string ("XY", "XYZ") or a tuple of chunk
+        dimensions. Chunk size will be limited to MAX_CHUNK_SIZE and adjusted
+        to divide evenly into dimensions.
         """
         if not input_chunks:
             _logger.debug("No chunk size specified, using ZYX.")
             chunks = [1, 1, self.z, self.y, self.x]
         elif isinstance(input_chunks, tuple):
+            if not len(input_chunks) == 5:
+                raise ValueError(
+                    "Input chunks must be a tuple of 5 dimensions, got "
+                    f"{len(input_chunks)} dimensions."
+                )
             chunks = list(input_chunks)
         elif isinstance(input_chunks, str):
             if input_chunks.lower() == "xy":
@@ -272,19 +263,20 @@ def _gen_chunks(self, input_chunks):
                 f"Chunk type {type(input_chunks)} is not supported."
             )
 
+        shape = (self.t, self.c, self.z, self.y, self.x)
+        original_chunks = chunks.copy()
+
+        # Limit chunk size to MAX_CHUNK_SIZE by halving Z
         bytes_per_pixel = np.dtype(self.reader.dtype).itemsize
-        while (
-            chunks[-3] > 1
-            and np.prod(chunks, dtype=np.int64) * bytes_per_pixel
-            > MAX_CHUNK_SIZE
-        ):
-            chunks[-3] = int(np.ceil(chunks[-3] / 2))
+        chunk_zyx_shape = _limit_zyx_chunk_size(
+            shape, bytes_per_pixel, MAX_CHUNK_SIZE, chunks=chunks
+        )
+        chunks[-3:] = list(chunk_zyx_shape)
 
-        data_dims = [self.t, self.c, self.z, self.y, self.x]
-        original_chunks = chunks.copy()
-        chunks = _adjust_chunks_for_divisibility(chunks, data_dims)
+        # Adjust chunks to divide evenly into dimensions
+        chunks = _adjust_chunks_for_divisibility(shape, chunks)
         for i, (orig, adj, dim) in enumerate(
-            zip(original_chunks, chunks, data_dims)
+            zip(original_chunks, chunks, shape)
         ):
             if orig != adj:
                 _logger.warning(

iohub/ngff/nodes.py

@@ -207,8 +207,7 @@ def __setitem__(self, key, value):
     def __delitem__(self, key):
         """.. Warning: this does NOT clean up metadata!"""
         key = normalize_path(str(key))
-        if key in self._member_names:
-            del self[key]
+        del self._group[key]
 
     def __contains__(self, key):
         key = normalize_path(str(key))
@@ -1068,6 +1067,14 @@ def initialize_pyramid(self, levels: int) -> None:
                 _scale_integers(array.chunks, factor), len(shape)
             )
 
+            if array.shards is not None:
+                shards = array.shards[:-3] + _scale_integers(
+                    array.shards[-3:], factor
+                )
+                shards_ratio = tuple(s // c for c, s in zip(chunks, shards))
+            else:
+                shards_ratio = None
+
             transforms = deepcopy(
                 self.metadata.multiscales[0]
                 .datasets[0]
@@ -1083,9 +1090,113 @@ def initialize_pyramid(self, levels: int) -> None:
                 shape=shape,
                 dtype=array.dtype,
                 chunks=chunks,
+                shards_ratio=shards_ratio,
                 transform=transforms,
             )
 
+    def compute_pyramid(
+        self,
+        levels: int | None = None,
+        method: str = "mean",
+    ) -> None:
+        """Compute pyramid by downsampling from source level.
+
+        Creates pyramid structure if none exists. Uses cascade downsampling
+        where each level is derived from the previous level to prevent
+        aliasing artifacts and chunk boundary issues.
+
+        Parameters
+        ----------
+        levels : int, optional
+            Number of pyramid levels. If None, uses existing pyramid structure.
+        method : str, optional
+            Downsampling method: "mean", "median", "mode", "min", "max",
+            "stride". By default "mean".
+
+        Raises
+        ------
+        ValueError
+            If level 0 array doesn't exist, pyramid structure is invalid,
+            or if a pyramid already exists with a different number of levels.
+
+        Examples
+        --------
+        >>> # Create and compute pyramid with 4 levels
+        >>> pos.compute_pyramid(levels=4, method="mean")
+
+        >>> # Recompute existing pyramid structure
+        >>> pos.compute_pyramid(method="median")
+
+        >>> # Change pyramid levels (must delete first)
+        >>> pos.delete_pyramid()
+        >>> pos.compute_pyramid(levels=3, method="mean")
+        """
+        from iohub.ngff.utils import _downsample_tensorstore
+
+        num_arrays = len(self.array_keys())
+        if num_arrays == 0:
+            raise ValueError(
+                "No level 0 array exists. Create base array before computing "
+                "pyramid."
+            )
+
+        if levels is None:
+            if num_arrays == 1:
+                raise ValueError(
+                    "Pyramid structure doesn't exist and levels=None. "
+                    "Specify 'levels' parameter to create pyramid."
+                )
+            levels = num_arrays
+
+        if num_arrays == 1:
+            self.initialize_pyramid(levels=levels)
+        elif num_arrays != levels:
+            raise ValueError(
+                f"Pyramid structure exists with {num_arrays} levels but "
+                f"{levels} requested. Call delete_pyramid() first to remove "
+                "existing pyramid."
+            )
+
+        # Compute pyramid data via cascade downsampling
+        for level in range(1, levels):
+            previous_level_array = self[str(level - 1)]
+            current_level_array = self[str(level)]
+
+            previous_ts = previous_level_array.tensorstore()
+            current_ts = current_level_array.tensorstore()
+
+            current_scale = self.get_effective_scale(str(level))
+            previous_scale = self.get_effective_scale(str(level - 1))
+
+            downsample_factors = [
+                int(round(current_scale[i] / previous_scale[i]))
+                for i in range(len(current_scale))
+            ]
+
+            _downsample_tensorstore(
+                source_ts=previous_ts,
+                target_ts=current_ts,
+                downsample_factors=downsample_factors,
+                method=method,
+            )
+
+    def delete_pyramid(self) -> None:
+        """Delete all dataset pyramid levels except the base (level 0) array.
+
+        Use this before calling compute_pyramid() with different levels
+        on a position that already has a pyramid structure.
+
+        This method removes both the zarr arrays and updates the OME-NGFF
+        multiscales metadata to reflect the deletion.
+        """
+        _multiscale = self.metadata.multiscales[0]
+
+        for dataset in _multiscale.datasets[1:]:
+            del self[dataset.path]
+
+        _multiscale.datasets = _multiscale.datasets[:1]
+        self.dump_meta()
+
     @property
     def scale(self) -> list[float]:
         """