如何加速44年ERA5格点风速百分位数计算代码？

问题背景 #

需要计算NetCDF文件中某一经纬度格点44年的风速百分位数，现有528个100m u分量和528个100m v分量文件（单个文件1.4G），原代码运行极慢，需求优化。

原代码：

import os
import numpy as np
import xarray as xr

data_folder1 = 'K:\\reanalysis\ERA5\\100m_speed\\100m u\\1979-2022'
data_folder2 = 'K:\\reanalysis\ERA5\\100m_speed\\100m v\\1979-2022'

file_names_u = os.listdir(data_folder1)
file_names_v = os.listdir(data_folder2)

def read_wind(file_path_u, file_path_v, lat_value, lon_value):
    with xr.open_dataset(file_path_u) as ds:
        u_wind = ds['u100'].sel(latitude=lat_value, longitude=lon_value).values
    with xr.open_dataset(file_path_v) as ds:
        v_wind = ds['v100'].sel(latitude=lat_value, longitude=lon_value).values
    return u_wind, v_wind

lat_value = 0
lon_value = 90
all_u_winds = []
all_v_winds = []

for file_name_u, file_name_v in zip(file_names_u, file_names_v):
    file_path_u = os.path.join(data_folder1, file_name_u)
    file_path_v = os.path.join(data_folder2, file_name_v)
    u_wind, v_wind = read_wind(file_path_u, file_path_v, lat_value, lon_value)
    all_u_winds.extend(u_wind)
    all_v_winds.extend(v_wind)

all_u_winds = np.array(all_u_winds)
all_v_winds = np.array(all_v_winds)

# 计算风速
all_windspeeds = np.sqrt(np.square(all_u_winds) + np.square(all_v_winds))

percentile_50 = np.percentile(all_windspeeds, 50)
print(f"The 50% percentile of the wind speed values is: {percentile_50}")

优化方案 #

1. 批量读取文件，削减IO开销 #

原代码循环打开单个文件导致频繁磁盘IO，改用xarray.open_mfdataset批量加载所有文件，直接提取目标格点数据：

import os
import numpy as np
import xarray as xr

data_folder_u = 'K:\\reanalysis\\ERA5\\100m_speed\\100m u\\1979-2022'
data_folder_v = 'K:\\reanalysis\\ERA5\\100m_speed\\100m v\\1979-2022'

# 批量读取u分量，分块加载避免内存溢出
ds_u = xr.open_mfdataset(os.path.join(data_folder_u, '*.nc'), chunks={'time': 1000})
u_data = ds_u['u100'].sel(latitude=0, longitude=90).values.flatten()

# 批量读取v分量
ds_v = xr.open_mfdataset(os.path.join(data_folder_v, '*.nc'), chunks={'time': 1000})
v_data = ds_v['v100'].sel(latitude=0, longitude=90).values.flatten()

# 计算风速与百分位数
windspeeds = np.sqrt(u_data**2 + v_data**2)
p50 = np.percentile(windspeeds, 50)
print(f"50%风速百分位数: {p50}")

2. 分块并行计算，利用多CPU资源 #

通过dask分块处理数据，让xarray自动并行加载和计算，进一步提升效率：

import os
import xarray as xr

data_folder_u = 'K:\\reanalysis\\ERA5\\100m_speed\\100m u\\1979-2022'
data_folder_v = 'K:\\reanalysis\\ERA5\\100m_speed\\100m v\\1979-2022'

# 批量读取并设置分块
ds_u = xr.open_mfdataset(os.path.join(data_folder_u, '*.nc'), chunks={'time': 2000})
ds_v = xr.open_mfdataset(os.path.join(data_folder_v, '*.nc'), chunks={'time': 2000})

# 直接在数据集上计算目标格点风速
windspeed = xr.ufuncs.sqrt(ds_u['u100'].sel(latitude=0, longitude=90)**2 + ds_v['v100'].sel(latitude=0, longitude=90)**2)

# 计算百分位数（自动用dask并行处理）
p50 = windspeed.quantile(0.5).values
print(f"50%风速百分位数: {p50}")

3. 核心优化点说明 #

• 减少IO次数：批量读取替代循环单文件操作，大幅降低磁盘IO开销；
• 精准加载数据：提前用sel提取目标经纬度，避免加载全局数据，减少内存占用；
• 分块并行：通过chunks参数让xarray调用dask分块处理，利用多CPU核心并行计算；
• 避免内存冗余：直接从数据集提取numpy数组，省去列表扩展再转换的额外开销。

额外建议 #

• 确保u、v分量文件的时间维度完全对齐，避免时间序列混乱；
• 若磁盘IO是瓶颈，将文件迁移至SSD存储；
• 可手动启动dask客户端调整并行资源：

  from dask.distributed import Client
  client = Client(n_workers=4)  # 根据CPU核心数设置
  

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