ngio 0.4.0a4__py3-none-any.whl → 0.4.1__py3-none-any.whl

ngio/io_pipes/_io_pipes_masked.py

@@ -16,8 +16,8 @@ from ngio.io_pipes._io_pipes import (
     NumpySetter,
 )
 from ngio.io_pipes._io_pipes_roi import roi_to_slicing_dict
-from ngio.io_pipes._io_pipes_utils import SlicingInputType
 from ngio.io_pipes._match_shape import dask_match_shape, numpy_match_shape
+from ngio.io_pipes._ops_slices import SlicingInputType
 from ngio.io_pipes._ops_transforms import TransformProtocol
 from ngio.io_pipes._zoom_transform import BaseZoomTransform
 
@@ -112,6 +112,7 @@ def _setup_numpy_getters(
 class NumpyGetterMasked(DataGetter[np.ndarray]):
     def __init__(
         self,
+        *,
         zarr_array: zarr.Array,
         dimensions: Dimensions,
         label_zarr_array: zarr.Array,
@@ -152,6 +153,7 @@ class NumpyGetterMasked(DataGetter[np.ndarray]):
             slicing_ops=self._data_getter.slicing_ops,
             axes_ops=self._data_getter.axes_ops,
             transforms=self._data_getter.transforms,
+            roi=roi,
         )
 
     @property
@@ -167,8 +169,8 @@ class NumpyGetterMasked(DataGetter[np.ndarray]):
             label_data=label_data,
             label=self.label_id,
             data_shape=data.shape,
-            label_axes=self._label_data_getter.axes_ops.in_memory_axes,
-            data_axes=self._data_getter.axes_ops.in_memory_axes,
+            label_axes=self._label_data_getter.axes_ops.output_axes,
+            data_axes=self._data_getter.axes_ops.output_axes,
             allow_rescaling=self._allow_rescaling,
         )
         if bool_mask.shape != data.shape:
@@ -180,6 +182,7 @@ class NumpyGetterMasked(DataGetter[np.ndarray]):
 class NumpySetterMasked(DataSetter[np.ndarray]):
     def __init__(
         self,
+        *,
         zarr_array: zarr.Array,
         dimensions: Dimensions,
         label_zarr_array: zarr.Array,
@@ -226,6 +229,7 @@ class NumpySetterMasked(DataSetter[np.ndarray]):
             slicing_ops=self._data_setter.slicing_ops,
             axes_ops=self._data_setter.axes_ops,
             transforms=self._data_setter.transforms,
+            roi=roi,
         )
 
     @property
@@ -235,14 +239,13 @@ class NumpySetterMasked(DataSetter[np.ndarray]):
     def set(self, patch: np.ndarray) -> None:
         data = self._data_getter()
         label_data = self._label_data_getter()
-        print(data.shape, label_data.shape)
 
         bool_mask = _numpy_label_to_bool_mask(
             label_data=label_data,
             label=self.label_id,
             data_shape=data.shape,
-            label_axes=self._label_data_getter.axes_ops.in_memory_axes,
-            data_axes=self._data_getter.axes_ops.in_memory_axes,
+            label_axes=self._label_data_getter.axes_ops.output_axes,
+            data_axes=self._data_getter.axes_ops.output_axes,
             allow_rescaling=self._allow_rescaling,
         )
         if bool_mask.shape != data.shape:
@@ -342,6 +345,7 @@ def _setup_dask_getters(
 class DaskGetterMasked(DataGetter[DaskArray]):
     def __init__(
         self,
+        *,
         zarr_array: zarr.Array,
         dimensions: Dimensions,
         label_zarr_array: zarr.Array,
@@ -381,6 +385,7 @@ class DaskGetterMasked(DataGetter[DaskArray]):
             slicing_ops=self._data_getter.slicing_ops,
             axes_ops=self._data_getter.axes_ops,
             transforms=self._data_getter.transforms,
+            roi=roi,
         )
 
     @property
@@ -395,8 +400,8 @@ class DaskGetterMasked(DataGetter[DaskArray]):
             label_data=label_data,
             label=self.label_id,
             data_shape=data_shape,
-            label_axes=self._label_data_getter.axes_ops.in_memory_axes,
-            data_axes=self._data_getter.axes_ops.in_memory_axes,
+            label_axes=self._label_data_getter.axes_ops.output_axes,
+            data_axes=self._data_getter.axes_ops.output_axes,
             allow_rescaling=self._allow_rescaling,
         )
         if bool_mask.shape != data.shape:
