Replace xarray-ome with iohub native xarray support

iohub@main has Position.to_xarray() which returns a dask-backed
DataArray with proper coordinates and scales. This eliminates the
xarray-ome dependency and the brittle zarr store path extraction.

Dims are now lowercase (t,c,z,y,x) per iohub convention.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

demo_utils.py

@@ -21,7 +21,7 @@ import numpy as np
 import torch
 import xarray as xr
 from numpy.typing import NDArray
-from xarray_ome import open_ome_dataset
+from iohub import open_ome_zarr
 
 from waveorder import util
 from waveorder.models import isotropic_thin_3d
@@ -129,9 +129,7 @@ def load_fov_from_plate(
     plate, row: str, column: str, field: str, resolution: int = 0
 ) -> xr.DataArray:
     """
-    Load a specific FOV from HCS plate using hybrid iohub + xarray-ome approach.
-
-    Uses iohub for navigation, then xarray-ome for fast data loading.
+    Load a specific FOV from HCS plate using iohub.
 
     Parameters
     ----------
@@ -149,28 +147,11 @@ def load_fov_from_plate(
     Returns
     -------
     xr.DataArray
-        Image data with labeled dimensions (T, C, Z, Y, X)
+        Image data with labeled dimensions (t, c, z, y, x)
     """
-    # Navigate to position using iohub (fast)
     position_key = f"{row}/{column}/{field}"
     position = plate[position_key]
-
-    # Get full zarr path from position (handle both Zarr V2 and V3)
-    store = position.zgroup.store
-    if hasattr(store, 'path'):
-        base_path = Path(store.path)  # Zarr V2 (DirectoryStore)
-    elif hasattr(store, 'root'):
-        base_path = Path(store.root)  # Zarr V3 (LocalStore)
-    else:
-        raise RuntimeError(f"Unknown store type: {type(store)}")
-
-    position_path = base_path / position.zgroup.path
-
-    # Load with xarray-ome (fast and reliable)
-    fov_dataset = open_ome_dataset(position_path, resolution=resolution, validate=False)
-    data_xr = fov_dataset["image"]
-
-    return data_xr
+    return position.to_xarray()
 
 
 # === Data Loading ===
@@ -192,7 +173,7 @@ def load_ome_zarr_fov(
     Returns
     -------
     xr.DataArray
-        Image data with labeled dimensions (T, C, Z, Y, X)
+        Image data with labeled dimensions (t, c, z, y, x)
     """
     zarr_path = Path(zarr_path)
     fov_path = Path(fov_path)
@@ -200,13 +181,8 @@ def load_ome_zarr_fov(
     print(f"Loading zarr store from: {zarr_path}")
     print(f"Accessing FOV: {fov_path}")
 
-    # Load as xarray Dataset
-    fov_dataset: xr.Dataset = open_ome_dataset(
-        zarr_path / fov_path, resolution=resolution, validate=False
-    )
-
-    # Extract the image DataArray
-    data_xr = fov_dataset["image"]
+    position = open_ome_zarr(str(zarr_path / fov_path), mode="r")
+    data_xr = position.to_xarray()
 
     print(f"Loaded data shape: {dict(data_xr.sizes)}")
     print(f"Dimensions: {list(data_xr.dims)}")
@@ -316,11 +292,11 @@ def extract_2d_slice(
     # Build selection dictionary for indexed dimensions
     sel_dict = {}
     if t is not None:
-        sel_dict["T"] = int(t)
+        sel_dict["t"] = int(t)
     if c is not None:
-        sel_dict["C"] = int(c)
+        sel_dict["c"] = int(c)
     if z is not None:
-        sel_dict["Z"] = int(z)
+        sel_dict["z"] = int(z)
 
     # Extract slice using xarray's labeled indexing
     slice_xr = data_xr.isel(**sel_dict) if sel_dict else data_xr
@@ -455,7 +431,10 @@ def print_data_summary(data_xr: xr.DataArray) -> None:
     for dim in info["dims"]:
         coords = info["coords"][dim]
         if len(coords) > 0:
-            print(f"  {dim}: [{coords[0]:.2f} ... {coords[-1]:.2f}] (n={len(coords)})")
+            try:
+                print(f"  {dim}: [{coords[0]:.2f} ... {coords[-1]:.2f}] (n={len(coords)})")
+            except (ValueError, TypeError):
+                print(f"  {dim}: [{coords[0]} ... {coords[-1]}] (n={len(coords)})")
 
     # Print memory size estimate
     total_elements = np.prod(list(info["sizes"].values()))

optimize_demo.py

@@ -159,7 +159,7 @@ def load_selected_fov(field: str, current_z: int, plate_metadata):
         new_pixel_scales = (Config.PIXEL_SIZE_Z, Config.PIXEL_SIZE_YX, Config.PIXEL_SIZE_YX)
 
         # Update Z slider
-        z_max = new_data_xr.sizes["Z"] - 1
+        z_max = new_data_xr.sizes["z"] - 1
         new_z = min(current_z, z_max)
 
         print(f"✅ Loaded: {dict(new_data_xr.sizes)}")
@@ -232,7 +232,7 @@ def run_reconstruction_ui(
     Uses slider parameters directly for a single fast reconstruction.
     """
     # Extract full Z-stack for timepoint 0 (for reconstruction)
-    zyx_stack = data_xr_state.isel(T=0, C=Config.CHANNEL).values
+    zyx_stack = data_xr_state.isel(t=0, c=Config.CHANNEL).values
 
     # Get current Z-slice for comparison (left side of ImageSlider)
     original_normalized = extract_2d_slice(
@@ -276,7 +276,7 @@ def run_optimization_ui(
     iteration history, iteration slider, and SLIDER UPDATES.
     """
     # Extract full Z-stack for timepoint 0 (for reconstruction)
-    zyx_stack = data_xr_state.isel(T=0, C=Config.CHANNEL).values
+    zyx_stack = data_xr_state.isel(t=0, c=Config.CHANNEL).values
 
     # Get current Z-slice for comparison (left side of ImageSlider)
     original_normalized = extract_2d_slice(
@@ -529,7 +529,7 @@ def create_gradio_interface(plate_metadata, default_fields, data_xr, pixel_scale
                     data_xr,
                     t=0,
                     c=Config.CHANNEL,
-                    z=data_xr.sizes["Z"] // 2,
+                    z=data_xr.sizes["z"] // 2,
                     normalize=True,
                     verbose=False,
                 )
@@ -547,8 +547,8 @@ def create_gradio_interface(plate_metadata, default_fields, data_xr, pixel_scale
                 gr.Markdown("### 🎛️ Navigation")
                 z_slider = gr.Slider(
                     minimum=0,
-                    maximum=data_xr.sizes["Z"] - 1,
-                    value=data_xr.sizes["Z"] // 2,
+                    maximum=data_xr.sizes["z"] - 1,
+                    value=data_xr.sizes["z"] // 2,
                     step=1,
                     label="Z-slice",
                     scale=1,

requirements.txt

@@ -3,8 +3,7 @@
 
 # Install waveorder from gradio-demo branch
 git+https://github.com/mehta-lab/waveorder.git@gradio-demo
-git+https://github.com/ianhi/xarray-ome.git@main
-git+https://github.com/czbiohub-sf/iohub.git@v0.3.0a2
+git+https://github.com/czbiohub-sf/iohub.git@main
 
 # Gradio for web interface
 gradio>=6.0.2