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

ngio/io_pipes/_ops_slices.py

@@ -1,4 +1,5 @@
 import math
+from collections.abc import Mapping, Sequence
 from typing import TypeAlias, assert_never
 from warnings import warn
 
@@ -7,10 +8,20 @@ import numpy as np
 import zarr
 from pydantic import BaseModel, ConfigDict
 
+from ngio.common._dimensions import Dimensions
+from ngio.io_pipes._ops_slices_utils import compute_slice_chunks
+from ngio.ome_zarr_meta.ngio_specs import Axis
 from ngio.utils import NgioValueError
 
+SlicingInputType: TypeAlias = slice | Sequence[int] | int | None
 SlicingType: TypeAlias = slice | tuple[int, ...] | int
 
+##############################################################
+#
+# "SlicingOps" model
+#
+##############################################################
+
 
 def _int_boundary_check(value: int, shape: int) -> int:
     """Ensure that the integer value is within the boundaries of the array shape."""
@@ -63,26 +74,41 @@ class SlicingOps(BaseModel):
 
     on_disk_axes: tuple[str, ...]
     on_disk_shape: tuple[int, ...]
-    slicing_tuple: tuple[SlicingType, ...] | None = None
+    on_disk_chunks: tuple[int, ...]
+    slicing_tuple: tuple[SlicingType, ...]
     model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True)
 
     @property
-    def normalized_slicing_tuple(self) -> None | tuple[SlicingType, ...]:
+    def normalized_slicing_tuple(self) -> 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
+        return _slicing_tuple_boundary_check(
+            slicing_tuple=self.slicing_tuple,
+            array_shape=self.on_disk_shape,
+        )
+
+    @property
+    def slice_axes(self) -> tuple[str, ...]:
+        """The axes after slicing."""
+        in_memory_axes = []
+        for ax, sl in zip(self.on_disk_axes, self.slicing_tuple, strict=True):
+            if isinstance(sl, int):
+                continue
+            in_memory_axes.append(ax)
+        return tuple(in_memory_axes)
+
+    def slice_chunks(self) -> set[tuple[int, ...]]:
+        """The required to read or write the slice."""
+        return compute_slice_chunks(
+            shape=self.on_disk_shape,
+            chunks=self.on_disk_chunks,
+            slicing_tuple=self.normalized_slicing_tuple,
+        )
 
     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)
@@ -126,12 +152,16 @@ def _check_tuple_in_slicing_tuple(
     return ax, first_tuple
 
 
+##############################################################
+#
+# Slicing implementations
+#
+##############################################################
+
+
 def get_slice_as_numpy(zarr_array: zarr.Array, slicing_ops: SlicingOps) -> np.ndarray:
+    """Get a slice of a zarr array as a numpy array."""
     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)
@@ -149,12 +179,9 @@ def get_slice_as_numpy(zarr_array: zarr.Array, slicing_ops: SlicingOps) -> np.nd
 
 
 def get_slice_as_dask(zarr_array: zarr.Array, slicing_ops: SlicingOps) -> da.Array:
