ngio 0.2.0a2__py3-none-any.whl → 0.5.0b4__py3-none-any.whl

ngio/ome_zarr_meta/ngio_specs/_axes.py

@@ -1,18 +1,15 @@
 """Fractal internal module for axes handling."""
 
-from collections.abc import Collection
+from collections.abc import Sequence
 from enum import Enum
-from logging import Logger
-from typing import TypeVar
+from typing import Literal, TypeAlias, TypeVar
 
-import numpy as np
 from pydantic import BaseModel, ConfigDict, Field
 
 from ngio.utils import NgioValidationError, NgioValueError
 
-logger = Logger(__name__)
-
 T = TypeVar("T")
+SlicingType: TypeAlias = slice | tuple[int, ...] | int
 
 ################################################################################################
 #
@@ -32,98 +29,94 @@ class AxisType(str, Enum):
     space = "space"
 
 
-class SpaceUnits(str, Enum):
-    """Allowed space units."""
-
-    nanometer = "nanometer"
-    nm = "nm"
-    micrometer = "micrometer"
-    um = "um"
-    millimeter = "millimeter"
-    mm = "mm"
-    centimeter = "centimeter"
-    cm = "cm"
-
-    @classmethod
-    def default(cls) -> "SpaceUnits":
-        return SpaceUnits.um
-
-
-class TimeUnits(str, Enum):
-    """Allowed time units."""
-
-    seconds = "seconds"
-    s = "s"
-
-    @classmethod
-    def default(cls) -> "TimeUnits":
-        return TimeUnits.s
+SpaceUnits = Literal[
+    "micrometer",
+    "nanometer",
+    "angstrom",
+    "picometer",
+    "millimeter",
+    "centimeter",
+    "decimeter",
+    "meter",
+    "inch",
+    "foot",
+    "yard",
+    "mile",
+    "kilometer",
+    "hectometer",
+    "megameter",
+    "gigameter",
+    "terameter",
+    "petameter",
+    "exameter",
+    "parsec",
+    "femtometer",
+    "attometer",
+    "zeptometer",
+    "yoctometer",
+    "zettameter",
+    "yottameter",
+]
+DefaultSpaceUnit = "micrometer"
+
+TimeUnits = Literal[
+    "attosecond",
+    "centisecond",
+    "day",
+    "decisecond",
+    "exasecond",
+    "femtosecond",
+    "gigasecond",
+    "hectosecond",
+    "hour",
+    "kilosecond",
+    "megasecond",
+    "microsecond",
+    "millisecond",
+    "minute",
+    "nanosecond",
+    "petasecond",
+    "picosecond",
+    "second",
+    "terasecond",
+    "yoctosecond",
+    "yottasecond",
+    "zeptosecond",
+    "zettasecond",
+]
+DefaultTimeUnit = "second"
 
 
 class Axis(BaseModel):
     """Axis infos model."""
 
-    on_disk_name: str
-    unit: SpaceUnits | TimeUnits | None = None
+    name: str
+    unit: str | None = None
     axis_type: AxisType | None = None
 
     model_config = ConfigDict(extra="forbid", frozen=True)
 
     def implicit_type_cast(self, cast_type: AxisType) -> "Axis":
