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

ngio/io_pipes/_ops_axes.py

@@ -0,0 +1,146 @@
+from collections.abc import Sequence
+from typing import TypeVar
+
+import dask.array as da
+import numpy as np
+
+from ngio.ome_zarr_meta.ngio_specs._axes import AxesOps
+
+
+def _apply_numpy_axes_ops(
+    array: np.ndarray,
+    squeeze_axes: tuple[int, ...] | None = None,
+    transpose_axes: tuple[int, ...] | None = None,
+    expand_axes: tuple[int, ...] | None = None,
+) -> np.ndarray:
+    """Apply axes operations to a numpy array."""
+    if squeeze_axes is not None:
+        array = np.squeeze(array, axis=squeeze_axes)
+    if transpose_axes is not None:
+        array = np.transpose(array, axes=transpose_axes)
+    if expand_axes is not None:
+        array = np.expand_dims(array, axis=expand_axes)
+    return array
+
+
+def _apply_dask_axes_ops(
+    array: da.Array,
+    squeeze_axes: tuple[int, ...] | None = None,
+    transpose_axes: tuple[int, ...] | None = None,
+    expand_axes: tuple[int, ...] | None = None,
+) -> da.Array:
+    """Apply axes operations to a dask array."""
+    if squeeze_axes is not None:
+        array = da.squeeze(array, axis=squeeze_axes)
+    if transpose_axes is not None:
+        array = da.transpose(array, axes=transpose_axes)
+    if expand_axes is not None:
+        array = da.expand_dims(array, axis=expand_axes)
+    return array
+
+
+T = TypeVar("T")
+
+
+def _apply_sequence_axes_ops(
+    input_: Sequence[T],
+    default: T,
+    squeeze_axes: tuple[int, ...] | None = None,
+    transpose_axes: tuple[int, ...] | None = None,
+    expand_axes: tuple[int, ...] | None = None,
+) -> list[T]:
+    input_list = list(input_)
+    if squeeze_axes is not None:
+        for offset, ax in enumerate(squeeze_axes):
+            input_list.pop(ax - offset)
+
+    if transpose_axes is not None:
+        input_list = [input_list[i] for i in transpose_axes]
+
+    if expand_axes is not None:
+        for ax in expand_axes:
+            input_list.insert(ax, default)
+
+    return input_list
+
+
+def get_as_numpy_axes_ops(
+    array: np.ndarray,
+    axes_ops: AxesOps,
+) -> np.ndarray:
+    """Apply axes operations to a numpy array."""
+    return _apply_numpy_axes_ops(
+        array,
+        squeeze_axes=axes_ops.get_squeeze_op,
+        transpose_axes=axes_ops.get_transpose_op,
+        expand_axes=axes_ops.get_expand_op,
+    )
+
+
+def get_as_dask_axes_ops(
+    array: da.Array,
+    axes_ops: AxesOps,
+) -> da.Array:
+    """Apply axes operations to a dask array."""
+    return _apply_dask_axes_ops(
+        array,
+        squeeze_axes=axes_ops.get_squeeze_op,
+        transpose_axes=axes_ops.get_transpose_op,
+        expand_axes=axes_ops.get_expand_op,
+    )
+
+
+def get_as_sequence_axes_ops(
+    input_: Sequence[T],
+    axes_ops: AxesOps,
+    default: T,
+) -> list[T]:
+    """Apply axes operations to a sequence."""
+    return _apply_sequence_axes_ops(
+        input_,
+        default=default,
+        squeeze_axes=axes_ops.get_squeeze_op,
+        transpose_axes=axes_ops.get_transpose_op,
+        expand_axes=axes_ops.get_expand_op,
+    )
+
+
+def set_as_numpy_axes_ops(
+    array: np.ndarray,
+    axes_ops: AxesOps,
+) -> np.ndarray:
+    """Apply inverse axes operations to a numpy array."""
+    return _apply_numpy_axes_ops(
+        array,
+        squeeze_axes=axes_ops.set_squeeze_op,
+        transpose_axes=axes_ops.set_transpose_op,
+        expand_axes=axes_ops.set_expand_op,
+    )
+
+
+def set_as_dask_axes_ops(
+    array: da.Array,
+    axes_ops: AxesOps,
+) -> da.Array:
+    """Apply inverse axes operations to a dask array."""
+    return _apply_dask_axes_ops(
+        array,
+        squeeze_axes=axes_ops.set_squeeze_op,
+        transpose_axes=axes_ops.set_transpose_op,
+        expand_axes=axes_ops.set_expand_op,
+    )
+
+
+def set_as_sequence_axes_ops(
+    input_: Sequence[T],
+    axes_ops: AxesOps,
+    default: T,
+) -> list[T]:
+    """Apply inverse axes operations to a sequence."""
+    return _apply_sequence_axes_ops(
+        input_,
+        default=default,
+        squeeze_axes=axes_ops.set_squeeze_op,
+        transpose_axes=axes_ops.set_transpose_op,
+        expand_axes=axes_ops.set_expand_op,
+    )

