ngio 0.4.0a2__py3-none-any.whl → 0.4.0a4__py3-none-any.whl

ngio/io_pipes/_io_pipes.py

@@ -0,0 +1,286 @@
+from abc import ABC, abstractmethod
+from collections.abc import Sequence
+from typing import Generic, TypeVar
+
+import numpy as np
+import zarr
+from dask.array import Array as DaskArray
+
+from ngio.common import Dimensions
+from ngio.io_pipes._io_pipes_utils import SlicingInputType, setup_io_pipe
+from ngio.io_pipes._ops_axes import (
+    get_as_dask_axes_ops,
+    get_as_numpy_axes_ops,
+    set_as_dask_axes_ops,
+    set_as_numpy_axes_ops,
+)
+from ngio.io_pipes._ops_slices import (
+    SlicingOps,
+    get_slice_as_dask,
+    get_slice_as_numpy,
+    set_slice_as_dask,
+    set_slice_as_numpy,
+)
+from ngio.io_pipes._ops_transforms import (
+    TransformProtocol,
+    get_as_dask_transform,
+    get_as_numpy_transform,
+    set_as_dask_transform,
+    set_as_numpy_transform,
+)
+from ngio.ome_zarr_meta.ngio_specs._axes import AxesOps
+
+##############################################################
+#
+# "From Disk" Pipes
+#
+##############################################################
+
+ArrayType = TypeVar("ArrayType", np.ndarray, DaskArray)
+
+
+class DataGetter(ABC, Generic[ArrayType]):
+    def __init__(
+        self,
+        zarr_array: zarr.Array,
+        slicing_ops: SlicingOps,
+        axes_ops: AxesOps,
+        transforms: Sequence[TransformProtocol] | None = None,
+    ) -> None:
+        self._zarr_array = zarr_array
+        self._slicing_ops = slicing_ops
+        self._axes_ops = axes_ops
+        self._transforms = transforms
+
+    @property
+    def zarr_array(self) -> zarr.Array:
+        return self._zarr_array
+
+    @property
+    def slicing_ops(self) -> SlicingOps:
+        return self._slicing_ops
+
+    @property
+    def axes_ops(self) -> AxesOps:
+        return self._axes_ops
+
+    @property
+    def in_memory_axes(self) -> tuple[str, ...]:
+        return self._axes_ops.in_memory_axes
+
+    @property
+    def transforms(self) -> Sequence[TransformProtocol] | None:
+        return self._transforms
+
+    def __call__(self) -> ArrayType:
+        return self.get()
+
+    @abstractmethod
+    def get(self) -> ArrayType:
+        pass
+
+
+class DataSetter(ABC, Generic[ArrayType]):
+    def __init__(
+        self,
+        zarr_array: zarr.Array,
+        slicing_ops: SlicingOps,
+        axes_ops: AxesOps,
+        transforms: Sequence[TransformProtocol] | None = None,
+    ) -> None:
+        self._zarr_array = zarr_array
+        self._slicing_ops = slicing_ops
+        self._axes_ops = axes_ops
+        self._transforms = transforms
+
+    @property
+    def zarr_array(self) -> zarr.Array:
+        return self._zarr_array
+
+    @property
+    def slicing_ops(self) -> SlicingOps:
+        return self._slicing_ops
+
+    @property
+    def axes_ops(self) -> AxesOps:
+        return self._axes_ops
+
+    @property
+    def in_memory_axes(self) -> tuple[str, ...]:
+        return self._axes_ops.in_memory_axes
+
+    @property
+    def transforms(self) -> Sequence[TransformProtocol] | None:
+        return self._transforms
+
+    def __call__(self, patch: ArrayType) -> None:
+        return self.set(patch)
+
+    @abstractmethod
+    def set(self, patch: ArrayType) -> None:
+        pass
+
+
+class NumpyGetter(DataGetter[np.ndarray]):
+    def __init__(
+        self,
+        zarr_array: zarr.Array,
+        dimensions: Dimensions,
+        axes_order: Sequence[str] | None = None,
+        transforms: Sequence[TransformProtocol] | None = None,
+        slicing_dict: dict[str, SlicingInputType] | None = None,
+        remove_channel_selection: bool = False,
+    ) -> None:
+        """Build a pipe to get a numpy or dask array from a zarr array."""
+        slicing_ops, axes_ops = setup_io_pipe(
+            dimensions=dimensions,
+            slicing_dict=slicing_dict,
+            axes_order=axes_order,
+            remove_channel_selection=remove_channel_selection,
+        )
+        super().__init__(
+            zarr_array=zarr_array,
+            slicing_ops=slicing_ops,
+            axes_ops=axes_ops,
+            transforms=transforms,
+        )
+
+    def get(self) -> np.ndarray:
+        array = get_slice_as_numpy(self._zarr_array, slicing_ops=self._slicing_ops)
+        array = get_as_numpy_axes_ops(array, axes_ops=self._axes_ops)
+        array = get_as_numpy_transform(
+            array,
+            slicing_ops=self._slicing_ops,
+            axes_ops=self._axes_ops,
+            transforms=self._transforms,
+        )
+        return array
+
+
+class DaskGetter(DataGetter[DaskArray]):
+    def __init__(
+        self,
+        zarr_array: zarr.Array,
+        dimensions: Dimensions,
+        axes_order: Sequence[str] | None = None,
+        transforms: Sequence[TransformProtocol] | None = None,
+        slicing_dict: dict[str, SlicingInputType] | None = None,
+        remove_channel_selection: bool = False,
+    ) -> None:
+        """Build a pipe to get a numpy or dask array from a zarr array."""
+        slicing_ops, axes_ops = setup_io_pipe(
+            dimensions=dimensions,
+            slicing_dict=slicing_dict,
+            axes_order=axes_order,
+            remove_channel_selection=remove_channel_selection,
+        )
+        super().__init__(
+            zarr_array=zarr_array,
+            slicing_ops=slicing_ops,
+            axes_ops=axes_ops,
+            transforms=transforms,
+        )
+
+    def get(self) -> DaskArray:
+        array = get_slice_as_dask(self._zarr_array, slicing_ops=self._slicing_ops)
+        array = get_as_dask_axes_ops(array, axes_ops=self._axes_ops)
+        array = get_as_dask_transform(
+            array,
+            slicing_ops=self._slicing_ops,
+            axes_ops=self._axes_ops,
+            transforms=self._transforms,
+        )
+        return array
+
+
+##############################################################
+#
+# "To Disk" Pipes
+#
+##############################################################
+
+
+class NumpySetter(DataSetter[np.ndarray]):
+    def __init__(
+        self,
+        zarr_array: zarr.Array,
+        dimensions: Dimensions,
+        axes_order: Sequence[str] | None = None,
+        transforms: Sequence[TransformProtocol] | None = None,
+        slicing_dict: dict[str, SlicingInputType] | None = None,
+        remove_channel_selection: bool = False,
+    ) -> None:
+        """Build a pipe to get a numpy or dask array from a zarr array."""
+        slicing_ops, axes_ops = setup_io_pipe(
+            dimensions=dimensions,
+            slicing_dict=slicing_dict,
+            axes_order=axes_order,
+            remove_channel_selection=remove_channel_selection,
+        )
+        super().__init__(
+            zarr_array=zarr_array,
+            slicing_ops=slicing_ops,
+            axes_ops=axes_ops,
+            transforms=transforms,
+        )
+
+    def set(self, patch: np.ndarray) -> None:
+        """Get a numpy array from the zarr array with ops."""
+        patch = set_as_numpy_transform(
+            array=patch,
+            slicing_ops=self._slicing_ops,
+            axes_ops=self._axes_ops,
+            transforms=self._transforms,
+        )
+        patch = set_as_numpy_axes_ops(
+            array=patch,
+            axes_ops=self._axes_ops,
+        )
+        set_slice_as_numpy(
+            zarr_array=self._zarr_array,
+            patch=patch,
+            slicing_ops=self._slicing_ops,
+        )
+
+
+class DaskSetter(DataSetter[DaskArray]):
+    def __init__(
+        self,
+        zarr_array: zarr.Array,
+        dimensions: Dimensions,
+        axes_order: Sequence[str] | None = None,
+        transforms: Sequence[TransformProtocol] | None = None,
+        slicing_dict: dict[str, SlicingInputType] | None = None,
+        remove_channel_selection: bool = False,
+    ) -> None:
+        """Build a pipe to get a numpy or dask array from a zarr array."""
+        slicing_ops, axes_ops = setup_io_pipe(
+            dimensions=dimensions,
+            slicing_dict=slicing_dict,
+            axes_order=axes_order,
+            remove_channel_selection=remove_channel_selection,
+        )
+        super().__init__(
+            zarr_array=zarr_array,
+            slicing_ops=slicing_ops,
+            axes_ops=axes_ops,
+            transforms=transforms,
+        )
+
+    def set(self, patch: DaskArray) -> None:
+        """Get a dask array from the zarr array with ops."""
+        patch = set_as_dask_transform(
+            array=patch,
+            slicing_ops=self._slicing_ops,
+            axes_ops=self._axes_ops,
+            transforms=self._transforms,
+        )
+        patch = set_as_dask_axes_ops(
+            array=patch,
+            axes_ops=self._axes_ops,
+        )
+        set_slice_as_dask(
+            zarr_array=self._zarr_array,
+            patch=patch,
+            slicing_ops=self._slicing_ops,
+        )