ENH: add zonal_stats

.github/workflows/release.yaml

@@ -11,10 +11,10 @@ jobs:
     runs-on: ubuntu-18.04
 
     steps:
-      - uses: actions/checkout@master
+      - uses: actions/checkout@v4
 
       - name: Set up Python
-        uses: actions/setup-python@v2
+        uses: actions/setup-python@v4
         with:
           python-version: "3.x"
 
@@ -24,7 +24,7 @@ jobs:
           python -m build
 
       - name: Publish distribution to PyPI
-        uses: pypa/gh-action-pypi-publish@master
+        uses: pypa/gh-action-pypi-publish@v1
         with:
           user: __token__
           password: ${{ secrets.pypi_password }}

.github/workflows/tests.yaml

@@ -17,12 +17,14 @@ jobs:
     strategy:
       matrix:
         os: [macos-latest, ubuntu-latest, windows-latest]
-        environment-file: [ci/latest.yaml]
+        environment-file: [ci/312.yaml]
         include:
           - environment-file: ci/39.yaml
             os: ubuntu-latest
           - environment-file: ci/310.yaml
             os: ubuntu-latest
+          - environment-file: ci/311.yaml
+            os: ubuntu-latest
           - environment-file: ci/dev.yaml
             os: ubuntu-latest
     defaults:
@@ -31,14 +33,13 @@ jobs:
 
     steps:
       - name: checkout repo
-        uses: actions/checkout@v2
+        uses: actions/checkout@v4
 
       - name: setup micromamba
-        uses: mamba-org/provision-with-micromamba@main
+        uses: mamba-org/setup-micromamba@v1
         with:
           environment-file: ${{ matrix.environment-file }}
           micromamba-version: "latest"
-          channel-priority: "flexible"
 
       - name: Install xvec
         run: pip install .

.gitignore

@@ -144,3 +144,5 @@ doc/source/generated
 .ruff_cache
 doc/source/cube.joblib.compressed
 doc/source/cube.pickle
+
+cache/

.pre-commit-config.yaml

@@ -1,14 +1,10 @@
 files: 'xvec\/'
 repos:
-  - repo: https://github.com/psf/black
-    rev: 23.7.0
-    hooks:
-      - id: black
-        language_version: python3
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: "v0.0.280"
+    rev: v0.1.2
     hooks:
       - id: ruff
+      - id: ruff-format
 
 ci:
   autofix_prs: false

ci/310.yaml

@@ -4,12 +4,19 @@ channels:
 dependencies:
   - python=3.10
   # required
-  - shapely=2
+  - shapely >=2
   - xarray
+  - rioxarray
+  - joblib
+  - rasterio
+  - tqdm
   - pyproj
   # testing
   - pytest
   - pytest-cov
   - pytest-xdist
   - pytest-reportlog
   - geopandas-base
+  - geodatasets
+  - pyogrio
+

ci/latest.yaml → ci/311.yaml

@@ -2,14 +2,21 @@ name: xvec
 channels:
   - conda-forge
 dependencies:
-  - python
+  - python=3.11
   # required
-  - shapely=2
+  - shapely >=2
   - xarray
+  - rioxarray
+  - joblib
+  - rasterio
+  - tqdm
   - pyproj
   # testing
   - pytest
   - pytest-cov
   - pytest-xdist
   - pytest-reportlog
   - geopandas-base
+  - geodatasets
+  - pyogrio
+

ci/312.yaml

@@ -0,0 +1,22 @@
+name: xvec
+channels:
+  - conda-forge
+dependencies:
+  - python=3.12
+  # required
+  - shapely >=2
+  - xarray
+  - rioxarray
+  - joblib
+  - rasterio
+  - tqdm
+  - pyproj
+  # testing
+  - pytest
+  - pytest-cov
+  - pytest-xdist
+  - pytest-reportlog
+  - geopandas-base
+  - geodatasets
+  - pyogrio
+

ci/39.yaml

@@ -4,12 +4,18 @@ channels:
 dependencies:
   - python=3.9
   # required
-  - shapely=2
+  - shapely >=2
   - xarray
