Flexible coordinate transform

xarray/__init__.py

@@ -30,6 +30,7 @@
     where,
 )
 from xarray.core.concat import concat
+from xarray.core.coordinate_transform import CoordinateTransform
 from xarray.core.coordinates import Coordinates
 from xarray.core.dataarray import DataArray
 from xarray.core.dataset import Dataset
@@ -109,6 +110,7 @@
     "CFTimeIndex",
     "Context",
     "Coordinates",
+    "CoordinateTransform",
     "DataArray",
     "Dataset",
     "DataTree",

xarray/core/coordinate_transform.py

@@ -0,0 +1,78 @@
+from collections.abc import Hashable, Iterable, Mapping
+from typing import Any
+
+import numpy as np
+
+
+class CoordinateTransform:
+    """Abstract coordinate transform with dimension & coordinate names."""
+
+    coord_names: tuple[Hashable, ...]
+    dims: tuple[str, ...]
+    dim_size: dict[str, int]
+    dtype: Any
+
+    def __init__(
+        self,
+        coord_names: Iterable[Hashable],
+        dim_size: Mapping[str, int],
+        dtype: Any = None,
+    ):
+        self.coord_names = tuple(coord_names)
+        self.dims = tuple(dim_size)
+        self.dim_size = dict(dim_size)
+
+        if dtype is None:
+            dtype = np.dtype(np.float64)
+        self.dtype = dtype
+
+    def forward(self, dim_positions: dict[str, Any]) -> dict[Hashable, Any]:
+        """Perform grid -> world coordinate transformation.
+
+        Parameters
+        ----------
+        dim_positions : dict
+            Grid location(s) along each dimension (axis).
+
+        Returns
+        -------
+        coord_labels : dict
+            World coordinate labels.
+
+        """
+        # TODO: cache the results in order to avoid re-computing
+        # all labels when accessing the values of each coordinate one at a time
+        raise NotImplementedError
+
+    def reverse(self, coord_labels: dict[Hashable, Any]) -> dict[str, Any]:
+        """Perform world -> grid coordinate reverse transformation.
+
+        Parameters
+        ----------
+        labels : dict
+            World coordinate labels.
+
+        Returns
+        -------
+        dim_positions : dict
+            Grid relative location(s) along each dimension (axis).
+
+        """
+        raise NotImplementedError
+
+    def equals(self, other: "CoordinateTransform") -> bool:
+        """Check equality with another CoordinateTransform of the same kind."""
+        raise NotImplementedError
+
+    def generate_coords(self, dims: tuple[str] | None = None) -> dict[Hashable, Any]:
+        """Compute all coordinate labels at once."""
+        if dims is None:
+            dims = self.dims
+
+        positions = np.meshgrid(
+            *[np.arange(self.dim_size[d]) for d in dims],
+            indexing="ij",
+        )
+        dim_positions = {dim: positions[i] for i, dim in enumerate(dims)}
+
+        return self.forward(dim_positions)

xarray/core/indexes.py

@@ -10,7 +10,9 @@
 import pandas as pd
 
 from xarray.core import formatting, nputils, utils
+from xarray.core.coordinate_transform import CoordinateTransform
 from xarray.core.indexing import (
+    CoordinateTransformIndexingAdapter,
     IndexSelResult,
     PandasIndexingAdapter,
     PandasMultiIndexingAdapter,
@@ -24,6 +26,7 @@
 )
 
 if TYPE_CHECKING:
+    from xarray.core.coordinate import Coordinates
     from xarray.core.types import ErrorOptions, JoinOptions, Self
     from xarray.core.variable import Variable
 
@@ -1372,6 +1375,132 @@ def rename(self, name_dict, dims_dict):
         )
 
 
