Remove rioxarray dependency

Author: dcherian

pyproject.toml

@@ -54,7 +54,6 @@ docs = [
   "pooch",
   "dask-geopandas",
   "rasterio",
-  "rioxarray",
   "exactextract",
   "sparse",
   "netCDF4",

src/rasterix/raster_index.py

@@ -21,6 +21,7 @@
 
 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")
 
@@ -35,6 +36,10 @@
 def assign_index(obj: T_Xarray, *, x_dim: str | None = None, y_dim: str | None = None) -> 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 ``"spatial_ref"`` variable. If not present, then the affine is determine from 1D coordinate
+    variables named ``x_dim`` and ``y_dim`` provided to this function.
+
     Parameters
     ----------
     obj : xarray.DataArray or xarray.Dataset
@@ -52,19 +57,19 @@ def assign_index(obj: T_Xarray, *, x_dim: str | None = None, y_dim: str | None =
     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)
+
     index = RasterIndex.from_transform(
-        obj.rio.transform(), width=obj.sizes[x_dim], height=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
     )
     coords = Coordinates.from_xindex(index)
     return obj.assign_coords(coords)

src/rasterix/rasterize/rasterio.py

@@ -15,7 +15,8 @@
 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):

src/rasterix/utils.py

@@ -0,0 +1,39 @@
+import xarray as xr
+from affine import Affine
+
+
+def get_affine(obj: xr.Dataset | xr.DataArray, *, x_dim="x", y_dim="y") -> 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.
+
+    Returns
+    -------
+    affine: Affine
+    """
+    spatial_ref = obj.coords.get("spatial_ref")
+    if spatial_ref is not None and "GeoTransform" in spatial_ref.attrs:
+        return Affine.from_gdal(*map(float, spatial_ref.attrs["GeoTransform"].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)

tests/test_exact.py

@@ -9,14 +9,15 @@
 import sparse
 import xarray as xr
 import xarray.testing as xrt
+import xproj  # noqa
 from exactextract import exact_extract
 from hypothesis import example, given, settings
 from xarray.tests import raise_if_dask_computes
 
 from rasterix.rasterize.exact import CoverageWeights, coverage, xy_to_raster_source
 
 dataset = xr.tutorial.open_dataset("eraint_uvz").rename({"latitude": "y", "longitude": "x"})
-dataset = dataset.rio.write_crs("epsg:4326")
+dataset = dataset.proj.assign_crs(spatial_ref="epsg:4326")
 world = gpd.read_file(geodatasets.get_path("naturalearth land"))
 XSIZE = dataset.x.size
 YSIZE = dataset.y.size

tests/test_rasterize.py

@@ -4,6 +4,7 @@
 import pytest
 import rioxarray  # noqa
 import xarray as xr
+import xproj  # noqa
 from xarray.tests import raise_if_dask_computes
 
 from rasterix.rasterize.rasterio import geometry_mask, rasterize
@@ -20,7 +21,7 @@ def test_rasterize(clip):
         snapshot = snapshot.sel(latitude=slice(83.25, None))
 
     ds = xr.tutorial.open_dataset("eraint_uvz")
-    ds = ds.rio.write_crs("epsg:4326")
+    ds = ds.proj.assign_crs(spatial_ref="epsg:4326")
     world = gpd.read_file(geodatasets.get_path("naturalearth land"))
 
     kwargs = dict(xdim="longitude", ydim="latitude", clip=clip)
@@ -54,7 +55,7 @@ def test_geometry_mask(clip, invert):
         snapshot = ~snapshot
 
     ds = xr.tutorial.open_dataset("eraint_uvz")
-    ds = ds.rio.write_crs("epsg:4326")
+    ds = ds.proj.assign_crs(spatial_ref="epsg:4326")
     world = gpd.read_file(geodatasets.get_path("naturalearth land"))
 
     kwargs = dict(xdim="longitude", ydim="latitude", clip=clip, invert=invert)