ngio 0.3.5__py3-none-any.whl → 0.4.0a2__py3-none-any.whl

ngio/ome_zarr_meta/ngio_specs/_axes.py

@@ -1,15 +1,16 @@
 """Fractal internal module for axes handling."""
 
-from collections.abc import Collection
+from collections.abc import Sequence
 from enum import Enum
-from typing import Literal, TypeVar
+from typing import Literal, TypeAlias, TypeVar
 
 import numpy as np
 from pydantic import BaseModel, ConfigDict, Field
 
-from ngio.utils import NgioValidationError, NgioValueError, ngio_logger
+from ngio.utils import NgioValidationError, NgioValueError
 
 T = TypeVar("T")
+SlicingType: TypeAlias = slice | tuple[int, ...] | int
 
 ################################################################################################
 #
@@ -98,24 +99,10 @@ class Axis(BaseModel):
 
     def implicit_type_cast(self, cast_type: AxisType) -> "Axis":
         unit = self.unit
-        if self.axis_type != cast_type:
-            ngio_logger.warning(
-                f"Axis {self.on_disk_name} has type {self.axis_type}. "
-                f"Casting to {cast_type}."
-            )
-
         if cast_type == AxisType.time and unit is None:
-            ngio_logger.warning(
-                f"Time axis {self.on_disk_name} has unit {self.unit}. "
-                f"Casting to {DefaultSpaceUnit}."
-            )
             unit = DefaultTimeUnit
 
         if cast_type == AxisType.space and unit is None:
-            ngio_logger.warning(
-                f"Space axis {self.on_disk_name} has unit {unit}. "
-                f"Casting to {DefaultSpaceUnit}."
-            )
             unit = DefaultSpaceUnit
 
         return Axis(on_disk_name=self.on_disk_name, axis_type=cast_type, unit=unit)
@@ -170,8 +157,28 @@ class AxesSetup(BaseModel):
 
     model_config = ConfigDict(extra="forbid", frozen=True)
 
-
-def _check_unique_names(axes: Collection[Axis]):
+    def canonical_map(self) -> dict[str, str]:
+        """Get the canonical map of axes."""
+        return {
+            "t": self.t,
+            "c": self.c,
+            "z": self.z,
+            "y": self.y,
+            "x": self.x,
+        }
+
+    def inverse_canonical_map(self) -> dict[str, str]:
+        """Get the on disk map of axes."""
+        return {
+            self.t: "t",
+            self.c: "c",
+            self.z: "z",
+            self.y: "y",
+            self.x: "x",
+        }
+
+
+def _check_unique_names(axes: Sequence[Axis]):
     """Check if all axes on disk have unique names."""
     names = [ax.on_disk_name for ax in axes]
     if len(set(names)) != len(names):
@@ -190,7 +197,7 @@ def _check_non_canonical_axes(axes_setup: AxesSetup, allow_non_canonical_axes: b
         )
 
 
-def _check_axes_validity(axes: Collection[Axis], axes_setup: AxesSetup):
+def _check_axes_validity(axes: Sequence[Axis], axes_setup: AxesSetup):
     """Check if all axes are valid."""
     _axes_setup = axes_setup.model_dump(exclude={"others"})
     _all_known_axes = [*_axes_setup.values(), *axes_setup.others]
@@ -204,7 +211,7 @@ def _check_axes_validity(axes: Collection[Axis], axes_setup: AxesSetup):
 
 
 def _check_canonical_order(
-    axes: Collection[Axis], axes_setup: AxesSetup, strict_canonical_order: bool
+    axes: Sequence[Axis], axes_setup: AxesSetup, strict_canonical_order: bool
 ):
     """Check if the axes are in the canonical order."""
     if not strict_canonical_order:
