如何用map函数与多进程简化栅格矢量处理代码并提速

核心需求 #

• 现有m个栅格文件和n个矢量文件，需为每个栅格遍历所有矢量执行掩膜+分区统计操作；当前为每个矢量单独编写重复for循环，代码冗余，需优化为类似R中map/map2df的批量处理逻辑
• 当前单脚本仅占用单个处理器，运行效率低下，需通过多进程技术缩短处理时间

现有冗余代码 #

# filenames_dat[i] 对应栅格输入文件
df1 = gpd.read_file("input_path")
df2 = gpd.read_file("input_path")

# 处理df1的循环
for i in range(len(raster_path)):
    array_name, trans_name = mask(filenames_dat[i], shapes=df1.geometry, crop=True, nodata=np.nan)
    zs= zonal_stats(df1, array_name[0], affine=trans_name, stats=['mean','sum'], nodata=np.nan, all_touched=True)
    df1['amg'+str(filenames[i])] = [x[('mean')] for x in zs]
    df1['mpg'+str(filenames[i])] = [x[('sum')] for x in zs]
    print(i)
    
df1csv = pd.DataFrame(df1)
df1csv.to_csv(cwd+'/rasteroutput/df1.csv', index = False)

# 处理df2的循环，逻辑与上面完全重复
for i in range(len(raster_path)):
    array_name, trans_name = mask(filenames_dat[i], shapes=df2.geometry, crop=True, nodata=np.nan)
    zs= zonal_stats(df2, array_name[0], affine=trans_name, stats=['mean','sum'], nodata=np.nan, all_touched=True)
    df2['amg'+str(filenames[i])] = [x[('mean')] for x in zs]
    df2['mpg'+str(filenames[i])] = [x[('sum')] for x in zs]
    print(i)

df2csv = pd.DataFrame(df2)
df2csv.to_csv(cwd+'/rasteroutput/df2.csv', index = False)

优化尝试中的问题 #

已编写poly2函数，但无法通过map正确传递多参数（期望实现类似R的map2df功能），调用list(map(poly2,lst,df))时报错：

def poly2(i,df):## i对应栅格索引
    df = df
    array_name, trans_name = mask(filenames_dat[i], shapes=df.geometry, crop=True, nodata=np.nan)
    zs= zonal_stats(df, array_name[0], affine=trans_name, stats=['mean','sum'], nodata=np.nan, all_touched=True)
    df['amg'+str(filenames[i])] = [x[('mean')] for x in zs]
    df['mpg'+str(filenames[i])] = [x[('sum')] for x in zs]
    print(i)

lst=[]    
for i in range(len(raster_path)):
    lst.append(i)
poly2(i=1, df=df) # 单独调用可行，但map调用报错

list(map(poly2,lst,df)) ## 此处报错

解决方案 #

1. 多参数map实现（解决代码冗余） #

map函数要求所有可迭代对象长度匹配，且按顺序传递参数。你之前的错误是直接传入GeoDataFrame，map会遍历其每一行而非传入整个对象，需用itertools.repeat或functools.partial处理：

方法1：用itertools.repeat固定矢量参数

from itertools import repeat

# 把所有矢量文件存入列表，支持扩展更多矢量
vector_dfs = [df1, df2]

def process_raster_for_vector(raster_idx, vector_df):
    raster = filenames_dat[raster_idx]
    array_name, trans_name = mask(raster, shapes=vector_df.geometry, crop=True, nodata=np.nan)
    zs = zonal_stats(vector_df, array_name[0], affine=trans_name, stats=['mean','sum'], nodata=np.nan, all_touched=True)
    raster_name = filenames[raster_idx]
    vector_df[f'amg{raster_name}'] = [x['mean'] for x in zs]
    vector_df[f'mpg{raster_name}'] = [x['sum'] for x in zs]
    print(f"栅格{raster_idx}处理完成")
    return vector_df

# 遍历每个矢量，用map批量处理所有栅格
for idx, vec_df in enumerate(vector_dfs):
    # 用repeat将同一个矢量重复len(raster_path)次，与栅格列表匹配
    list(map(process_raster_for_vector, lst, repeat(vec_df)))
    vec_df.to_csv(f"{cwd}/rasteroutput/df{idx+1}.csv", index=False)

方法2：用functools.partial固定矢量参数

from functools import partial

for idx, vec_df in enumerate(vector_dfs):
    # 固定vector_df参数，生成仅接收raster_idx的函数
    partial_func = partial(process_raster_for_vector, vector_df=vec_df)
    list(map(partial_func, lst))
    vec_df.to_csv(f"{cwd}/rasteroutput/df{idx+1}.csv", index=False)

2. 多进程实现（解决单进程效率问题） #

多进程中避免传递GeoDataFrame引用（序列化易出错），建议直接传递文件路径，在子进程中独立读取：

import multiprocessing as mp

# 重构函数：接收栅格列表、矢量路径、输出路径
def process_single_vector(raster_list, vector_path, output_path):
    vec_df = gpd.read_file(vector_path)
    for raster_idx in raster_list:
        raster = filenames_dat[raster_idx]
        array_name, trans_name = mask(raster, shapes=vec_df.geometry, crop=True, nodata=np.nan)
        zs = zonal_stats(vec_df, array_name[0], affine=trans_name, stats=['mean','sum'], nodata=np.nan, all_touched=True)
        raster_name = filenames[raster_idx]
        vec_df[f'amg{raster_name}'] = [x['mean'] for x in zs]
        vec_df[f'mpg{raster_name}'] = [x['sum'] for x in zs]
        print(f"矢量{vector_path}完成栅格{raster_idx}处理")
    vec_df.to_csv(output_path, index=False)

# 主进程启动多进程
if __name__ == "__main__":
    # 替换为实际矢量文件路径和输出路径
    vector_paths = ["input_path1", "input_path2"]
    output_paths = [f"{cwd}/rasteroutput/df{idx+1}.csv" for idx in range(len(vector_paths))]
    
    # 进程数设为CPU核心数
    pool = mp.Pool(processes=mp.cpu_count())
    # 为每个矢量分配独立进程，处理其所有栅格任务
    pool.starmap(process_single_vector, [(lst, vp, op) for vp, op in zip(vector_paths, output_paths)])
    pool.close()
    pool.join()

关键注意事项 #

• map的可迭代对象长度必须一致，否则会提前终止
• 多进程中禁止共享可变对象（如已加载的GeoDataFrame），避免序列化冲突
• 若栅格文件过大，建议在子进程中单独读取，减少内存占用

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