ngio 0.1.5__py3-none-any.whl → 0.2.0a1__py3-none-any.whl

ngio/__init__.py

@@ -2,14 +2,40 @@
 
 from importlib.metadata import PackageNotFoundError, version
 
-from ngio.core import Image, Label, NgffImage
-
-__all__ = ["Image", "Label", "NgffImage"]
-
-
 try:
     __version__ = version("ngio")
 except PackageNotFoundError:  # pragma: no cover
     __version__ = "uninstalled"
 __author__ = "Lorenzo Cerrone"
 __email__ = "lorenzo.cerrone@uzh.ch"
+
+from ngio.common import ArrayLike, Dimensions
+from ngio.hcs import OmeZarrPlate, OmeZarrWell, open_omezarr_plate, open_omezarr_well
+from ngio.images import (
+    Image,
+    Label,
+    OmeZarrContainer,
+    create_empty_omezarr,
+    create_omezarr_from_array,
+    open_image,
+    open_omezarr_container,
+)
+from ngio.ome_zarr_meta.ngio_specs import AxesSetup, PixelSize
+
+__all__ = [
+    "ArrayLike",
+    "AxesSetup",
+    "Dimensions",
+    "Image",
+    "Label",
+    "OmeZarrContainer",
+    "OmeZarrPlate",
+    "OmeZarrWell",
+    "PixelSize",
+    "create_empty_omezarr",
+    "create_omezarr_from_array",
+    "open_image",
+    "open_omezarr_container",
+    "open_omezarr_plate",
+    "open_omezarr_well",
+]

ngio/common/__init__.py

@@ -0,0 +1,44 @@
+"""Common classes and functions that are used across the package."""
+
+from ngio.common._array_pipe import get_pipe, set_pipe
+from ngio.common._axes_transforms import (
+    transform_dask_array,
+    transform_list,
+    transform_numpy_array,
+)
+from ngio.common._common_types import ArrayLike
+from ngio.common._dimensions import Dimensions
+from ngio.common._pyramid import consolidate_pyramid, init_empty_pyramid, on_disk_zoom
+from ngio.common._roi import RasterCooROI, WorldCooROI
+from ngio.common._slicer import (
+    SliceTransform,
+    compute_and_slices,
+    dask_get_slice,
+    dask_set_slice,
+    numpy_get_slice,
+    numpy_set_slice,
+)
+from ngio.common._zoom import dask_zoom, numpy_zoom
+
+__all__ = [
+    "ArrayLike",
+    "Dimensions",
+    "RasterCooROI",
+    "SliceTransform",
+    "WorldCooROI",
+    "compute_and_slices",
+    "consolidate_pyramid",
+    "dask_get_slice",
+    "dask_set_slice",
+    "dask_zoom",
+    "get_pipe",
+    "init_empty_pyramid",
+    "numpy_get_slice",
+    "numpy_set_slice",
+    "numpy_zoom",
+    "on_disk_zoom",
+    "set_pipe",
+    "transform_dask_array",
+    "transform_list",
+    "transform_numpy_array",
+]

ngio/common/_array_pipe.py

