rasterix 0.1a3__py3-none-any.whl → 0.1a4__py3-none-any.whl

rasterix/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.1a3'
-__version_tuple__ = version_tuple = (0, 1, 'a3')
+__version__ = version = '0.1a4'
+__version_tuple__ = version_tuple = (0, 1, 'a4')

rasterix/raster_index.py

@@ -3,59 +3,97 @@ from __future__ import annotations
 import math
 import textwrap
 from collections.abc import Hashable, Iterable, Mapping, Sequence
-from typing import Any, Self, TypeVar
+from typing import Any, Self, TypeVar, cast
 
 import numpy as np
 import pandas as pd
+import xproj
 from affine import Affine
-from xarray import Coordinates, DataArray, Dataset, Index, Variable
+from pyproj import CRS
+from xarray import Coordinates, DataArray, Dataset, Index, Variable, get_options
 from xarray.core.coordinate_transform import CoordinateTransform
 
 # TODO: import from public API once it is available
-from xarray.core.indexes import CoordinateTransformIndex, PandasIndex
+from xarray.core.indexes import CoordinateTransformIndex
 from xarray.core.indexing import IndexSelResult, merge_sel_results
+from xarray.core.types import JoinOptions
+from xproj.typing import CRSAwareIndex
 
 from rasterix.odc_compat import BoundingBox, bbox_intersection, bbox_union, maybe_int, snap_grid
 from rasterix.rioxarray_compat import guess_dims
+from rasterix.utils import get_affine
 
 T_Xarray = TypeVar("T_Xarray", "DataArray", "Dataset")
 
 __all__ = ["assign_index", "RasterIndex"]
 
 
-def assign_index(obj: T_Xarray, *, x_dim: str | None = None, y_dim: str | None = None) -> T_Xarray:
+# X/Y axis order conventions used for public and internal attributes
+XAXIS = 0
+YAXIS = 1
+
+
+def assign_index(
+    obj: T_Xarray, *, x_dim: str | None = None, y_dim: str | None = None, crs: bool = True
+) -> T_Xarray:
     """Assign a RasterIndex to an Xarray DataArray or Dataset.
 
+    By default, the affine transform is guessed by first looking for a ``GeoTransform`` attribute
+    on a CF "grid mapping" variable (commonly ``"spatial_ref"``). If not present, then the affine is determined from 1D coordinate
+    variables named ``x_dim`` and ``y_dim`` provided to this function.
+
     Parameters
     ----------
     obj : xarray.DataArray or xarray.Dataset
-        The object to assign the index to. Must have a rio accessor with a transform.
+        The object to assign the index to.
     x_dim : str, optional
         Name of the x dimension. If None, will be automatically detected.
     y_dim : str, optional
         Name of the y dimension. If None, will be automatically detected.
+    crs: bool, optional
+       Auto-detect CRS using xproj?
 
     Returns
     -------
     xarray.DataArray or xarray.Dataset
         The input object with RasterIndex coordinates assigned.
 
+    Notes
+    -----
+    The "grid mapping" variable is determined following the CF conventions:
+
+    - If a DataArray is provided, we look for an attribute named ``"grid_mapping"``.
+    - For a Dataset, we pull the first detected ``"grid_mapping"`` attribute when iterating over data variables.
+
+    The value of this attribute is a variable name containing projection information (commonly ``"spatial_ref"``).
+    We then look for a ``"GeoTransform"`` attribute on this variable (following GDAL convention).
+
+    References
+    ----------
+    - `CF conventions document <http://cfconventions.org/Data/cf-conventions/cf-conventions-1.11/cf-conventions.html#grid-mappings-and-projections>`_.
+    - `GDAL docs on GeoTransform <https://gdal.org/en/stable/tutorials/geotransforms_tut.html>`_.
+
     Examples
     --------
     >>> import xarray as xr
-    >>> import rioxarray  # Required for rio accessor
+    >>> import rioxarray  # Required for reading TIFF
     >>> da = xr.open_dataset("path/to/raster.tif", engine="rasterio")
     >>> indexed_da = assign_index(da)
     """
-    import rioxarray  # noqa
-
     if x_dim is None or y_dim is None:
         guessed_x, guessed_y = guess_dims(obj)
     x_dim = x_dim or guessed_x
     y_dim = y_dim or guessed_y
 
+    affine = get_affine(obj, x_dim=x_dim, y_dim=y_dim, clear_transform=True)
+
     index = RasterIndex.from_transform(
-        obj.rio.transform(), obj.sizes[x_dim], obj.sizes[y_dim], x_dim=x_dim, y_dim=y_dim
+        affine,
+        width=obj.sizes[x_dim],
+        height=obj.sizes[y_dim],
+        x_dim=x_dim,
+        y_dim=y_dim,
+        crs=obj.proj.crs if crs else None,
     )
     coords = Coordinates.from_xindex(index)
     return obj.assign_coords(coords)
@@ -91,8 +129,8 @@ class AffineTransform(CoordinateTransform):
         self.affine = affine
 
         # array dimensions in reverse order (y = rows, x = cols)
-        self.xy_dims = self.dims[0], self.dims[1]
-        self.dims = self.dims[1], self.dims[0]
+        self.xy_dims = self.dims[XAXIS], self.dims[YAXIS]
+        self.dims = self.dims[YAXIS], self.dims[XAXIS]
 
     def forward(self, dim_positions):
         positions = tuple(dim_positions[dim] for dim in self.xy_dims)
@@ -114,13 +152,17 @@ class AffineTransform(CoordinateTransform):
 
         return results
 
-    def equals(self, other, *, exclude):
+    def equals(self, other: CoordinateTransform, *, exclude: frozenset[Hashable] | None = None) -> bool:
         if exclude is not None:
             raise NotImplementedError
         if not isinstance(other, AffineTransform):
             return False
         return self.affine == other.affine and self.dim_size == other.dim_size
 
