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

ngio/ome_zarr_meta/ngio_specs/_axes.py

@@ -4,7 +4,6 @@ from collections.abc import Sequence
 from enum import Enum
 from typing import Literal, TypeAlias, TypeVar
 
-import numpy as np
 from pydantic import BaseModel, ConfigDict, Field
 
 from ngio.utils import NgioValidationError, NgioValueError
@@ -91,7 +90,7 @@ DefaultTimeUnit = "second"
 class Axis(BaseModel):
     """Axis infos model."""
 
-    on_disk_name: str
+    name: str
     unit: str | None = None
     axis_type: AxisType | None = None
 
@@ -105,7 +104,7 @@ class Axis(BaseModel):
         if cast_type == AxisType.space and unit is None:
             unit = DefaultSpaceUnit
 
-        return Axis(on_disk_name=self.on_disk_name, axis_type=cast_type, unit=unit)
+        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."""
@@ -167,6 +166,11 @@ class AxesSetup(BaseModel):
             "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 {
@@ -177,10 +181,15 @@ class AxesSetup(BaseModel):
             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(
@@ -202,10 +211,10 @@ def _check_axes_validity(axes: Sequence[Axis], axes_setup: AxesSetup):
     _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}"
             )
 
@@ -216,17 +225,17 @@ def _check_canonical_order(
     """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}"
         )
 
 
@@ -253,32 +262,20 @@ def validate_axes(
     )
 
 
-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)
-
-    @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:
-    """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: Sequence[Axis],
+        axes: Sequence[Axis],
         # user defined args
         axes_setup: AxesSetup | None = None,
         allow_non_canonical_axes: bool = False,
@@ -287,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
@@ -296,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,
@@ -307,7 +304,7 @@ class AxesMapper:
 
         self._canonical_order = canonical_axes_order()
 
-        self._on_disk_axes = on_disk_axes
+        self._axes = axes
         self._axes_setup = axes_setup
 
         self._index_mapping = self._compute_index_mapping()
@@ -339,9 +336,9 @@ class AxesMapper:
             _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]
+        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
@@ -351,11 +348,11 @@ class AxesMapper:
 
     @property
     def axes(self) -> tuple[Axis, ...]:
-        return tuple(self._on_disk_axes)
+        return tuple(self._axes)
 
     @property
     def axes_names(self) -> tuple[str, ...]:
-        return tuple(ax.on_disk_name for ax in self._on_disk_axes)
+        return tuple(ax.name for ax in self._axes)
 
     @property
     def allow_non_canonical_axes(self) -> bool:
@@ -367,10 +364,82 @@ class AxesMapper:
         """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."""
         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)
@@ -392,105 +461,18 @@ class AxesMapper:
                     break
             else:
                 new_axes.append(axes)
-        self._on_disk_axes = new_axes
-
-    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 not isinstance(name, str):
-                raise NgioValueError(
-                    f"Invalid axis name '{name}'. Axis names must be strings."
-                )
-        inv_canonical_map = self.axes_setup.inverse_canonical_map()
+        self._axes = new_axes
 
-        # 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):
-            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: Sequence[str]) -> SlicingOps:
-        """Get the new order of the axes."""
-        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: Sequence[str]) -> SlicingOps:
-        """Get the new order of the axes."""
-        axes_to_squeeze, transposition_order, axes_to_expand = self._reorder_axes(names)
-        # Inverse transpose is just the transpose with the inverse 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) -> SlicingOps:
-        """Get the new order of the axes."""
-        other = self._axes_setup.others
-        return self.to_order(other + list(self._canonical_order))
-
-    def from_canonical(self) -> SlicingOps:
-        """Get the new order of the axes."""
-        other = self._axes_setup.others
-        return self.from_order(other + list(self._canonical_order))
-
-
-def canonical_axes(
+def build_canonical_axes_handler(
     axes_names: Sequence[str],
     space_units: SpaceUnits | str | None = DefaultSpaceUnit,
     time_units: TimeUnits | str | None = DefaultTimeUnit,
-) -> list[Axis]:
+    # 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:
@@ -502,25 +484,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,
+    )

ngio/ome_zarr_meta/ngio_specs/_dataset.py

@@ -3,85 +3,49 @@
 from collections.abc import Sequence
 
 from ngio.ome_zarr_meta.ngio_specs._axes import (
-    AxesMapper,
-    AxesSetup,
-    Axis,
-    AxisType,
-    DefaultSpaceUnit,
-    DefaultTimeUnit,
-    SpaceUnits,
-    TimeUnits,
+    AxesHandler,
 )
 from ngio.ome_zarr_meta.ngio_specs._pixel_size import PixelSize
 from ngio.utils import NgioValidationError
 
 
 class Dataset:
-    """Model for a dataset in the multiscale.
-
-    To initialize the Dataset object, the path, the axes, scale, and translation list
-    can be provided with on_disk order.
-    """
+    """Model for a dataset in the multiscale."""
 
     def __init__(
         self,
         *,
         # args coming from ngff specs
         path: str,
-        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,
-        strict_canonical_order: bool = False,
+        axes_handler: AxesHandler,
+        scale: Sequence[float],
+        translation: Sequence[float] | None = None,
     ):
         """Initialize the Dataset object.
 
         Args:
             path (str): The path of the dataset.
-            on_disk_axes (list[Axis]): The list of axes in the multiscale.
-            on_disk_scale (list[float]): The list of scale transformation.
+            axes_handler (AxesHandler): The axes handler object.
+            scale (list[float]): The list of scale transformation.
                 The scale transformation must have the same length as the axes.
