diff --git a/src/rasterstats/main.py b/src/rasterstats/main.py index fb099c7..3dd00da 100644 --- a/src/rasterstats/main.py +++ b/src/rasterstats/main.py @@ -34,7 +34,8 @@ def raster_stats(*args, **kwargs): def zonal_stats(vectors, raster, layer_num=0, band_num=1, nodata_value=None, global_src_extent=False, categorical=False, stats=None, - copy_properties=False, all_touched=False, transform=None): + copy_properties=False, all_touched=False, transform=None, + add_stats=[None], add_func=[None]): """Summary statistics of a raster, broken out by vector geometries. Attributes @@ -78,6 +79,9 @@ def zonal_stats(vectors, raster, layer_num=0, band_num=1, nodata_value=None, transform : list of float, optional GDAL-style geotransform coordinates when `raster` is an ndarray. Required when `raster` is an ndarray, otherwise ignored. + add_stats : List with name additional statistics to compute, optional + Each name in add_stats is a string with the name of a function + that takes Numpy (masked) matrices and computes and returns a statistic Returns ------- @@ -296,6 +300,19 @@ def zonal_stats(vectors, raster, layer_num=0, band_num=1, nodata_value=None, except KeyError: rmax = float(masked.max()) feature_stats['range'] = rmax - rmin + if add_stats: + if add_func: + try: + exec('from '+add_func+' import *') + except: + print('Could not load module %s') % add_func + break + for mystat in range(len(add_stats)): + adstat=add_stats[mystat] + try: + exec('feature_stats["'+adstat+'"] = '+adstat+'(masked)') + except: + print('Statistic %s in add_stats not defined in add_func') % adstat # Use the enumerated id as __fid__ feature_stats['__fid__'] = i diff --git a/tests/myfunc.py b/tests/myfunc.py new file mode 100755 index 0000000..0018cf5 --- /dev/null +++ b/tests/myfunc.py @@ -0,0 +1,7 @@ +#!/usr/bin/env python +# Additional functions to be used in raster stat computation +from __future__ import division +import numpy as np + +def mymean(x): + return np.ma.mean(x) diff --git a/tests/test_zonal.py b/tests/test_zonal.py index 5e41aa5..f528b48 100644 --- a/tests/test_zonal.py +++ b/tests/test_zonal.py @@ -8,6 +8,9 @@ OGRError, feature_to_geojson, bbox_to_pixel_offsets from shapely.geometry import shape, box import json +import sys + +sys.path.append(os.path.dirname(os.path.abspath(__file__))) DATA = os.path.join(os.path.dirname(os.path.abspath(__file__)), "data") raster = os.path.join(DATA, 'slope.tif') @@ -389,3 +392,12 @@ def test_featurecollection(): # geointerface featurecollection stats2 = zonal_stats(df, raster) assert stats == stats2 + +def test_add_stats(): + from geopandas import GeoDataFrame + polygons = os.path.join(DATA, 'polygons.shp') + df = GeoDataFrame.from_file(polygons) + mymean= lambda x: np.ma.mean(x) + stats = zonal_stats(df.geometry, raster, add_stats=['mymean'], add_func='myfunc') + for i in range(len(stats)): + assert stats[i]['mean'] == stats[i]['mymean']