[WIP] test wraparound indexing

Closes #26

Author: dcherian

src/rasterix/raster_index.py

@@ -288,6 +288,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]
 
@@ -520,8 +522,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/rasterize/utils.py

@@ -6,7 +6,6 @@
 import geopandas as gpd
 import numpy as np
 import xarray as xr
-from affine import Affine
 
 if TYPE_CHECKING:
     import dask.array
@@ -34,20 +33,6 @@ def clip_to_bbox(
     return obj
 
 
-def get_affine(obj: xr.Dataset | xr.DataArray, *, xdim="x", ydim="y") -> Affine:
-    spatial_ref = obj.coords["spatial_ref"]
-    if "GeoTransform" in spatial_ref.attrs:
-        return Affine.from_gdal(*map(float, spatial_ref.attrs["GeoTransform"].split(" ")))
-    else:
-        x = obj.coords[xdim]
-        y = obj.coords[ydim]
-        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)
-
-
 def is_in_memory(*, obj, geometries) -> bool:
     return not obj.chunks and isinstance(geometries, gpd.GeoDataFrame)
 

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
 
@@ -363,3 +363,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