-        if self.axis_type != cast_type:
-            logger.warning(
-                f"Axis {self.on_disk_name} has type {self.axis_type}. "
-                f"Casting to {cast_type}."
-            )
-        new_axis = Axis(
-            on_disk_name=self.on_disk_name, axis_type=cast_type, unit=self.unit
-        )
-        if cast_type == AxisType.time and not isinstance(self.unit, TimeUnits):
-            logger.warning(
-                f"Time axis {self.on_disk_name} has unit {self.unit}. "
-                f"Casting to {TimeUnits.default()}."
-            )
-            new_axis = Axis(
-                on_disk_name=self.on_disk_name,
-                axis_type=AxisType.time,
-                unit=TimeUnits.default(),
-            )
-        elif cast_type == AxisType.space and not isinstance(self.unit, SpaceUnits):
-            logger.warning(
-                f"Space axis {self.on_disk_name} has unit {self.unit}. "
-                f"Casting to {SpaceUnits.default()}."
-            )
-            new_axis = Axis(
-                on_disk_name=self.on_disk_name,
-                axis_type=AxisType.space,
-                unit=SpaceUnits.default(),
-            )
-        elif cast_type == AxisType.channel and self.unit is not None:
-            logger.warning(
-                f"Channel axis {self.on_disk_name} has unit {self.unit}. Removing unit."
-            )
-            new_axis = Axis(
-                on_disk_name=self.on_disk_name,
-                axis_type=AxisType.channel,
-                unit=None,
-            )
-        return new_axis
+        unit = self.unit
+        if cast_type == AxisType.time and unit is None:
+            unit = DefaultTimeUnit
+
+        if cast_type == AxisType.space and unit is None:
+            unit = DefaultSpaceUnit
+
+        return Axis(name=self.name, axis_type=cast_type, unit=unit)
 
     def canonical_axis_cast(self, canonical_name: str) -> "Axis":
         """Cast the implicit axis to the correct type."""
         match canonical_name:
             case "t":
-                if self.axis_type != AxisType.time or not isinstance(
-                    self.unit, TimeUnits
-                ):
+                if self.axis_type != AxisType.time or self.unit is None:
                     return self.implicit_type_cast(AxisType.time)
             case "c":
-                if self.axis_type != AxisType.channel or self.unit is not None:
+                if self.axis_type != AxisType.channel:
                     return self.implicit_type_cast(AxisType.channel)
             case "z" | "y" | "x":
-                if self.axis_type != AxisType.space or not isinstance(
-                    self.unit, SpaceUnits
-                ):
+                if self.axis_type != AxisType.space or self.unit is None:
                     return self.implicit_type_cast(AxisType.space)
         return self
 
@@ -163,10 +156,40 @@ 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 get_on_disk_name(self, canonical_name: str) -> str | None:
+        """Get the on disk name of the axis by its canonical name."""
+        canonical_map = self.canonical_map()
+        return canonical_map.get(canonical_name, None)
+
+    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 get_canonical_name(self, on_disk_name: str) -> str | None:
+        """Get the canonical name of the axis by its on disk name."""
+        inv_map = self.inverse_canonical_map()
+        return inv_map.get(on_disk_name, None)
+
+
+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]
+    names = [ax.name for ax in axes]
     if len(set(names)) != len(names):
         duplicates = {item for item in names if names.count(item) > 1}
         raise NgioValidationError(
@@ -183,41 +206,41 @@ 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]
     for ax in axes:
-        if ax.on_disk_name not in _all_known_axes:
+        if ax.name not in _all_known_axes:
             raise NgioValidationError(
-                f"Invalid axis name '{ax.on_disk_name}'. "
-                f"Please correct map `{ax.on_disk_name}` "
+                f"Invalid axis name '{ax.name}'. "
+                f"Please correct map `{ax.name}` "
                 f"using the AxesSetup model {axes_setup}"
             )
 
 
 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:
         return
-    _on_disk_names = [ax.on_disk_name for ax in axes]
+    _names = [ax.name for ax in axes]
     _canonical_order = []
     for name in canonical_axes_order():
         mapped_name = getattr(axes_setup, name)
-        if mapped_name in _on_disk_names:
+        if mapped_name in _names:
             _canonical_order.append(mapped_name)
 