ngio/io_pipes/_ops_slices.py

@@ -0,0 +1,218 @@
+import math
+from typing import TypeAlias, assert_never
+from warnings import warn
+
+import dask.array as da
+import numpy as np
+import zarr
+from pydantic import BaseModel, ConfigDict
+
+from ngio.utils import NgioValueError
+
+SlicingType: TypeAlias = slice | tuple[int, ...] | int
+
+
+def _int_boundary_check(value: int, shape: int) -> int:
+    """Ensure that the integer value is within the boundaries of the array shape."""
+    if value < 0 or value >= shape:
+        raise NgioValueError(
+            f"Invalid index {value}. Index is out of bounds for axis with size {shape}."
+        )
+    return value
+
+
+def _slicing_tuple_boundary_check(
+    slicing_tuple: tuple[SlicingType, ...],
+    array_shape: tuple[int, ...],
+) -> tuple[SlicingType, ...]:
+    """Ensure that the slicing tuple is within the boundaries of the array shape.
+
+    This function normalizes the slicing tuple to ensure that the selection
+    is within the boundaries of the array shape.
+    """
+    if len(slicing_tuple) != len(array_shape):
+        raise NgioValueError(
+            f"Invalid slicing tuple {slicing_tuple}. "
+            f"Length {len(slicing_tuple)} does not match array shape {array_shape}."
+        )
+    out_slicing_tuple = []
+    for sl, sh in zip(slicing_tuple, array_shape, strict=True):
+        if isinstance(sl, slice):
+            start, stop, step = sl.start, sl.stop, sl.step
+            if start is not None:
+                start = math.floor(start)
+                start = max(0, min(start, sh))
+            if stop is not None:
+                stop = math.ceil(stop)
+                stop = max(0, min(stop, sh))
+            out_slicing_tuple.append(slice(start, stop, step))
+        elif isinstance(sl, int):
+            _int_boundary_check(sl, shape=sh)
+            out_slicing_tuple.append(sl)
+        elif isinstance(sl, tuple):
+            [_int_boundary_check(i, shape=sh) for i in sl]
+            out_slicing_tuple.append(sl)
+        else:
+            assert_never(sl)
+
+    return tuple(out_slicing_tuple)
+
+
+class SlicingOps(BaseModel):
+    """Class to hold slicing operations."""
+
+    on_disk_axes: tuple[str, ...]
+    on_disk_shape: tuple[int, ...]
+    slicing_tuple: tuple[SlicingType, ...] | None = None
+    model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True)
+
+    @property
+    def normalized_slicing_tuple(self) -> None | tuple[SlicingType, ...]:
+        """Normalize the slicing tuple to be within the array shape boundaries."""
+        if self.slicing_tuple is not None:
+            return _slicing_tuple_boundary_check(
+                slicing_tuple=self.slicing_tuple,
+                array_shape=self.on_disk_shape,
+            )
+        return None
+
+    def get(self, ax_name: str, normalize: bool = False) -> SlicingType:
+        """Get the slicing tuple."""
+        slicing_tuple = (
+            self.slicing_tuple if not normalize else self.normalized_slicing_tuple
+        )
+        if slicing_tuple is None:
+            return slice(None)
+        if ax_name not in self.on_disk_axes:
+            return slice(None)
+        ax_index = self.on_disk_axes.index(ax_name)
+        return slicing_tuple[ax_index]
+
+
+def _check_tuple_in_slicing_tuple(
+    slicing_tuple: tuple[SlicingType, ...],
+) -> tuple[None, None] | tuple[int, tuple[int, ...]]:
+    """Check if there are any tuple in the slicing tuple.
+
+    The zarr python api only supports int or slices, not tuples.
+    Ngio support a single tuple in the slicing tuple to allow non-contiguous
+    selection (main use case: selecting multiple channels).
+    """
+    # Find if the is any tuple in the slicing tuple
+    # If there is one we need to handle it differently
+    tuple_in_slice = [
+        (i, s) for i, s in enumerate(slicing_tuple) if isinstance(s, tuple)
+    ]
+    if not tuple_in_slice:
+        # No tuple in the slicing tuple
+        return None, None
+
+    if len(tuple_in_slice) > 1:
+        raise NotImplementedError(
+            "Slicing with multiple non-contiguous tuples/lists "
+            "is not supported yet in Ngio. Use directly the "
+            "zarr.Array api to get the correct array slice."
+        )
+    # Complex case, we have exactly one tuple in the slicing tuple
+    ax, first_tuple = tuple_in_slice[0]
+    if len(first_tuple) > 100:
+        warn(
+            "Performance warning: "
+            "Non-contiguous slicing with a tuple/list with more than 100 elements is "
+            "not natively supported by zarr. This is implemented by Ngio by performing "
+            "multiple reads and stacking the result.",
+            stacklevel=2,
+        )
+    return ax, first_tuple
+
+
+def get_slice_as_numpy(zarr_array: zarr.Array, slicing_ops: SlicingOps) -> np.ndarray:
+    slicing_tuple = slicing_ops.normalized_slicing_tuple
+    if slicing_tuple is None:
+        # Base case, no slicing, return the full array
+        return zarr_array[...]
+
+    # Find if the is any tuple in the slicing tuple
+    # If there is one we need to handle it differently
+    ax, first_tuple = _check_tuple_in_slicing_tuple(slicing_tuple)
+    if ax is None:
+        # Simple case, no tuple in the slicing tuple
+        return zarr_array[slicing_tuple]
+
+    assert first_tuple is not None
+    slices = [
+        zarr_array[(*slicing_tuple[:ax], idx, *slicing_tuple[ax + 1 :])]
+        for idx in first_tuple
+    ]
+    out_array = np.stack(slices, axis=ax)
+    return out_array
+
+
+def get_slice_as_dask(zarr_array: zarr.Array, slicing_ops: SlicingOps) -> da.Array:
+    da_array = da.from_zarr(zarr_array)
+    slicing_tuple = slicing_ops.normalized_slicing_tuple
+    if slicing_tuple is None:
+        # Base case, no slicing, return the full array
+        return da_array[...]
+
+    # Find if the is any tuple in the slicing tuple
+    # If there is one we need to handle it differently
+    ax, first_tuple = _check_tuple_in_slicing_tuple(slicing_tuple)
+    if ax is None:
+        # Base case, no tuple in the slicing tuple
+        return da_array[slicing_tuple]
+
+    assert first_tuple is not None
+    slices = [
+        da_array[(*slicing_tuple[:ax], idx, *slicing_tuple[ax + 1 :])]
+        for idx in first_tuple
+    ]
+    out_array = da.stack(slices, axis=ax)
+    return out_array
+
+
+def set_slice_as_numpy(
+    zarr_array: zarr.Array,
+    patch: np.ndarray,
+    slicing_ops: SlicingOps,
+) -> None:
+    slice_tuple = slicing_ops.normalized_slicing_tuple
+    if slice_tuple is None:
+        # Base case, no slicing, write the full array
+        zarr_array[...] = patch
+        return
+
+    ax, first_tuple = _check_tuple_in_slicing_tuple(slice_tuple)
+    if ax is None:
+        # Base case, no tuple in the slicing tuple
+        zarr_array[slice_tuple] = patch
+        return
+
+    # Complex case, we have exactly one tuple in the slicing tuple
+    assert first_tuple is not None
+    for i, idx in enumerate(first_tuple):
+        _sub_slice = (*slice_tuple[:ax], idx, *slice_tuple[ax + 1 :])
+        zarr_array[_sub_slice] = np.take(patch, indices=i, axis=ax)
+
+
+def set_slice_as_dask(
+    zarr_array: zarr.Array, patch: da.Array, slicing_ops: SlicingOps
+) -> None:
+    slice_tuple = slicing_ops.normalized_slicing_tuple
+    if slice_tuple is None:
+        # Base case, no slicing, write the full array
+        da.to_zarr(arr=patch, url=zarr_array)
+        return
+    ax, first_tuple = _check_tuple_in_slicing_tuple(slice_tuple)
+    if ax is None:
+        # Base case, no tuple in the slicing tuple
+        da.to_zarr(arr=patch, url=zarr_array, region=slice_tuple)
+        return
+
+    # Complex case, we have exactly one tuple in the slicing tuple
+    assert first_tuple is not None
+    for i, idx in enumerate(first_tuple):
+        _sub_slice = (*slice_tuple[:ax], slice(idx, idx + 1), *slice_tuple[ax + 1 :])
+        sub_patch = da.take(patch, indices=i, axis=ax)
+        sub_patch = da.expand_dims(sub_patch, axis=ax)
+        da.to_zarr(arr=sub_patch, url=zarr_array, region=_sub_slice)

