diff --git a/docs/manual.rst b/docs/manual.rst index 15225a4..2a05496 100644 --- a/docs/manual.rst +++ b/docs/manual.rst @@ -174,6 +174,20 @@ then use it in your ``zonal_stats`` call like so:: ... add_stats={'mymean':mymean}) [{'count': 75, 'mymean': 14.660084635416666}, {'count': 50, 'mymean': 56.605761718750003}] +To have access to geometry properties, a dictionnary can be passed to the user-defined function:: + + >>> def mymean_prop(x,prop): + ... return np.ma.mean(x) * prop['id'] + +then use it in your ``zonal_stats`` call like so:: + + >>> zonal_stats("tests/data/polygons.shp", + ... "tests/data/slope.tif", + ... stats="count", + ... add_stats={'mymean_prop':mymean_prop}, + ... properties=['id']) + [{'count': 75, 'mymean_prop': 14.660084635416666}, {'count': 50, 'mymean_prop': 113.2115234375}] + GeoJSON output ^^^^^^^^^^^^^^ diff --git a/src/rasterstats/main.py b/src/rasterstats/main.py index 3719663..5a57a67 100644 --- a/src/rasterstats/main.py +++ b/src/rasterstats/main.py @@ -112,6 +112,7 @@ def gen_zonal_stats( with zonal stats appended as additional properties. Use with `prefix` to ensure unique and meaningful property names. + Returns ------- generator of dicts (if geojson_out is False) @@ -258,7 +259,11 @@ def gen_zonal_stats( if add_stats is not None: for stat_name, stat_func in add_stats.items(): - feature_stats[stat_name] = stat_func(masked) + try: + feature_stats[stat_name] = stat_func(masked, feat['properties']) + except TypeError: + # backwards compatible with single-argument function + feature_stats[stat_name] = stat_func(masked) if raster_out: feature_stats['mini_raster_array'] = masked diff --git a/tests/test_zonal.py b/tests/test_zonal.py index b5c3014..898b8eb 100644 --- a/tests/test_zonal.py +++ b/tests/test_zonal.py @@ -289,6 +289,17 @@ def mymean(x): assert stats[i]['mean'] == stats[i]['mymean'] +def test_add_stats_prop(): + polygons = os.path.join(DATA, 'polygons.shp') + + def mymean_prop(x, prop): + return np.ma.mean(x) * prop['id'] + + stats = zonal_stats(polygons, raster, add_stats={'mymean_prop': mymean_prop}) + for i in range(len(stats)): + assert stats[i]['mymean_prop'] == stats[i]['mean'] * (i+1) + + def test_mini_raster(): polygons = os.path.join(DATA, 'polygons.shp') stats = zonal_stats(polygons, raster, raster_out=True)