+class CoordinateTransformIndex(Index):
+    """Helper class for creating Xarray indexes based on coordinate transforms.
+
+    - wraps a :py:class:`CoordinateTransform` instance
+    - takes care of creating the index (lazy) coordinates
+    - supports point-wise label-based selection
+    - supports exact alignment only, by comparing indexes based on their transform
+      (not on their explicit coordinate labels)
+
+    """
+
+    transform: CoordinateTransform
+
+    def __init__(
+        self,
+        transform: CoordinateTransform,
+    ):
+        self.transform = transform
+
+    def create_variables(
+        self, variables: Mapping[Any, Variable] | None = None
+    ) -> IndexVars:
+        from xarray.core.variable import Variable
+
+        new_variables = {}
+
+        for name in self.transform.coord_names:
+            # copy attributes, if any
+            attrs: Mapping[Hashable, Any] | None
+
+            if variables is not None and name in variables:
+                var = variables[name]
+                attrs = var.attrs
+            else:
+                attrs = None
+
+            data = CoordinateTransformIndexingAdapter(self.transform, name)
+            new_variables[name] = Variable(self.transform.dims, data, attrs=attrs)
+
+        return new_variables
+
+    def create_coordinates(self) -> Coordinates:
+        # TODO: move this in xarray.Index base class?
+        from xarray.core.coordinates import Coordinates
+
+        variables = self.create_variables()
+        indexes = {name: self for name in variables}
+        return Coordinates(coords=variables, indexes=indexes)
+
+    def isel(
+        self, indexers: Mapping[Any, int | slice | np.ndarray | Variable]
+    ) -> Self | None:
+        # TODO: support returning a new index (e.g., possible to re-calculate the
+        # the transform or calculate another transform on a reduced dimension space)
+        return None
+
+    def sel(
+        self, labels: dict[Any, Any], method=None, tolerance=None
+    ) -> IndexSelResult:
+        from xarray.core.dataarray import DataArray
+        from xarray.core.variable import Variable
+
+        if method != "nearest":
+            raise ValueError(
+                "CoordinateTransformIndex only supports selection with method='nearest'"
+            )
+
+        labels_set = set(labels)
+        coord_names_set = set(self.transform.coord_names)
+
+        missing_labels = coord_names_set - labels_set
+        if missing_labels:
+            missing_labels_str = ",".join([f"{name}" for name in missing_labels])
+            raise ValueError(f"missing labels for coordinate(s): {missing_labels_str}.")
+
+        label0_obj = next(iter(labels.values()))
+        dim_size0 = getattr(label0_obj, "sizes", None)
+
+        is_xr_obj = [
+            isinstance(label, DataArray | Variable) for label in labels.values()
+        ]
+        if not all(is_xr_obj):
+            raise TypeError(
+                "CoordinateTransformIndex only supports advanced (point-wise) indexing "
+                "with either xarray.DataArray or xarray.Variable objects."
+            )
+        dim_size = [getattr(label, "sizes", None) for label in labels.values()]
+        if any([ds != dim_size0 for ds in dim_size]):
+            raise ValueError(
+                "CoordinateTransformIndex only supports advanced (point-wise) indexing "
+                "with xarray.DataArray or xarray.Variable objects of macthing dimensions."
+            )
+
+        coord_labels = {
+            name: labels[name].values for name in self.transform.coord_names
+        }
+        dim_positions = self.transform.reverse(coord_labels)
+
+        results = {}
+        dims0 = tuple(dim_size0)
+        for dim, pos in dim_positions.items():
+            # TODO: rounding the decimal positions is not always the behavior we expect
+            # (there are different ways to represent implicit intervals)
+            # we should probably make this customizable.
+            pos = np.round(pos).astype("int")
+            if isinstance(label0_obj, Variable):
+                xr_pos = Variable(dims0, pos)
+            else:
+                # dataarray
+                xr_pos = DataArray(pos, dims=dims0)
+            results[dim] = xr_pos
+
+        return IndexSelResult(results)
+
+    def equals(self, other: Self) -> bool:
+        return self.transform.equals(other.transform)
+
+    def rename(
+        self,
+        name_dict: Mapping[Any, Hashable],
+        dims_dict: Mapping[Any, Hashable],
+    ) -> Self:
+        # TODO: maybe update self.transform coord_names, dim_size and dims attributes
+        return self
+
+
 def create_default_index_implicit(
     dim_variable: Variable,
     all_variables: Mapping | Iterable[Hashable] | None = None,