-    if _on_disk_names != _canonical_order:
+    if _names != _canonical_order:
         raise NgioValidationError(
             f"Invalid axes order. The axes must be in the canonical order. "
-            f"Expected {_canonical_order}, but found {_on_disk_names}"
+            f"Expected {_canonical_order}, but found {_names}"
         )
 
 
 def validate_axes(
-    axes: Collection[Axis],
+    axes: Sequence[Axis],
     axes_setup: AxesSetup,
     allow_non_canonical_axes: bool = False,
     strict_canonical_order: bool = False,
@@ -239,33 +262,20 @@ 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 AxesSqueeze(AxesTransformation):
-    axes: tuple[int, ...]
-
-
-class AxesMapper:
-    """Map on disk axes to canonical axes.
-
-    This class is used to map the on disk axes to the canonical axes.
+class AxesHandler:
+    """This class is used to handle and operate on OME-Zarr axes.
 
+    The class also provides:
+        - methods to reorder, squeeze and expand axes.
+        - methods to validate the axes.
+        - methods to get axis by name or index.
+        - methods to operate on the axes.
     """
 
     def __init__(
         self,
         # spec dictated args
-        on_disk_axes: Collection[Axis],
+        axes: Sequence[Axis],
         # user defined args
         axes_setup: AxesSetup | None = None,
         allow_non_canonical_axes: bool = False,
@@ -274,7 +284,7 @@ class AxesMapper:
         """Create a new AxesMapper object.
 
         Args:
-            on_disk_axes (list[Axis]): The axes on disk.
+            axes (list[Axis]): The axes on disk.
             axes_setup (AxesSetup, optional): The axis setup. Defaults to None.
             allow_non_canonical_axes (bool, optional): Allow non canonical axes.
             strict_canonical_order (bool, optional): Check if the axes are in the
@@ -283,7 +293,7 @@ class AxesMapper:
         axes_setup = axes_setup if axes_setup is not None else AxesSetup()
 
         validate_axes(
-            axes=on_disk_axes,
+            axes=axes,
             axes_setup=axes_setup,
             allow_non_canonical_axes=allow_non_canonical_axes,
             strict_canonical_order=strict_canonical_order,
@@ -293,163 +303,181 @@ 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 = 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_name, on_disk_name in canonical_map.items():
+            if on_disk_name in _index_mapping.keys():
+                _index_mapping[canonical_name] = _index_mapping[on_disk_name]
         return _index_mapping
 
     @property
-    def on_disk_axes(self) -> list[Axis]:
-        return list(self._on_disk_axes)
+    def axes_setup(self) -> AxesSetup:
+        """Return the axes setup."""
+        return self._axes_setup
+
+    @property
+    def axes(self) -> tuple[Axis, ...]:
+        return tuple(self._axes)
+
+    @property
+    def axes_names(self) -> tuple[str, ...]:
+        """On disk axes names."""
+        return tuple(ax.name for ax in self._axes)
 
     @property
-    def on_disk_axes_names(self) -> list[str]:
-        return [ax.on_disk_name for ax in self._on_disk_axes]
+    def allow_non_canonical_axes(self) -> bool:
+        """Return if non canonical axes are allowed."""
+        return self._allow_non_canonical_axes
+
+    @property
+    def strict_canonical_order(self) -> bool:
+        """Return if strict canonical order is enforced."""
+        return self._strict_canonical_order
+
+    @property
+    def space_unit(self) -> str | None:
+        """Return the space unit for a given axis."""
+        x_axis = self.get_axis("x")
+        y_axis = self.get_axis("y")
+
+        if x_axis is None or y_axis is None:
+            raise NgioValidationError(
+                "The dataset must have x and y axes to determine the space unit."
+            )
+
+        if x_axis.unit == y_axis.unit:
+            return x_axis.unit
+        else:
+            raise NgioValidationError(
+                "Inconsistent space units. "
+                f"x={x_axis.unit} and y={y_axis.unit} should have the same unit."
+            )
+
+    @property
+    def time_unit(self) -> str | None:
+        """Return the time unit for a given axis."""
+        t_axis = self.get_axis("t")
+        if t_axis is None:
+            return None
+        return t_axis.unit
+
+    def to_units(
+        self,
+        *,
+        space_unit: SpaceUnits = DefaultSpaceUnit,
+        time_unit: TimeUnits = DefaultTimeUnit,
+    ) -> "AxesHandler":
+        """Convert the pixel size to the given units.
+
+        Args:
+            space_unit(str): The space unit to convert to.
+            time_unit(str): The time unit to convert to.
+        """
+        new_axes = []
+        for ax in self.axes:
+            if ax.axis_type == AxisType.space:
+                new_ax = Axis(
+                    name=ax.name,
+                    axis_type=ax.axis_type,
+                    unit=space_unit,
+                )
+                new_axes.append(new_ax)
+            elif ax.axis_type == AxisType.time:
+                new_ax = Axis(name=ax.name, axis_type=ax.axis_type, unit=time_unit)
+                new_axes.append(new_ax)
+            else:
+                new_axes.append(ax)
+
+        return AxesHandler(
+            axes=new_axes,
+            axes_setup=self.axes_setup,
+            allow_non_canonical_axes=self.allow_non_canonical_axes,
+            strict_canonical_order=self.strict_canonical_order,
+        )
 
     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 has_axis(self, axis_name: str) -> bool:
+        """Return whether the axis exists."""
+        index = self.get_index(axis_name)
+        if index is None:
+            return False
+        return True
+
+    def get_canonical_name(self, name: str) -> str | None:
+        """Get the canonical name of the axis by name."""
+        return self._axes_setup.get_canonical_name(name)
 
     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))
                     break
             else:
                 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()