+    def __repr__(self) -> str:
+        params = ", ".join(f"{pn}={getattr(self.affine, pn):.4g}" for pn in "abcdef")
+        return f"{type(self).__name__}({params})"
+
 
 class AxisAffineTransform(CoordinateTransform):
     """Axis-independent wrapper of an affine 2D transform with no skew/rotation."""
@@ -169,7 +211,7 @@ class AxisAffineTransform(CoordinateTransform):
 
         return {self.dim: positions}
 
-    def equals(self, other, *, exclude):
+    def equals(self, other: CoordinateTransform, *, exclude: frozenset[Hashable] | None = None) -> bool:
         if not isinstance(other, AxisAffineTransform):
             return False
         if exclude is not None:
@@ -188,14 +230,10 @@ class AxisAffineTransform(CoordinateTransform):
         return self.forward({self.dim: np.arange(self.size)})
 
     def slice(self, slice: slice) -> AxisAffineTransform:
-        start = max(slice.start or 0, 0)
-        stop = min(slice.stop or self.size, self.size)
-        step = slice.step or 1
-
-        # TODO: support reverse transform (i.e., start > stop)?
-        assert start < stop
-
-        size = (stop - start) // step
+        newrange = range(self.size)[slice]
+        start = newrange.start
+        step = newrange.step or 1
+        size = len(newrange)
         scale = float(step)
 
         if self.is_xaxis:
@@ -212,6 +250,10 @@ class AxisAffineTransform(CoordinateTransform):
             dtype=self.dtype,
         )
 
+    def __repr__(self) -> str:
+        params = ", ".join(f"{pn}={getattr(self.affine, pn):.4g}" for pn in "abcdef")
+        return f"{type(self).__name__}({params}, axis={'X' if self.is_xaxis else 'Y'}, dim={self.dim!r})"
+
 
 class AxisAffineTransformIndex(CoordinateTransformIndex):
     """Axis-independent Xarray Index for an affine 2D transform with no
@@ -225,9 +267,6 @@ class AxisAffineTransformIndex(CoordinateTransformIndex):
     - Data slicing computes a new affine transform and returns a new
       `AxisAffineTransformIndex` object
 
-    - Otherwise data selection creates and returns a new Xarray
-      `PandasIndex` object for non-scalar indexers
-
     - The index can be converted to a `pandas.Index` object (useful for Xarray
       operations that don't work with Xarray indexes yet).
 
@@ -244,7 +283,7 @@ class AxisAffineTransformIndex(CoordinateTransformIndex):
 
     def isel(  # type: ignore[override]
         self, indexers: Mapping[Any, int | slice | np.ndarray | Variable]
-    ) -> AxisAffineTransformIndex | PandasIndex | None:
+    ) -> AxisAffineTransformIndex | None:
         idxer = indexers[self.dim]
 
         # generate a new index with updated transform if a slice is given
@@ -256,14 +295,17 @@ class AxisAffineTransformIndex(CoordinateTransformIndex):
         # no index for scalar value
         elif np.ndim(idxer) == 0:
             return None
-        # otherwise return a PandasIndex with values computed by forward transformation
+        # otherwise drop the index
+        # (TODO: return a PandasIndex with values computed by forward transformation when it
+        #  will be possible to auto-convert RasterIndex to PandasIndex for x and/or y axis)
         else:
-            values = self.axis_transform.forward({self.dim: idxer})[self.axis_transform.coord_name]
-            if isinstance(idxer, Variable):
-                new_dim = idxer.dims[0]
-            else:
-                new_dim = self.dim
-            return PandasIndex(values, new_dim, coord_dtype=values.dtype)
+            return None
+            # values = self.axis_transform.forward({self.dim: idxer})[self.axis_transform.coord_name]
+            # if isinstance(idxer, Variable):
+            #     new_dim = idxer.dims[0]
+            # else:
+            #     new_dim = self.dim
+            # return PandasIndex(values, new_dim, coord_dtype=values.dtype)
 
     def sel(self, labels, method=None, tolerance=None):
         coord_name = self.axis_transform.coord_name
@@ -311,22 +353,10 @@ class AxisAffineTransformIndex(CoordinateTransformIndex):
         return pd.Index(values[self.dim])
 
 
-# The types of Xarray indexes that may be wrapped by RasterIndex
-WrappedIndex = AxisAffineTransformIndex | PandasIndex | CoordinateTransformIndex
-WrappedIndexCoords = Hashable | tuple[Hashable, Hashable]
-
-
-def _filter_dim_indexers(index: WrappedIndex, indexers: Mapping) -> Mapping:
-    if isinstance(index, CoordinateTransformIndex):
-        dims = index.transform.dims
-    else:
-        # PandasIndex
-        dims = (str(index.dim),)
+WrappedIndex = tuple[AxisAffineTransformIndex, AxisAffineTransformIndex] | CoordinateTransformIndex
 
-    return {dim: indexers[dim] for dim in dims if dim in indexers}
 
-
-class RasterIndex(Index):
+class RasterIndex(Index, xproj.ProjIndexMixin):
     """Xarray index for raster coordinate indexing and spatial operations.
 
     RasterIndex provides spatial indexing capabilities for raster data by wrapping
@@ -341,11 +371,13 @@ class RasterIndex(Index):
     - **Rectilinear grids**: Uses separate 1D indexes for independent x/y axes,
       enabling more efficient slicing operations.
 
-    Parameters
-    ----------
-    indexes : Mapping[WrappedIndexCoords, WrappedIndex]
-        Dictionary mapping coordinate names to their corresponding index objects.
-        Keys are either single coordinate names or tuples for coupled coordinates.
+    RasterIndex is CRS-aware, i.e., it has a ``crs`` property that is used for
+    checking equality or compatibility with other RasterIndex instances. CRS is
+    optional.
+
+    Do not use :py:meth:`~rasterix.RasterIndex.__init__` directly. Instead use
+    :py:meth:`~rasterix.RasterIndex.from_transform` or
+    :py:func:`~rasterix.assign_index`.
 
     Attributes
     ----------
