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

ngio/io_pipes/_io_pipes_utils.py

@@ -0,0 +1,299 @@
+from collections.abc import Mapping, Sequence
+from typing import TypeAlias
+
+from ngio.common._dimensions import Dimensions
+from ngio.io_pipes._ops_slices import SlicingOps, SlicingType
+from ngio.ome_zarr_meta.ngio_specs import Axis
+from ngio.ome_zarr_meta.ngio_specs._axes import AxesOps
+from ngio.utils import NgioValueError
+
+SlicingInputType: TypeAlias = slice | Sequence[int] | int | None
+
+
+def _try_to_slice(value: Sequence[int]) -> slice | tuple[int, ...]:
+    """Try to convert a list of integers into a slice if they are contiguous.
+
+    - If the input is empty, return an empty tuple.
+    - If the input is sorted, and contains contiguous integers,
+      return a slice from the minimum to the maximum integer.
+    - Otherwise, return the input as a tuple.
+
+    This is useful for optimizing array slicing operations
+    by allowing the use of slices when possible, which can be more efficient.
+    """
+    if not value:
+        raise NgioValueError("Ngio does not support empty sequences as slice input.")
+
+    if not all(isinstance(i, int) for i in value):
+        _value = []
+        for i in value:
+            try:
+                _value.append(int(i))
+            except Exception as e:
+                raise NgioValueError(
+                    f"Invalid value {i} of type {type(i)} in sequence {value}"
+                ) from e
+        value = _value
+    # If the input is not sorted, return it as a tuple
+    max_input = max(value)
+    min_input = min(value)
+    assert min_input >= 0, "Input must contain non-negative integers"
+
+    if sorted(value) == list(range(min_input, max_input + 1)):
+        return slice(min_input, max_input + 1)
+
+    return tuple(value)
+
+
+def _remove_channel_slicing(
+    slicing_dict: dict[str, SlicingInputType],
+    dimensions: Dimensions,
+) -> dict[str, SlicingInputType]:
+    """This utility function removes the channel selection from the slice kwargs.
+
+    if ignore_channel_selection is True, it will remove the channel selection
+    regardless of the dimensions. If the ignore_channel_selection is False
+    it will fail.
+    """
+    if dimensions.is_multi_channels:
+        return slicing_dict
+
+    if "c" in slicing_dict:
+        slicing_dict.pop("c", None)
+    return slicing_dict
+
+
+def _check_slicing_virtual_axes(slice_: SlicingInputType) -> bool:
+    """Check if the slice_ is compatible with virtual axes.
+
+    Virtual axes are axes that are not present in the actual data,
+    such as time or channel axes in some datasets.
+    So the only valid slices for virtual axes are:
+    - None: means all data along the axis
+    - 0: means the first element along the axis
+    - slice([0, None], [1, None])
+    """
+    if slice_ is None or slice_ == 0:
+        return True
+    if isinstance(slice_, slice):
+        if slice_.start is None and slice_.stop is None:
+            return True
+        if slice_.start == 0 and slice_.stop is None:
+            return True
+        if slice_.start is None and slice_.stop == 0:
+            return True
+        if slice_.start == 0 and slice_.stop == 1:
+            return True
+    if isinstance(slice_, Sequence):
+        if len(slice_) == 1 and slice_[0] == 0:
+            return True
+    return False
+
+
+def _clean_slicing_dict(
+    dimensions: Dimensions,
+    slicing_dict: Mapping[str, SlicingInputType],
+    remove_channel_selection: bool = False,
+) -> dict[str, SlicingInputType]:
+    """Clean the slicing dict.
+
+    This function will:
+        - Validate that the axes in the slicing_dict are present in the dimensions.
+        - Make sure that the slicing_dict uses the on-disk axis names.
+        - Check for duplicate axis names in the slicing_dict.
+        - Clean up channel selection if the dimensions
+    """
+    clean_slicing_dict: dict[str, SlicingInputType] = {}
+    for axis_name, slice_ in slicing_dict.items():
+        axis = dimensions.axes_handler.get_axis(axis_name)
+        if axis is None:
+            # Virtual axes should be allowed to be selected
+            # Common use case is still allowing channel_selection
+            # When the zarr has not channel axis.
+            if not _check_slicing_virtual_axes(slice_):
+                raise NgioValueError(
+                    f"Invalid axis selection:{axis_name}={slice_}. "
+                    f"Not found on the on-disk axes {dimensions.axes}."
+                )
+            # Virtual axes can be safely ignored
+            continue
+        if axis.name in clean_slicing_dict:
+            raise NgioValueError(
+                f"Duplicate axis {axis.name} in slice kwargs. "
+                "Please provide unique axis names."
+            )
+        clean_slicing_dict[axis.name] = slice_
+
+    if remove_channel_selection:
+        clean_slicing_dict = _remove_channel_slicing(
+            slicing_dict=clean_slicing_dict, dimensions=dimensions
+        )
+    return clean_slicing_dict
+
+
+def _normalize_axes_order(
+    dimensions: Dimensions,
+    axes_order: Sequence[str],
+) -> list[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 new_axes_order
+
+
+def _normalize_slicing_tuple(
+    axis: Axis,
+    slicing_dict: dict[str, SlicingInputType],
+    no_axes_ops: bool,
+    axes_order: list[str],
+) -> tuple[SlicingType, str | None]:
+    """Normalize the slicing dict to tuple.
+
+    Since the slicing dict can contain different types of values
+    We need to normalize them to more predictable types.
+    The output types are:
+    - slice
+    - int
+    - tuple of int (for non-contiguous selection)
+    """
+    axis_name = axis.name
+    if axis_name not in slicing_dict:
+        # If no slice is provided for the axis, use a full slice
+        return slice(None), None
+
+    value = slicing_dict[axis_name]
+    if value is None:
+        return slice(None), None
+
+    if isinstance(value, slice):
+        return value, None
+    elif isinstance(value, int):
+        # If axes ops are requested, we need to preserve the dimension
+        # When we slice because the axes ops will be applied later
+        # If no axes ops are requested, we can safely keep the integer
+        # which will remove the dimension
+        if (not no_axes_ops) or (axis_name in axes_order):
+            # Axes ops require all dimensions to be preserved
+            value = slice(value, value + 1)
+            return value, None
+        return value, axis_name
+    elif isinstance(value, Sequence):
+        # If a contiguous sequence of integers is provided,
+        # convert it to a slice for efficiency
+        # Alternatively, it will be converted to a tuple of ints
+        return _try_to_slice(value), None
+
+    raise NgioValueError(
+        f"Invalid slice definition {value} of type {type(value)}. "
+        "Allowed types are: int, slice, sequence of int or None."
+    )
+
+
+def _build_slicing_tuple(
+    *,
+    dimensions: Dimensions,
+    slicing_dict: dict[str, SlicingInputType],
+    axes_order: list[str] | None = None,
+    no_axes_ops: bool = False,
+    remove_channel_selection: bool = False,
+) -> tuple[tuple[SlicingType, ...] | None, list[str]]:
+    """Assemble slices to be used to query the array."""
+    if len(slicing_dict) == 0:
+        # Skip unnecessary computation if no slicing is requested
+        return None, []
+    _axes_order = (
+        _normalize_axes_order(dimensions=dimensions, axes_order=axes_order)
+        if axes_order is not None
+        else []
+    )
+    _slicing_dict = _clean_slicing_dict(
+        dimensions=dimensions,
+        slicing_dict=slicing_dict,
+        remove_channel_selection=remove_channel_selection,
+    )
+
+    slicing_tuple = []
+    axes_to_remove = []
+    for axis in dimensions.axes_handler.axes:
+        sl, ax_to_remove = _normalize_slicing_tuple(
+            axis=axis,
+            slicing_dict=_slicing_dict,
+            no_axes_ops=no_axes_ops,
+            axes_order=_axes_order,
+        )
+        slicing_tuple.append(sl)
+        if ax_to_remove is not None:
+            axes_to_remove.append(ax_to_remove)
+    slicing_tuple = tuple(slicing_tuple)
+    # Slicing tuple can have only one element of type tuple
+    # If multiple tuple are present it will lead to errors
+    # when querying the array
+    if sum(isinstance(s, tuple) for s in slicing_tuple) > 1:
+        raise NgioValueError(
+            f"Invalid slicing tuple {slicing_tuple}. Ngio does not support "
+            "multiple non-contiguous selections (tuples) in the slicing tuple. "
+            "Please use slices or single integer selections instead."
+        )
+    return slicing_tuple, axes_to_remove
+
+
+def _build_axes_ops(
+    *,
+    axes_order: Sequence[str] | None,
+    dimensions: Dimensions,
+) -> tuple[list[str] | None, AxesOps]:
+    if axes_order is None:
+        return None, AxesOps(
+            on_disk_axes=dimensions.axes_handler.axes_names,
+            in_memory_axes=dimensions.axes_handler.axes_names,
+        )
+
+    axes_order = _normalize_axes_order(dimensions=dimensions, axes_order=axes_order)
+    axes_ops = dimensions.axes_handler.get_axes_ops(axes_order)
+    return axes_order, axes_ops
+
+
+def setup_io_pipe(
+    *,
+    dimensions: Dimensions,
+    slicing_dict: dict[str, SlicingInputType] | None = None,
+    axes_order: Sequence[str] | None = None,
+    remove_channel_selection: bool = False,
+) -> tuple[SlicingOps, AxesOps]:
+    """Setup the slicing tuple and axes ops for an IO pipe."""
+    slicing_dict = slicing_dict or {}
+    axes_order, axes_ops = _build_axes_ops(
+        axes_order=axes_order,
+        dimensions=dimensions,
+    )
+
+    slicing_tuple, axes_to_remove = _build_slicing_tuple(
+        dimensions=dimensions,
+        slicing_dict=slicing_dict,
+        axes_order=axes_order,
+        no_axes_ops=axes_ops.is_no_op,
+        remove_channel_selection=remove_channel_selection,
+    )
+
+    if axes_to_remove:
+        in_memory_axes = tuple(
+            ax for ax in axes_ops.in_memory_axes if ax not in axes_to_remove
+        )
+        axes_ops = AxesOps(
+            on_disk_axes=axes_ops.on_disk_axes,
+            in_memory_axes=in_memory_axes,
+        )
+    slicing_ops = SlicingOps(
+        on_disk_axes=dimensions.axes_handler.axes_names,
+        slicing_tuple=slicing_tuple,
+        on_disk_shape=dimensions.shape,
+    )
+    return slicing_ops, axes_ops

ngio/io_pipes/_match_shape.py

@@ -0,0 +1,376 @@
+from collections.abc import Sequence
+from enum import Enum
+
+import dask.array as da
+import numpy as np
+
+from ngio.utils import NgioValueError, ngio_logger
+
+
+class Action(str, Enum):
+    NONE = "none"
+    PAD = "pad"
+    TRIM = "trim"
+    RESCALING = "rescaling"
+
+
+def _compute_pad_widths(
+    array_shape: tuple[int, ...],
+    actions: list[Action],
+    target_shape: tuple[int, ...],
+) -> tuple[tuple[int, int], ...]:
+    pad_def = []
+    for act, s, ts in zip(actions, array_shape, target_shape, strict=True):
+        if act == Action.PAD:
+            total_pad = ts - s
+            before = total_pad // 2
+            after = total_pad - before
+            pad_def.append((before, after))
+        else:
+            pad_def.append((0, 0))
+    ngio_logger.warning(
+        f"Images have a different shape ({array_shape} vs {target_shape}). "
+        f"Resolving by padding: {pad_def}",
+        stacklevel=2,
+    )
+    return tuple(pad_def)
+
+
+def _numpy_pad(
+    array: np.ndarray,
+    actions: list[Action],
+    target_shape: tuple[int, ...],
+    pad_mode: str = "constant",
+    constant_values: int | float = 0,
+) -> np.ndarray:
+    if all(act != Action.PAD for act in actions):
+        return array
+    pad_widths = _compute_pad_widths(array.shape, actions, target_shape)
+    return np.pad(array, pad_widths, mode=pad_mode, constant_values=constant_values)  # type: ignore
+
+
+def _dask_pad(
+    array: da.Array,
+    actions: list[Action],
+    target_shape: tuple[int, ...],
+    pad_mode: str = "constant",
+    constant_values: int | float = 0,
+) -> da.Array:
+    if all(act != Action.PAD for act in actions):
+        return array
+    shape = tuple(int(s) for s in array.shape)
+    pad_widths = _compute_pad_widths(shape, actions, target_shape)
+    return da.pad(array, pad_widths, mode=pad_mode, constant_values=constant_values)
+
+
+def _compute_trim_slices(
+    array_shape: tuple[int, ...],
+    actions: list[Action],
+    target_shape: tuple[int, ...],
+) -> tuple[slice, ...]:
+    slices = []
+    for act, s, ts in zip(actions, array_shape, target_shape, strict=True):
+        if act == Action.TRIM:
+            slices.append(slice(0, ts))
+        else:
+            slices.append(slice(0, s))
+
+    ngio_logger.warning(
+        f"Images have a different shape ({array_shape} vs {target_shape}). "
+        f"Resolving by trimming: {slices}",
+        stacklevel=2,
+    )
+    return tuple(slices)
+
+
+def _numpy_trim(
+    array: np.ndarray, actions: list[Action], target_shape: tuple[int, ...]
+) -> np.ndarray:
+    if all(act != Action.TRIM for act in actions):
+        return array
+    slices = _compute_trim_slices(array.shape, actions, target_shape)
+    return array[tuple(slices)]
+
+
+def _dask_trim(
+    array: da.Array, actions: list[Action], target_shape: tuple[int, ...]
+) -> da.Array:
+    if all(act != Action.TRIM for act in actions):
+        return array
+    shape = tuple(int(s) for s in array.shape)
+    slices = _compute_trim_slices(shape, actions, target_shape)
+    return array[tuple(slices)]
+
+
+def _compute_rescaling_shape(
+    array_shape: tuple[int, ...],
+    actions: list[Action],
+    target_shape: tuple[int, ...],
+) -> tuple[int, ...]:
+    rescaling_shape = []
+    factor = []
+    for act, s, ts in zip(actions, array_shape, target_shape, strict=True):
+        if act == Action.RESCALING:
+            rescaling_shape.append(ts)
+            factor.append(ts / s)
+        else:
+            rescaling_shape.append(s)
+            factor.append(1.0)
+
+    ngio_logger.warning(
+        f"Images have a different shape ({array_shape} vs {target_shape}). "
+        f"Resolving by scaling with factors {factor}.",
+        stacklevel=2,
+    )
+    return tuple(rescaling_shape)
+
+
+def _numpy_rescaling(
+    array: np.ndarray, actions: list[Action], target_shape: tuple[int, ...]
+) -> np.ndarray:
+    if all(act != Action.RESCALING for act in actions):
+        return array
+    from ngio.common._zoom import numpy_zoom
+
+    rescaling_shape = _compute_rescaling_shape(array.shape, actions, target_shape)
+    return numpy_zoom(source_array=array, target_shape=rescaling_shape, order="nearest")
+
+
+def _dask_rescaling(
+    array: da.Array, actions: list[Action], target_shape: tuple[int, ...]
+) -> da.Array:
+    if all(act != Action.RESCALING for act in actions):
+        return array
+    from ngio.common._zoom import dask_zoom
+
+    shape = tuple(int(s) for s in array.shape)
+    rescaling_shape = _compute_rescaling_shape(shape, actions, target_shape)
+    return dask_zoom(source_array=array, target_shape=rescaling_shape, order="nearest")
+
+
+def _check_axes(array_shape, reference_shape, array_axes, reference_axes):
+    if len(array_shape) != len(array_axes):
+        raise NgioValueError(
+            f"Array shape {array_shape} and reference axes {array_axes} "
+            "must have the same number of dimensions."
+        )
+    if len(reference_shape) != len(reference_axes):
+        raise NgioValueError(
+            f"Reference shape {reference_shape} and reference axes {reference_axes} "
+            "must have the same number of dimensions."
+        )
+
+    # Check if the array axes are a subset of the target axes
+    diff = set(array_axes) - set(reference_axes)
+    if diff:
+        raise NgioValueError(
+            f"Array axes {array_axes} are not a subset "
+            f"of reference axes {reference_axes}"
+        )
+
+    # Array must be smaller or equal in number of dimensions
+    if len(array_axes) > len(reference_axes):
+        raise NgioValueError(
+            f"Array has more dimensions ({len(array_axes)}) "
+            f"than reference ({len(reference_axes)}). "
+            "Cannot match shapes if the array has more dimensions."
+        )
+
+
+def _compute_reshape_and_actions(
+    array_shape: tuple[int, ...],
+    reference_shape: tuple[int, ...],
+    array_axes: list[str],
+    reference_axes: list[str],
+    tolerance: int = 1,
+    allow_rescaling: bool = True,
+) -> tuple[tuple[int, ...], list[Action]]:
+    # Reshape array to match reference shape
+    # And determine actions to be taken
+    # to match the shapes
+    reshape_tuple = []
+    actions = []
+    errors = []
+    left_pointer = 0
+    for ref_ax, ref_shape in zip(reference_axes, reference_shape, strict=True):
+        if ref_ax not in array_axes:
+            reshape_tuple.append(1)
+            actions.append(Action.NONE)
+        elif ref_ax == array_axes[left_pointer]:
+            s2 = array_shape[left_pointer]
+            reshape_tuple.append(s2)
+            left_pointer += 1
+
+            if s2 == ref_shape or s2 == 1:
+                actions.append(Action.NONE)
+            elif s2 < ref_shape:
+                if (ref_shape - s2) <= tolerance:
+                    actions.append(Action.PAD)
+                elif allow_rescaling:
+                    actions.append(Action.RESCALING)
+                else:
+                    errors.append(
+                        f"Cannot pad axis={ref_ax}:{s2}->{ref_shape} "
+                        "because shape difference is outside tolerance "
+                        f"{tolerance}."
+                    )
+            elif s2 > ref_shape:
+                if (s2 - ref_shape) <= tolerance:
+                    actions.append(Action.TRIM)
+                elif allow_rescaling:
+                    actions.append(Action.RESCALING)
+                else:
+                    errors.append(
+                        f"Cannot trim axis={ref_ax}:{s2}->{ref_shape} "
+                        "because shape difference is outside tolerance "
+                        f"{tolerance}."
+                    )
+            else:
+                raise RuntimeError("Unreachable code reached.")
+        else:
+            raise NgioValueError(
+                f"Axes order mismatch {array_axes} -> {reference_axes}. "
+                "Cannot match shapes if the order is different."
+            )
+    if errors:
+        raise NgioValueError(
+            "Array shape cannot be matched to reference shape:\n\n".join(errors)
+        )
+    return tuple(reshape_tuple), actions
+
+
+def numpy_match_shape(
+    array: np.ndarray,
+    reference_shape: tuple[int, ...],
+    array_axes: Sequence[str],
+    reference_axes: Sequence[str],
+    tolerance: int = 1,
+    pad_mode: str = "constant",
+    pad_values: int | float = 0,
+    allow_rescaling: bool = True,
+):
+    """Match the shape of a numpy array to a reference shape.
+
+    This function will reshape, pad, trim and broadcast the input array
+    to match the reference shape. If the shapes cannot be matched within
+    the specified tolerance, an error is raised.
+
+    The reference axes must be a superset of the array axes, and the order
+    of the axes must be the same.
+
+    Args:
+        array (np.ndarray): The input array to be reshaped.
+        reference_shape (tuple[int, ...]): The target shape to match.
+        array_axes (Sequence[str]): The axes names of the input array.
+        reference_axes (Sequence[str]): The axes names of the reference shape.
+        tolerance (int): The maximum number of pixels by which dimensions
+            can differ when matching shapes.
+        allow_broadcast (bool): If True, allow broadcasting new dimensions to
+            match the reference shape. If False, single-dimension axes will
+            be left as is.
+        pad_mode (str): The mode to use for padding. See numpy.pad for options.
+        pad_values (int | float): The constant value to use for padding if
+            pad_mode is 'constant'.
+        allow_rescaling (bool): If True, when the array differs more than the
+            tolerance, it will be rescalingd to the reference shape. If False,
+            an error will be raised.
+    """
+    _check_axes(
+        array_shape=array.shape,
+        reference_shape=reference_shape,
+        array_axes=array_axes,
+        reference_axes=reference_axes,
+    )
+    if array.shape == reference_shape:
+        # Shapes already match
+        return array
+
+    array_axes = list(array_axes)
+    reference_axes = list(reference_axes)
+
+    reshape_tuple, actions = _compute_reshape_and_actions(
+        array_shape=array.shape,
+        reference_shape=reference_shape,
+        array_axes=array_axes,
+        reference_axes=reference_axes,
+        tolerance=tolerance,
+        allow_rescaling=allow_rescaling,
+    )
+    array = array.reshape(reshape_tuple)
+    array = _numpy_rescaling(array=array, actions=actions, target_shape=reference_shape)
+    array = _numpy_pad(
+        array=array,
+        actions=actions,
+        target_shape=reference_shape,
+        pad_mode=pad_mode,
+        constant_values=pad_values,
+    )
+    array = _numpy_trim(array=array, actions=actions, target_shape=reference_shape)
+    return array
+
+
+def dask_match_shape(
+    array: da.Array,
+    reference_shape: tuple[int, ...],
+    array_axes: Sequence[str],
+    reference_axes: Sequence[str],
+    tolerance: int = 1,
+    pad_mode: str = "constant",
+    pad_values: int | float = 0,
+    allow_rescaling: bool = True,
+) -> da.Array:
+    """Match the shape of a dask array to a reference shape.
+
+    This function will reshape, pad, trim and broadcast the input array
+    to match the reference shape. If the shapes cannot be matched within
+    the specified tolerance, an error is raised.
+
+    The reference axes must be a superset of the array axes, and the order
+    of the axes must be the same.
+
+    Args:
+        array (da.Array): The input array to be reshaped.
+        reference_shape (tuple[int, ...]): The target shape to match.
+        array_axes (Sequence[str]): The axes names of the input array.
+        reference_axes (Sequence[str]): The axes names of the reference shape.
+        tolerance (int): The maximum number of pixels by which dimensions
+            can differ when matching shapes.
+        pad_mode (str): The mode to use for padding. See numpy.pad for options.
+        pad_values (int | float): The constant value to use for padding if
+            pad_mode is 'constant'.
+        allow_rescaling (bool): If True, when the array differs more than the
+            tolerance, it will be rescalingd to the reference shape. If False,
+            an error will be raised.
+    """
+    array_shape = tuple(int(s) for s in array.shape)
+    _check_axes(
+        array_shape=array_shape,
+        reference_shape=reference_shape,
+        array_axes=array_axes,
+        reference_axes=reference_axes,
+    )
+    if array_shape == reference_shape:
+        # Shapes already match
+        return array
+    array_axes = list(array_axes)
+    reference_axes = list(reference_axes)
+
+    reshape_tuple, actions = _compute_reshape_and_actions(
+        array_shape=tuple(int(s) for s in array.shape),
+        reference_shape=reference_shape,
+        array_axes=array_axes,
+        reference_axes=reference_axes,
+        tolerance=tolerance,
+        allow_rescaling=allow_rescaling,
+    )
+    array = da.reshape(array, reshape_tuple)
+    array = _dask_rescaling(array=array, actions=actions, target_shape=reference_shape)
+    array = _dask_pad(
+        array=array,
+        actions=actions,
+        target_shape=reference_shape,
+        pad_mode=pad_mode,
+        constant_values=pad_values,
+    )
+    array = _dask_trim(array=array, actions=actions, target_shape=reference_shape)
+    return array