Support wraparound indexing

Author: dcherian

src/rasterix/raster_index.py

@@ -303,6 +303,8 @@ def isel(  # type: ignore[override]
             # return PandasIndex(values, new_dim, coord_dtype=values.dtype)
 
     def sel(self, labels, method=None, tolerance=None):
+        # CoordinateTransformIndex only supports "nearest"
+        method = method or "nearest"
         coord_name = self.axis_transform.coord_name
         label = labels[coord_name]
 
@@ -535,8 +537,8 @@ def from_transform(
         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)
+            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),

src/rasterix/rioxarray_compat.py

@@ -204,6 +204,8 @@ def guess_dims(obj: T_Xarray) -> tuple[str, str]:
         x_dim = "longitude"
         y_dim = "latitude"
     else:
+        x_dim = None
+        y_dim = None
         # look for coordinates with CF attributes
         for coord in obj.coords:
             # make sure to only look in 1D coordinates
@@ -213,12 +215,18 @@ def guess_dims(obj: T_Xarray) -> tuple[str, str]:
             if (obj.coords[coord].attrs.get("axis", "").upper() == "X") or (
                 obj.coords[coord].attrs.get("standard_name", "").lower()
                 in ("longitude", "projection_x_coordinate")
+                or obj.coords[coord].attrs.get("units", "").lower() in ("degrees_east",)
             ):
                 x_dim = coord
             elif (obj.coords[coord].attrs.get("axis", "").upper() == "Y") or (
                 obj.coords[coord].attrs.get("standard_name", "").lower()
                 in ("latitude", "projection_y_coordinate")
+                or obj.coords[coord].attrs.get("units", "").lower() in ("degrees_north",)
             ):
                 y_dim = coord
 
+    if not x_dim or not y_dim:
+        raise ValueError(
+            "Could not guess names of x, y coordinate variables. Please explicitly pass `x_dim` and `y_dim`."
+        )
     return x_dim, y_dim

tests/test_raster_index.py

@@ -6,7 +6,7 @@
 import rioxarray  # noqa
 import xarray as xr
 from affine import Affine
-from xarray.testing import assert_identical
+from xarray.testing import assert_equal, assert_identical
 
 from rasterix import RasterIndex, assign_index
 from rasterix.utils import get_grid_mapping_var
@@ -389,3 +389,29 @@ def test_repr() -> None:
 
     index3 = RasterIndex.from_transform(Affine.identity(), width=12, height=10, crs="epsg:31370")
     assert repr(index3).startswith("RasterIndex(crs=EPSG:31370)")
+
+
+def test_assign_index_cant_guess_error():
+    ds = xr.Dataset(
+        {"sst": (("time", "lat", "lon"), np.ones((1, 89, 180)))},
+        coords={"lat": np.arange(88, -89, -2), "lon": np.arange(0, 360, 2)},
+    )
+    with pytest.raises(ValueError, match="guess"):
+        assign_index(ds)
+
+
+def test_wraparound_indexing_longitude():
+    ds = xr.Dataset(
+        {"sst": (("time", "lat", "lon"), np.random.random((1, 89, 180)))},
+        coords={
+            "lat": ("lat", np.arange(88, -89, -2), {"axis": "Y"}),
+            "lon": ("lon", np.arange(0, 360, 2), {"axis": "X"}),
+        },
+    )
+    indexed = ds.pipe(assign_index)
+    # We lose existing attrs when calling ``assign_index``.
+    assert_equal(ds.sel(lon=[220, 240]), indexed.sel(lon=[-140, -120]))
+    assert_equal(ds.sel(lon=220), indexed.sel(lon=-140))
+    assert_equal(ds.sel(lon=slice(220, 240)), indexed.sel(lon=slice(-140, -120)))
+    assert_equal(ds.sel(lon=slice(240, 220)), indexed.sel(lon=slice(-120, -140)))
+    # assert_equal(ds.sel(lon=[220]), indexed.sel(lon=[-140]))  # FIXME