ngio/io_pipes/_ops_transforms.py

@@ -0,0 +1,104 @@
+from collections.abc import Sequence
+from typing import Protocol
+
+import dask.array as da
+import numpy as np
+
+from ngio.io_pipes._ops_slices import SlicingOps
+from ngio.ome_zarr_meta.ngio_specs._axes import AxesOps
+
+
+class TransformProtocol(Protocol):
+    """Protocol for a generic transform."""
+
+    def get_as_numpy_transform(
+        self, array: np.ndarray, slicing_ops: SlicingOps, axes_ops: AxesOps
+    ) -> np.ndarray:
+        """A transformation to be applied after loading a numpy array."""
+        ...
+
+    def get_as_dask_transform(
+        self, array: da.Array, slicing_ops: SlicingOps, axes_ops: AxesOps
+    ) -> da.Array:
+        """A transformation to be applied after loading a dask array."""
+        ...
+
+    def set_as_numpy_transform(
+        self, array: np.ndarray, slicing_ops: SlicingOps, axes_ops: AxesOps
+    ) -> np.ndarray:
+        """A transformation to be applied before writing a numpy array."""
+        ...
+
+    def set_as_dask_transform(
+        self, array: da.Array, slicing_ops: SlicingOps, axes_ops: AxesOps
+    ) -> da.Array:
+        """A transformation to be applied before writing a dask array."""
+        ...
+
+
+def get_as_numpy_transform(
+    array: np.ndarray,
+    slicing_ops: SlicingOps,
+    axes_ops: AxesOps,
+    transforms: Sequence[TransformProtocol] | None = None,
+) -> np.ndarray:
+    """Apply a numpy transform to an array."""
+    if transforms is None:
+        return array
+
+    for transform in transforms:
+        array = transform.get_as_numpy_transform(
+            array, slicing_ops=slicing_ops, axes_ops=axes_ops
+        )
+    return array
+
+
+def get_as_dask_transform(
+    array: da.Array,
+    slicing_ops: SlicingOps,
+    axes_ops: AxesOps,
+    transforms: Sequence[TransformProtocol] | None = None,
+) -> da.Array:
+    """Apply a dask transform to an array."""
+    if transforms is None:
+        return array
+
+    for transform in transforms:
+        array = transform.get_as_dask_transform(
+            array, slicing_ops=slicing_ops, axes_ops=axes_ops
+        )
+    return array
+
+
+def set_as_numpy_transform(
+    array: np.ndarray,
+    slicing_ops: SlicingOps,
+    axes_ops: AxesOps,
+    transforms: Sequence[TransformProtocol] | None = None,
+) -> np.ndarray:
+    """Apply inverse numpy transforms to an array."""
+    if transforms is None:
+        return array
+
+    for transform in transforms:
+        array = transform.set_as_numpy_transform(
+            array, slicing_ops=slicing_ops, axes_ops=axes_ops
+        )
+    return array
+
+
+def set_as_dask_transform(
+    array: da.Array,
+    slicing_ops: SlicingOps,
+    axes_ops: AxesOps,
+    transforms: Sequence[TransformProtocol] | None = None,
+) -> da.Array:
+    """Apply inverse dask transforms to an array."""
+    if transforms is None:
+        return array
+
+    for transform in transforms:
+        array = transform.set_as_dask_transform(
+            array, slicing_ops=slicing_ops, axes_ops=axes_ops
+        )
+    return array