@@ -224,7 +231,7 @@ def _check_canonical_order(
 
 
 def validate_axes(
-    axes: Collection[Axis],
+    axes: Sequence[Axis],
     axes_setup: AxesSetup,
     allow_non_canonical_axes: bool = False,
     strict_canonical_order: bool = False,
@@ -246,20 +253,19 @@ def validate_axes(
     )
 
 
-class AxesTransformation(BaseModel):
-    model_config = ConfigDict(extra="forbid", frozen=True, arbitrary_types_allowed=True)
-
-
-class AxesTranspose(AxesTransformation):
-    axes: tuple[int, ...]
-
-
-class AxesExpand(AxesTransformation):
-    axes: tuple[int, ...]
+class SlicingOps(BaseModel):
+    slice_tuple: tuple[SlicingType, ...] | None = None
+    transpose_axes: tuple[int, ...] | None = None
+    expand_axes: tuple[int, ...] | None = None
+    squeeze_axes: tuple[int, ...] | None = None
+    model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True)
 
-
-class AxesSqueeze(AxesTransformation):
-    axes: tuple[int, ...]
+    @property
+    def requires_axes_ops(self) -> bool:
+        """Check if the slicing operations require axes operations."""
+        if self.expand_axes or self.transpose_axes or self.squeeze_axes:
+            return True
+        return False
 
 
 class AxesMapper:
@@ -272,7 +278,7 @@ class AxesMapper:
     def __init__(
         self,
         # spec dictated args
-        on_disk_axes: Collection[Axis],
+        on_disk_axes: Sequence[Axis],
         # user defined args
         axes_setup: AxesSetup | None = None,
         allow_non_canonical_axes: bool = False,
@@ -300,56 +306,42 @@ class AxesMapper:
         self._strict_canonical_order = strict_canonical_order
 
         self._canonical_order = canonical_axes_order()
-        self._extended_canonical_order = [*axes_setup.others, *self._canonical_order]
 
         self._on_disk_axes = on_disk_axes
         self._axes_setup = axes_setup
 
-        self._name_mapping = self._compute_name_mapping()
         self._index_mapping = self._compute_index_mapping()
 
         # Validate the axes type and cast them if necessary
         # This needs to be done after the name mapping is computed
-        self.validate_axex_type()
-
-    def _compute_name_mapping(self):
-        """Compute the name mapping.
-
-        The name mapping is a dictionary with keys the canonical axes names
-        and values the on disk axes names.
-        """
-        _name_mapping = {}
-        axis_setup_dict = self._axes_setup.model_dump(exclude={"others"})
-        _on_disk_names = self.on_disk_axes_names
-        for canonical_key, on_disk_value in axis_setup_dict.items():
-            if on_disk_value in _on_disk_names:
-                _name_mapping[canonical_key] = on_disk_value
-            else:
-                _name_mapping[canonical_key] = None
-
-        for on_disk_name in _on_disk_names:
-            if on_disk_name not in _name_mapping.keys():
-                _name_mapping[on_disk_name] = on_disk_name
-
-        for other in self._axes_setup.others:
-            if other not in _name_mapping.keys():
-                _name_mapping[other] = None
-        return _name_mapping
+        self.validate_axes_type()
 
     def _compute_index_mapping(self):
         """Compute the index mapping.
 
         The index mapping is a dictionary with keys the canonical axes names
         and values the on disk axes index.
+
+        Example:
+            If the on disk axes are ['channel', 't', 'z', 'y', 'x'],
+            the index mapping will be:
+            {
+                'c': 0,
+                'channel': 0,
+                't': 1,
+                'z': 2,
+                'y': 3,
+                'x': 4,
+            }
         """
         _index_mapping = {}
-        for canonical_key, on_disk_value in self._name_mapping.items():
-            if on_disk_value is not None:
-                _index_mapping[canonical_key] = self.on_disk_axes_names.index(
-                    on_disk_value
-                )
-            else:
-                _index_mapping[canonical_key] = None
+        for i, ax in enumerate(self.axes_names):
+            _index_mapping[ax] = i
+        # If the axis is not in the canonical order we also set it.
+        canonical_map = self._axes_setup.canonical_map()
+        for canonical_key, on_disk_value in canonical_map.items():
+            if on_disk_value in _index_mapping.keys():
+                _index_mapping[canonical_key] = _index_mapping[on_disk_value]
         return _index_mapping
 
     @property
@@ -358,12 +350,12 @@ class AxesMapper:
         return self._axes_setup
 
     @property
-    def on_disk_axes(self) -> list[Axis]:
-        return list(self._on_disk_axes)
+    def axes(self) -> tuple[Axis, ...]:
+        return tuple(self._on_disk_axes)
 
     @property
-    def on_disk_axes_names(self) -> list[str]:
-        return [ax.on_disk_name for ax in self._on_disk_axes]
+    def axes_names(self) -> tuple[str, ...]:
+        return tuple(ax.on_disk_name for ax in self._on_disk_axes)
 
     @property
     def allow_non_canonical_axes(self) -> bool:
@@ -377,28 +369,23 @@ class AxesMapper:
 
     def get_index(self, name: str) -> int | None:
         """Get the index of the axis by name."""
-        if name not in self._index_mapping.keys():
-            raise NgioValueError(
-                f"Invalid axis name '{name}'. "
-                f"Possible values are {self._index_mapping.keys()}"
-            )
-        return self._index_mapping[name]
+        return self._index_mapping.get(name, None)
 
     def get_axis(self, name: str) -> Axis | None:
         """Get the axis object by name."""
         index = self.get_index(name)
         if index is None:
             return None
-        return self.on_disk_axes[index]
+        return self.axes[index]
 
-    def validate_axex_type(self):
+    def validate_axes_type(self):
         """Validate the axes type.
 
         If the axes type is not correct, a warning is issued.
         and the axis is implicitly cast to the correct type.
         """
         new_axes = []
-        for axes in self.on_disk_axes:
+        for axes in self.axes:
             for name in self._canonical_order:
                 if axes == self.get_axis(name):
                     new_axes.append(axes.canonical_axis_cast(name))
@@ -407,71 +394,107 @@ class AxesMapper:
                 new_axes.append(axes)
         self._on_disk_axes = new_axes
 
-    def _change_order(
-        self, names: Collection[str]
-    ) -> tuple[tuple[int, ...], tuple[int, ...]]:
-        unique_names = set()
+    def _reorder_axes(
+        self, names: Sequence[str]
+    ) -> tuple[tuple[int, ...] | None, tuple[int, ...] | None, tuple[int, ...] | None]:
+        """Change the order of the axes."""
+        # Validate the names
+        unique_names = set(names)
+        if len(unique_names) != len(names):
+            raise NgioValueError(
+                "Duplicate axis names found. Please provide unique names for each axis."
+            )
         for name in names:
-            if name not in self._index_mapping.keys():
+            if not isinstance(name, str):
                 raise NgioValueError(
-                    f"Invalid axis name '{name}'. "
-                    f"Possible values are {self._index_mapping.keys()}"
+                    f"Invalid axis name '{name}'. Axis names must be strings."
                 )
-            _unique_name = self._name_mapping.get(name)
-            if _unique_name is None:
-                continue
-            if _unique_name in unique_names:
-                raise NgioValueError(
-                    f"Duplicate axis name, two or more '{_unique_name}' were found. "
-                    f"Please provide unique names."
-                )
-            unique_names.add(_unique_name)
-
-        if len(self.on_disk_axes_names) > len(unique_names):
-            missing_names = set(self.on_disk_axes_names) - unique_names
-            raise NgioValueError(
-                f"Some axes where not queried. "
-                f"Please provide the following missing axes {missing_names}"
-            )
-        _indices, _insert = [], []
+        inv_canonical_map = self.axes_setup.inverse_canonical_map()
+
+        # Step 1: Check find squeeze axes
+        _axes_to_squeeze: list[int] = []
+        axes_names_after_squeeze = []
+        for i, ax in enumerate(self.axes_names):
+            # If the axis is not in the names, it means we need to squeeze it
+            ax_canonical = inv_canonical_map.get(ax, None)
+            if ax not in names and ax_canonical not in names:
+                _axes_to_squeeze.append(i)
+            elif ax in names:
+                axes_names_after_squeeze.append(ax)
+            elif ax_canonical in names:
+                # If the axis is in the canonical map, we add it to the names
+                axes_names_after_squeeze.append(ax_canonical)
+
+        axes_to_squeeze = tuple(_axes_to_squeeze) if len(_axes_to_squeeze) > 0 else None
+
+        # Step 2: Find the transposition order
+        _transposition_order: list[int] = []
+        axes_names_after_transpose = []
+        for ax in names:
+            if ax in axes_names_after_squeeze:
+                _transposition_order.append(axes_names_after_squeeze.index(ax))
+                axes_names_after_transpose.append(ax)
+
+        if np.allclose(_transposition_order, range(len(_transposition_order))):
+            # If the transposition order is the identity, we don't need to transpose
+            transposition_order = None
+        else:
+            transposition_order = tuple(_transposition_order)
+
+        # Step 3: Find axes to expand
+        _axes_to_expand: list[int] = []
         for i, name in enumerate(names):
-            _index = self._index_mapping[name]
-            if _index is None:
-                _insert.append(i)
-            else:
-                _indices.append(self._index_mapping[name])
-        return tuple(_indices), tuple(_insert)
+            if name not in axes_names_after_transpose:
+                # If the axis is not in the mapping, it means we need to expand it
+                _axes_to_expand.append(i)
+
+        axes_to_expand = tuple(_axes_to_expand) if len(_axes_to_expand) > 0 else None
+        return axes_to_squeeze, transposition_order, axes_to_expand
 
-    def to_order(self, names: Collection[str]) -> tuple[AxesTransformation, ...]:
+    def to_order(self, names: Sequence[str]) -> SlicingOps:
         """Get the new order of the axes."""
-        _indices, _insert = self._change_order(names)
-        return AxesTranspose(axes=_indices), AxesExpand(axes=_insert)
+        axes_to_squeeze, transposition_order, axes_to_expand = self._reorder_axes(names)
+        return SlicingOps(
+            transpose_axes=transposition_order,
+            expand_axes=axes_to_expand,
+            squeeze_axes=axes_to_squeeze,
+        )
 
-    def from_order(self, names: Collection[str]) -> tuple[AxesTransformation, ...]:
+    def from_order(self, names: Sequence[str]) -> SlicingOps:
         """Get the new order of the axes."""
-        _indices, _insert = self._change_order(names)
+        axes_to_squeeze, transposition_order, axes_to_expand = self._reorder_axes(names)
         # Inverse transpose is just the transpose with the inverse indices
-        _reverse_indices = tuple(np.argsort(_indices))
-        return AxesSqueeze(axes=_insert), AxesTranspose(axes=_reverse_indices)
+        if transposition_order is None:
+            _reverse_indices = None
+        else:
+            _reverse_indices = tuple(np.argsort(transposition_order))
+
+        return SlicingOps(
+            transpose_axes=_reverse_indices,
+            expand_axes=axes_to_squeeze,
+            squeeze_axes=axes_to_expand,
+        )
 
-    def to_canonical(self) -> tuple[AxesTransformation, ...]:
+    def to_canonical(self) -> SlicingOps:
         """Get the new order of the axes."""
-        return self.to_order(self._extended_canonical_order)
+        other = self._axes_setup.others
+        return self.to_order(other + list(self._canonical_order))
 
-    def from_canonical(self) -> tuple[AxesTransformation, ...]:
+    def from_canonical(self) -> SlicingOps:
         """Get the new order of the axes."""
-        return self.from_order(self._extended_canonical_order)
+        other = self._axes_setup.others
+        return self.from_order(other + list(self._canonical_order))
 
 
 def canonical_axes(
-    axes_names: Collection[str],
-    space_units: SpaceUnits | None = DefaultSpaceUnit,
-    time_units: TimeUnits | None = DefaultTimeUnit,
+    axes_names: Sequence[str],
+    space_units: SpaceUnits | str | None = DefaultSpaceUnit,
+    time_units: TimeUnits | str | None = DefaultTimeUnit,
 ) -> list[Axis]:
     """Create a new canonical axes mapper.
 
     Args:
-        axes_names (Collection[str] | int): The axes names on disk.
+        axes_names (Sequence[str] | int): The axes names on disk.
             - The axes should be in ['t', 'c', 'z', 'y', 'x']
             - The axes should be in strict canonical order.
             - If an integer is provided, the axes are created from the last axis

ngio/ome_zarr_meta/ngio_specs/_channels.py

@@ -3,7 +3,7 @@
 Stores the same information as the Omero section of the ngff 0.4 metadata.
 """
 
-from collections.abc import Collection
+from collections.abc import Sequence
 from difflib import SequenceMatcher
 from enum import Enum
 from typing import Any, TypeVar
@@ -77,7 +77,8 @@ class NgioColors(str, Enum):
             # try to match the color to the channel name
             similarity[color] = SequenceMatcher(None, channel_name, color).ratio()
         # Get the color with the highest similarity
-        color_str = max(similarity, key=similarity.get)  # type: ignore
+        color_str = max(similarity, key=similarity.get)  # type: ignore (max type overload fails to infer type)
+        assert isinstance(color_str, str), "Color name must be a string."
         return NgioColors.__members__[color_str]
 
 
@@ -287,7 +288,7 @@ class Channel(BaseModel):
 T = TypeVar("T")
 
 
-def _check_elements(elements: Collection[T], expected_type: Any) -> Collection[T]:
+def _check_elements(elements: Sequence[T], expected_type: Any) -> Sequence[T]:
     """Check that the elements are of the same type."""
     if len(elements) == 0:
         raise NgioValidationError("At least one element must be provided.")
@@ -301,7 +302,7 @@ def _check_elements(elements: Collection[T], expected_type: Any) -> Collection[T
     return elements
 
 
-def _check_unique(elements: Collection[T]) -> Collection[T]:
+def _check_unique(elements: Sequence[T]) -> Sequence[T]:
     """Check that the elements are unique."""
     if len(set(elements)) != len(elements):
         raise NgioValidationError("All elements must be unique.")
@@ -329,35 +330,35 @@ class ChannelsMeta(BaseModel):
     @classmethod
     def default_init(
         cls,
-        labels: Collection[str | None] | int,
-        wavelength_id: Collection[str | None] | None = None,
-        colors: Collection[str | NgioColors | None] | None = None,
-        start: Collection[int | float | None] | int | float | None = None,
-        end: Collection[int | float | None] | int | float | None = None,
-        active: Collection[bool | None] | None = None,
+        labels: Sequence[str | None] | int,
+        wavelength_id: Sequence[str | None] | None = None,
+        colors: Sequence[str | NgioColors | None] | None = None,
+        start: Sequence[int | float | None] | int | float | None = None,
+        end: Sequence[int | float | None] | int | float | None = None,
+        active: Sequence[bool | None] | None = None,
         data_type: Any = np.uint16,
         **omero_kwargs: dict,
     ) -> "ChannelsMeta":
         """Create a ChannelsMeta object with the default unit.
 
         Args:
-            labels(Collection[str | None] | int): The list of channels names
+            labels(Sequence[str | None] | int): The list of channels names
                 in the image. If an integer is provided, the channels will be
                 named "channel_i".
-            wavelength_id(Collection[str | None] | None): The wavelength ID of the
+            wavelength_id(Sequence[str | None] | None): The wavelength ID of the
                 channel. If None, the wavelength ID will be the same as the
                 channel name.
-            colors(Collection[str | NgioColors | None] | None): The list of
+            colors(Sequence[str | NgioColors | None] | None): The list of
                 colors for the channels. If None, the colors will be random.
-            start(Collection[int | float | None] | int | float | None): The start
+            start(Sequence[int | float | None] | int | float | None): The start
                 value of the channel. If None, the start value will be the
                 minimum value of the data type.
-            end(Collection[int | float | None] | int | float | None): The end
+            end(Sequence[int | float | None] | int | float | None): The end
                 value of the channel. If None, the end value will be the
                 maximum value of the data type.
             data_type(Any): The data type of the channel. Will be used to set the
                 min and max values of the channel.
-            active (Collection[bool | None] | None): Whether the channel should
+            active (Sequence[bool | None] | None): Whether the channel should
                 be shown by default.
             omero_kwargs(dict): Extra fields to store in the omero attributes.
         """
@@ -367,9 +368,9 @@ class ChannelsMeta(BaseModel):
         labels = _check_elements(labels, str)
         labels = _check_unique(labels)
 
-        _wavelength_id: Collection[str | None] = [None] * len(labels)
+        _wavelength_id: Sequence[str | None] = [None] * len(labels)
         if wavelength_id is None:
-            _wavelength_id: Collection[str | None] = [None] * len(labels)
+            _wavelength_id: Sequence[str | None] = [None] * len(labels)
         else:
             _wavelength_id = _check_elements(wavelength_id, str)
             _wavelength_id = _check_unique(wavelength_id)
@@ -425,3 +426,39 @@ class ChannelsMeta(BaseModel):
                 )
             )
         return cls(channels=channels, **omero_kwargs)
+
+    @property
+    def channel_labels(self) -> list[str]:
+        """Get the labels of the channels in the image."""
+        return [channel.label for channel in self.channels]
+
+    @property
+    def channel_wavelength_ids(self) -> list[str | None]:
+        """Get the wavelength IDs of the channels in the image."""
+        return [channel.wavelength_id for channel in self.channels]
+
+    def get_channel_idx(
+        self, channel_label: str | None = None, wavelength_id: str | None = None
+    ) -> int:
+        """Get the index of a channel by its label or wavelength ID."""
+        # Only one of the arguments must be provided
+        if channel_label is not None and wavelength_id is not None:
+            raise NgioValueError(
+                "get_channel_idx must receive either label or wavelength_id, not both."
+            )
+
+        if channel_label is not None:
+            if channel_label not in self.channel_labels:
+                raise NgioValueError(f"Channel with label {channel_label} not found.")
+            return self.channel_labels.index(channel_label)
+
+        if wavelength_id is not None:
+            if wavelength_id not in self.channel_wavelength_ids:
+                raise NgioValueError(
+                    f"Channel with wavelength ID {wavelength_id} not found."
+                )
+            return self.channel_wavelength_ids.index(wavelength_id)
+
+        raise NgioValueError(
+            "get_channel_idx must receive either label or wavelength_id"
+        )

ngio/ome_zarr_meta/ngio_specs/_dataset.py

@@ -1,6 +1,6 @@
 """Fractal internal module for dataset metadata handling."""
 
-from collections.abc import Collection
+from collections.abc import Sequence
 
 from ngio.ome_zarr_meta.ngio_specs._axes import (
     AxesMapper,
@@ -28,9 +28,9 @@ class Dataset:
         *,
         # args coming from ngff specs
         path: str,
-        on_disk_axes: Collection[Axis],
-        on_disk_scale: Collection[float],
-        on_disk_translation: Collection[float] | None = None,
+        on_disk_axes: Sequence[Axis],
+        on_disk_scale: Sequence[float],
+        on_disk_translation: Sequence[float] | None = None,
         # user defined args
         axes_setup: AxesSetup | None = None,
         allow_non_canonical_axes: bool = False,
@@ -123,8 +123,8 @@ class Dataset:
             y=self.get_scale("y"),
             z=self.get_scale("z"),
             t=self.get_scale("t"),
-            space_unit=self.space_unit,  # type: ignore
-            time_unit=self.time_unit,  # type: ignore
+            space_unit=self.space_unit,
+            time_unit=self.time_unit,
         )
 
     @property
@@ -145,7 +145,7 @@ class Dataset:
             time_unit(str): The time unit to convert to.
         """
         new_axes = []
-        for ax in self.axes_mapper.on_disk_axes:
+        for ax in self.axes_mapper.axes:
             if ax.axis_type == AxisType.space:
                 new_ax = Axis(
                     on_disk_name=ax.on_disk_name,

ngio/ome_zarr_meta/ngio_specs/_ngio_hcs.py

@@ -51,11 +51,11 @@ class CustomWellImage(WellImage04):
 
 
 class CustomWellMeta(WellMeta04):
-    images: list[CustomWellImage]  # type: ignore[valid-type]
+    images: list[CustomWellImage]  # type: ignore (override of WellMeta04.images)
 
 
 class CustomWellAttrs(WellAttrs04):
-    well: CustomWellMeta  # type: ignore[valid-type]
+    well: CustomWellMeta  # type: ignore (override of WellAttrs04.well)
 
 
 class NgioWellMeta(CustomWellAttrs):
@@ -531,7 +531,7 @@ class NgioPlateMeta(HCSAttrs):
             acquisitions = None
 
         if version is None:
-            version = self.plate.version  # type: ignore[assignment]
+            version = self.plate.version  # type: ignore (version is NgffVersions or None)
 
         return NgioPlateMeta(
             plate=Plate(