@@ -376,38 +408,95 @@ class RasterIndex(Index):
     For rectilinear grids without rotation, RasterIndex creates separate 1D indexes
     for x and y coordinates, which enables efficient slicing operations. For grids
     with rotation or skew, it uses a coupled 2D transform.
+
     """
 
-    _wrapped_indexes: dict[WrappedIndexCoords, WrappedIndex]
-    _shape: dict[str, int]
+    _index: WrappedIndex
+    _axis_independent: bool
+    _crs: CRS | None
+    _xy_shape: tuple[int, int]
+    _xy_dims: tuple[str, str]
+    _xy_coord_names: tuple[Hashable, Hashable]
+
+    def __init__(self, index: WrappedIndex, crs: CRS | Any | None = None):
+        if isinstance(index, CoordinateTransformIndex) and isinstance(index.transform, AffineTransform):
+            self._axis_independent = False
+            xtransform = cast(AffineTransform, index.transform)
+            dim_size = xtransform.dim_size
+            xy_dims = xtransform.xy_dims
+            self._xy_shape = (dim_size[xy_dims[XAXIS]], dim_size[xy_dims[YAXIS]])
+            self._xy_dims = xtransform.xy_dims
+            self._xy_coord_names = (xtransform.coord_names[XAXIS], xtransform.coord_names[YAXIS])
+        elif (
+            isinstance(index, tuple)
+            and len(index) == 2
+            and isinstance(index[XAXIS], AxisAffineTransformIndex)
+            and isinstance(index[YAXIS], AxisAffineTransformIndex)
+        ):
+            self._axis_independent = True
+            self._xy_shape = (index[XAXIS].axis_transform.size, index[YAXIS].axis_transform.size)
+            self._xy_dims = (index[XAXIS].axis_transform.dim, index[YAXIS].axis_transform.dim)
+            self._xy_coord_names = (
+                index[XAXIS].axis_transform.coord_name,
+                index[YAXIS].axis_transform.coord_name,
+            )
+        else:
+            raise ValueError(f"Could not create RasterIndex. Received invalid index {index!r}")
 
-    def __init__(self, indexes: Mapping[WrappedIndexCoords, WrappedIndex]):
-        idx_keys = list(indexes)
-        idx_vals = list(indexes.values())
+        self._index = index
 
-        # either one or the other configuration (dependent vs. independent x/y axes)
-        axis_dependent = (
-            len(indexes) == 1
-            and isinstance(idx_keys[0], tuple)
-            and isinstance(idx_vals[0], CoordinateTransformIndex)
-        )
-        axis_independent = len(indexes) in (1, 2) and all(
-            isinstance(idx, AxisAffineTransformIndex | PandasIndex) for idx in idx_vals
-        )
-        assert axis_dependent ^ axis_independent
-
-        self._wrapped_indexes = dict(indexes)
-        if axis_independent:
-            self._shape = {
-                "x": self._wrapped_indexes["x"].axis_transform.size,
-                "y": self._wrapped_indexes["y"].axis_transform.size,
-            }
+        if crs is not None:
+            crs = CRS.from_user_input(crs)
+        self._crs = crs
+
+    @property
+    def _wrapped_indexes(self) -> tuple[CoordinateTransformIndex | AxisAffineTransformIndex, ...]:
+        """Returns the wrapped index objects as a tuple."""
+        if not isinstance(self._index, tuple):
+            return (self._index,)
         else:
-            self._shape = next(iter(self._wrapped_indexes.values())).dim_size
+            return self._index
+
+    @property
+    def _xy_indexes(self) -> tuple[AxisAffineTransformIndex, AxisAffineTransformIndex]:
+        assert self._axis_independent
+        return cast(tuple[AxisAffineTransformIndex, AxisAffineTransformIndex], self._index)
+
+    @property
+    def _xxyy_index(self) -> CoordinateTransformIndex:
+        assert not self._axis_independent
+        return cast(CoordinateTransformIndex, self._index)
+
+    @property
+    def xy_shape(self) -> tuple[int, int]:
+        """Return the dimension size of the X and Y axis, respectively."""
+        return self._xy_shape
+
+    @property
+    def xy_dims(self) -> tuple[str, str]:
+        """Return the dimension name of the X and Y axis, respectively."""
+        return self._xy_dims
+
+    @property
+    def xy_coord_names(self) -> tuple[Hashable, Hashable]:
+        """Return the name of the coordinate variables representing labels on
+        the X and Y axis, respectively.
+
+        """
+        return self._xy_coord_names
 
     @classmethod
     def from_transform(
-        cls, affine: Affine, width: int, height: int, x_dim: str = "x", y_dim: str = "y"
+        cls,
+        affine: Affine,
+        *,
+        width: int,
+        height: int,
+        x_dim: str = "x",
+        y_dim: str = "y",
+        x_coord_name: str = "xc",
+        y_coord_name: str = "yc",
+        crs: CRS | Any | None = None,
     ) -> RasterIndex:
         """Create a RasterIndex from an affine transform and raster dimensions.
 
@@ -420,10 +509,17 @@ class RasterIndex(Index):
             Number of pixels in the x direction.
         height : int
             Number of pixels in the y direction.
-        x_dim : str, default "x"
+        x_dim : str, optional
             Name for the x dimension.
-        y_dim : str, default "y"
+        y_dim : str, optional
             Name for the y dimension.
+        x_coord_name : str, optional
+            Name for the x dimension. For non-rectilinear transforms only.
+        y_coord_name : str, optional
+            Name for the y dimension. For non-rectilinear transforms only.
+        crs : :class:`pyproj.crs.CRS` or any, optional
+            The coordinate reference system. Any value accepted by
+            :meth:`pyproj.crs.CRS.from_user_input`.
 
         Returns
         -------
