Add xarray integration to Position (#370)

* Add Position.to_xarray() method

* Add Position.write_xarray() method

* Add tests for to_xarray and write_xarray

* better channel names handling

* lowercase dimension names

* nest xarray atributes under "iohub" key

* rng fixture

* revise tests to use hypothesis

* use oindex instead of double loop

Author: talonchandler

iohub/ngff/nodes.py

@@ -15,6 +15,7 @@
 from typing import Generator, Literal, Sequence, Type, TypeAlias, overload
 
 import numpy as np
+import xarray as xr
 import zarr.codecs
 from numpy.typing import ArrayLike, DTypeLike, NDArray
 from pydantic import ValidationError
@@ -984,18 +985,12 @@ def initialize_pyramid(self, levels: int) -> None:
             chunks = _pad_shape(_scale_integers(array.chunks, factor), len(shape))
 
             if array.shards is not None:
-                shards = array.shards[:-3] + _scale_integers(
-                    array.shards[-3:], factor
-                )
+                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]
-                .coordinate_transformations
-            )
+            transforms = deepcopy(self.metadata.multiscales[0].datasets[0].coordinate_transformations)
             for tr in transforms:
                 if tr.type == "scale":
                     for i in range(len(tr.scale))[-3:]:
@@ -1051,16 +1046,12 @@ def compute_pyramid(
 
         num_arrays = len(self.array_keys())
         if num_arrays == 0:
-            raise ValueError(
-                "No level 0 array exists. Create base array before computing "
-                "pyramid."
-            )
+            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."
+                    "Pyramid structure doesn't exist and levels=None. Specify 'levels' parameter to create pyramid."
                 )
             levels = num_arrays
 