@@ -408,6 +413,7 @@ class DaskGetterMasked(DataGetter[DaskArray]):
 class DaskSetterMasked(DataSetter[DaskArray]):
     def __init__(
         self,
+        *,
         zarr_array: zarr.Array,
         dimensions: Dimensions,
         label_zarr_array: zarr.Array,
@@ -456,6 +462,7 @@ class DaskSetterMasked(DataSetter[DaskArray]):
             slicing_ops=self._data_setter.slicing_ops,
             axes_ops=self._data_setter.axes_ops,
             transforms=self._data_setter.transforms,
+            roi=roi,
         )
 
     @property
@@ -471,8 +478,8 @@ class DaskSetterMasked(DataSetter[DaskArray]):
             label_data=label_data,
             label=self.label_id,
             data_shape=data_shape,
-            label_axes=self._label_data_getter.axes_ops.in_memory_axes,
-            data_axes=self._data_getter.axes_ops.in_memory_axes,
+            label_axes=self._label_data_getter.axes_ops.output_axes,
+            data_axes=self._data_getter.axes_ops.output_axes,
             allow_rescaling=self._allow_rescaling,
         )
         if bool_mask.shape != data.shape:

ngio/io_pipes/_io_pipes_roi.py

@@ -10,13 +10,14 @@ from ngio.io_pipes._io_pipes import (
     NumpyGetter,
     NumpySetter,
 )