-        for name in names:
-            if name not in self._index_mapping.keys():
-                raise NgioValueError(
-                    f"Invalid axis name '{name}'. "
-                    f"Possible values are {self._index_mapping.keys()}"
-                )
-            _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)
+        self._axes = new_axes
 
-        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 = [], []
-        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)
-
-    def to_order(self, names: Collection[str]) -> tuple[AxesTransformation, ...]:
-        """Get the new order of the axes."""
-        _indices, _insert = self._change_order(names)
-        return AxesTranspose(axes=_indices), AxesExpand(axes=_insert)
-
-    def from_order(self, names: Collection[str]) -> tuple[AxesTransformation, ...]:
-        """Get the new order of the axes."""
-        _indices, _insert = self._change_order(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)
-
-    def to_canonical(self) -> tuple[AxesTransformation, ...]:
-        """Get the new order of the axes."""
-        return self.to_order(self._extended_canonical_order)
-
-    def from_canonical(self) -> tuple[AxesTransformation, ...]:
-        """Get the new order of the axes."""
-        return self.from_order(self._extended_canonical_order)
-
-
-def canonical_axes(
-    axes_names: Collection[str],
-    space_units: SpaceUnits | None = None,
-    time_units: TimeUnits | None = None,
-) -> list[Axis]:
+
+def build_canonical_axes_handler(
+    axes_names: Sequence[str],
+    space_units: SpaceUnits | str | None = DefaultSpaceUnit,
+    time_units: TimeUnits | str | None = DefaultTimeUnit,
+    # user defined args
+    axes_setup: AxesSetup | None = None,
+    allow_non_canonical_axes: bool = False,
+    strict_canonical_order: bool = False,
+) -> AxesHandler:
     """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
@@ -457,25 +485,31 @@ def canonical_axes(
                 e.g. 3 -> ["z", "y", "x"]
         space_units (SpaceUnits, optional): The space units. Defaults to None.
         time_units (TimeUnits, optional): The time units. Defaults to None.
+        axes_setup (AxesSetup, optional): The axis setup. Defaults to None.
+        allow_non_canonical_axes (bool, optional): Allow non canonical axes.
+            Defaults to False.
+        strict_canonical_order (bool, optional): Check if the axes are in the
+            canonical order. Defaults to False.
 
     """
     axes = []
     for name in axes_names:
         match name:
             case "t":
-                axes.append(
-                    Axis(on_disk_name=name, axis_type=AxisType.time, unit=time_units)
-                )
+                axes.append(Axis(name=name, axis_type=AxisType.time, unit=time_units))
             case "c":
-                axes.append(Axis(on_disk_name=name, axis_type=AxisType.channel))
+                axes.append(Axis(name=name, axis_type=AxisType.channel))
             case "z" | "y" | "x":
-                axes.append(
-                    Axis(on_disk_name=name, axis_type=AxisType.space, unit=space_units)
-                )
+                axes.append(Axis(name=name, axis_type=AxisType.space, unit=space_units))
             case _:
                 raise NgioValueError(
                     f"Invalid axis name '{name}'. "
                     "Only 't', 'c', 'z', 'y', 'x' are allowed."
                 )
 
-    return axes
+    return AxesHandler(
+        axes=axes,
+        axes_setup=axes_setup,
+        allow_non_canonical_axes=allow_non_canonical_axes,
+        strict_canonical_order=strict_canonical_order,
+    )