@@ -443,24 +539,38 @@ class RasterIndex(Index):
         >>> from affine import Affine
         >>> transform = Affine(1.0, 0.0, 0.0, 0.0, -1.0, 100.0)
         >>> index = RasterIndex.from_transform(transform, width=100, height=100)
+
+        Create a non-rectilinear index:
+
+        >>> index = RasterIndex.from_transform(
+        ...     Affine.rotation(45), width=100, height=100, x_coord_name="x1", y_coord_name="x2"
+        ... )
         """
-        indexes: dict[WrappedIndexCoords, AxisAffineTransformIndex | CoordinateTransformIndex]
+        index: WrappedIndex
 
         # pixel centered coordinates
         affine = affine * Affine.translation(0.5, 0.5)
 
         if affine.is_rectilinear and affine.b == affine.d == 0:
-            x_transform = AxisAffineTransform(affine, width, "x", x_dim, is_xaxis=True)
-            y_transform = AxisAffineTransform(affine, height, "y", y_dim, is_xaxis=False)
-            indexes = {
-                "x": AxisAffineTransformIndex(x_transform),
-                "y": AxisAffineTransformIndex(y_transform),
-            }
+            x_transform = AxisAffineTransform(affine, width, x_dim, x_dim, is_xaxis=True)
+            y_transform = AxisAffineTransform(affine, height, y_dim, y_dim, is_xaxis=False)
+            index = (
+                AxisAffineTransformIndex(x_transform),
+                AxisAffineTransformIndex(y_transform),
+            )
         else:
-            xy_transform = AffineTransform(affine, width, height, x_dim=x_dim, y_dim=y_dim)
-            indexes = {("x", "y"): CoordinateTransformIndex(xy_transform)}
+            xy_transform = AffineTransform(
+                affine,
+                width,
+                height,
+                x_dim=x_dim,
+                y_dim=y_dim,
+                x_coord_name=x_coord_name,
+                y_coord_name=y_coord_name,
+            )
+            index = CoordinateTransformIndex(xy_transform)
 
-        return cls(indexes)
+        return cls(index, crs=crs)
 
     @classmethod
     def from_variables(
@@ -475,92 +585,116 @@ class RasterIndex(Index):
     def create_variables(self, variables: Mapping[Any, Variable] | None = None) -> dict[Hashable, Variable]:
         new_variables: dict[Hashable, Variable] = {}
 
-        for index in self._wrapped_indexes.values():
+        for index in self._wrapped_indexes:
             new_variables.update(index.create_variables())
 
+        if self.crs is not None:
+            xname, yname = self.xy_coord_names
+            xattrs, yattrs = self.crs.cs_to_cf()
+            if "axis" in xattrs and "axis" in yattrs:
+                # The axis order is defined by the projection
+                # So we have to figure which is which.
+                # This is an ugly hack that works for common cases.
+                if xattrs["axis"] == "Y":
+                    xattrs, yattrs = yattrs, xattrs
+                new_variables[xname].attrs = xattrs
+                new_variables[yname].attrs = yattrs
+
         return new_variables
 
+    @property
+    def crs(self) -> CRS | None:
+        """Returns the coordinate reference system (CRS) of the index as a
+        :class:`pyproj.crs.CRS` object, or ``None`` if CRS is undefined.
+        """
+        return self._crs
+
+    def _proj_set_crs(self: RasterIndex, spatial_ref: Hashable, crs: CRS) -> RasterIndex:
+        # Returns a raster index shallow copy with a replaced CRS
+        # (XProj integration via xproj.ProjIndexMixin)
+        # Note: XProj already handles the case of overriding any existing CRS
+        return RasterIndex(self._wrapped_indexes, crs=crs)
+
     def isel(self, indexers: Mapping[Any, int | slice | np.ndarray | Variable]) -> RasterIndex | None:
-        new_indexes: dict[WrappedIndexCoords, WrappedIndex] = {}
-
-        for coord_names, index in self._wrapped_indexes.items():
-            index_indexers = _filter_dim_indexers(index, indexers)
-            if not index_indexers:
-                # no selection to perform: simply propagate the index
-                # TODO: uncomment when https://github.com/pydata/xarray/issues/10063 is fixed
-                # new_indexes[coord_names] = index
-                ...
-            else:
-                new_index = index.isel(index_indexers)
-                if new_index is not None:
-                    new_indexes[coord_names] = new_index
-
-        if new_indexes:
-            # TODO: if there's only a single PandasIndex can we just return it?
-            # (maybe better to keep it wrapped if we plan to later make RasterIndex CRS-aware)
-            return RasterIndex(new_indexes)
-        else:
+        if not self._axis_independent:
+            # preserve RasterIndex is not supported in the case of coupled x/y 2D coordinates
             return None
+        else:
+            new_indexes = []
+
+            for index in self._xy_indexes:
+                dim = index.axis_transform.dim
+
+                if dim not in indexers:
+                    # simply propagate the index
+                    new_indexes.append(index)
+                else:
+                    new_index = index.isel({dim: indexers[dim]})
+                    if new_index is not None:
+                        new_indexes.append(new_index)
+
+            # TODO: if/when supported in Xarray, return PandasIndex instances for either the
+            # x or the y axis (or both) instead of returning None (drop the index)
+            if len(new_indexes) == 2:
+                return RasterIndex(tuple(new_indexes), crs=self.crs)
+            else:
+                return None
 
     def sel(self, labels: dict[Any, Any], method=None, tolerance=None) -> IndexSelResult:
-        results = []
+        if not self._axis_independent:
+            return self._xxyy_index.sel(labels, method=method, tolerance=tolerance)
+        else:
+            results: list[IndexSelResult] = []
 
-        for coord_names, index in self._wrapped_indexes.items():
-            if not isinstance(coord_names, tuple):
-                coord_names = (coord_names,)
-            index_labels = {k: v for k, v in labels.items() if k in coord_names}
-            if index_labels:
-                results.append(index.sel(index_labels, method=method, tolerance=tolerance))
+            for index in self._xy_indexes:
+                coord_name = index.axis_transform.coord_name
+                if coord_name in labels:
+                    res = index.sel({coord_name: labels[coord_name]}, method=method, tolerance=tolerance)
+                    results.append(res)
 
-        return merge_sel_results(results)
+            return merge_sel_results(results)
 
     def equals(self, other: Index, *, exclude=None) -> bool:
         if exclude is None:
             exclude = {}
+
         if not isinstance(other, RasterIndex):
             return False
-        if set(self._wrapped_indexes) != set(other._wrapped_indexes):
+        if not self._proj_crs_equals(cast(CRSAwareIndex, other), allow_none=True):
+            return False
+        if self._axis_independent != other._axis_independent:
             return False
 
-        return all(
-            index.equals(other._wrapped_indexes[k])  # type: ignore[arg-type]
-            for k, index in self._wrapped_indexes.items()
-            if k not in exclude
-        )
-
-    def to_pandas_index(self) -> pd.Index:
-        # conversion is possible only if this raster index encapsulates
-        # exactly one AxisAffineTransformIndex or a PandasIndex associated
-        # to either the x or y axis (1-dimensional) coordinate.
-        if len(self._wrapped_indexes) == 1:
-            index = next(iter(self._wrapped_indexes.values()))
-            if isinstance(index, AxisAffineTransformIndex | PandasIndex):
-                return index.to_pandas_index()
-
-        raise ValueError("Cannot convert RasterIndex to pandas.Index")
-
-    def __repr__(self):
-        items: list[str] = []
-
-        for coord_names, index in self._wrapped_indexes.items():
-            items += [repr(coord_names) + ":", textwrap.indent(repr(index), "    ")]
-
-        return "RasterIndex\n" + "\n".join(items)
+        if self._axis_independent:
+            return all(
+                idx.equals(other_idx)
+                for idx, other_idx in zip(self._xy_indexes, other._xy_indexes)
+                if idx.axis_transform.dim not in exclude
+            )
+        else:
+            return self._xxyy_index.equals(other._xxyy_index)
 
     def transform(self) -> Affine:
         """Affine transform for top-left corners."""
         return self.center_transform() * Affine.translation(-0.5, -0.5)
 
+    def as_geotransform(self, *, decimals: int | None = None) -> str:
+        """Convert the affine transform to a string suitable for saving as the GeoTransform attribute."""
+        gt = self.transform().to_gdal()
+        if decimals is not None:
+            fmt = f".{decimals}f"
+            return " ".join(f"{num:{fmt}}" for num in gt)
+        else:
+            return " ".join(map(str, gt))
+
     def center_transform(self) -> Affine:
         """Affine transform for cell centers."""
-        if len(self._wrapped_indexes) > 1:
-            x = self._wrapped_indexes["x"].axis_transform.affine
-            y = self._wrapped_indexes["y"].axis_transform.affine
-            aff = Affine(x.a, x.b, x.c, y.d, y.e, y.f)
+        if not self._axis_independent:
+            return cast(AffineTransform, self._xxyy_index.transform).affine
         else:
-            index = next(iter(self._wrapped_indexes.values()))
-            aff = index.affine
-        return aff
+            x = self._xy_indexes[XAXIS].axis_transform.affine
+            y = self._xy_indexes[YAXIS].axis_transform.affine
+            return Affine(x.a, x.b, x.c, y.d, y.e, y.f)
 
     @property
     def bbox(self) -> BoundingBox:
@@ -570,10 +704,8 @@ class RasterIndex(Index):
         -------
         BoundingBox
         """