+    """Get a slice of a zarr array as a dask 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)
@@ -177,11 +204,6 @@ def set_slice_as_numpy(
     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
@@ -195,17 +217,31 @@ def set_slice_as_numpy(
         zarr_array[_sub_slice] = np.take(patch, indices=i, axis=ax)
 
 
+def handle_int_set_as_dask(
+    patch: da.Array,
+    slicing_tuple: tuple[SlicingType, ...],
+) -> tuple[da.Array, tuple[SlicingType, ...]]:
+    """Handle the case where the slicing tuple contains integers.
+
+    In this case we need to expand the patch array to match the slicing tuple.
+    """
+    new_slicing_tuple = list(slicing_tuple)
+    for i, sl in enumerate(slicing_tuple):
+        if isinstance(sl, int):
+            patch = da.expand_dims(patch, axis=i)
+            new_slicing_tuple[i] = slice(sl, sl + 1)
+    return patch, tuple(new_slicing_tuple)
+
+
 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)
+    patch, slice_tuple = handle_int_set_as_dask(patch, slice_tuple)
     if ax is None:
         # Base case, no tuple in the slicing tuple
+        # assert False
         da.to_zarr(arr=patch, url=zarr_array, region=slice_tuple)
         return
 
@@ -216,3 +252,195 @@ def set_slice_as_dask(
         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)
+
+
+##############################################################
+#
+# Builder functions
+#
+##############################################################
+
+
+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_slicing_tuple(
+    axis: Axis,
+    slicing_dict: dict[str, SlicingInputType],
+) -> SlicingType:
+    """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)
+
+    value = slicing_dict[axis_name]
+    if value is None:
+        return slice(None)
+    if isinstance(value, slice) or isinstance(value, int):
+        return value
+    elif isinstance(value, Sequence):
+        # If a contiguous sequence of integers is provided,
+        # convert it to a slice for simplicity.
+        # Alternatively, it will be converted to a tuple of ints
+        return _try_to_slice(value)
+
+    raise NgioValueError(
+        f"Invalid slice definition {value} of type {type(value)}. "
+        "Allowed types are: int, slice, sequence of int or None."
+    )
+
+
+def build_slicing_ops(
+    *,
+    dimensions: Dimensions,
+    slicing_dict: dict[str, SlicingInputType] | None,
+    remove_channel_selection: bool = False,
+) -> SlicingOps:
+    """Assemble slices to be used to query the array."""
+    slicing_dict = slicing_dict or {}
+    _slicing_dict = _clean_slicing_dict(
+        dimensions=dimensions,
+        slicing_dict=slicing_dict,
+        remove_channel_selection=remove_channel_selection,
+    )
+
+    slicing_tuple = tuple(
+        _normalize_slicing_tuple(
+            axis=axis,
+            slicing_dict=_slicing_dict,
+        )
+        for axis in dimensions.axes_handler.axes
+    )
+    return SlicingOps(
+        on_disk_axes=dimensions.axes_handler.axes_names,
+        on_disk_shape=dimensions.shape,
+        on_disk_chunks=dimensions.chunks,
+        slicing_tuple=slicing_tuple,
+    )

ngio/io_pipes/_ops_slices_utils.py

