variable zonal stats using exactextract and automatic detection of method

Author: martinfleis

xvec/accessor.py

@@ -16,7 +16,9 @@
 from .plotting import _plot
 from .utils import transform_geom
 from .zonal import (
+    _get_method,
     _variable_zonal,
+    _variable_zonal_exactextract,
     _zonal_stats_exactextract,
     _zonal_stats_iterative,
     _zonal_stats_rasterize,
@@ -983,7 +985,7 @@ def zonal_stats(
         stats: str | Callable | Sequence[str | Callable | tuple] = "mean",
         name: str = "geometry",
         index: bool | None = None,
-        method: str = "rasterize",
+        method: str | None = None,
         all_touched: bool = False,
         n_jobs: int = -1,
         nodata: Any = None,
@@ -1005,10 +1007,10 @@ def zonal_stats(
 
         Parameters
         ----------
-        geometry : Sequence[shapely.Geometry]
+        geometry : Sequence[shapely.Geometry] | xr.DataArray
             An arrray-like (1-D) of shapely geometries, like a numpy array or
-            :class:`geopandas.GeoSeries`. Polygon and LineString geometry types are
-            supported.
+            :class:`geopandas.GeoSeries` or xr.DataArray holding variable geometry.
+            Polygon and LineString geometry types are supported.
         x_coords : Hashable
             name of the coordinates containing ``x`` coordinates (i.e. the first value
             in the coordinate pair encoding the vertex of the polygon)
@@ -1045,7 +1047,8 @@ def zonal_stats(
 
             ``"rasterize"``
                 uses :func:`rasterio.features.rasterize` and is faster, but can lead to
-                loss of information in case of small polygons or lines.
+                loss of information in case of small polygons or lines. Not supported
+                for zonal stats using variable geometry (n-D array of geometry).
 
             ``"iterate"``
                 iterates over geometries and uses
@@ -1057,7 +1060,8 @@ def zonal_stats(
                 that is covered by the polygon and uses
                 :func:`exactextract.exact_extract`.
 
-            The default is ``"rasterize"``.
+            The default is selected based on the availability of engines in the order
+            of priority 1. ``"exactextract"``, 2. ``"rasterize"`` 3. ``"iterate"``.
         all_touched : bool, optional
             If True, all pixels touched by geometries will be considered. If False, only
             pixels whose center is within the polygon or that are selected by
@@ -1156,17 +1160,40 @@ def zonal_stats(
         --------
         extract_points : extraction of values for the raster object for points
         """
+
         if isinstance(geometry, xr.DataArray) and len(geometry.dims) > 1:
-            return _variable_zonal(
-                self,
-                variable_geometry=geometry,
-                x_coords=x_coords,
-                y_coords=y_coords,
-                stats=stats,
-                all_touched=all_touched,
-                nodata=nodata,
+            if method is None:
+                method = _get_method(variable=True)
+
+            if method == "iterate":
+                return _variable_zonal(
+                    self,
+                    variable_geometry=geometry,
+                    x_coords=x_coords,
+                    y_coords=y_coords,
+                    stats=stats,
+                    all_touched=all_touched,
+                    n_jobs=n_jobs,
+                    nodata=nodata,
+                )
+            if method == "exactextract":
+                return _variable_zonal_exactextract(
+                    self,
+                    geometry=geometry,
+                    x_coords=x_coords,
+                    y_coords=y_coords,
+                    stats=stats,
+                    nodata=nodata,
+                    strategy=strategy,
+                )
+            raise ValueError(
+                f"Method '{method}' is not supported for zonal statistics based on "
+                "variable geometry. Use one of `exactextract` or `iterate`."
             )
 
+        if method is None:
+            method = _get_method(variable=False)
+
         if method == "rasterize":
             result = _zonal_stats_rasterize(
                 self,

xvec/tests/test_zonal_stats.py

@@ -30,7 +30,7 @@ def glaciers():
     return glaciers, sentinel_2
 
 
-@pytest.mark.parametrize("method", ["rasterize", "iterate", "exactextract"])
+@pytest.mark.parametrize("method", [None, "rasterize", "iterate", "exactextract"])
 def test_structure(method):
     da = xr.DataArray(
         np.ones((10, 10, 5)),
@@ -45,7 +45,7 @@ def test_structure(method):
     polygon2 = shapely.geometry.Polygon([(6, 22), (9, 22), (9, 29), (6, 26)])
     polygons = gpd.GeoSeries([polygon1, polygon2], crs="EPSG:4326")
 
-    if method == "exactextract":
+    if method in ["exactextract", None]:
         expected = xr.DataArray(
             np.array([[12.0] * 5, [16.5] * 5]),
             coords={
@@ -117,13 +117,13 @@ def test_match():
     xr.testing.assert_allclose(rasterize, iterate)
 
 
-@pytest.mark.parametrize("method", ["rasterize", "iterate", "exactextract"])
+@pytest.mark.parametrize("method", [None, "rasterize", "iterate", "exactextract"])
 def test_dataset(method):
     ds = xr.tutorial.open_dataset("eraint_uvz")
     world = gpd.read_file(geodatasets.get_path("naturalearth land"))
     result = ds.xvec.zonal_stats(world.geometry, "longitude", "latitude", method=method)
 
-    if method == "exactextract":
+    if method in ["exactextract", None]:
         xr.testing.assert_allclose(
             xr.Dataset(
                 {
@@ -147,7 +147,7 @@ def test_dataset(method):
         )
 
 
-@pytest.mark.parametrize("method", ["rasterize", "iterate", "exactextract"])
+@pytest.mark.parametrize("method", [None, "rasterize", "iterate", "exactextract"])
 def test_dataarray(method):
     ds = xr.tutorial.open_dataset("eraint_uvz")
     world = gpd.read_file(geodatasets.get_path("naturalearth land"))
@@ -157,13 +157,13 @@ def test_dataarray(method):
 
     assert result.shape == (127, 2, 3)
     assert result.dims == ("geometry", "month", "level")
-    if method == "exactextract":
+    if method in ["exactextract", None]:
         assert result.mean() == pytest.approx(61625.53438858)
     else:
         assert result.mean() == pytest.approx(61367.76185577)
 
 
-@pytest.mark.parametrize("method", ["rasterize", "iterate", "exactextract"])
+@pytest.mark.parametrize("method", [None, "rasterize", "iterate", "exactextract"])
 def test_stat(method):
     ds = xr.tutorial.open_dataset("eraint_uvz")
     world = gpd.read_file(geodatasets.get_path("naturalearth land"))
@@ -174,7 +174,7 @@ def test_stat(method):
     median_ = ds.z.xvec.zonal_stats(
         world.geometry, "longitude", "latitude", method=method, stats="median"
     )
-    if method == "exactextract":
+    if method in ["exactextract", None]:
         quantile_ = ds.z.xvec.zonal_stats(
             world.geometry,
             "longitude",
@@ -192,7 +192,7 @@ def test_stat(method):
             q=0.2,
         )
 
-    if method == "exactextract":
+    if method in ["exactextract", None]:
         assert mean_.mean() == pytest.approx(61625.53438858)
         assert median_.mean() == pytest.approx(61628.67168691)
         assert quantile_.mean() == pytest.approx(61540.75632235)
@@ -308,11 +308,11 @@ def test_callable(method):
     xr.testing.assert_identical(da_agg, da_std)
 
 
-@pytest.mark.parametrize("method", ["rasterize", "iterate", "exactextract"])
+@pytest.mark.parametrize("method", [None, "rasterize", "iterate", "exactextract"])
 def test_multiple(method):
     ds = xr.tutorial.open_dataset("eraint_uvz")
     world = gpd.read_file(geodatasets.get_path("naturalearth land"))
-    if method == "exactextract":
+    if method in ["exactextract", None]:
         result = ds.xvec.zonal_stats(
             world.geometry[:10].boundary,
             "longitude",
@@ -366,7 +366,7 @@ def test_multiple(method):
         ).all()
 
 
-@pytest.mark.parametrize("method", ["rasterize", "iterate", "exactextract"])
+@pytest.mark.parametrize("method", [None, "rasterize", "iterate", "exactextract"])
 def test_invalid(method):
     ds = xr.tutorial.open_dataset("eraint_uvz")
     world = gpd.read_file(geodatasets.get_path("naturalearth land"))
@@ -394,7 +394,25 @@ def test_invalid(method):
         )
 
 
-def test_variable_geometry_multiple(glaciers):
+@pytest.mark.parametrize("method", [None, "iterate", "exactextract"])
+def test_variable_geometry_multiple(glaciers, method):
+    da, sentinel_2 = glaciers
+
+    result = sentinel_2.xvec.zonal_stats(
+        da.geometry,
+        x_coords="x",
+        y_coords="y",
+        stats=[
+            "mean",
+            "sum",
+        ],
+        method=method,
+    )
+
+    assert result.sizes == {"year": 3, "name": 5, "zonal_statistics": 2, "band": 11}
+
+
+def test_variable_geometry_iterate_custom(glaciers):
     da, sentinel_2 = glaciers
 
     result = sentinel_2.xvec.zonal_stats(
@@ -407,24 +425,31 @@ def test_variable_geometry_multiple(glaciers):
             ("numpymean", np.nanmean),
             np.nanmean,
         ],
+        method="iterate",
     )
 
     assert result.sizes == {"year": 3, "name": 5, "zonal_statistics": 4, "band": 11}
-    assert result.statistics.mean() == 17067828
+    assert result.mean() == 17067828
 
 
-def test_variable_geometry_single(glaciers):
+@pytest.mark.parametrize("method", [None, "iterate", "exactextract"])
+def test_variable_geometry_single(glaciers, method):
     da, sentinel_2 = glaciers
 
     result = sentinel_2.xvec.zonal_stats(
         da.geometry,
         x_coords="x",
         y_coords="y",
         stats="mean",
+        method=method,
     )
 
     assert result.sizes == {"year": 3, "name": 5, "band": 11}
-    assert result.statistics.mean() == 13168.585
+
+    if method in ("exactextract", None):
+        assert result.mean() == pytest.approx(13168.076)
+    else:
+        assert result.mean() == 13168.585
 
 
 def test_exactextract_strategy():
@@ -458,7 +483,7 @@ def test_exactextract_strategy():
         )
 
 
-@pytest.mark.parametrize("method", ["rasterize", "iterate", "exactextract"])
+@pytest.mark.parametrize("method", [None, "rasterize", "iterate", "exactextract"])
 def test_nodata(method):
     ds = xr.tutorial.open_dataset("eraint_uvz")
     world = gpd.read_file(geodatasets.get_path("naturalearth land"))