基于Dask Array用Shapefile掩蔽并平均大型NetCDF数据：性能优化咨询

我正在处理ERA5-Land的40年全球逐小时降水数据，需要用流域Shapefile对数据掩蔽，计算每个时间步的流域平均值，最后为每个流域保存一个NetCDF文件。现有代码可以正常运行，但处理速度极慢，预估完成全部5000个流域的处理需要数月时间，求优化方案。

原代码步骤 #

1. 创建文件列表 #

# create lists of ERA5 and catchment boundary files
ERA5_files = [os.path.join(dirpath,f) for (dirpath, dirnames, filenames) in os.walk("path_to_all_files") for f in filenames] # 约500个文件

catchment_files = [os.path.join(dp, f) for dp, dn, filenames in os.walk("path_to_shapefiles") for f in filenames if os.path.splitext(f)[1] == '.shp'] # 约5000个Shapefile

2. 合并Shapefile为GeoDataFrame #

gdfs = []
for f in catchment_files:
    gdf = gpd.read_file(f)
    gdfs.append(gdf)

gdf = gpd.GeoDataFrame(pd.concat(gdfs))
# 过滤无效流域
gdf = gdf[gdf['gauge_id'].isin(passed)]

3. 数据处理函数 #

# 预处理函数：转换单位、调整经度格式、设置CRS
def intconv(nc_stack):
    nc_stack['tp'] = nc_stack['tp']*100
    nc_stack.astype('uint16')

    # 将ERA5经度从0-360转换为-180-180
    lon_switch = nc_stack.assign_coords(lon=(((nc_stack.lon + 180) % 360) - 180)).sortby('lon')
    lon_switch = lon_switch.rio.set_spatial_dims("lon", "lat", inplace=True)
    lon_switch = lon_switch.rio.write_crs("EPSG:4326")

    return lon_switch

# 掩蔽、计算均值并保存NetCDF
def nc_clip_return(nc_stack, gdf, a):
    shape = gdf[gdf.gauge_id == a]
    nc_stack = nc_stack.load()
    da_clip = nc_stack.rio.clip(
        shape.geometry
    )

    # 计算空间均值并保存
    grid_mean = da_clip.mean(dim=["lon", "lat"],skipna=True)
    name = '/save_path/' + a + '_0200_era5.nc'
    grid_mean.to_netcdf(name, mode='w',format='NETCDF4_CLASSIC')

4. 读取数据集并启动Dask #

from dask.distributed import Client

client = Client()
client

start = time.time()
test = xr.open_mfdataset(ERA5_files[0:200], chunks={"time": 12}, parallel=True)
# test = xr.decode_cf(test)
end = time.time()
print(end - start)

5. 循环处理每个流域 #

# 遍历流域ID，计算均值并保存
for a in gdf.gauge_id.unique():
    start = time.time()
    test = test.rio.set_spatial_dims("lon", "lat", inplace=True)
    nc_clip_return(a = a, nc_stack = test, gdf = gdf) 
    
    end = time.time()
    print(end - start)

优化方案 #

1. 移除重复操作与内存加载 #

• 将test.rio.set_spatial_dims移到循环外，只执行一次
• 删除nc_clip_return中的nc_stack.load()：Dask的核心优势是延迟计算，提前加载全量数据会耗尽内存并拖慢速度，让掩蔽、均值计算都基于Dask数组执行，仅在保存时触发计算

2. 预缓存流域几何数据 #

提前将流域几何按ID整理为字典，避免每次循环重复筛选：

# 预生成流域ID与几何的映射字典
gauge_geom_map = {row.gauge_id: row.geometry for _, row in gdf.iterrows()}

修改nc_clip_return，直接传入几何对象：

def nc_clip_return(nc_stack, geom, gauge_id):
    da_clip = nc_stack.rio.clip([geom])
    grid_mean = da_clip.mean(dim=["lon", "lat"], skipna=True)
    name = f'/save_path/{gauge_id}_0200_era5.nc'
    # 启用压缩减少文件大小与IO时间
    grid_mean.to_netcdf(name, mode='w', format='NETCDF4_CLASSIC', 
                        encoding={'tp': {'zlib': True, 'complevel': 5}})

3. 优化Dask分块与集群配置 #

• 调整时间分块大小：原chunks={"time":12}过小，会导致任务调度开销过大，建议改为chunks={"time": 24*7}（一周数据），让每个分块大小保持在100-200MB左右
• 启动Dask时指定资源参数，充分利用CPU：

client = Client(n_workers=4, threads_per_worker=2, memory_limit='8GB')

4. 并行处理流域任务 #

用Dask的delayed将串行循环改为并行执行：

from dask import delayed

# 预处理数据集（仅执行一次）
test = test.rio.set_spatial_dims("lon", "lat", inplace=False)
test = test.rio.write_crs("EPSG:4326")

# 生成所有延迟任务
tasks = []
for gauge_id, geom in gauge_geom_map.items():
    task = delayed(nc_clip_return)(nc_stack=test, geom=geom, gauge_id=gauge_id)
    tasks.append(task)

# 批量执行任务
dask.compute(tasks)

5. 修复预处理函数的无效操作 #

原intconv中nc_stack.astype('uint16')未赋值，导致类型转换不生效，修改后将预处理嵌入文件读取流程：

def intconv(nc_ds):
    # 转换单位并修改数据类型
    nc_ds['tp'] = (nc_ds['tp'] * 100).astype('uint16')
    # 调整经度格式
    lon = (((nc_ds.lon + 180) % 360) - 180)
    nc_ds = nc_ds.assign_coords(lon=lon).sortby('lon')
    nc_ds = nc_ds.rio.set_spatial_dims("lon", "lat")
    nc_ds = nc_ds.rio.write_crs("EPSG:4326")
    return nc_ds

# 读取文件时直接执行预处理
test = xr.open_mfdataset(ERA5_files[0:200], chunks={"time": 24*7}, 
                         parallel=True, preprocess=intconv)

6. 提升IO性能 #

• 将保存路径换成SSD，大幅降低写入延迟
• 给GeoDataFrame建立空间索引，加速几何检索：

gdf = gdf.set_index('gauge_id')
gdf.sindex

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