@@ -0,0 +1,196 @@
+from collections.abc import Iterable, Iterator
+from itertools import product
+from typing import TypeAlias, TypeVar
+
+from ngio.utils import NgioValueError, ngio_logger
+
+T = TypeVar("T")
+
+##############################################################
+#
+# Check slice overlaps
+#
+##############################################################
+
+
+def _pairs_stream(iterable: Iterable[T]) -> Iterator[tuple[T, T]]:
+    # Same as combinations but yields pairs as soon as they are generated
+    seen: list[T] = []
+    for a in iterable:
+        for b in seen:
+            yield b, a
+        seen.append(a)
+
+
+SlicingType: TypeAlias = slice | tuple[int, ...] | int
+
+
+def check_elem_intersection(s1: SlicingType, s2: SlicingType) -> bool:
+    """Compare if two SlicingType elements intersect.
+
+    If they are a slice, check if they overlap.
+    If they are integers, check if they are equal.
+    If they are tuples, check if they have any common elements.
+    """
+    if not isinstance(s1, type(s2)):
+        raise NgioValueError(
+            f"Slices must be of the same type. Got {type(s1)} and {type(s2)}"
+        )
+
+    if isinstance(s1, slice) and isinstance(s2, slice):
+        # Handle slice objects
+        start1, stop1, step1 = s1.start or 0, s1.stop or float("inf"), s1.step or 1
+        start2, stop2, step2 = s2.start or 0, s2.stop or float("inf"), s2.step or 1
+
+        if step1 is not None and step2 != 1:
+            raise NotImplementedError(
+                "Intersection for slices with step != 1 is not implemented"
+            )
+
+        if step2 is not None and step1 != 1:
+            raise NotImplementedError(
+                "Intersection for slices with step != 1 is not implemented"
+            )
+
+        return not (stop1 <= start2 or stop2 <= start1)
+    elif isinstance(s1, int) and isinstance(s2, int):
+        # Handle integer indices
+        return s1 == s2
+    elif isinstance(s1, tuple) and isinstance(s2, tuple):
+        if set(s1) & set(s2):
+            return True
+        return False
+    else:
+        raise TypeError("Unsupported slice type")
+
+
+def check_slicing_tuple_intersection(
+    s1: tuple[SlicingType, ...], s2: tuple[SlicingType, ...]
+) -> bool:
+    """For a tuple of SlicingType, check if all elements intersect."""
+    if len(s1) != len(s2):
+        raise NgioValueError("Slices must have the same length")
+    return all(check_elem_intersection(a, b) for a, b in zip(s1, s2, strict=True))
+
+
+def check_if_regions_overlap(slices: Iterable[tuple[SlicingType, ...]]) -> bool:
+    """Check for overlaps in a list of slicing tuples using brute-force method.
+
+    This is O(n^2) and not efficient for large lists.
+    Returns True if any overlaps are found.
+    """
+    for it, (si, sj) in enumerate(_pairs_stream(slices)):
+        overalap = check_slicing_tuple_intersection(si, sj)
+        if overalap:
+            return True
+
+        if it == 10_000:
+            ngio_logger.warning(
+                "Performance Warning check_for_overlaps is O(n^2) and may be slow for "
+                "large numbers of regions."
+            )
+    return False
+
+
+##############################################################
+#
+# Check chunk overlaps
+#
+##############################################################
+
+
+def _normalize_slice(slc: slice, size: int) -> tuple[int, int]:
+    if slc.step not in (None, 1):
+        raise NgioValueError(f"Only step=1 slices supported, got step={slc.step}")
+    start = 0 if slc.start is None else slc.start
+    stop = size if slc.stop is None else slc.stop
+    if start < 0 or stop < 0:
+        raise NgioValueError("Negative slice bounds are not supported")
+    # clamp to [0, size]
+    start = min(start, size)
+    stop = min(stop, size)
+    if start > stop:
+        # empty selection
+        return (0, 0)
+    return start, stop
+
+
+def _chunk_indices_for_axis(sel: SlicingType, size: int, csize: int) -> list[int]:
+    """From a selection for a single axis, return the list chunk indices touched."""
+    if isinstance(sel, slice):
+        start, stop = _normalize_slice(sel, size)
+        if start >= stop:  # empty
+            return []
+        first = start // csize
+        last = (stop - 1) // csize
+        return list(range(first, last + 1))
+
+    if isinstance(sel, int):
+        if sel < 0 or sel >= size:
+            raise IndexError(f"index {sel} out of bounds for axis of size {size}")
+        return [sel // csize]
+
+    if isinstance(sel, tuple):
+        if not sel:
+            return []
+        chunks_hit = {}
+        for v in sel:
+            if not isinstance(v, int):
+                raise TypeError("Only integers allowed inside tuple selections")
+            if v < 0 or v >= size:
+                raise IndexError(f"index {v} out of bounds for axis of size {size}")
+            chunks_hit[v // csize] = None
+        return sorted(chunks_hit.keys())
+
+    raise TypeError(f"Unsupported index type: {type(sel)!r}")
+
+
+def compute_slice_chunks(
+    shape: tuple[int, ...],
+    chunks: tuple[int, ...],
+    slicing_tuple: tuple[SlicingType, ...],
+) -> set[tuple[int, ...]]:
+    """Compute the set of chunk coordinates touched by `slicing_tuple`.
+
+    Args:
+        shape: overall array shape (s1, s2, ...)
+        chunks: chunk shape (c1, c2, ...)
+        slicing_tuple: tuple of slices, ints, or tuples of ints
+    """
+    if len(slicing_tuple) != len(shape):
+        raise NgioValueError(
+            f"key must have {len(shape)} items, got {len(slicing_tuple)}"
+        )
+
+    per_axis_chunks: list[list[int]] = [
+        _chunk_indices_for_axis(sel, size, csize)
+        for sel, size, csize in zip(slicing_tuple, shape, chunks, strict=True)
+    ]
+
+    # If any axis yields no chunks, the overall selection is empty.
+    if any(len(ax) == 0 for ax in per_axis_chunks):
+        return set()
+
+    return {tuple(idx) for idx in product(*per_axis_chunks)}
+
+
+def check_if_chunks_overlap(
+    slices: Iterable[tuple[SlicingType, ...]],
+    shape: tuple[int, ...],
+    chunks: tuple[int, ...],
+) -> bool:
+    """Check for overlaps in a list of slicing tuples using brute-force method.
+
+    This is O(n^2) and not efficient for large lists.
+    Returns True if any overlaps are found.
+    """
+    slices_chunks = (compute_slice_chunks(shape, chunks, si) for si in slices)
+    for it, (si, sj) in enumerate(_pairs_stream(slices_chunks)):
+        if si & sj:
+            return True
+        if it == 10_000:
+            ngio_logger.warning(
+                "Performance Warning check_for_chunks_overlaps is O(n^2) and may be "
+                "slow for large numbers of regions."
+            )
+    return False

ngio/io_pipes/_ops_transforms.py

@@ -4,8 +4,8 @@ from typing import Protocol
 import dask.array as da
 import numpy as np
 
+from ngio.io_pipes._ops_axes import AxesOps
 from ngio.io_pipes._ops_slices import SlicingOps
-from ngio.ome_zarr_meta.ngio_specs._axes import AxesOps
 
 
 class TransformProtocol(Protocol):

ngio/io_pipes/_zoom_transform.py

@@ -10,8 +10,8 @@ from ngio.common._zoom import (
     dask_zoom,
     numpy_zoom,
 )
+from ngio.io_pipes._ops_axes import AxesOps
 from ngio.io_pipes._ops_slices import SlicingOps
-from ngio.ome_zarr_meta import AxesOps
 
 
 class BaseZoomTransform:
@@ -55,10 +55,10 @@ class BaseZoomTransform:
         axes_ops: AxesOps,
         slicing_ops: SlicingOps,
     ) -> tuple[int, ...]:
-        assert len(array_shape) == len(axes_ops.in_memory_axes)
+        assert len(array_shape) == len(axes_ops.output_axes)
 
         target_shape = []
-        for shape, ax_name in zip(array_shape, axes_ops.in_memory_axes, strict=True):
+        for shape, ax_name in zip(array_shape, axes_ops.output_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
@@ -81,10 +81,10 @@ class BaseZoomTransform:
         axes_ops: AxesOps,
         slicing_ops: SlicingOps,
     ) -> tuple[int, ...]:
-        assert len(array_shape) == len(axes_ops.in_memory_axes)
+        assert len(array_shape) == len(axes_ops.output_axes)
 
         target_shape = []
-        for shape, ax_name in zip(array_shape, axes_ops.in_memory_axes, strict=True):
+        for shape, ax_name in zip(array_shape, axes_ops.output_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

ngio/ome_zarr_meta/__init__.py

@@ -14,7 +14,6 @@ from ngio.ome_zarr_meta._meta_handlers import (
 )
 from ngio.ome_zarr_meta.ngio_specs import (
     AxesHandler,
-    AxesOps,
     Dataset,
     ImageInWellPath,
     NgffVersions,
@@ -29,7 +28,6 @@ from ngio.ome_zarr_meta.ngio_specs import (
 
 __all__ = [
     "AxesHandler",
-    "AxesOps",
     "Dataset",
     "ImageInWellPath",
     "ImageMetaHandler",

ngio/ome_zarr_meta/ngio_specs/__init__.py

@@ -8,7 +8,6 @@ This models can be tr
 
 from ngio.ome_zarr_meta.ngio_specs._axes import (
     AxesHandler,
-    AxesOps,
     AxesSetup,
     Axis,
     AxisType,
@@ -46,7 +45,6 @@ from ngio.ome_zarr_meta.ngio_specs._pixel_size import PixelSize
 
 __all__ = [
     "AxesHandler",
-    "AxesOps",
     "AxesSetup",
     "Axis",
     "AxisType",