-        return BoundingBox.from_transform(
-            shape=tuple(self._shape[k] for k in ("y", "x")),
-            transform=self.transform(),
-        )
+        yx_shape = (self.xy_shape[YAXIS], self.xy_shape[XAXIS])
+        return BoundingBox.from_transform(shape=yx_shape, transform=self.transform())
 
     @classmethod
     def concat(
@@ -581,7 +713,7 @@ class RasterIndex(Index):
         indexes: Sequence[Self],
         dim: Hashable,
         positions: Iterable[Iterable[int]] | None = None,
-    ) -> Self:
+    ) -> RasterIndex:
         if len(indexes) == 1:
             return next(iter(indexes))
 
@@ -593,7 +725,7 @@ class RasterIndex(Index):
         new_bbox = bbox_union(as_compatible_bboxes(*indexes, concat_dim=dim))
         return indexes[0]._new_with_bbox(new_bbox)
 
-    def _new_with_bbox(self, bbox: BoundingBox) -> Self:
+    def _new_with_bbox(self, bbox: BoundingBox) -> RasterIndex:
         affine = self.transform()
         new_affine, Nx, Ny = bbox_to_affine(bbox, rx=affine.a, ry=affine.e)
         # TODO: set xdim, ydim explicitly
@@ -601,8 +733,14 @@ class RasterIndex(Index):
         assert new_index.bbox == bbox
         return new_index
 