@@ -1084,10 +1075,7 @@ def compute_pyramid(
             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_factors = [int(round(current_scale[i] / previous_scale[i])) for i in range(len(current_scale))]
 
             _downsample_tensorstore(
                 source_ts=previous_ts,
@@ -1328,6 +1316,159 @@ def set_contrast_limits(self, channel_name: str, window: WindowDict):
         self.metadata.omero.channels[channel_index].window = window
         self.dump_meta()
 
+    def to_xarray(self) -> xr.DataArray:
+        """Export full Position data as a labeled xarray.DataArray (tczyx).
+
+        The DataArray is backed by a dask array (lazy, no data loaded
+        until ``.values`` or ``.compute()`` is called).
+
+        Coordinate units follow CF conventions: each coordinate carries
+        its own ``attrs["units"]`` (e.g. ``xa.coords["z"].attrs["units"]
+        == "micrometer"``). ``xa.attrs`` is reserved for value-level
+        metadata (e.g. ``xa.attrs["units"] = "nanometer"`` for ret).
+
+        Returns
+        -------
+        xr.DataArray
+            5D labeled array with coordinates derived from
+            channel names and physical scales/units.
+        """
+        all_channel_names = self.channel_names
+        scale = self.scale
+        translation = self.get_effective_translation(self.metadata.multiscales[0].datasets[0].path)
+
+        data = self.data.dask_array()
+        T, C, Z, Y, X = data.shape
+
+        # Build axis unit lookup from OME metadata
+        axis_units = {}
+        for axis in self.axes:
+            unit = getattr(axis, "unit", None)
+            if unit is not None:
+                axis_units[axis.name.lower()] = unit
+
+        # CF convention: units live in per-coordinate attrs
+        physical = {"t": (T, 0), "z": (Z, 2), "y": (Y, 3), "x": (X, 4)}
+        coords = {"c": ("c", all_channel_names)}
+        for dim, (size, idx) in physical.items():
+            values = np.arange(size) * scale[idx] + translation[idx]
+            attrs = {"units": axis_units[dim]} if dim in axis_units else {}
+            coords[dim] = (dim, values, attrs)
+
+        # Restore any previously saved DataArray attrs from zarr
+        iohub_dict = self.zattrs.get("iohub", {})
+        saved_attrs = dict(iohub_dict.get("xarray_attrs", {}))
+
+        return xr.DataArray(
+            data,
+            dims=("t", "c", "z", "y", "x"),
+            coords=coords,
+            attrs=saved_attrs,
+        )
+
+    def write_xarray(self, data_array: xr.DataArray, image: str = "0") -> None:
+        """Write an xarray.DataArray into this Position.
+
+        Supports writing a subset of channels and/or timepoints.
+        The image array is created on first call; subsequent calls
+        write into the existing array at the correct indices.
+
+        Scales, translations, axis units, and DataArray attrs are
+        set from the first write and updated on subsequent writes.
+
+        Parameters
+        ----------
+        data_array : xr.DataArray
+            5D labeled array with tczyx dimensions.
+            The "c" coordinate must be a subset of this Position's
+            channel names. "t" coordinates are mapped to time indices
+            via the scale and translation.
+        image : str, optional
+            Name of the image array to write to, by default "0".
+        """
+        if tuple(data_array.dims) != ("t", "c", "z", "y", "x"):
+            raise ValueError(f"DataArray dims must be ('t', 'c', 'z', 'y', 'x'), got {data_array.dims}")
+
+        # Validate channels are a subset
+        xa_channels = list(data_array.coords["c"].values)
+        for ch in xa_channels:
+            if ch not in self.channel_names:
+                raise ValueError(f"Channel '{ch}' not in this Position's channel names {self.channel_names}")
+
+        # Infer scales and translations from coordinates
+        def _coord_scale(coord_values):
+            if len(coord_values) < 2:
+                return 1.0
+            return float(coord_values[1] - coord_values[0])
+
+        t_scale = _coord_scale(data_array.coords["t"].values)
+        z_scale = _coord_scale(data_array.coords["z"].values)
+        y_scale = _coord_scale(data_array.coords["y"].values)
+        x_scale = _coord_scale(data_array.coords["x"].values)
+
+        t_trans = float(data_array.coords["t"].values[0])
+        z_trans = float(data_array.coords["z"].values[0])
+        y_trans = float(data_array.coords["y"].values[0])
+        x_trans = float(data_array.coords["x"].values[0])
+
+        # Read coordinate units from per-coordinate attrs (CF convention)
+        def _coord_unit(dim, default):
+            return data_array.coords[dim].attrs.get("units", default)
+
+        self.axes = [
+            TimeAxisMeta(name="T", unit=_coord_unit("t", "second")),
+            ChannelAxisMeta(name="C"),
+            SpaceAxisMeta(name="Z", unit=_coord_unit("z", "micrometer")),
+            SpaceAxisMeta(name="Y", unit=_coord_unit("y", "micrometer")),
+            SpaceAxisMeta(name="X", unit=_coord_unit("x", "micrometer")),
+        ]
+
+        transforms = [
+            TransformationMeta(
+                type="scale",
+                scale=[t_scale, 1.0, z_scale, y_scale, x_scale],
+            )
+        ]
+        if any(v != 0.0 for v in [t_trans, z_trans, y_trans, x_trans]):
+            transforms.append(
+                TransformationMeta(
+                    type="translation",
+                    translation=[t_trans, 0.0, z_trans, y_trans, x_trans],
+                )
+            )
+
+        np_data = data_array.values
+
+        # Create image array if it doesn't exist yet
+        if image not in self:
+            T_full = len(data_array.coords["t"])
+            _, _, Z, Y, X = np_data.shape
+            full_shape = (T_full, len(self.channel_names), Z, Y, X)
+            self.create_zeros(
+                image,
+                shape=full_shape,
+                dtype=np_data.dtype,
+                transform=transforms,
+            )
+
+        # Map channel names to indices
+        c_indices = [self.get_channel_index(ch) for ch in xa_channels]
+
+        # Map T coordinates to indices using scale and translation
+        scale = self.get_effective_scale(image)
+        translation = self.get_effective_translation(image)
+        t_coords = data_array.coords["t"].values
+        t_indices = np.round((t_coords - translation[0]) / scale[0]).astype(int)
+
+        arr = self[image]
+        arr.oindex[t_indices, c_indices] = np_data
+
+        # Persist DataArray attrs to zarr for round-tripping
+        if data_array.attrs:
+            iohub_dict = dict(self.zattrs.get("iohub", {}))
+            iohub_dict["xarray_attrs"] = dict(data_array.attrs)
+            self.zattrs["iohub"] = iohub_dict
+
 
 class TiledPosition(Position):
     """Variant of the NGFF position node

tests/conftest.py

@@ -9,6 +9,11 @@
 from wget import download
 
 
+@pytest.fixture
+def rng():
+    return np.random.default_rng(42)
+
+
 def _download_ndtiff_v3_labeled_positions(test_data: Path) -> None:
     ghfs = fsspec.filesystem(
         "github",