ngio/io_pipes/_zoom_transform.py

@@ -0,0 +1,175 @@
+import math
+from collections.abc import Sequence
+
+import dask.array as da
+import numpy as np
+
+from ngio.common._dimensions import Dimensions
+from ngio.common._zoom import (
+    InterpolationOrder,
+    dask_zoom,
+    numpy_zoom,
+)
+from ngio.io_pipes._ops_slices import SlicingOps
+from ngio.ome_zarr_meta import AxesOps
+
+
+class BaseZoomTransform:
+    def __init__(
+        self,
+        input_dimensions: Dimensions,
+        target_dimensions: Dimensions,
+        order: InterpolationOrder = "nearest",
+    ) -> None:
+        self._input_dimensions = input_dimensions
+        self._target_dimensions = target_dimensions
+        self._input_pixel_size = input_dimensions.pixel_size
+        self._target_pixel_size = target_dimensions.pixel_size
+        self._order: InterpolationOrder = order
+
+    def _normalize_shape(
+        self, slice_: slice | int | tuple, scale: float, max_dim: int
+    ) -> int:
+        if isinstance(slice_, slice):
+            _start = slice_.start or 0
+            _start_int = math.floor(_start * scale)
+            if slice_.stop is not None:
+                _stop = slice_.stop * scale
+                _stop = min(_stop, max_dim)
+            else:
+                _stop = max_dim
+            _stop_int = math.ceil(_stop)
+            target_shape = _stop_int - _start_int
+
+        elif isinstance(slice_, int):
+            target_shape = 1
+        elif isinstance(slice_, tuple):
+            target_shape = len(slice_) * scale
+        else:
+            raise ValueError(f"Unsupported slice type: {type(slice_)}")
+        return math.ceil(target_shape)
+
+    def _compute_zoom_shape(
+        self,
+        array_shape: Sequence[int],
+        axes_ops: AxesOps,
+        slicing_ops: SlicingOps,
+    ) -> tuple[int, ...]:
+        assert len(array_shape) == len(axes_ops.in_memory_axes)
+
+        target_shape = []
+        for shape, ax_name in zip(array_shape, axes_ops.in_memory_axes, strict=True):
+            ax_type = self._input_dimensions.axes_handler.get_axis(ax_name)
+            if ax_type is not None and ax_type.axis_type == "channel":
+                # Do not scale channel axis
+                target_shape.append(shape)
+                continue
+            t_dim = self._target_dimensions.get(ax_name, default=1)
+            in_pix = self._input_pixel_size.get(ax_name, default=1.0)
+            t_pix = self._target_pixel_size.get(ax_name, default=1.0)
+            slice_ = slicing_ops.get(ax_name, normalize=False)
+            scale = in_pix / t_pix
+            _target_shape = self._normalize_shape(
+                slice_=slice_, scale=scale, max_dim=t_dim
+            )
+            target_shape.append(_target_shape)
+        return tuple(target_shape)
+
+    def _compute_inverse_zoom_shape(
+        self,
+        array_shape: Sequence[int],
+        axes_ops: AxesOps,
+        slicing_ops: SlicingOps,
+    ) -> tuple[int, ...]:
+        assert len(array_shape) == len(axes_ops.in_memory_axes)
+
+        target_shape = []
+        for shape, ax_name in zip(array_shape, axes_ops.in_memory_axes, strict=True):
+            ax_type = self._input_dimensions.axes_handler.get_axis(ax_name)
+            if ax_type is not None and ax_type.axis_type == "channel":
+                # Do not scale channel axis
+                target_shape.append(shape)
+                continue
+            in_dim = self._input_dimensions.get(ax_name, default=1)
+            slice_ = slicing_ops.get(ax_name=ax_name, normalize=True)
+            target_shape.append(
+                self._normalize_shape(slice_=slice_, scale=1, max_dim=in_dim)
+            )
+
+        # Since we are basing the rescaling on the slice, we need to ensure
+        # that the input image we got is roughly the right size.
+        # This is a safeguard against user errors.
+        expected_shape = self._compute_zoom_shape(
+            array_shape=target_shape, axes_ops=axes_ops, slicing_ops=slicing_ops
+        )
+        if any(
+            abs(es - s) > 1 for es, s in zip(expected_shape, array_shape, strict=True)
+        ):
+            raise ValueError(
+                f"Input array shape {array_shape} is not compatible with the expected "
+                f"shape {expected_shape} based on the zoom transform.\n"
+            )
+        return tuple(target_shape)
+
+    def _numpy_zoom(
+        self, array: np.ndarray, target_shape: tuple[int, ...]
+    ) -> np.ndarray:
+        if array.shape == target_shape:
+            return array
+        return numpy_zoom(
+            source_array=array, target_shape=target_shape, order=self._order
+        )
+
+    def _dask_zoom(
+        self,
+        array: da.Array,
+        array_shape: tuple[int, ...],
+        target_shape: tuple[int, ...],
+    ) -> da.Array:
+        if array_shape == target_shape:
+            return array
+        return dask_zoom(
+            source_array=array, target_shape=target_shape, order=self._order
+        )
+
+    def get_as_numpy_transform(
+        self, array: np.ndarray, slicing_ops: SlicingOps, axes_ops: AxesOps
+    ) -> np.ndarray:
+        """Apply the scaling transformation to a numpy array."""
+        target_shape = self._compute_zoom_shape(
+            array_shape=array.shape, axes_ops=axes_ops, slicing_ops=slicing_ops
+        )
+        return self._numpy_zoom(array=array, target_shape=target_shape)
+
+    def get_as_dask_transform(
+        self, array: da.Array, slicing_ops: SlicingOps, axes_ops: AxesOps
+    ) -> da.Array:
+        """Apply the scaling transformation to a dask array."""
+        array_shape = tuple(int(s) for s in array.shape)
+        target_shape = self._compute_zoom_shape(
+            array_shape=array_shape, axes_ops=axes_ops, slicing_ops=slicing_ops
+        )
+        return self._dask_zoom(
+            array=array, array_shape=array_shape, target_shape=target_shape
+        )
+
+    def set_as_numpy_transform(
+        self, array: np.ndarray, slicing_ops: SlicingOps, axes_ops: AxesOps
+    ) -> np.ndarray:
+        """Apply the inverse scaling transformation to a numpy array."""
+        target_shape = self._compute_inverse_zoom_shape(
+            array_shape=array.shape, axes_ops=axes_ops, slicing_ops=slicing_ops
+        )
+        return self._numpy_zoom(array=array, target_shape=target_shape)
+
+    def set_as_dask_transform(
+        self, array: da.Array, slicing_ops: SlicingOps, axes_ops: AxesOps
+    ) -> da.Array:
+        """Apply the inverse scaling transformation to a dask array."""
+        array_shape = tuple(int(s) for s in array.shape)
+        target_shape = self._compute_inverse_zoom_shape(
+            array_shape=array_shape, axes_ops=axes_ops, slicing_ops=slicing_ops
+        )
+        return self._dask_zoom(
+            array=array, array_shape=array_shape, target_shape=target_shape
+        )

ngio/ome_zarr_meta/__init__.py

@@ -13,7 +13,8 @@ from ngio.ome_zarr_meta._meta_handlers import (
     get_well_meta_handler,
 )
 from ngio.ome_zarr_meta.ngio_specs import (
-    AxesMapper,
+    AxesHandler,
+    AxesOps,
     Dataset,
     ImageInWellPath,
     NgffVersions,
@@ -22,11 +23,13 @@ from ngio.ome_zarr_meta.ngio_specs import (
     NgioPlateMeta,
     NgioWellMeta,
     PixelSize,
+    build_canonical_axes_handler,
     path_in_well_validation,
 )
 
 __all__ = [
-    "AxesMapper",
+    "AxesHandler",
+    "AxesOps",
     "Dataset",
     "ImageInWellPath",
     "ImageMetaHandler",
@@ -40,6 +43,7 @@ __all__ = [
     "NgioPlateMeta",
     "NgioWellMeta",
     "PixelSize",
+    "build_canonical_axes_handler",
     "find_image_meta_handler",
     "find_label_meta_handler",
     "find_plate_meta_handler",