+  - rioxarray
+  - joblib
+  - rasterio
+  - tqdm
   - pyproj
   # testing
   - pytest
   - pytest-cov
   - pytest-xdist
   - pytest-reportlog
   - geopandas-base
+  - geodatasets
+  - pyogrio

ci/dev.yaml

@@ -7,12 +7,18 @@ dependencies:
   - cython
   - geos
   - proj
+  - rioxarray
+  - joblib
+  - rasterio
+  - tqdm
   # testing
   - pytest
   - pytest-cov
   - pytest-xdist
   - pytest-reportlog
   - geopandas-base
+  - geodatasets
+  - pyogrio
   - pip
   - pip:
       - git+https://github.com/shapely/shapely.git@main

doc/source/api.rst

@@ -55,6 +55,7 @@ Methods
     Dataset.xvec.to_geodataframe
     Dataset.xvec.to_geopandas
     Dataset.xvec.extract_points
+    Dataset.xvec.zonal_stats
 
 
 DataArray.xvec
@@ -89,4 +90,5 @@ Methods
     DataArray.xvec.query
     DataArray.xvec.to_geodataframe
     DataArray.xvec.to_geopandas
-    DataArray.xvec.extract_points
+    DataArray.xvec.extract_points
+    DataArray.xvec.zonal_stats

doc/source/conf.py

@@ -44,6 +44,8 @@
     "xarray": ("https://docs.xarray.dev/en/latest/", None),
     "geopandas": ("https://geopandas.org/en/latest", None),
     "pandas": ("https://pandas.pydata.org/docs", None),
+    "rasterio": ("https://rasterio.readthedocs.io/en/latest/", None),
+
 }
 
 # -- Options for HTML output -------------------------------------------------

pyproject.toml

@@ -56,3 +56,4 @@ line-length = 88
 select = ["E", "F", "W", "I", "UP", "B", "A", "C4", "Q"]
 exclude = ["doc"]
 target-version = "py39"
+ignore = ['E501', 'Q000', 'Q001', 'Q002', 'Q003', 'W191']

xvec/accessor.py

@@ -11,6 +11,7 @@
 from pyproj import CRS, Transformer
 
 from .index import GeometryIndex
+from .zonal import _zonal_stats_iterative, _zonal_stats_rasterize
 
 
 @xr.register_dataarray_accessor("xvec")
@@ -918,6 +919,119 @@ def to_geodataframe(
         )
         return df
 