@@ -0,0 +1,160 @@
+from collections.abc import Collection, Iterable
+from typing import Literal
+
+import dask
+import dask.delayed
+import numpy as np
+import zarr
+
+from ngio.common._axes_transforms import transform_dask_array, transform_numpy_array
+from ngio.common._common_types import ArrayLike
+from ngio.common._dimensions import Dimensions
+from ngio.common._slicer import (
+    SliceTransform,
+    compute_and_slices,
+    dask_get_slice,
+    dask_set_slice,
+    numpy_get_slice,
+    numpy_set_slice,
+)
+from ngio.ome_zarr_meta.ngio_specs import AxesTransformation
+from ngio.utils import NgioValueError
+
+
+def _compute_from_disk_transforms(
+    *,
+    dimensions: Dimensions,
+    axes_order: Collection[str] | None = None,
+    **slice_kwargs: slice | int | Iterable[int],
+) -> tuple[SliceTransform, tuple[AxesTransformation, ...]]:
+    slices = compute_and_slices(dimensions=dimensions, **slice_kwargs)
+
+    if axes_order is None:
+        return slices, ()
+
+    additional_transformations = dimensions._axes_mapper.to_order(axes_order)
+    return slices, additional_transformations
+
+
+def _compute_to_disk_transforms(
+    *,
+    dimensions: Dimensions,
+    axes_order: Collection[str] | None = None,
+    **slice_kwargs: slice | int | Iterable[int],
+) -> tuple[SliceTransform, tuple[AxesTransformation, ...]]:
+    slices = compute_and_slices(dimensions=dimensions, **slice_kwargs)
+    if axes_order is None:
+        return slices, ()
+
+    additional_transformations = dimensions._axes_mapper.from_order(axes_order)
+    return slices, additional_transformations
+
+
+def _numpy_get_pipe(
+    array: zarr.Array,
+    slices: SliceTransform,
+    transformations: tuple[AxesTransformation, ...],
+) -> np.ndarray:
+    array = numpy_get_slice(array, slices)
+    return transform_numpy_array(array, transformations)
+
+
+def _delayed_numpy_get_pipe(
+    array: zarr.Array,
+    slices: SliceTransform,
+    transformations: tuple[AxesTransformation, ...],
+) -> dask.delayed:
+    array = dask.delayed(numpy_get_slice)(array, slices)
+    return dask.delayed(transform_numpy_array)(array, transformations)
+
+
+def _dask_get_pipe(
+    array: zarr.Array,
+    slices: SliceTransform,
+    transformations: tuple[AxesTransformation, ...],
+) -> dask.array:
+    array = dask_get_slice(array, slices)
+    return transform_dask_array(array, transformations)
+
+
+def _numpy_set_pipe(
+    array: zarr.Array,
+    patch: np.ndarray,
+    slices: SliceTransform,
+    transformations: tuple[AxesTransformation, ...],
+) -> None:
+    patch = transform_numpy_array(patch, transformations)
+    numpy_set_slice(array, patch, slices)
+
+
+def _dask_set_pipe(
+    array: zarr.Array,
+    patch: np.ndarray,
+    slices: SliceTransform,
+    transformations: tuple[AxesTransformation, ...],
+) -> None:
+    patch = transform_dask_array(patch, transformations)
+    dask_set_slice(array, patch, slices)
+
+
+def _delayed_numpy_set_pipe(
+    array: zarr.Array,
+    patch: np.ndarray,
+    slices: SliceTransform,
+    transformations: tuple[AxesTransformation, ...],
+) -> dask.delayed:
+    patch = dask.delayed(transform_numpy_array)(patch, transformations)
+    return dask.delayed(numpy_set_slice)(array, patch, slices)
+
+
+def get_pipe(
+    array: zarr.Array,
+    *,
+    dimensions: Dimensions,
+    axes_order: Collection[str] | None = None,
+    mode: Literal["numpy", "dask", "delayed"] = "numpy",
+    **slice_kwargs: slice | int | Iterable[int],
+):
+    slices, transformations = _compute_from_disk_transforms(
+        dimensions=dimensions, axes_order=axes_order, **slice_kwargs
+    )
+    match mode:
+        case "numpy":
+            return _numpy_get_pipe(array, slices, transformations)
+        case "dask":
+            return _dask_get_pipe(array, slices, transformations)
+
+        case "delayed_numpy":
+            return _delayed_numpy_get_pipe(array, slices, transformations)
+
+        case _:
+            raise NgioValueError(
+                f"Unknown get pipe mode {mode}, expected 'numpy', 'dask' or 'delayed'."
+            )
+
+
+def set_pipe(
+    array: zarr.Array,
+    patch: ArrayLike,
+    *,
+    dimensions: Dimensions,
+    axes_order: Collection[str] | None = None,
+    **slice_kwargs: slice | int | Iterable[int],
+):
+    slices, transformations = _compute_to_disk_transforms(
+        dimensions=dimensions, axes_order=axes_order, **slice_kwargs
+    )
+    if isinstance(patch, dask.array.Array):
+        _dask_set_pipe(
+            array=array, patch=patch, slices=slices, transformations=transformations
+        )
+    elif isinstance(patch, np.ndarray):
+        _numpy_set_pipe(
+            array=array, patch=patch, slices=slices, transformations=transformations
+        )
+    elif isinstance(patch, dask.delayed.Delayed):
+        _delayed_numpy_set_pipe(
+            array=array, patch=patch, slices=slices, transformations=transformations
+        )
+    else:
+        raise NgioValueError("Unknown patch type, expected numpy, dask or delayed.")

ngio/common/_axes_transforms.py

