Zonal functions accept vector zones directly (#999)

* Accept vector zones in zonal functions (#998)

stats(), crosstab(), apply(), and crop() now accept GeoDataFrames or
list-of-(geometry, value) pairs as the zones argument. When vector
input is detected, rasterize() is called internally using the values
raster as the template. Adds column and rasterize_kw parameters.

* Add vector zones docs and notebook examples (#998)

Document column and rasterize_kw parameters in stats(), crosstab(),
apply(), and crop() docstrings. Add vector zones section to the
3_Zonal user guide notebook with GeoDataFrame, list-of-pairs,
crosstab, and rasterize_kw examples.

Author: brendancol

examples/user_guide/3_Zonal.ipynb

@@ -3,18 +3,7 @@
   {
    "cell_type": "markdown",
    "metadata": {},
-   "source": [
-    "# Xarray-spatial\n",
-    "### User Guide: Zonal\n",
-    "-----\n",
-    "\n",
-    "Xarray-spatial's zonal functions provide an easy way to generate statistics for zones within a raster aggregate. It's set up with a default set of calculations, or you can input any set of custom calculations you'd like to perform.\n",
-    "\n",
-    "[Generate terrain](#Generate-Terrain-Data)\n",
-    "[Zonal Statistics](#Zonal-Statistics)\n",
-    "\n",
-    "-----------\n"
-   ]
+   "source": "# Xarray-spatial\n### User Guide: Zonal\n-----\n\nXarray-spatial's zonal functions provide an easy way to generate statistics for zones within a raster aggregate. It's set up with a default set of calculations, or you can input any set of custom calculations you'd like to perform.\n\n[Generate terrain](#Generate-Terrain-Data)\n[Zonal Statistics](#Zonal-Statistics)\n[Zonal Crosstab](#Zonal-crosstab)\n[Vector Zones](#Vector-zones)\n\n-----------"
   },
   {
    "cell_type": "markdown",
@@ -239,11 +228,71 @@
    "source": "# Calculate crosstab: rows are trail segments, columns are tree species (1=Oak, 2=Pine, 3=Maple)\nresult = tree_species_agg.xrs.zonal_crosstab(zones_agg, zone_ids=[11, 12, 13, 14, 15, 16], cat_ids=[1, 2, 3])\nresult.columns = ['Trail Segment', 'Oak', 'Pine', 'Maple']\nresult"
   },
   {
-   "cell_type": "code",
+   "cell_type": "markdown",
    "execution_count": null,
    "metadata": {},
    "outputs": [],
-   "source": []
+   "source": "## Vector zones\n\nZonal functions can also accept vector geometries directly -- a GeoDataFrame or a list of `(geometry, value)` pairs. The rasterization to match the values grid happens automatically inside the function call, so there's no separate `rasterize()` step.\n\nThis works with `zonal_stats`, `zonal_crosstab`, `zonal_apply`, and `crop`."
+  },
+  {
+   "cell_type": "markdown",
+   "source": "### GeoDataFrame zones\n\nLet's define three rectangular regions as polygons in a GeoDataFrame and compute elevation statistics directly.",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "source": "import geopandas as gpd\nfrom shapely.geometry import box\nfrom xrspatial.zonal import stats\n\n# Three rectangular regions across the terrain\nzones_gdf = gpd.GeoDataFrame(\n    {'region_id': [1.0, 2.0, 3.0],\n     'name': ['West valley', 'Central ridge', 'East slope']},\n    geometry=[\n        box(-20, -20, -5, 20),   # left third\n        box(-5, -20, 5, 20),     # middle third\n        box(5, -20, 20, 20),     # right third\n    ],\n)\n\n# Pass the GeoDataFrame directly -- no manual rasterize() needed\nresult = stats(zones_gdf, terrain, column='region_id',\n               stats_funcs=['mean', 'min', 'max', 'std', 'count'])\nresult",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "source": "The same thing works with the accessor. The values raster is used as the template for rasterization, so the zones automatically align to the grid.",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "source": "# Accessor style -- same result\nterrain.xrs.zonal_stats(zones_gdf, column='region_id',\n                        stats_funcs=['mean', 'min', 'max', 'std', 'count'])",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "source": "### List-of-pairs zones\n\nIf you don't have a GeoDataFrame, you can pass a list of `(geometry, zone_id)` tuples. No `column` argument needed here since the zone ID is the second element of each pair.",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "source": "pairs = [\n    (box(-15, -15, 0, 0), 10.0),   # southwest quadrant\n    (box(0, 0, 15, 15), 20.0),     # northeast quadrant\n]\n\nstats(pairs, terrain, stats_funcs=['mean', 'count'])",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "source": "### Crosstab with vector zones\n\nVector zones work with `zonal_crosstab` too. Here we'll use the tree species data from earlier with our GeoDataFrame regions.",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "source": "from xrspatial.zonal import crosstab\n\nct = crosstab(zones_gdf, tree_species_agg, column='region_id', cat_ids=[1, 2, 3])\nct.columns = ['Region', 'Oak', 'Pine', 'Maple']\nct",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "source": "### Forwarding rasterize options\n\nYou can pass extra options to the internal `rasterize()` call via `rasterize_kw`. For example, `all_touched=True` will include pixels that touch the geometry boundary, not just those whose centre falls inside.",
+   "metadata": {}
+  },
+  {
+   "cell_type": "code",
+   "source": "# Small polygon -- compare default vs all_touched\nsmall_zone = gpd.GeoDataFrame(\n    {'zone': [1.0]}, geometry=[box(-2, -2, 2, 2)]\n)\n\ndefault = stats(small_zone, terrain, column='zone', stats_funcs=['count'])\ntouched = stats(small_zone, terrain, column='zone', stats_funcs=['count'],\n                rasterize_kw={'all_touched': True})\n\nprint(f\"Default pixel count:     {int(default['count'].iloc[0])}\")\nprint(f\"all_touched pixel count: {int(touched['count'].iloc[0])}\")",
+   "metadata": {},
+   "execution_count": null,
+   "outputs": []
   }
  ],
  "metadata": {

xrspatial/tests/test_zonal.py

@@ -1686,3 +1686,145 @@ def _guarded_unique(a, *args, **kwargs):
         result = result.compute()
     assert isinstance(result, pd.DataFrame)
     assert len(result) > 0
+
+
+# ---------------------------------------------------------------------------
+# Vector zones (GeoDataFrame / list-of-pairs) — implicit rasterization
+# ---------------------------------------------------------------------------
+
+try:
+    from shapely.geometry import box as shapely_box
+    import geopandas as gpd
+    _has_vector = True
+except ImportError:
+    _has_vector = False
+
+skip_no_vector = pytest.mark.skipif(
+    not _has_vector, reason="shapely/geopandas not installed"
+)
+
+
+def _make_grid(width=20, height=20, bounds=(0, 0, 100, 100)):
+    """Build a template DataArray with pixel-centre coords."""
+    xmin, ymin, xmax, ymax = bounds
+    px = (xmax - xmin) / width
+    py = (ymax - ymin) / height
+    x = np.linspace(xmin + px / 2, xmax - px / 2, width)
+    y = np.linspace(ymax - py / 2, ymin + py / 2, height)
+    return xr.DataArray(
+        np.ones((height, width), dtype=np.float64),
+        dims=['y', 'x'],
+        coords={'y': y, 'x': x},
+    )
+
+
+@skip_no_vector
+class TestVectorZones:
+    """Verify that zonal functions accept vector zones directly."""
+
+    def _zones_raster_and_gdf(self):
+        """Two non-overlapping boxes covering different quadrants."""
+        from xrspatial.rasterize import rasterize
+
+        values = _make_grid()
+        gdf = gpd.GeoDataFrame(
+            {'zone_id': [1.0, 2.0]},
+            geometry=[
+                shapely_box(0, 50, 50, 100),   # top-left
+                shapely_box(50, 0, 100, 50),    # bottom-right
+            ],
+        )
+        zones_raster = rasterize(gdf, like=values, column='zone_id')
+        return values, gdf, zones_raster
+
+    # -- stats --
+
+    def test_stats_gdf_matches_raster(self):
+        values, gdf, zones_raster = self._zones_raster_and_gdf()
+        expected = stats(zones_raster, values, stats_funcs=['mean', 'count'])
+        result = stats(gdf, values, stats_funcs=['mean', 'count'],
+                       column='zone_id')
+        pd.testing.assert_frame_equal(result, expected)
+
+    def test_stats_list_of_pairs(self):
+        values, gdf, zones_raster = self._zones_raster_and_gdf()
+        pairs = [
+            (shapely_box(0, 50, 50, 100), 1.0),
+            (shapely_box(50, 0, 100, 50), 2.0),
+        ]
+        expected = stats(zones_raster, values, stats_funcs=['mean', 'count'])
+        result = stats(pairs, values, stats_funcs=['mean', 'count'])
+        pd.testing.assert_frame_equal(result, expected)
+
+    def test_stats_gdf_missing_column_raises(self):
+        values = _make_grid()
+        gdf = gpd.GeoDataFrame(
+            {'zone_id': [1.0]},
+            geometry=[shapely_box(0, 0, 50, 50)],
+        )
+        with pytest.raises(ValueError, match="column is required"):
+            stats(gdf, values)
+
+    def test_stats_pairs_with_column_raises(self):
+        values = _make_grid()
+        pairs = [(shapely_box(0, 0, 50, 50), 1.0)]
+        with pytest.raises(ValueError, match="column should not be set"):
+            stats(pairs, values, column='zone_id')
+
+    def test_stats_accessor(self):
+        import xrspatial  # noqa: F401  — registers accessor
+        values, gdf, zones_raster = self._zones_raster_and_gdf()
+        expected = stats(zones_raster, values, stats_funcs=['mean', 'count'])
+        result = values.xrs.zonal_stats(gdf, stats_funcs=['mean', 'count'],
+                                        column='zone_id')
+        pd.testing.assert_frame_equal(result, expected)
+
+    # -- crosstab --
+
+    def test_crosstab_gdf(self):
+        values, gdf, zones_raster = self._zones_raster_and_gdf()
+        expected = crosstab(zones_raster, values)
+        result = crosstab(gdf, values, column='zone_id')
+        pd.testing.assert_frame_equal(result, expected)
+
+    # -- apply --
+
+    def test_apply_gdf(self):
+        values, gdf, zones_raster = self._zones_raster_and_gdf()
+        # rasterize produces float zones; apply needs int zones
+        zones_int = zones_raster.copy(data=zones_raster.values.astype(int))
+        fn = lambda x: x * 2
+        expected = apply(zones_int, values, fn)
+        result = apply(gdf, values, fn, column='zone_id',
+                       rasterize_kw={'dtype': int})
+        xr.testing.assert_identical(result, expected)
+
+    # -- crop --
+
+    def test_crop_gdf(self):
+        values, gdf, zones_raster = self._zones_raster_and_gdf()
+        expected = crop(zones_raster, values, zones_ids=[1.0])
+        result = crop(gdf, values, zones_ids=[1.0], column='zone_id')
+        xr.testing.assert_identical(result, expected)
+
+    # -- rasterize_kw forwarding --
+
+    def test_stats_rasterize_kw_all_touched(self):
+        values, gdf, _ = self._zones_raster_and_gdf()
+        from xrspatial.rasterize import rasterize
+        zones_at = rasterize(gdf, like=values, column='zone_id',
+                             all_touched=True)
+        expected = stats(zones_at, values, stats_funcs=['count'])
+        result = stats(gdf, values, stats_funcs=['count'],
+                       column='zone_id',
+                       rasterize_kw={'all_touched': True})
+        pd.testing.assert_frame_equal(result, expected)
+
+    # -- raster zones still work --
+
+    def test_raster_zones_unchanged(self):
+        """Passing a DataArray directly should still work as before."""
+        values, _, zones_raster = self._zones_raster_and_gdf()
+        result = stats(zones_raster, values, stats_funcs=['count'])
+        assert isinstance(result, pd.DataFrame)
+        assert len(result) > 0