-            on_disk_translation (list[float] | None): The list of translation.
-            axes_setup (AxesSetup): The axes setup object
-            allow_non_canonical_axes (bool): Allow non-canonical axes.
-            strict_canonical_order (bool): Strict canonical order.
+            translation (list[float] | None): The list of translation.
+                The translation must have the same length as the axes.
         """
         self._path = path
-        self._axes_mapper = AxesMapper(
-            on_disk_axes=on_disk_axes,
-            axes_setup=axes_setup,
-            allow_non_canonical_axes=allow_non_canonical_axes,
-            strict_canonical_order=strict_canonical_order,
-        )
+        self._axes_handler = axes_handler
 
-        if len(on_disk_scale) != len(on_disk_axes):
+        if len(scale) != len(axes_handler.axes):
             raise NgioValidationError(
                 "The length of the scale transformation must be the same as the axes."
             )
-        self._on_disk_scale = list(on_disk_scale)
+        self._scale = list(scale)
 
-        on_disk_translation = on_disk_translation or [0.0] * len(on_disk_axes)
-        if len(on_disk_translation) != len(on_disk_axes):
+        translation = translation or [0.0] * len(axes_handler.axes)
+        if len(translation) != len(axes_handler.axes):
             raise NgioValidationError(
                 "The length of the translation must be the same as the axes."
             )
-        self._on_disk_translation = list(on_disk_translation)
-
-    def get_scale(self, axis_name: str) -> float:
-        """Return the scale for a given axis."""
-        idx = self._axes_mapper.get_index(axis_name)
-        if idx is None:
-            return 1.0
-        return self._on_disk_scale[idx]
-
-    def get_translation(self, axis_name: str) -> float:
-        """Return the translation for a given axis."""
-        idx = self._axes_mapper.get_index(axis_name)
-        if idx is None:
-            return 0.0
-        return self._on_disk_translation[idx]
+        self._translation = list(translation)
 
     @property
     def path(self) -> str:
@@ -89,84 +53,46 @@ class Dataset:
         return self._path
 
     @property
-    def space_unit(self) -> str | None:
-        """Return the space unit for a given axis."""
-        x_axis = self._axes_mapper.get_axis("x")
-        y_axis = self._axes_mapper.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._axes_mapper.get_axis("t")
-        if t_axis is None:
-            return None
-        return t_axis.unit
+    def axes_handler(self) -> AxesHandler:
+        """Return the axes handler object."""
+        return self._axes_handler
 
     @property
     def pixel_size(self) -> PixelSize:
         """Return the pixel size for the dataset."""
         return PixelSize(
-            x=self.get_scale("x"),
-            y=self.get_scale("y"),
-            z=self.get_scale("z"),
-            t=self.get_scale("t"),
-            space_unit=self.space_unit,
-            time_unit=self.time_unit,
+            x=self.get_scale("x", default=1.0),
+            y=self.get_scale("y", default=1.0),
+            z=self.get_scale("z", default=1.0),
+            t=self.get_scale("t", default=1.0),
+            space_unit=self.axes_handler.space_unit,
+            time_unit=self.axes_handler.time_unit,
         )
 
     @property
-    def axes_mapper(self) -> AxesMapper:
-        """Return the axes mapper object."""
-        return self._axes_mapper
+    def scale(self) -> tuple[float, ...]:
+        """Return the scale transformation as a tuple."""
+        return tuple(self._scale)
 
-    def to_units(
-        self,
-        *,
-        space_unit: SpaceUnits = DefaultSpaceUnit,
-        time_unit: TimeUnits = DefaultTimeUnit,
-    ) -> "Dataset":
-        """Convert the pixel size to the given units.
+    @property
+    def translation(self) -> tuple[float, ...]:
+        """Return the translation as a tuple."""
+        return tuple(self._translation)
 
-        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_mapper.axes:
-            if ax.axis_type == AxisType.space:
-                new_ax = Axis(
-                    on_disk_name=ax.on_disk_name,
-                    axis_type=ax.axis_type,
-                    unit=space_unit,
-                )
-                new_axes.append(new_ax)
-            elif ax.axis_type == AxisType.time:
-                new_ax = Axis(
-                    on_disk_name=ax.on_disk_name, axis_type=ax.axis_type, unit=time_unit
-                )
-                new_axes.append(new_ax)
-            else:
-                new_axes.append(ax)
+    def get_scale(self, axis_name: str, default: float | None = None) -> float:
+        """Return the scale for a given axis."""
+        idx = self.axes_handler.get_index(axis_name)
+        if idx is None:
+            if default is not None:
+                return default
+            raise ValueError(f"Axis {axis_name} not found in axes {self.axes_handler}.")
+        return self._scale[idx]
 
-        return Dataset(
-            path=self.path,
-            on_disk_axes=new_axes,
-            on_disk_scale=self._on_disk_scale,
-            on_disk_translation=self._on_disk_translation,
-            axes_setup=self.axes_mapper.axes_setup,
-            allow_non_canonical_axes=self.axes_mapper.allow_non_canonical_axes,
-            strict_canonical_order=self.axes_mapper.strict_canonical_order,
-        )
+    def get_translation(self, axis_name: str, default: float | None = None) -> float:
+        """Return the translation for a given axis."""
+        idx = self.axes_handler.get_index(axis_name)
+        if idx is None:
+            if default is not None:
+                return default
+            raise ValueError(f"Axis {axis_name} not found in axes {self.axes_handler}.")
+        return self._translation[idx]