@@ -0,0 +1,63 @@
+from typing import TypeVar
+
+import dask.array as da
+import numpy as np
+
+from ngio.ome_zarr_meta.ngio_specs._axes import (
+    AxesExpand,
+    AxesSqueeze,
+    AxesTransformation,
+    AxesTranspose,
+)
+
+T = TypeVar("T")
+
+
+def transform_list(
+    input_list: list[T], default: T, operations: tuple[AxesTransformation, ...]
+) -> list[T]:
+    if isinstance(input_list, tuple):
+        input_list = list(input_list)
+
+    for operation in operations:
+        if isinstance(operation, AxesTranspose):
+            input_list = [input_list[i] for i in operation.axes]
+
+        if isinstance(operation, AxesExpand):
+            for ax in operation.axes:
+                input_list.insert(ax, default)
+        elif isinstance(operation, AxesSqueeze):
+            for offset, ax in enumerate(operation.axes):
+                input_list.pop(ax - offset)
+
+    return input_list
+
+
+def transform_numpy_array(
+    array: np.ndarray, operations: tuple[AxesTransformation, ...]
+) -> np.ndarray:
+    for operation in operations:
+        if isinstance(operation, AxesTranspose):
+            array = np.transpose(array, operation.axes)
+        elif isinstance(operation, AxesExpand):
+            array = np.expand_dims(array, axis=operation.axes)
+        elif isinstance(operation, AxesSqueeze):
+            array = np.squeeze(array, axis=operation.axes)
+        else:
+            raise ValueError(f"Unknown operation {operation}")
+    return array
+
+
+def transform_dask_array(
+    array: da.Array, operations: tuple[AxesTransformation, ...]
+) -> da.Array:
+    for operation in operations:
+        if isinstance(operation, AxesTranspose):
+            array = da.transpose(array, axes=operation.axes)
+        elif isinstance(operation, AxesExpand):
+            array = da.expand_dims(array, axis=operation.axes)
+        elif isinstance(operation, AxesSqueeze):
+            array = da.squeeze(array, axis=operation.axes)
+        else:
+            raise ValueError(f"Unknown operation {operation}")
+    return array

ngio/common/_common_types.py

@@ -0,0 +1,5 @@
+import numpy as np
+import zarr
+from dask import array as da
+
+ArrayLike = np.ndarray | da.core.Array | zarr.Array  # type: ignore

ngio/common/_dimensions.py

@@ -0,0 +1,113 @@
+"""Dimension metadata.
+
+This is not related to the NGFF metadata,
+but it is based on the actual metadata of the image data.
+"""
+
+from collections.abc import Collection
+
+from ngio.common._axes_transforms import transform_list
+from ngio.ome_zarr_meta import AxesMapper
+from ngio.utils import NgioValidationError, NgioValueError
+
+
+class Dimensions:
+    """Dimension metadata."""
+
+    def __init__(
+        self,
+        shape: tuple[int, ...],
+        axes_mapper: AxesMapper,
+    ) -> None:
+        """Create a Dimension object from a Zarr array.
+
+        Args:
+            shape: The shape of the Zarr array.
+            axes_mapper: The axes mapper object.
+        """
+        self._shape = shape
+        self._axes_mapper = axes_mapper
+
+        if len(self._shape) != len(self._axes_mapper.on_disk_axes):
+            raise NgioValidationError(
+                "The number of dimensions must match the number of axes. "
+                f"Expected Axis {self._axes_mapper.on_disk_axes_names} but got shape "
+                f"{self._shape}."
+            )
+
+    def __str__(self) -> str:
+        """Return the string representation of the object."""
+        dims = ", ".join(
+            f"{ax.on_disk_name}: {s}"
+            for ax, s in zip(self._axes_mapper.on_disk_axes, self._shape, strict=True)
+        )
+        return f"Dimensions({dims})"
+
+    def get(self, axis_name: str, strict: bool = True) -> int:
+        """Return the dimension of the given axis name.
+
+        Args:
+            axis_name: The name of the axis (either canonical or non-canonical).
+            strict: If True, raise an error if the axis does not exist.
+        """
+        index = self._axes_mapper.get_index(axis_name)
+        if index is None and strict:
+            raise NgioValueError(f"Axis {axis_name} does not exist.")
+        elif index is None:
+            return 1
+        return self._shape[index]
+
+    def get_shape(self, axes_order: Collection[str]) -> tuple[int, ...]:
+        """Return the shape in the given axes order."""
+        transforms = self._axes_mapper.to_order(axes_order)
+        return tuple(transform_list(list(self._shape), 1, transforms))
+
+    def get_canonical_shape(self) -> tuple[int, ...]:
+        """Return the shape in the canonical order."""
+        transforms = self._axes_mapper.to_canonical()
+        return tuple(transform_list(list(self._shape), 1, transforms))
+
+    def __repr__(self) -> str:
+        """Return the string representation of the object."""
+        return str(self)
+
+    @property
+    def on_disk_shape(self) -> tuple[int, ...]:
+        """Return the shape as a tuple."""
+        return tuple(self._shape)
+
+    @property
+    def is_time_series(self) -> bool:
+        """Return whether the data is a time series."""
+        if self.get("t", strict=False) == 1:
+            return False
+        return True
+
+    @property
+    def is_2d(self) -> bool:
+        """Return whether the data is 2D."""
+        if self.get("z", strict=False) != 1:
+            return False
+        return True
+
+    @property
+    def is_2d_time_series(self) -> bool:
+        """Return whether the data is a 2D time series."""
+        return self.is_2d and self.is_time_series
+
+    @property
+    def is_3d(self) -> bool:
+        """Return whether the data is 3D."""
+        return not self.is_2d
+
+    @property
+    def is_3d_time_series(self) -> bool:
+        """Return whether the data is a 3D time series."""
+        return self.is_3d and self.is_time_series
+
+    @property
+    def is_multi_channels(self) -> bool:
+        """Return whether the data has multiple channels."""
+        if self.get("c", strict=False) == 1:
+            return False
+        return True