-from ngio.io_pipes._io_pipes_utils import SlicingInputType
+from ngio.io_pipes._ops_slices import SlicingInputType
 from ngio.io_pipes._ops_transforms import TransformProtocol
 from ngio.ome_zarr_meta.ngio_specs._pixel_size import PixelSize
 from ngio.utils import NgioValueError
 
 
 def roi_to_slicing_dict(
+    *,
     roi: Roi | RoiPixels,
     pixel_size: PixelSize | None = None,
     slicing_dict: dict[str, SlicingInputType] | None = None,
@@ -42,6 +43,7 @@ def roi_to_slicing_dict(
 class NumpyRoiGetter(NumpyGetter):
     def __init__(
         self,
+        *,
         zarr_array: zarr.Array,
         dimensions: Dimensions,
         roi: Roi | RoiPixels,
@@ -62,12 +64,14 @@ class NumpyRoiGetter(NumpyGetter):
             transforms=transforms,
             slicing_dict=input_slice_kwargs,
             remove_channel_selection=remove_channel_selection,
+            roi=roi,
         )
 
 
 class DaskRoiGetter(DaskGetter):
     def __init__(
         self,
+        *,
         zarr_array: zarr.Array,
         dimensions: Dimensions,
         roi: Roi | RoiPixels,
@@ -88,12 +92,14 @@ class DaskRoiGetter(DaskGetter):
             transforms=transforms,
             slicing_dict=input_slice_kwargs,
             remove_channel_selection=remove_channel_selection,
+            roi=roi,
         )
 
 
 class NumpyRoiSetter(NumpySetter):
     def __init__(
         self,
+        *,
         zarr_array: zarr.Array,
         dimensions: Dimensions,
         roi: Roi | RoiPixels,
@@ -114,12 +120,14 @@ class NumpyRoiSetter(NumpySetter):
             transforms=transforms,
             slicing_dict=input_slice_kwargs,
             remove_channel_selection=remove_channel_selection,
+            roi=roi,
         )
 
 
 class DaskRoiSetter(DaskSetter):
     def __init__(
         self,
+        *,
         zarr_array: zarr.Array,
         dimensions: Dimensions,
         roi: Roi | RoiPixels,
@@ -140,4 +148,5 @@ class DaskRoiSetter(DaskSetter):
             transforms=transforms,
             slicing_dict=input_slice_kwargs,
             remove_channel_selection=remove_channel_selection,
+            roi=roi,
         )

ngio/io_pipes/_io_pipes_types.py

@@ -0,0 +1,56 @@
+from collections.abc import Sequence
+from typing import Protocol, TypeVar
+
+import zarr
+
+from ngio.common._roi import Roi, RoiPixels
+from ngio.io_pipes._ops_axes import AxesOps
+from ngio.io_pipes._ops_slices import SlicingOps
+from ngio.io_pipes._ops_transforms import TransformProtocol
+
+GetterDataType = TypeVar("GetterDataType", covariant=True)
+SetterDataType = TypeVar("SetterDataType", contravariant=True)
+
+
+class DataGetterProtocol(Protocol[GetterDataType]):
+    @property
+    def zarr_array(self) -> zarr.Array: ...
+
+    @property
+    def slicing_ops(self) -> SlicingOps: ...
+
+    @property
+    def axes_ops(self) -> AxesOps: ...
+
+    @property
+    def transforms(self) -> Sequence[TransformProtocol] | None: ...
+
+    @property
+    def roi(self) -> Roi | RoiPixels: ...
+
+    def __call__(self) -> GetterDataType:
+        return self.get()
+
+    def get(self) -> GetterDataType: ...
+
+
+class DataSetterProtocol(Protocol[SetterDataType]):
+    @property
+    def zarr_array(self) -> zarr.Array: ...
+
+    @property
+    def slicing_ops(self) -> SlicingOps: ...
+
+    @property
+    def axes_ops(self) -> AxesOps: ...
+
+    @property
+    def transforms(self) -> Sequence[TransformProtocol] | None: ...
+
+    @property
+    def roi(self) -> Roi | RoiPixels: ...
+
+    def __call__(self, patch: SetterDataType) -> None:
+        return self.set(patch)
+
+    def set(self, patch: SetterDataType) -> None: ...

ngio/io_pipes/_ops_axes.py

@@ -3,8 +3,80 @@ from typing import TypeVar
 
 import dask.array as da
 import numpy as np
+from pydantic import BaseModel, ConfigDict
 
-from ngio.ome_zarr_meta.ngio_specs._axes import AxesOps
+from ngio.common._dimensions import Dimensions
+from ngio.utils import NgioValueError
+
+##############################################################
+#
+# "AxesOps" Model
+#
+##############################################################
+
+
+class AxesOps(BaseModel):
+    """Model to represent axes operations.
+
+    This model will be used to transform objects from on disk axes to in memory axes.
+    """
+
+    input_axes: tuple[str, ...]
+    output_axes: tuple[str, ...]
+    transpose_op: tuple[int, ...] | None = None
+    expand_op: tuple[int, ...] | None = None
+    squeeze_op: tuple[int, ...] | None = None
+    model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True)
+
+    @property
+    def is_no_op(self) -> bool:
+        """Check if all operations are no ops."""
+        if (
+            self.transpose_op is None
+            and self.expand_op is None
+            and self.squeeze_op is None
+        ):
+            return True
+        return False
+
+    @property
+    def get_transpose_op(self) -> tuple[int, ...] | None:
+        """Get the transpose axes."""
+        return self.transpose_op
+
+    @property
+    def get_expand_op(self) -> tuple[int, ...] | None:
+        """Get the expand axes."""
+        return self.expand_op
+
+    @property
+    def get_squeeze_op(self) -> tuple[int, ...] | None:
+        """Get the squeeze axes."""
+        return self.squeeze_op
+
+    @property
+    def set_transpose_op(self) -> tuple[int, ...] | None:
+        """Set the transpose axes."""
+        if self.transpose_op is None:
+            return None
+        return tuple(np.argsort(self.transpose_op))
+
+    @property
+    def set_expand_op(self) -> tuple[int, ...] | None:
+        """Set the expand axes."""
+        return self.squeeze_op
+
+    @property
+    def set_squeeze_op(self) -> tuple[int, ...] | None:
+        """Set the squeeze axes."""
+        return self.expand_op
+
+
+##############################################################
+#
+# Axes Operations implementations
+#
+##############################################################
 
 
 def _apply_numpy_axes_ops(
@@ -144,3 +216,129 @@ def set_as_sequence_axes_ops(
         transpose_axes=axes_ops.set_transpose_op,
         expand_axes=axes_ops.set_expand_op,
     )
+
+
+##############################################################
+#
+# Builder functions
+#
+##############################################################
+
+
+def _check_output_axes(axes: Sequence[str]) -> None:
+    """Check that the input axes are valid."""
+    unique_names = set(axes)
+    if len(unique_names) != len(axes):
+        raise NgioValueError(
+            "Duplicate axis names found. Please provide unique names for each axis."
+        )
+    for name in axes:
+        if not isinstance(name, str):
+            raise NgioValueError(
+                f"Invalid axis name '{name}'. Axis names must be strings."
+            )
+
+
+def _build_squeeze_tuple(
+    input_axes: tuple[str, ...], output_axes: tuple[str, ...]
+) -> tuple[tuple[int, ...], tuple[str, ...]]:
+    """Build a tuple of axes to squeeze."""
+    axes_to_squeeze = []
+    axes_after_squeeze = []
+    for i, ax in enumerate(input_axes):
+        if ax not in output_axes:
+            axes_to_squeeze.append(i)
+        else:
+            axes_after_squeeze.append(ax)
+    return tuple(axes_to_squeeze), tuple(axes_after_squeeze)
+
+
+def _build_transpose_tuple(
+    input_axes: tuple[str, ...], output_axes: tuple[str, ...]
+) -> tuple[tuple[int, ...], tuple[str, ...]]:
+    """Build a tuple of axes to transpose."""
+    transposition_order = []
+    axes_names_after_transpose = []
+    for ax in output_axes:
+        if ax in input_axes:
+            transposition_order.append(input_axes.index(ax))
+            axes_names_after_transpose.append(ax)
+    return tuple(transposition_order), tuple(axes_names_after_transpose)
+
+
+def _build_expand_tuple(
+    input_axes: tuple[str, ...], output_axes: tuple[str, ...]
+) -> tuple[int, ...]:
+    """Build a tuple of axes to expand."""
+    axes_to_expand = []
+    for i, ax in enumerate(output_axes):
+        if ax not in input_axes:
+            axes_to_expand.append(i)
+    return tuple(axes_to_expand)
+
+
+def _build_axes_ops(
+    input_axes: tuple[str, ...], output_axes: tuple[str, ...]
+) -> AxesOps:
+    """Change the order of the axes."""
+    # Validate the names
+    _check_output_axes(output_axes)
+    # Step 1: Check find squeeze axes
+    axes_to_squeeze, input_axes = _build_squeeze_tuple(input_axes, output_axes)
+    # Step 2: Find the transposition order
+    transposition_order, input_axes = _build_transpose_tuple(input_axes, output_axes)
+    # Step 3: Find axes to expand
+    axes_to_expand = _build_expand_tuple(input_axes, output_axes)
+
+    # If the operations are empty, make them None
+    if len(axes_to_squeeze) == 0:
+        axes_to_squeeze = None
+
+    if np.allclose(transposition_order, np.arange(len(transposition_order))):
+        # If the transposition order is the identity, we don't need to transpose
+        transposition_order = None
+    if len(axes_to_expand) == 0:
+        axes_to_expand = None
+
+    return AxesOps(
+        input_axes=input_axes,
+        output_axes=output_axes,
+        transpose_op=transposition_order,
+        expand_op=axes_to_expand,
+        squeeze_op=axes_to_squeeze,
+    )
+
+
+def _normalize_axes_order(
+    dimensions: Dimensions,
+    axes_order: Sequence[str],
+) -> tuple[str, ...]:
+    """Convert axes order to the on-disk axes names.
+
+    In this way there is not unambiguity in the axes order.
+    """
+    new_axes_order = []
+    for axis_name in axes_order:
+        axis = dimensions.axes_handler.get_axis(axis_name)
+        if axis is None:
+            new_axes_order.append(axis_name)
+        else:
+            new_axes_order.append(axis.name)
+    return tuple(new_axes_order)
+
+
+def build_axes_ops(
+    *,
+    dimensions: Dimensions,
+    input_axes: tuple[str, ...],
+    axes_order: Sequence[str] | None,
+) -> AxesOps:
+    if axes_order is None:
+        return AxesOps(
+            input_axes=input_axes,
+            output_axes=input_axes,
+        )
+    output_axes = _normalize_axes_order(dimensions=dimensions, axes_order=axes_order)
+
+    axes_ops = _build_axes_ops(input_axes=input_axes, output_axes=output_axes)
+    return axes_ops