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

ngio/ome_zarr_meta/ngio_specs/__init__.py

@@ -7,15 +7,16 @@ This models can be tr
 """
 
 from ngio.ome_zarr_meta.ngio_specs._axes import (
-    AxesMapper,
+    AxesHandler,
+    AxesOps,
     AxesSetup,
     Axis,
     AxisType,
     DefaultSpaceUnit,
     DefaultTimeUnit,
-    SlicingOps,
     SpaceUnits,
     TimeUnits,
+    build_canonical_axes_handler,
     canonical_axes_order,
     canonical_label_axes_order,
 )
@@ -44,7 +45,8 @@ from ngio.ome_zarr_meta.ngio_specs._ngio_image import (
 from ngio.ome_zarr_meta.ngio_specs._pixel_size import PixelSize
 
 __all__ = [
-    "AxesMapper",
+    "AxesHandler",
+    "AxesOps",
     "AxesSetup",
     "Axis",
     "AxisType",
@@ -65,9 +67,9 @@ __all__ = [
     "NgioPlateMeta",
     "NgioWellMeta",
     "PixelSize",
-    "SlicingOps",
     "SpaceUnits",
     "TimeUnits",
+    "build_canonical_axes_handler",
     "canonical_axes_order",
     "canonical_label_axes_order",
     "default_channel_name",

ngio/ome_zarr_meta/ngio_specs/_axes.py

@@ -91,7 +91,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 +105,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 +167,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 +182,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 +212,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 +226,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 +263,77 @@ 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
+class AxesOps(BaseModel):
+    """Model to represent axes operations.
+
+    This model will be used to transform objects from on disk axes to in memory axes.
+    """
+
+    on_disk_axes: tuple[str, ...]
+    in_memory_axes: tuple[str, ...]
+    transpose_op: tuple[int, ...] | None = None
+    expand_op: tuple[int, ...] | None = None
+    squeeze_op: 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:
+    def is_no_op(self) -> bool:
+        """Check if all operations are no ops."""
+        if (
+            self.transpose_op is None
+            and self.expand_op is None
+            and self.squeeze_op is None
+        ):
             return True
         return False
 
+    @property
+    def get_transpose_op(self) -> tuple[int, ...] | None:
+        """Get the transpose axes."""
+        return self.transpose_op
+
+    @property
+    def get_expand_op(self) -> tuple[int, ...] | None:
+        """Get the expand axes."""
+        return self.expand_op
+
+    @property
+    def get_squeeze_op(self) -> tuple[int, ...] | None:
+        """Get the squeeze axes."""
+        return self.squeeze_op
+
+    @property
+    def set_transpose_op(self) -> tuple[int, ...] | None:
+        """Set the transpose axes."""
+        if self.transpose_op is None:
+            return None
+        return tuple(np.argsort(self.transpose_op))
+
+    @property
+    def set_expand_op(self) -> tuple[int, ...] | None:
+        """Set the expand axes."""
+        return self.squeeze_op
+
+    @property
+    def set_squeeze_op(self) -> tuple[int, ...] | None:
+        """Set the squeeze axes."""
+        return self.expand_op
 
-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 +342,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 +351,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 +362,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 +394,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 +406,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 +422,75 @@ 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 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,7 +512,7 @@ class AxesMapper:
                     break
             else:
                 new_axes.append(axes)
-        self._on_disk_axes = new_axes
+        self._axes = new_axes
 
     def _reorder_axes(
         self, names: Sequence[str]
@@ -451,46 +571,32 @@ class AxesMapper:
         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."""
+    def get_axes_ops(self, names: Sequence[str]) -> AxesOps:
+        """Get the axes operations to go from on-disk to in-memory 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,
+        return AxesOps(
+            on_disk_axes=self.axes_names,
+            in_memory_axes=tuple(names),
+            transpose_op=transposition_order,
+            expand_op=axes_to_expand,
+            squeeze_op=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."""
+    def get_canonical_axes_ops(self) -> AxesOps:
+        """Get the axes operations to go from on-disk to canonical in-memory axes."""
         other = self._axes_setup.others
-        return self.to_order(other + list(self._canonical_order))
+        return self.get_axes_ops(list(self._canonical_order) + other)
 
-    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 +608,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]