基于Xarray数据集区域求和为GeoDataFrame加列：求优化建议

我发现了geocube这个出色的区域统计库，完美解决了问题。该教程提供了针对DataArray分析的优秀示例；以下是我针对Dataset中多栅格分析编写的代码。

再次恳请各位指出优化空间！谢谢。

Alex

import geopandas as gpd
from rioxarray import open_rasterio
import xarray as xr

ssurgo_data = gpd.read_file("../../test/test_data/input/soil_data_group.geojson")
ssurgo_data = ssurgo_data.loc[ssurgo_data.hzdept_r==0]
ssurgo_data["mukey"] = ssurgo_data.mukey.astype(int)
tif_dict = {2000: 'path_to_2000.tif', 2001: 'path_to_2001.tif'}
xda_dict = {i: open_rasterio(tif_dict[i], mask_and_scale=True) for i in tif_dict.keys()}
elevation = xr.Dataset(data_vars=xda_dict)
elevation = elevation.rio.clip(
    ssurgo_data.geometry.values, ssurgo_data.crs, from_disk=True
).sel(band=1).drop("band")
elevation.name = "elevation"
out_grid = make_geocube(
    vector_data=ssurgo_data,
    measurements=["mukey"],
    like=elevation, # ensure the data are on the same grid
)
for j in tif_dict.keys():
    out_grid[j] = (elevation[j].dims, elevation[j].values, elevation[j].attrs, elevation[j].encoding)
out_grid

以下代码虽然能运行，但速度极慢。我正在寻找优化方法，欢迎各位提出建议！谢谢！

我接触xarray时间不长，推测自己只是遗漏了简单的方法。我有一个包含20-30个栅格的xarray.Dataset，每个文件对应同一区域不同类别（年份、降雨、温度等）的数据。

同时我有一个覆盖大致同一时空范围的geopandas.GeoDataFrame，所有几何对象均为shapely.Polygon。我希望为该GeoDataFrame添加一列，对应Dataset中的一个文件，列值为每个Polygon覆盖区域内数据的总和。

Alex

import geopandas as gpd
from rioxarray import open_rasterio
import xarray as xr


rasters = {
     2000: open_rasterio('path_to_2000_data.tif'), 
     2001: open_rasterio('path_to_2001_data.tif')
# ... include data across other years, categories, etc.
}
xds = xr.Dataset(rasters)

gdf = gpd.read_parquet('path_to_geometry.geoparquet')

def helper_xr_pop(
        poly: Polygon,
        xds: xr.Dataset
) -> dict:
    try:
        # clip the dataset by the polygon
        out = xds.rio.clip([poly])
        # convert the sums to a dictionary ({2000: 10, 2001: 20})
        out = out.sum().out.to_dict()
    except:
        out = {}
    return out

XR_DATA = 'xr_data'
gdf[XR_DATA] = dask_gdf[GEO_COL].apply(lambda x: helper_xr_pop(x, xds), axis=1)
s = gdf[[GEO_COL]].apply(lambda x: helper_gpd_pop(*x, xds), axis=1)
pop_df = pd.json_normalize(s)
gdf = gdf.drop(columns=[XR_DATA])
out_gdf = pd.concat(objs=[gdf, pop_df], axis=1)

我尝试用pandas/geopandas、dask/dask_geopandas和polars实现，但耗时都远超预期。我直觉xarray有对应的方法，但尚未找到。我查阅了几天文档，学到了很多，但没找到问题的答案。

内容的提问来源于stack exchange，提问作者Alex Smith