Port to iohub 0.3.2 ImageArray API (#407)

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Co-authored-by: Alexandr Kalinin <alxndrkalinin@users.noreply.github.com>

Author: alxndrkalinin

applications/airtable/pyproject.toml

@@ -8,14 +8,13 @@ description = "Interface to the Computational Imaging Airtable database"
 keywords = [ "airtable", "metadata", "microscopy", "zarr" ]
 license = "BSD-3-Clause"
 authors = [ { name = "Biohub", email = "compmicro@czbiohub.org" } ]
-requires-python = ">=3.11"
+requires-python = ">=3.12"
 classifiers = [
   "Development Status :: 3 - Alpha",
   "Intended Audience :: Science/Research",
   "License :: OSI Approved :: BSD License",
   "Operating System :: OS Independent",
   "Programming Language :: Python :: 3 :: Only",
-  "Programming Language :: Python :: 3.11",
   "Programming Language :: Python :: 3.12",
   "Programming Language :: Python :: 3.13",
   "Programming Language :: Python :: 3.14",

applications/cytoland/pyproject.toml

@@ -14,14 +14,13 @@ keywords = [
 ]
 license = "BSD-3-Clause"
 authors = [ { name = "Biohub", email = "compmicro@czbiohub.org" } ]
-requires-python = ">=3.11"
+requires-python = ">=3.12"
 classifiers = [
   "Development Status :: 4 - Beta",
   "Intended Audience :: Science/Research",
   "License :: OSI Approved :: BSD License",
   "Operating System :: OS Independent",
   "Programming Language :: Python :: 3 :: Only",
-  "Programming Language :: Python :: 3.11",
   "Programming Language :: Python :: 3.12",
   "Programming Language :: Python :: 3.13",
   "Programming Language :: Python :: 3.14",

applications/dynacell/pyproject.toml

@@ -14,14 +14,13 @@ keywords = [
 ]
 license = "BSD-3-Clause"
 authors = [ { name = "Biohub", email = "compmicro@czbiohub.org" } ]
-requires-python = ">=3.11"
+requires-python = ">=3.12"
 classifiers = [
   "Development Status :: 3 - Alpha",
   "Intended Audience :: Science/Research",
   "License :: OSI Approved :: BSD License",
   "Operating System :: OS Independent",
   "Programming Language :: Python :: 3 :: Only",
-  "Programming Language :: Python :: 3.11",
   "Programming Language :: Python :: 3.12",
   "Programming Language :: Python :: 3.13",
   "Programming Language :: Python :: 3.14",

applications/dynaclr/pyproject.toml

@@ -15,14 +15,13 @@ keywords = [
 ]
 license = "BSD-3-Clause"
 authors = [ { name = "Biohub", email = "compmicro@czbiohub.org" } ]
-requires-python = ">=3.11"
+requires-python = ">=3.12"
 classifiers = [
   "Development Status :: 4 - Beta",
   "Intended Audience :: Science/Research",
   "License :: OSI Approved :: BSD License",
   "Operating System :: OS Independent",
   "Programming Language :: Python :: 3 :: Only",
-  "Programming Language :: Python :: 3.11",
   "Programming Language :: Python :: 3.12",
   "Programming Language :: Python :: 3.13",
   "Programming Language :: Python :: 3.14",

applications/dynaclr/src/dynaclr/data/datamodule.py

@@ -12,9 +12,11 @@
 from __future__ import annotations
 
 import logging
+import os
 
 import numpy as np
 import pandas as pd
+from iohub.core.config import TensorStoreConfig
 from lightning.pytorch import LightningDataModule
 from monai.data.thread_buffer import ThreadDataLoader
 from monai.transforms import Compose, MapTransform
@@ -238,6 +240,12 @@ def __init__(
         # Loss hyperparameters (informational)
         # Other
         self.cache_pool_bytes = cache_pool_bytes
+        cpus = os.environ.get("SLURM_CPUS_PER_TASK")
+        cpus = int(cpus) if cpus is not None else (os.cpu_count() or 4)
+        self.tensorstore_config = TensorStoreConfig(
+            data_copy_concurrency=cpus,
+            cache_pool_bytes=cache_pool_bytes or None,
+        )
         self.seed = seed
         self.include_wells = include_wells
         self.exclude_fovs = exclude_fovs
@@ -361,13 +369,13 @@ def _setup_experiment_split(self, registry: ExperimentRegistry) -> None:
             positive_cell_source=self.positive_cell_source,
             positive_match_columns=self.positive_match_columns,
             max_border_shift=self.max_border_shift,
+            tensorstore_config=self.tensorstore_config,
         )
         self.train_dataset = MultiExperimentTripletDataset(
             index=train_index,
             fit=True,
             tau_range_hours=self.tau_range,
             tau_decay_rate=self.tau_decay_rate,
-            cache_pool_bytes=self.cache_pool_bytes,
             channels_per_sample=self.channels_per_sample,
             positive_cell_source=self.positive_cell_source,
             positive_match_columns=self.positive_match_columns,
@@ -388,13 +396,13 @@ def _setup_experiment_split(self, registry: ExperimentRegistry) -> None:
                 positive_cell_source=self.positive_cell_source,
                 positive_match_columns=self.positive_match_columns,
                 max_border_shift=self.max_border_shift,
+                tensorstore_config=self.tensorstore_config,
             )
             self.val_dataset = MultiExperimentTripletDataset(
                 index=val_index,
                 fit=True,
                 tau_range_hours=self.tau_range,
                 tau_decay_rate=self.tau_decay_rate,
-                cache_pool_bytes=self.cache_pool_bytes,
                 channels_per_sample=self.channels_per_sample,
                 positive_cell_source=self.positive_cell_source,
                 positive_match_columns=self.positive_match_columns,
@@ -419,6 +427,7 @@ def _setup_fov_split(self, registry: ExperimentRegistry) -> None:
             num_workers=self.num_workers_index,
             positive_cell_source=self.positive_cell_source,
             positive_match_columns=self.positive_match_columns,
+            tensorstore_config=self.tensorstore_config,
         )
 
         rng = np.random.default_rng(self.seed)
@@ -458,7 +467,6 @@ def _setup_fov_split(self, registry: ExperimentRegistry) -> None:
             fit=True,
             tau_range_hours=self.tau_range,
             tau_decay_rate=self.tau_decay_rate,
-            cache_pool_bytes=self.cache_pool_bytes,
             channels_per_sample=self.channels_per_sample,
             positive_cell_source=self.positive_cell_source,
             positive_match_columns=self.positive_match_columns,
@@ -477,7 +485,6 @@ def _setup_fov_split(self, registry: ExperimentRegistry) -> None:
                 fit=True,
                 tau_range_hours=self.tau_range,
                 tau_decay_rate=self.tau_decay_rate,
-                cache_pool_bytes=self.cache_pool_bytes,
                 channels_per_sample=self.channels_per_sample,
                 positive_cell_source=self.positive_cell_source,
                 positive_match_columns=self.positive_match_columns,

applications/dynaclr/src/dynaclr/data/dataset.py

@@ -17,7 +17,6 @@
 from __future__ import annotations
 
 import logging
-import os
 from collections import defaultdict
 
 import numpy as np
@@ -112,8 +111,6 @@ class MultiExperimentTripletDataset(Dataset):
     return_negative : bool
         Reserved for future use.  Currently unused (NTXentLoss uses
         in-batch negatives).
-    cache_pool_bytes : int
-        Tensorstore cache pool size in bytes.
     channels_per_sample : int | list[str] | None
         Controls how many source channels to read per sample.
         ``None`` (default) — read all source channels, output ``(B, C, Z, Y, X)``.
@@ -148,7 +145,6 @@ def __init__(
         tau_range_hours: tuple[float, float] = (0.5, 2.0),
         tau_decay_rate: float = 2.0,
         return_negative: bool = False,
-        cache_pool_bytes: int = 0,
         channels_per_sample: int | list[str] | None = None,
         positive_cell_source: str = "lookup",
         positive_match_columns: list[str] | None = None,
@@ -205,25 +201,9 @@ def __init__(
                 _logger.info("Label encoder '%s' (%s): %d classes", batch_key, col, len(encoder))
 
         self._rng = np.random.default_rng()
-        self._setup_tensorstore_context(cache_pool_bytes)
-        self._build_match_lookup()
-
-    # ------------------------------------------------------------------
-    # Initialization helpers
-    # ------------------------------------------------------------------
-
-    def _setup_tensorstore_context(self, cache_pool_bytes: int) -> None:
-        """Configure tensorstore context with CPU limits based on SLURM env."""
-        cpus = os.environ.get("SLURM_CPUS_PER_TASK")
-        cpus = int(cpus) if cpus is not None else (os.cpu_count() or 4)
-        self._ts_context = ts.Context(
-            {
-                "data_copy_concurrency": {"limit": cpus},
-                "cache_pool": {"total_bytes_limit": cache_pool_bytes},
-            }
-        )
         self._tensorstores: dict[str, ts.TensorStore] = {}
         self._norm_meta_cache: dict[str, NormMeta | None] = {}
+        self._build_match_lookup()
 
     def _build_match_lookup(self) -> None:
         """Build lookup structures for O(1) positive candidate lookup.
@@ -531,10 +511,7 @@ def _get_tensorstore(self, position, fov_name: str) -> "ts.TensorStore":
         ts.TensorStore
         """
         if fov_name not in self._tensorstores:
-            self._tensorstores[fov_name] = position["0"].tensorstore(
-                context=self._ts_context,
-                recheck_cached_data="open",
-            )
+            self._tensorstores[fov_name] = position["0"].native
         return self._tensorstores[fov_name]
 
     def _slice_patch(

applications/dynaclr/src/dynaclr/data/index.py

@@ -14,6 +14,7 @@
 
 import numpy as np
 import pandas as pd
+from iohub.core.config import TensorStoreConfig
 from iohub.ngff import Plate, Position, open_ome_zarr
 
 from dynaclr.data.experiment import ExperimentRegistry
@@ -188,6 +189,7 @@ def __init__(
         positive_cell_source: str = "lookup",
         positive_match_columns: list[str] | None = None,
         max_border_shift: int = -1,
+        tensorstore_config: TensorStoreConfig | None = None,
     ) -> None:
         self.registry = registry
         self.yx_patch_size = yx_patch_size
@@ -197,6 +199,11 @@ def __init__(
         if max_border_shift < 0:
             max_border_shift = max(yx_patch_size[0] // 4, yx_patch_size[1] // 4)
         self.max_border_shift = max_border_shift
+        # Plates cached here feed Position objects whose arrays the dataset reads
+        # via ``position["0"].native`` (tensorstore handle). The tensorstore impl
+        # must be configured at open-time — default zarr would return a
+        # ``zarr.Array`` that has no ``.read().result()`` method.
+        self.tensorstore_config = tensorstore_config or TensorStoreConfig()
         self._store_cache: dict[str, Plate] = {}
 
         # Merge collection-level exclude_fovs with runtime exclude_fovs
@@ -362,7 +369,12 @@ def _resolve_positions_and_dims(self, tracks: pd.DataFrame) -> tuple[list[Positi
 
         for (store_path, well_name, fov_name), _group in tracks.groupby(["store_path", "well_name", "fov_name"]):
             if store_path not in self._store_cache:
-                self._store_cache[store_path] = open_ome_zarr(store_path, mode="r")
+                self._store_cache[store_path] = open_ome_zarr(
+                    store_path,
+                    mode="r",
+                    implementation="tensorstore",
+                    implementation_config=self.tensorstore_config,
+                )
             plate = self._store_cache[store_path]
             # fov_name may be just the FOV id (e.g. "000000") or the full
             # position path (e.g. "C/1/000000"). Prepend well_name when needed.

applications/qc/pyproject.toml

@@ -9,14 +9,13 @@ readme = "README.md"
 keywords = [ "microscopy", "quality control", "zarr" ]
 license = "BSD-3-Clause"
 authors = [ { name = "Biohub", email = "compmicro@czbiohub.org" } ]
-requires-python = ">=3.11"
+requires-python = ">=3.12"
 classifiers = [
   "Development Status :: 4 - Beta",
   "Intended Audience :: Science/Research",
   "License :: OSI Approved :: BSD License",
   "Operating System :: OS Independent",
   "Programming Language :: Python :: 3 :: Only",
-  "Programming Language :: Python :: 3.11",
   "Programming Language :: Python :: 3.12",
   "Programming Language :: Python :: 3.13",
   "Programming Language :: Python :: 3.14",

applications/qc/src/qc/focus.py

@@ -1,7 +1,6 @@
 """In-focus z-slice detection using midband spatial frequency power."""
 
 import numpy as np
-import tensorstore
 import torch
 from waveorder.focus import focus_from_transverse_band
 
@@ -46,17 +45,16 @@ def __init__(
         self.device = torch.device(device)
 
     def channels(self) -> list[str]:
+        """Return the channels this metric is configured for."""
         return self.channel_names
 
     def __call__(self, position, channel_name, channel_index, num_workers=4):
-        tzyx = (
-            position["0"]
-            .tensorstore(context=tensorstore.Context({"data_copy_concurrency": {"limit": num_workers}}))[
-                :, channel_index
-            ]
-            .read()
-            .result()
-        )
+        """Compute focus-slice index per timepoint for one channel of ``position``."""
+        # Tensorstore concurrency is configured on the plate at
+        # open-time (see qc_metrics.generate_qc_metadata); num_workers
+        # is retained here only to match the QCMetric abstract interface.
+        del num_workers
+        tzyx = position["0"].native[:, channel_index].read().result()
 
         T = tzyx.shape[0]
         focus_indices = np.empty(T, dtype=int)

applications/qc/src/qc/qc_metrics.py

@@ -4,6 +4,7 @@
 from abc import ABC, abstractmethod
 
 import iohub.ngff as ngff
+from iohub.core.config import TensorStoreConfig
 from tqdm import tqdm
 
 from viscy_utils.meta_utils import write_meta_field
@@ -83,42 +84,45 @@ def generate_qc_metadata(
     num_workers : int
         Number of workers for data loading.
     """
-    plate = ngff.open_ome_zarr(zarr_dir, mode="r+")
-    position_map = list(plate.positions())
+    with ngff.open_ome_zarr(
+        zarr_dir,
+        mode="r+",
+        implementation="tensorstore",
+        implementation_config=TensorStoreConfig(data_copy_concurrency=num_workers),
+    ) as plate:
+        position_map = list(plate.positions())
 
-    for metric in metrics:
-        channel_list = metric.channels()
+        for metric in metrics:
+            channel_list = metric.channels()
 
-        for channel_name in channel_list:
-            channel_index = plate.channel_names.index(channel_name)
-            _logger.info(f"Computing {metric.field_name} for channel '{channel_name}'")
+            for channel_name in channel_list:
+                channel_index = plate.channel_names.index(channel_name)
+                _logger.info(f"Computing {metric.field_name} for channel '{channel_name}'")
 
-            position_results = []
+                position_results = []
 
-            for _, pos in tqdm(position_map, desc="Positions"):
-                result = metric(pos, channel_name, channel_index, num_workers)
-                position_results.append((pos, result))
+                for _, pos in tqdm(position_map, desc="Positions"):
+                    result = metric(pos, channel_name, channel_index, num_workers)
+                    position_results.append((pos, result))
 
-            all_results = [r for _, r in position_results]
-            dataset_stats = metric.aggregate_dataset(all_results)
+                all_results = [r for _, r in position_results]
+                dataset_stats = metric.aggregate_dataset(all_results)
 
-            if dataset_stats:
-                write_meta_field(
-                    position=plate,
-                    metadata={"dataset_statistics": dataset_stats},
-                    field_name=metric.field_name,
-                    subfield_name=channel_name,
-                )
-
-            for pos, result in position_results:
-                metadata = {**result}
                 if dataset_stats:
-                    metadata["dataset_statistics"] = dataset_stats
-                write_meta_field(
-                    position=pos,
-                    metadata=metadata,
-                    field_name=metric.field_name,
-                    subfield_name=channel_name,
-                )
-
-    plate.close()
+                    write_meta_field(
+                        position=plate,
+                        metadata={"dataset_statistics": dataset_stats},
+                        field_name=metric.field_name,
+                        subfield_name=channel_name,
+                    )
+
+                for pos, result in position_results:
+                    metadata = {**result}
+                    if dataset_stats:
+                        metadata["dataset_statistics"] = dataset_stats
+                    write_meta_field(
+                        position=pos,
+                        metadata=metadata,
+                        field_name=metric.field_name,
+                        subfield_name=channel_name,
+                    )