-    def join(self, other: RasterIndex, how) -> RasterIndex:
-        if set(self._wrapped_indexes.keys()) != set(other._wrapped_indexes.keys()):
+    def join(self, other: RasterIndex, how: JoinOptions = "inner") -> RasterIndex:
+        if not self._proj_crs_equals(cast(CRSAwareIndex, other), allow_none=True):
+            raise ValueError(
+                "raster indexes on objects to align do not have the same CRS\n"
+                f"first index:\n{self!r}\n\nsecond index:\n{other!r}"
+            )
+
+        if len(self._wrapped_indexes) != len(other._wrapped_indexes):
             # TODO: better error message
             raise ValueError(
                 "Alignment is only supported between RasterIndexes, when both contain compatible transforms."
@@ -619,6 +757,7 @@ class RasterIndex(Index):
         return self._new_with_bbox(new_bbox)
 
     def reindex_like(self, other: Self, method=None, tolerance=None) -> dict[Hashable, Any]:
+        x_dim, y_dim = self.xy_dims
         affine = self.transform()
         ours, theirs = as_compatible_bboxes(self, other, concat_dim=None)
         inter = bbox_intersection([ours, theirs])
@@ -631,14 +770,35 @@ class RasterIndex(Index):
         tol: float = 0.01
 
         indexers = {}
-        indexers["x"] = get_indexer(
-            theirs.left, ours.left, inter.left, inter.right, spacing=dx, tol=tol, size=other._shape["x"]
+        indexers[x_dim] = get_indexer(
+            theirs.left, ours.left, inter.left, inter.right, spacing=dx, tol=tol, size=other.xy_shape[XAXIS]
         )
-        indexers["y"] = get_indexer(
-            theirs.top, ours.top, inter.top, inter.bottom, spacing=dy, tol=tol, size=other._shape["y"]
+        indexers[y_dim] = get_indexer(
+            theirs.top, ours.top, inter.top, inter.bottom, spacing=dy, tol=tol, size=other.xy_shape[YAXIS]
         )
         return indexers
 
+    def _repr_inline_(self, max_width: int) -> str:
+        # TODO: remove when fixed in Xarray (https://github.com/pydata/xarray/pull/10415)
+        if max_width is None:
+            max_width = get_options()["display_width"]
+
+        srs = xproj.format_crs(self.crs, max_width=max_width)
+        return f"{self.__class__.__name__} (crs={srs})"
+
+    def __repr__(self) -> str:
+        srs = xproj.format_crs(self.crs)
+
+        def index_repr(idx) -> str:
+            return textwrap.indent(f"{type(idx).__name__}({idx.transform!r})", "    ")
+
+        if self._axis_independent:
+            idx_repr = "\n".join(index_repr(idx) for idx in self._xy_indexes)
+        else:
+            idx_repr = index_repr(self._xxyy_index)
+
+        return f"RasterIndex(crs={srs})\n{idx_repr}"
+
 
 def get_indexer(off, our_off, start, stop, spacing, tol, size) -> np.ndarray:
     istart = math.ceil(maybe_int((start - off) / spacing, tol))
@@ -656,7 +816,7 @@ def get_indexer(off, our_off, start, stop, spacing, tol, size) -> np.ndarray:
     return idxr
 
 
-def bbox_to_affine(bbox: BoundingBox, rx, ry) -> Affine:
+def bbox_to_affine(bbox: BoundingBox, rx, ry) -> tuple[Affine, int, int]:
     # Fraction of a pixel that can be ignored, defaults to 1/100. Bounding box of the output
     # geobox is allowed to be smaller than supplied bounding box by that amount.
     # FIXME: translate user-provided `tolerance` to `tol`
@@ -675,10 +835,10 @@ def as_compatible_bboxes(*indexes: RasterIndex, concat_dim: Hashable | None) ->
     _assert_transforms_are_compatible(*transforms)
 
     expected_off_x = (transforms[0].c,) + tuple(
-        t.c + i._shape["x"] * t.a for i, t in zip(indexes[:-1], transforms[:-1])
+        t.c + i.xy_shape[XAXIS] * t.a for i, t in zip(indexes[:-1], transforms[:-1])
     )
     expected_off_y = (transforms[0].f,) + tuple(
-        t.f + i._shape["y"] * t.e for i, t in zip(indexes[:-1], transforms[:-1])
+        t.f + i.xy_shape[YAXIS] * t.e for i, t in zip(indexes[:-1], transforms[:-1])
     )
 
     off_x = tuple(t.c for t in transforms)
@@ -688,14 +848,17 @@ def as_compatible_bboxes(*indexes: RasterIndex, concat_dim: Hashable | None) ->
         if all(o == off_x[0] for o in off_x[1:]) and all(o == off_y[0] for o in off_y[1:]):
             raise ValueError("Attempting to concatenate arrays with same transform along X or Y.")
 
-    if concat_dim == "x":
+    # note: Xarray alignment already ensures that the indexes dimensions are compatible.
+    x_dim, y_dim = indexes[0].xy_dims
+
+    if concat_dim == x_dim:
         if any(off_y[0] != o for o in off_y[1:]):
             raise ValueError("offsets must be identical in X when concatenating along Y")
         if any(a != b for a, b in zip(off_x, expected_off_x)):
             raise ValueError(
                 f"X offsets are incompatible. Provided offsets {off_x}, expected offsets: {expected_off_x}"
             )
-    elif concat_dim == "y":
+    elif concat_dim == y_dim:
         if any(off_x[0] != o for o in off_x[1:]):
             raise ValueError("offsets must be identical in X when concatenating along Y")
 

rasterix/rasterize/rasterio.py

@@ -15,7 +15,8 @@ from rasterio.features import MergeAlg
 from rasterio.features import geometry_mask as geometry_mask_rio
 from rasterio.features import rasterize as rasterize_rio
 
-from .utils import XAXIS, YAXIS, clip_to_bbox, get_affine, is_in_memory, prepare_for_dask
+from ..utils import get_affine
+from .utils import XAXIS, YAXIS, clip_to_bbox, is_in_memory, prepare_for_dask
 
 F = TypeVar("F", bound=Callable[..., Any])
 
@@ -161,7 +162,7 @@ def rasterize(
         obj = clip_to_bbox(obj, geometries, xdim=xdim, ydim=ydim)
 
     rasterize_kwargs = dict(
-        all_touched=all_touched, merge_alg=merge_alg, affine=get_affine(obj, xdim=xdim, ydim=ydim), env=env
+        all_touched=all_touched, merge_alg=merge_alg, affine=get_affine(obj, x_dim=xdim, y_dim=ydim), env=env
     )
     # FIXME: box.crs == geometries.crs
 
@@ -325,7 +326,7 @@ def geometry_mask(
         obj = clip_to_bbox(obj, geometries, xdim=xdim, ydim=ydim)
 
     geometry_mask_kwargs = dict(
-        all_touched=all_touched, affine=get_affine(obj, xdim=xdim, ydim=ydim), env=env
+        all_touched=all_touched, affine=get_affine(obj, x_dim=xdim, y_dim=ydim), env=env
     )
 
     if is_in_memory(obj=obj, geometries=geometries):

rasterix/utils.py

@@ -0,0 +1,67 @@
+import xarray as xr
+from affine import Affine
+
+
+def get_grid_mapping_var(obj: xr.Dataset | xr.DataArray) -> xr.DataArray | None:
+    grid_mapping_var = None
+    if isinstance(obj, xr.DataArray):
+        if maybe := obj.attrs.get("grid_mapping", None):
+            if maybe in obj.coords:
+                grid_mapping_var = maybe
+    else:
+        # for datasets, grab the first one for simplicity
+        for var in obj.data_vars.values():
+            if maybe := var.attrs.get("grid_mapping"):
+                if maybe in obj.coords:
+                    # make sure it exists and is not an out-of-date attribute
+                    grid_mapping_var = maybe
+                    break
+    if grid_mapping_var is None and "spatial_ref" in obj.coords:
+        # hardcode this
+        grid_mapping_var = "spatial_ref"
+    if grid_mapping_var is not None:
+        return obj[grid_mapping_var]
+    return None
+
+
+def get_affine(
+    obj: xr.Dataset | xr.DataArray, *, x_dim="x", y_dim="y", clear_transform: bool = False
+) -> Affine:
+    """
+    Grabs an affine transform from an Xarray object.
+
+    This method will first look for the ``"GeoTransform"`` attribute on a variable named
+    ``"spatial_ref"``. If not, it will auto-guess the transform from the provided ``x_dim``,
+    and ``y_dim``.
+
+    Parameters
+    ----------
+    obj: xr.DataArray or xr.Dataset
+    x_dim: str, optional
+        Name of the X dimension coordinate variable.
+    y_dim: str, optional
+        Name of the Y dimension coordinate variable.
+    clear_transform: bool
+       Whether to delete the ``GeoTransform`` attribute if detected.
+
+    Returns
+    -------
+    affine: Affine
+    """
+    grid_mapping_var = get_grid_mapping_var(obj)
+    if grid_mapping_var is not None and (transform := grid_mapping_var.attrs.get("GeoTransform")):
+        if clear_transform:
+            del grid_mapping_var.attrs["GeoTransform"]
+        return Affine.from_gdal(*map(float, transform.split(" ")))
+    else:
+        x = obj.coords[x_dim]
+        y = obj.coords[y_dim]
+        if x.ndim != 1:
+            raise ValueError(f"Coordinate variable {x_dim=!r} must be 1D.")
+        if y.ndim != 1:
+            raise ValueError(f"Coordinate variable {y_dim=!r} must be 1D.")
+        dx = (x[1] - x[0]).item()
+        dy = (y[1] - y[0]).item()
+        return Affine.translation(
+            x[0].item() - dx / 2, (y[0] if dy < 0 else y[-1]).item() - dy / 2
+        ) * Affine.scale(dx, dy)

{rasterix-0.1a3.dist-info → rasterix-0.1a4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: rasterix
-Version: 0.1a3
+Version: 0.1a4
 Summary: Raster extensions for Xarray
 License-Expression: Apache-2.0
 License-File: LICENSE
@@ -17,6 +17,7 @@ Requires-Dist: affine
 Requires-Dist: numpy>=2
 Requires-Dist: pandas>=2
 Requires-Dist: xarray>=2025
+Requires-Dist: xproj>=0.2.0
 Provides-Extra: dask
 Requires-Dist: dask-geopandas; extra == 'dask'
 Provides-Extra: docs
@@ -52,57 +53,30 @@ Requires-Dist: hypothesis; extra == 'test'
 Requires-Dist: netcdf4; extra == 'test'
 Requires-Dist: pooch; extra == 'test'
 Requires-Dist: rasterio; extra == 'test'
-Requires-Dist: rioxarray; extra == 'test'
 Requires-Dist: sparse; extra == 'test'
 Description-Content-Type: text/markdown
 
 # rasterix: Raster tricks for Xarray
 
-[![GitHub Workflow CI Status](https://img.shields.io/github/actions/workflow/status/dcherian/rasterix/test.yml?branch=main&logo=github&style=flat)](https://github.com/dcherian/rasterix/actions)
+[![GitHub Workflow CI Status](https://img.shields.io/github/actions/workflow/status/xarray-contrib/rasterix/test.yml?branch=main&logo=github&style=flat)](https://github.com/xarray-contrib/rasterix/actions)
 [![Documentation Status](https://readthedocs.org/projects/rasterix/badge/?version=latest)](https://rasterix.readthedocs.io/en/latest/?badge=latest)
 [![PyPI](https://img.shields.io/pypi/v/rasterix.svg?style=flat)](https://pypi.org/project/rasterix/)
 [![Conda-forge](https://img.shields.io/conda/vn/conda-forge/rasterix.svg?style=flat)](https://anaconda.org/conda-forge/rasterix)
 
-<img src="rasterix.png" width="300">
+<img src="_static/rasterix.png" width="300">
 
 This WIP project contains tools to make it easier to analyze raster data with Xarray.
-
-The intent is to provide reusable building blocks for the many sub-ecosystems around: e.g. rioxarray, odc-geo, etc.
-
-## Contents
-
 It currently has two pieces.
 
-### 1. RasterIndex
-
-See `src/ rasterix/raster_index.py` and `notebooks/raster_index.ipynb` for a brief demo.
-
-### 2. Dask-aware rasterization wrappers
+1. `RasterIndex` for indexing using the affine transform recorded in GeoTIFFs.
+1. Dask-aware rasterization wrappers around `exactextract`, `rasterio.features.rasterize`, and `rasterio.features.geometry_mask`.
 
-See `src/rasterix/rasterize.py` for dask-aware wrappers around [`exactextract`](https://github.com/dcherian/rasterix/blob/ec3f51e60e25aa312e6f48c4b22f91bec70413ed/rasterize.py#L165), [`rasterio.features.rasterize`](https://github.com/dcherian/rasterix/blob/ec3f51e60e25aa312e6f48c4b22f91bec70413ed/rasterize.py#L307), and [`rasterio.features.geometry_mask`](https://github.com/dcherian/rasterix/blob/ec3f51e60e25aa312e6f48c4b22f91bec70413ed/rasterize.py#L472).
-
-This code is likely to move elsewhere!
+Our intent is to provide reusable building blocks for the many sub-ecosystems around: e.g. `rioxarray`, `odc.geo`, etc.
 
 ## Installing
 
-### PyPI
-
 `rasterix` alpha releases are available on pypi
 
 ```
 pip install rasterix
 ```
-
-## Developing
-
-1. Clone the repo
-   ```
-   git remote add upstream git@github.com:dcherian/rasterix.git
-   cd rasterix
-   ```
-1. [Install hatch](https://hatch.pypa.io/1.12/install/)
-1. Run the tests
-   ```
-   hatch env run --env test.py3.13 run-pytest  # Run the tests without coverage reports
-   hatch env run --env test.py3.13 run-coverage-html   # Run the tests with an html coverage report
-   ```

rasterix-0.1a4.dist-info/RECORD

@@ -0,0 +1,14 @@
+rasterix/__init__.py,sha256=iSC54A4b-7zgxtKV0CUok0O2-ApGtzVA-6hWG3a5AAA,283
+rasterix/_version.py,sha256=Rlam12WbDAQ2TzSYxzdnw4SLFERCr6sX8OCMwKNxEX4,514
+rasterix/odc_compat.py,sha256=MxjctCH0zV6VSSQymatJreHYUnUo3qkB1oIKbOxEK4A,19099
+rasterix/raster_index.py,sha256=uwEDjjmnESWpyqsJwAnZvpWrjnKDeg0m1qNNH9mYTWY,32808
+rasterix/rioxarray_compat.py,sha256=o32UBkWBpNhmC35ar0Ozn2QszpWdPvdDDXif3Mq2ZKg,12428
+rasterix/utils.py,sha256=y2N-3d9-D8tBKjeZXv4vcjY_fs_dbbx5SfBBkQibLNs,2441
+rasterix/rasterize/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+rasterix/rasterize/exact.py,sha256=UPB9zXpj7ZgkDy-x4vcNwY3J8DTrOTkRcIBMkhPQ04s,10558
+rasterix/rasterize/rasterio.py,sha256=gtdQxM9CiosalYk3rLKGkp0UmhZVO5parAVsR7IINoE,12883
+rasterix/rasterize/utils.py,sha256=iIuyTdigpFnts_-bUfas6FMsgix5lMqMQTxQaEBlyXM,3521
+rasterix-0.1a4.dist-info/METADATA,sha256=K5l2PDTNLvmdA6TVGvy2tg5-YRevdmhV5mYX_Os3Rbw,3402
+rasterix-0.1a4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+rasterix-0.1a4.dist-info/licenses/LICENSE,sha256=QFnsASMx8_yBNbrS7GVOhJ5CglGsLMj83Rn61uWyMs8,10265
+rasterix-0.1a4.dist-info/RECORD,,

rasterix-0.1a3.dist-info/RECORD

@@ -1,13 +0,0 @@
-rasterix/__init__.py,sha256=iSC54A4b-7zgxtKV0CUok0O2-ApGtzVA-6hWG3a5AAA,283
-rasterix/_version.py,sha256=hJd92hgvVD0aLSzXDQZBYe_xFmMQly9wq_1h2CuGX_g,514
-rasterix/odc_compat.py,sha256=MxjctCH0zV6VSSQymatJreHYUnUo3qkB1oIKbOxEK4A,19099
-rasterix/raster_index.py,sha256=yCq08-BvgEWe5ZbtknkMsXqKkZVqGyw85Otcx3_atf4,26130
-rasterix/rioxarray_compat.py,sha256=o32UBkWBpNhmC35ar0Ozn2QszpWdPvdDDXif3Mq2ZKg,12428
-rasterix/rasterize/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-rasterix/rasterize/exact.py,sha256=UPB9zXpj7ZgkDy-x4vcNwY3J8DTrOTkRcIBMkhPQ04s,10558
-rasterix/rasterize/rasterio.py,sha256=TdnDMGfsN3npP_LPEPXH2LV25AJEkmEuTOpke5GN8_8,12860
-rasterix/rasterize/utils.py,sha256=iIuyTdigpFnts_-bUfas6FMsgix5lMqMQTxQaEBlyXM,3521
-rasterix-0.1a3.dist-info/METADATA,sha256=05YXAANNx2f50SqN_bXO_N7gb1YR29S3GIMEmi1mEXQ,4262
-rasterix-0.1a3.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-rasterix-0.1a3.dist-info/licenses/LICENSE,sha256=QFnsASMx8_yBNbrS7GVOhJ5CglGsLMj83Rn61uWyMs8,10265
-rasterix-0.1a3.dist-info/RECORD,,