使用MetPy函数与多进程池并行计算CAPE时遇错误求解决

问题详情 #

尝试用multiprocessing.Pool的apply()方法并行化CAPE值估算，运行代码时出现单位错误，同时并行逻辑存在效率问题。

运行代码 #

import multiprocessing  
from multiprocessing import Pool
from pytictoc import TicToc # conda install pytictoc -c ecf

tims = 4 #FIXME: u.shape[0] USE 4 FOR 6 HOURLY DATA

if __name__ == '__main__': #PART OF THE SYNTAX
pool = Pool(processes=8) # set number of process


for tim in range(tims): # loop across 8 timestamps of a day
    print('processing time-step: ',tim)
    t = TicToc()
    t.tic()
    
    for lat in SP_test.latitude.values:
        for lon in SP_test.longitude.values:
            for tim in SP_test.time.values:
                Temp = SP_test.sel(time =tim, latitude=lat, longitude = lon).t
                #print(Temp)
                RH = SP_test.sel(time =tim, latitude=lat, longitude = lon).r
                #print(RH)
                TD = dewpoint_from_relative_humidity(Temp,RH)
                
                cape[tim] = [pool.apply(most_unstable_cape_cin,args = (p,Temp,TD))] # run parallel (across (lat,lon), but for the fixed current timestamp) processes; save all 91*151 results in 1 item of the list


    t.toc()

错误信息 #

ValueError: `most_unstable_cape_cin` given arguments with incorrect units: `dewpoint` requires "[temperature]" but given "dimensionless"

修改方案与替代实现 #

1. 修复单位错误 #

错误根源是传递给most_unstable_cape_cin的露点温度（TD）丢失了温度单位，大概率是xarray对象跨进程序列化时丢失了单位信息，有两种解决方式：

方法A：在函数内部补全单位

修改most_unstable_cape_cin函数，手动为输入参数添加正确单位（需导入metpy的units模块）：

from metpy.units import units

def most_unstable_cape_cin(p, temp, td):
    # 根据实际数据单位调整，比如degC或K
    temp = temp * units.degC
    td = td * units.degC
    # 原有函数逻辑...

方法B：传递带单位的数值

在调用前将xarray数据转换为带单位的数值对象，确保跨进程传递时单位不丢失：

from metpy.units import units

# 转换为带单位的量再传递
Temp_unit = Temp.values * units(Temp.units)
TD_unit = TD.values * units(TD.units)
pool.apply(most_unstable_cape_cin, args=(p, Temp_unit, TD_unit))

2. 优化并行逻辑（提升效率） #

原代码的嵌套循环+pool.apply()方式效率极低，因为apply()是阻塞式调用，每次仅执行一个任务，未真正利用多进程优势。建议改用pool.starmap()批量提交任务：

修正后的代码示例

import multiprocessing  
from multiprocessing import Pool
from pytictoc import TicToc
from metpy.units import units

def most_unstable_cape_cin(p, temp_val, td_val):
    # 补全单位
    temp = temp_val * units.degC
    td = td_val * units.degC
    # 原有函数逻辑，返回CAPE值
    ...

if __name__ == '__main__':
    pool = Pool(processes=8)
    t = TicToc()
    
    # 生成所有任务参数列表
    tasks = []
    # 避免变量名覆盖，外层循环用time_val替代tim
    for time_val in SP_test.time.values[:4]:  # 对应原tims=4的设置
        print(f'processing time-step: {time_val}')
        t.tic()
        
        for lat in SP_test.latitude.values:
            for lon in SP_test.longitude.values:
                Temp = SP_test.sel(time=time_val, latitude=lat, longitude=lon).t
                RH = SP_test.sel(time=time_val, latitude=lat, longitude=lon).r
                TD = dewpoint_from_relative_humidity(Temp, RH)
                # 提取数值存入任务列表
                tasks.append( (p, Temp.values, TD.values) )
    
    # 批量并行处理所有任务
    results = pool.starmap(most_unstable_cape_cin, tasks)
    
    # 将结果整理到cape数组（按维度顺序对应赋值）
    result_idx = 0
    for time_val in SP_test.time.values[:4]:
        for lat in SP_test.latitude.values:
            for lon in SP_test.longitude.values:
                cape.loc[dict(time=time_val, latitude=lat, longitude=lon)] = results[result_idx]
                result_idx += 1
    
    t.toc()
    pool.close()
    pool.join()

3. 更简洁的替代方案：xarray + Dask并行 #

如果使用xarray处理气象数据，推荐结合Dask实现并行，无需手动管理进程池，代码更简洁：

import dask
from dask.distributed import Client
import xarray as xr
from metpy.calc import dewpoint_from_relative_humidity, most_unstable_cape_cin
from metpy.units import units

# 启动Dask集群（8个工作进程）
client = Client(n_workers=8)

def compute_cape(temp, rh, p):
    # 计算露点
    td = dewpoint_from_relative_humidity(temp, rh)
    # 计算CAPE（忽略CIN值）
    cape_val, _ = most_unstable_cape_cin(p, temp, td)
    return cape_val

# 用xarray的apply_ufunc实现向量化并行
cape = xr.apply_ufunc(
    compute_cape,
    SP_test.t,
    SP_test.r,
    input_core_dims=[[], []],  # 根据数据实际维度调整
    kwargs={'p': p},
    dask='parallelized',
    output_dtypes=[float],
    vectorize=True
)

# 触发实际计算
cape = cape.compute()

# 关闭Dask集群
client.close()

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