+    def zonal_stats(
+        self,
+        polygons: Sequence[shapely.Geometry],
+        x_coords: Hashable,
+        y_coords: Hashable,
+        stats: str = "mean",
+        name: Hashable = "geometry",
+        index: bool = None,
+        method: str = "rasterize",
+        all_touched: bool = False,
+        n_jobs: int = -1,
+        **kwargs,
+    ):
+        """Extract the values from a dataset indexed by a set of geometries
+
+        The CRS of the raster and that of polygons need to be equal.
+        Xvec does not verify their equality.
+
+        Parameters
+        ----------
+        polygons : Sequence[shapely.Geometry]
+            An arrray-like (1-D) of shapely geometries, like a numpy array or
+            :class:`geopandas.GeoSeries`.
+        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)
+        y_coords : Hashable
+            name of the coordinates containing ``y`` coordinates (i.e. the second value
+            in the coordinate pair encoding the vertex of the polygon)
+        stats : string
+            Spatial aggregation statistic method, by default "mean". It supports the
+            following statistcs: ['mean', 'median', 'min', 'max', 'sum']
+        name : Hashable, optional
+            Name of the dimension that will hold the ``polygons``, by default "geometry"
+        index : bool, optional
+            If `polygons` is a GeoSeries, ``index=True`` will attach its index as another
+            coordinate to the geometry dimension in the resulting object. If
+            ``index=None``, the index will be stored if the `polygons.index` is a named
+            or non-default index. If ``index=False``, it will never be stored. This is
+            useful as an attribute link between the resulting array and the GeoPandas
+            object from which the polygons are sourced.
+        method : str, optional
+            The method of data extraction. The default is ``"rasterize"``, which uses
+            :func:`rasterio.features.rasterize` and is faster, but can lead to loss
+            of information in case of small polygons. Other option is ``"iterate"``, which
+            iterates over polygons and uses :func:`rasterio.features.geometry_mask`.
+        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
+            Bresenham’s line algorithm will be considered.
+        n_jobs : int, optional
+            Number of parallel threads to use. It is recommended to set this to the
+            number of physical cores of the CPU. ``-1`` uses all available cores. Applies
+            only if ``method="iterate"``.
+        **kwargs : optional
+            Keyword arguments to be passed to the aggregation function
+            (e.g., ``Dataset.mean(**kwargs)``).
+
+        Returns
+        -------
+        Dataset
+            A subset of the original object with N-1 dimensions indexed by
+            the the GeometryIndex.
+
+        """
+        # TODO: allow multiple stats at the same time (concat along a new axis),
+        # TODO: possibly as a list of tuples to include names?
+        # TODO: allow callable in stat (via .reduce())
+        if method == "rasterize":
+            result = _zonal_stats_rasterize(
+                self,
+                polygons=polygons,
+                x_coords=x_coords,
+                y_coords=y_coords,
+                stats=stats,
+                name=name,
+                all_touched=all_touched,
+                **kwargs,
+            )
+        elif method == "iterate":
+            result = _zonal_stats_iterative(
+                self,
+                polygons=polygons,
+                x_coords=x_coords,
+                y_coords=y_coords,
+                stats=stats,
+                name=name,
+                all_touched=all_touched,
+                n_jobs=n_jobs,
+                **kwargs,
+            )
+        else:
+            raise ValueError(
+                f"method '{method}' is not supported. Allowed options are 'rasterize' "
+                "and 'iterate'."
+            )
+
+        # save the index as a data variable
+        if isinstance(polygons, pd.Series):
+            if index is None:
+                if polygons.index.name is not None or not polygons.index.equals(
+                    pd.RangeIndex(0, len(polygons))
+                ):
+                    index = True
+            if index:
+                index_name = polygons.index.name if polygons.index.name else "index"
+                result = result.assign_coords({index_name: (name, polygons.index)})
+
+        # standardize the shape - each method comes with a different one
+        return result.transpose(
+            name, *tuple(d for d in self._obj.dims if d not in [x_coords, y_coords])
+        )
+
     def extract_points(
         self,
         points: Sequence[shapely.Geometry],
@@ -957,15 +1071,15 @@ def extract_points(
         crs : Any, optional
             Cordinate reference system of shapely geometries. If ``points`` have a
             ``.crs`` attribute (e.g. ``geopandas.GeoSeries`` or a ``DataArray`` with
-            ``"crs"`` in ``.attrs`), ``crs`` will be automatically inferred. For more
+            ``"crs"`` in ``.attrs``), ``crs`` will be automatically inferred. For more
             generic objects (numpy  array, list), CRS shall be specified manually.
         index : bool, optional
             If `points` is a GeoSeries, ``index=True`` will attach its index as another
             coordinate to the geometry dimension in the resulting object. If
             ``index=None``, the index will be stored if the `points.index` is a named
             or non-default index. If ``index=False``, it will never be stored. This is
             useful as an attribute link between the resulting array and the GeoPandas
-            object from which the points are sourced from.
+            object from which the points are sourced.
 
         Returns
         -------

xvec/index.py

@@ -35,7 +35,8 @@ def _get_common_crs(crs_set: set[CRS | None]):
             warnings.warn(  # noqa: B028
                 "CRS not set for some of the concatenation inputs. "
                 f"Setting output's CRS as {names[0]} "
-                "(the single non-null CRS provided)."
+                "(the single non-null CRS provided).",
+                stacklevel=3,
             )
         return crs_not_none[0]
 

xvec/tests/test_accessor.py

@@ -7,7 +7,7 @@
 from geopandas.testing import assert_geodataframe_equal
 from pandas.testing import assert_frame_equal
 
-import xvec  # noqa
+import xvec  # noqa: F401
 from xvec import GeometryIndex