加速计算最长连续热浪事件序列——30年格点数据集性能优化求助

问题背景 #

我有一个包含30年热浪观测日数据的掩码数据集，需要计算每个格点单元的最长热浪事件持续时长。现有代码处理单年份数据时耗时2-3分钟且运行正常，但处理30年全量数据时已耗时1小时仍未完成，急需优化方案。

数据集说明：ff是xarray.DataArray，维度为：time: 10958、latitude: 27、longitude: 21

原代码及问题分析 #

原代码通过手动遍历每个格点的时间序列实现逻辑，这种三层嵌套循环的方式效率极低——面对10958×27×21≈630万次的循环操作，Python单循环的性能瓶颈被无限放大，直接导致全量数据处理耗时剧增。

原代码如下：

fs = hw.where(hw.rolling(time=3).count()==3).sel(time=slice('1991','2020'))
ff = fs

is_heatwave = fs > 0

tim = 0
lati = 0
longi = 0

for t in range(1, ff.shape[0]*ff.shape[1]*ff.shape[2]):
    if np.any(ff[tim, lati, longi] > 1):
        ff[tim, lati, longi] = 1
    else:
        ff[tim, lati, longi] = 0

    longi = longi + 1
    if longi == 21:
        longi = 0
        lati = lati + 1
    
    if lati == 27:
        lati = 0
        tim = tim + 1
    
    if tim == 10959:
        break


tim = 0
lati = 0
longi = 0

for t in range(1, ff.shape[0]*ff.shape[1]*ff.shape[2]):
    if np.any(ff[tim, lati, longi] == 1):
        ff[tim, lati, longi] = ff[tim, lati, longi] + ff[tim-1, lati, longi]
    else:
        ff[tim, lati, longi] = ff[tim, lati, longi]

    longi = longi + 1
    if longi == 21:
        longi = 0
        lati = lati + 1
    
    if lati == 27:
        lati = 0
        tim = tim + 1
    
    if tim == 10959:
        break

ff_max = ff.max(axis=0)

优化方案 #

利用xarray和numpy的向量化操作替代循环，充分发挥数组运算的并行优势，能将处理时间压缩到几十秒级别。

完整优化代码 #

# 保留原逻辑：筛选连续3天的热浪事件
fs = hw.where(hw.rolling(time=3).count() == 3).sel(time=slice('1991','2020'))
is_heatwave = fs > 0

# 将布尔型热浪标识转为数值型，缺失值填充为0
ff = is_heatwave.astype(float).fillna(0)

# 计算连续热浪时长：遇到非热浪则重置计数
non_hw_mask = ff == 0
# 用非热浪点的累计数作为分组依据，实现分段累计
reset_points = non_hw_mask.cumsum(dim='time')
hw_duration = ff.groupby(reset_points).cumsum(dim='time')

# 获取每个格点的最长热浪持续时长
ff_max = hw_duration.max(dim='time')

优化核心逻辑 #

• 向量化运算：彻底抛弃Python循环，用底层C实现的数组运算处理全量数据，避免循环带来的性能损耗。
• 分组累计：通过cumsum生成分组标识，结合groupby实现“遇到非热浪则重置计数”的逻辑，替代原代码中对前一时刻值的逐个判断。
• xarray原生方法：利用xarray的维度管理能力，无需手动处理维度索引，代码简洁且不易出错。

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