ngio/common/_pyramid.py

@@ -0,0 +1,222 @@
+import math
+from collections.abc import Collection
+from typing import Literal
+
+import dask.array as da
+import numpy as np
+import zarr
+
+from ngio.common._zoom import _zoom_inputs_check, dask_zoom, numpy_zoom
+from ngio.utils import AccessModeLiteral, StoreOrGroup, open_group_wrapper
+
+
+def _on_disk_numpy_zoom(
+    source: zarr.Array,
+    target: zarr.Array,
+    order: Literal[0, 1, 2] = 1,
+) -> None:
+    target[...] = numpy_zoom(source[...], target_shape=target.shape, order=order)
+
+
+def _on_disk_dask_zoom(
+    source: zarr.Array,
+    target: zarr.Array,
+    order: Literal[0, 1, 2] = 1,
+) -> None:
+    source_array = da.from_zarr(source)
+    target_array = dask_zoom(source_array, target_shape=target.shape, order=order)
+
+    target_array = target_array.rechunk(target.chunks)
+    target_array.compute_chunk_sizes()
+    target_array.to_zarr(target)
+
+
+def _on_disk_coarsen(
+    source: zarr.Array,
+    target: zarr.Array,
+    _order: Literal[0, 1] = 1,
+    aggregation_function: np.ufunc | None = None,
+) -> None:
+    """Apply a coarsening operation from a source zarr array to a target zarr array.
+
+    Args:
+        source (zarr.Array): The source array to coarsen.
+        target (zarr.Array): The target array to save the coarsened result to.
+        _order (Literal[0, 1]): The order of interpolation is not really implemented
+            for coarsening, but it is kept for compatibility with the zoom function.
+            _order=1 -> linear interpolation ~ np.mean
+            _order=0 -> nearest interpolation ~ np.max
+        aggregation_function (np.ufunc): The aggregation function to use.
+    """
+    source_array = da.from_zarr(source)
+
+    _scale, _target_shape = _zoom_inputs_check(
+        source_array=source_array, scale=None, target_shape=target.shape
+    )
+
+    assert _target_shape == target.shape, (
+        "Target shape must match the target array shape"
+    )
+
+    if aggregation_function is None:
+        if _order == 1:
+            aggregation_function = np.mean
+        elif _order == 0:
+            aggregation_function = np.max
+        else:
+            raise ValueError(
+                f"Aggregation function must be provided for order {_order}"
+            )
+
+    coarsening_setup = {}
+    for i, s in enumerate(_scale):
+        factor = 1 / s
+        # This check is very strict, but it is necessary to avoid
+        # a few pixels shift in the coarsening
+        # We could add a tolerance
+        if factor.is_integer():
+            coarsening_setup[i] = int(factor)
+        else:
+            raise ValueError(
+                f"Coarsening factor must be an integer, got {factor} on axis {i}"
+            )
+
+    out_target = da.coarsen(
+        aggregation_function, source_array, coarsening_setup, trim_excess=True
+    )
+    out_target = out_target.rechunk(target.chunks)
+    out_target.to_zarr(target)
+
+
+def on_disk_zoom(
+    source: zarr.Array,
+    target: zarr.Array,
+    order: Literal[0, 1, 2] = 1,
+    mode: Literal["dask", "numpy", "coarsen"] = "dask",
+) -> None:
+    """Apply a zoom operation from a source zarr array to a target zarr array.
+
+    Args:
+        source (zarr.Array): The source array to zoom.
+        target (zarr.Array): The target array to save the zoomed result to.
+        order (Literal[0, 1, 2]): The order of interpolation. Defaults to 1.
+        mode (Literal["dask", "numpy", "coarsen"]): The mode to use. Defaults to "dask".
+    """
+    if not isinstance(source, zarr.Array):
+        raise ValueError("source must be a zarr array")
+
+    if not isinstance(target, zarr.Array):
+        raise ValueError("target must be a zarr array")
+
+    if source.dtype != target.dtype:
+        raise ValueError("source and target must have the same dtype")
+
+    match mode:
+        case "numpy":
+            return _on_disk_numpy_zoom(source, target, order)
+        case "dask":
+            return _on_disk_dask_zoom(source, target, order)
+        case "coarsen":
+            return _on_disk_coarsen(
+                source,
+                target,
+            )
+        case _:
+            raise ValueError("mode must be either 'dask', 'numpy' or 'coarsen'")
+
+
+def _find_closest_arrays(
+    processed: list[zarr.Array], to_be_processed: list[zarr.Array]
+) -> tuple[int, int]:
+    dist_matrix = np.zeros((len(processed), len(to_be_processed)))
+    for i, arr_to_proc in enumerate(to_be_processed):
+        for j, proc_arr in enumerate(processed):
+            dist_matrix[j, i] = np.sqrt(
+                np.sum(
+                    [
+                        (s1 - s2) ** 2
+                        for s1, s2 in zip(
+                            arr_to_proc.shape, proc_arr.shape, strict=False
+                        )
+                    ]
+                )
+            )
+
+    return np.unravel_index(dist_matrix.argmin(), dist_matrix.shape)
+
+
+def consolidate_pyramid(
+    source: zarr.Array,
+    targets: list[zarr.Array],
+    order: Literal[0, 1, 2] = 1,
+    mode: Literal["dask", "numpy", "coarsen"] = "dask",
+) -> None:
+    """Consolidate the Zarr array."""
+    processed = [source]
+    to_be_processed = targets
+
+    while to_be_processed:
+        source_id, target_id = _find_closest_arrays(processed, to_be_processed)
+
+        source_image = processed[source_id]
+        target_image = to_be_processed.pop(target_id)
+
+        on_disk_zoom(
+            source=source_image,
+            target=target_image,
+            mode=mode,
+            order=order,
+        )
+        processed.append(target_image)
+
+
+def init_empty_pyramid(
+    store: StoreOrGroup,
+    paths: list[str],
+    ref_shape: Collection[int],
+    scaling_factors: Collection[float],
+    chunks: Collection[int] | None = None,
+    dtype: str = "uint16",
+    mode: AccessModeLiteral = "a",
+) -> None:
+    # Return the an Image object
+    if chunks is not None and len(chunks) != len(ref_shape):
+        raise ValueError(
+            "The shape and chunks must have the same number of dimensions."
+        )
+
+    if len(ref_shape) != len(scaling_factors):
+        raise ValueError(
+            "The shape and scaling factor must have the same number of dimensions."
+        )
+
+    root_group, _ = open_group_wrapper(store, mode=mode)
+    for path in paths:
+        if any(s < 1 for s in ref_shape):
+            raise ValueError(
+                "Level shape must be at least 1 on all dimensions. "
+                f"Calculated shape: {ref_shape} at level {path}."
+            )
+        new_arr = root_group.zeros(
+            name=path,
+            shape=ref_shape,
+            dtype=dtype,
+            chunks=chunks,
+            dimension_separator="/",
+        )
+
+        # Todo redo this with when a proper build of pyramid is implemented
+        _shape = []
+        for s, sc in zip(ref_shape, scaling_factors, strict=True):
+            if math.floor(s / sc) % 2 == 0:
+                _shape.append(math.floor(s / sc))
+            else:
+                _shape.append(math.ceil(s / sc))
+        ref_shape = _shape
+
+        if chunks is None:
+            chunks = new_arr.chunks
+            if chunks is None:
+                raise ValueError("Something went wrong with the chunks")
+        chunks = [min(c, s) for c, s in zip(chunks, ref_shape, strict=True)]
+    return None