盘古大模型实战-EW帮帮网

微信交流群：https://732903.hlcode.cn/?id=NK1fWUR

0 前言

前一段时间，在学习人工智能的同时，也去了解了一下几乎是作为人工智能在气象上应用的一大里程碑式的研究成果-华为盘古气象大模型。正是盘古大模型的出现，促使天气预报的未来发展方向多了个除天气学方法、统计学方法、数值预报方法之外的AI方法。故写一篇博客来对之前的一部分实践内容进行总结。

1 盘古大模型

1.1初识盘古

盘古气象大模型是来自华为云的研究人员提出了一种新的高分辨率全球AI气象预报系统。盘古气象大模型是首个精度超过传统数值预报方法的AI方法，1小时-7天预测精度均高于传统数值方法（欧洲气象中心的operational IFS），同时预测速度提升10000倍，能够提供秒级的全球气象预报，包括位势、湿度、风速、温度、海平面气压等。盘古气象模型的水平空间分辨率达到0.25∘×0.25∘ ，时间分辨率为1小时，覆盖13层垂直高度，可以精准地预测细粒度气象特征。作为基础模型，盘古气象大模型还能够直接应用于多个下游场景。例如，在热带风暴预测任务中，盘古气象大模型的预测精度显著超过欧洲气象中心的高精度预报（ECMWF HRES Forecast）结果。

盘古大模型使用的数据，包括垂直高度上13个不同气压层，每层五种气象要素（温度、湿度、位势、经度和纬度方向的风速），以及地球表面的四种气象要素（2米温度、经度和纬度方向的10米风速、海平面气压）。

基于复杂多变的气象背景，研究者提出3D Earth-Specific Transformer。其主要思想是使用一个视觉transformer的3D变种来处理复杂的不均匀的气象要素。由于气象数据分辨率很大，因而相比于常见的vision transformer方法，研究人员将网络的encoder和decoder减少到2级（8个block），同时采用Swin transformer的滑窗注意力机制，以减少网络的计算量。

同时，气象要素数据对应的经纬度网格是不均匀的，而不同的要素在不同纬度、高度的分布也是不均匀的。对这些不均匀性的建模，有利于学习气象数据背后潜藏着的复杂物理规律，如科里奥利力等。为此，研究者在每一个transformer模块中引入和纬度、高度相关的绝对位置编码来学习每一次空间运算的不规则分量。这样改动后的transformer模块，被称为3D Earth-Specific Transformer。

为了缓解迭代误差，研究者提出一个简单而有效的策略。研究人员训练了4个不同预报间隔的模型，分别为1小时间隔、3小时间隔、6小时间隔、24小时间隔。进而，研究人员使用贪心算法调用这些模型，使得预测特定时间气象状况的迭代次数最小。例如，对于24小时预测，只需要调用一次24小时间隔的模型；而对于23小时预测，则需要调用三次6小时预报，一次3小时预报和两次1小时预报。通过使用多个不同时间间隔模型捕捉不同时序关系，盘古气象大模型不仅减少了迭代误差，并且避免了由递归训练带来的训练资源消耗。

1.2相关论文与博客

1.Nature

Accurate medium-range global weather forecasting with 3D neural networks, Nature, Volume 619, Pages 533–538, 2023.

2.arXiv 预印本

Pangu-Weather: A 3D High-Resolution Model for Fast and Accurate Global Weather Forecast, arXiv preprint: 2211.02556, 2022.

盘古天气：快速准确的全球天气预报的 3D 高分辨率模型

2 模型下载

盘古的开发者团队将参考代码与需要准备的库与说明都在开源的GitHub项目有所说明

GitHub链接：

GitHub - 198808xc/Pangu-Weather: An official implementation of Pangu-Weather

在这里，我们可以用git把项目clone下来。

Git地址：https://github.com/198808xc/Pangu-Weather.git

在Git中运行以下内容：

git clone https://github.com/198808xc/Pangu-Weather.git

同时，盘古提供的四个时段的预测模型（3h、6h、12h、24h）的模型也可以在GitHub上下载，或点击以下链接（Git中的链接）：

Please download the four pre-trained models (~1.1GB each) from Google drive or Baidu netdisk:

The 1-hour model (pangu_weather_1.onnx): Google drive/Baidu netdisk

The 3-hour model (pangu_weather_3.onnx): Google drive/Baidu netdisk

The 6-hour model (pangu_weather_6.onnx): Google drive/Baidu netdisk

The 24-hour model (pangu_weather_24.onnx): Google drive/Baidu netdisk

3 数据下载

3.1ERA5再分析数据

ECMWF（European Centre for Medium-Range Weather Forecasts）的再分析资料ECMWF Reanalysis v5（ERA5）数据。

ERA5再分析资料的下载地址：ERA5 hourly data on pressure levels from 1940 to present (copernicus.eu)

ERA5 hourly data on single levels from 1940 to present (copernicus.eu)

3.2ECMWF集合预报数据

TIGGE 数据集（国际全球大集合）由 13 个全球 NWP 中心自 2006 年 10 月起12 的集合预报数据组成。它是 THORPEX 交互式全球大集合的一部分，是世界气象组织于 2005 年 3 月启动的一项研究计划。

本篇使用TIGGE (The International Grand Global Ensemble)提供的集合预报数据，与盘古气象大模型的预测结果进行对比验证。

数据的下载网址如下：

https://apps.ecmwf.int/datasets/data/tigge/levtype=sfc/type=fc/

4 数据处理

在进行盘古模型的训练之前，我们首先需要知道输入数据的要求。

地表：输入数据分为4个，分别是海表面气压，10m经纬向风和2m温度

高度层：位势，比湿，温度，经纬向风

surface: 4 surface variables (MSLP, U10, V10, T2M)

upper:(Z, Q, T, U and V) 13levs

(1000hPa, 925hPa, 850hPa, 700hPa, 600hPa, 500hPa, 400hPa, 300hPa,250hPa, 200hPa, 150hPa, 100hPa and 50hPa）

721*1440 latitude*longitude

在官网选择好所需的ERA5数据之后，盘古的开发者团队在开源的GitHub项目上对数据类型有所强调，只能输入”.npy”格式数据，且在输入模型时要定义为float.32格式，所以我们要对下载好的数据进行进一步的处理。

代码

代码分为两个部分，一个是处理高度层（upper）数据为npy格式，一个是处理地表数据（surface）到npy格式。

第一部分

import xarray as xr
import numpy as np
import os
df=xr.open_dataset(r"D:\Pangu-Weather-ReadyToGo\Test-Data\mslp-10mU-10mV-2mT-2024.08.00.nc")
time_list=np.array(df.valid_time)
#(MSLP, U10, V10, T2M in the exact order)
save_dir="D:\Pangu-Weather-ReadyToGo\Test-Data\surface_data"
for i in range(len(time_list)):
    time = time_list[i]
    msl = np.expand_dims(df.msl.loc[time, :, :], 0)
    u10 = np.expand_dims(df.u10.loc[time, :, :], 0)
    v10 = np.expand_dims(df.v10.loc[time, :, :], 0)
    t2m = np.expand_dims(df.t2m.loc[time, :, :], 0)
    res_array = np.concatenate([msl,u10, v10, t2m], axis=0)
    strname=f'mslp-10mU-10mV-2mT-{str(time)[:-16]}'+'.npy'
    save_path = os.path.join(save_dir, strname)
    np.save(save_path, res_array)
    print(strname,"saved successfully")

第二部分

import xarray as xr
import numpy as np
import os
df=xr.open_dataset(r"D:\Pangu-Weather-ReadyToGo\Test-Data\Z-Q-T-U-V-13levs-2024.08.00.nc")
time_list=np.array(df.valid_time)
#(5,13,721,1440)
#(Z, Q, T, U and V in the exact order),
# 13 pressure levels (1000hPa, 925hPa, 850hPa, 700hPa, 600hPa, 500hPa,
# 400hPa, 300hPa, 250hPa, 200hPa, 150hPa, 100hPa and 50hPa in the exact order)
# print(df.pressure_level)
print(df)
save_dir=r"D:\Pangu-Weather-ReadyToGo\Test-Data\upper_data"
for i in range(len(time_list)):
    time = time_list[i]
    z = np.expand_dims(df.z.loc[time, :, :], 0)
    q = np.expand_dims(df.q.loc[time, :, :], 0)
    t = np.expand_dims(df.t.loc[time, :, :], 0)
    u = np.expand_dims(df.u.loc[time, :, :], 0)
    v = np.expand_dims(df.v.loc[time, :, :], 0)
    res_array = np.concatenate([z,q,t,u,v], axis=0)
    strname=f'z-q-t-u-v-{str(time)[:-16]}'+'.npy'
    save_path = os.path.join(save_dir, strname)
    np.save(save_path, res_array)
    print(strname,"saved successfully")

以上代码将下载的包含多个时间步的nc数据分割为逐个时间步的npy。

5 模型运行

这里以24h的预测模型为例。注意：需要写好下载的预测模型路径、数据读取路径和数据输出路径。

import os
import numpy as np
import onnx
import onnxruntime as ort
#-------------以下是数据要求-----------
'''
surface: 4 surface variables (MSLP, U10, V10, T2M in the exact order)
upper:(Z, Q, T, U and V) 13levs
(1000hPa, 925hPa, 850hPa, 700hPa, 600hPa, 500hPa, 400hPa, 300hPa,
 250hPa, 200hPa, 150hPa, 100hPa and 50hPa
721*1440 latitude longitude
initial fields of at 12:00UTC, 2018/09/27.
Note that ndarray (.astype(np.float32)), not in double precision.
'''
# The directory of your input and output data
input_data_dir1 = r"D:\Pangu-Weather-ReadyToGo\Test-Data\surface_data"
input_data_dir2 = r"D:\Pangu-Weather-ReadyToGo\Test-Data\upper_data"
output_data_dir = r"D:\Pangu-Weather-ReadyToGo\Test-Data\output_data"
model_24 = onnx.load("D:\Pangu-Weather-ReadyToGo\Pangu-Weather-main\models\pangu_weather_24.onnx")
file_name_surface='mslp-10mU-10mV-2mT-2024-08-06T00.npy'
file_name_surface_out='mslp-10mU-10mV-2mT-2024-08-07T00.npy'
file_name_upper="z-q-t-u-v-2024-08-06T00.npy"
file_name_upper_out='z-q-t-u-v-2024-08-07T00.npy'
# Set the behavier of onnxruntime
options = ort.SessionOptions()
options.enable_cpu_mem_arena=False
options.enable_mem_pattern = False
options.enable_mem_reuse = False
# Increase the number for faster inference and more memory consumption
options.intra_op_num_threads = 1

# Set the behavier of cuda provider
cuda_provider_options = {'arena_extend_strategy':'kSameAsRequested',}

# Initialize onnxruntime session for Pangu-Weather Models
ort_session_24 = ort.InferenceSession(r"D:\Pangu-Weather-ReadyToGo\Pangu-Weather-main\models\pangu_weather_24.onnx",
                                     sess_options=options, providers=[('CUDAExecutionProvider', cuda_provider_options)])

# Load the upper-air numpy arrays
input = np.load(os.path.join(input_data_dir2, file_name_upper)).astype(np.float32)
# Load the surface numpy arrays
input_surface = np.load(os.path.join(input_data_dir1, file_name_surface)).astype(np.float32)
# Run the inference session
output, output_surface = ort_session_24.run(None, {'input':input, 'input_surface':input_surface})
# Save the results
np.save(os.path.join(output_data_dir, f'{file_name_upper_out}'), output)
np.save(os.path.join(output_data_dir, f'{file_name_surface_out}'), output_surface)

最后再将模型的输出数据由npy转成nc，标注上输出各个数据的变量名称，完成模型运行。

同样的，可以使用循环的方式迭代调用盘古模型，或者以贪心算法组合调用不同预测时间步的模型，以下以24h迭代预测为例子：

# Load the upper-air numpy arrays
input = np.load(os.path.join(input_data_dir2, file_name_upper)).astype(np.float32)
# Load the surface numpy arrays
input_surface = np.load(os.path.join(input_data_dir1, file_name_surface)).astype(np.float32)
# Run the inference session
input_24, input_surface_24 = input, input_surface
for i in range(run_time):
    output, output_surface = ort_session_24.run(None, {'input':input_24, 'input_surface':input_surface_24})
    input_24, input_surface_24 = output, output_surface
    # Save the results
    file_name_upper_out = f'z-q-t-u-v-2024-08-0{i+2}T00.npy'
    file_name_surface_out = f'mslp-10mU-10mV-2mT-2024-08-0{i+2}T00.npy'
    np.save(os.path.join(output_data_dir, f'{file_name_upper_out}'), output)
    np.save(os.path.join(output_data_dir, f'{file_name_surface_out}'), output_surface)
    print(f'2024-08-0{i+2}T00 run successfully')
  # Your can save the results here

运行结果

用panoply简单挑出模型的其中一个运行结果（2m温度）来看，结果如下：

6 模型结果评估

6.1 简单评估

这里以广东范围的区域平均海平面气压来进行和EC预报数据的简单对比评估，代码如下：

import numpy as np
import xarray as xr
import matplotlib.pyplot as plt
import os
import glob
from scipy import stats

df=xr.open_dataset(r"D:\Pangu-Weather-ReadyToGo\Test-Data\mslp-10mU-10mV-2mT-2024.08.00.nc")
# u10 v10 t2m msl (valid_time, latitude, longitude) float32 71MB ...
var=df.msl.loc['2024-08-02':'2024-08-07',25.5:20,109.5:117.5]
var_mean=np.mean(var,axis=(1,2))/100
print(var_mean)
time=var.valid_time
#ecmwf的预报
file_path=r"D:\Pangu-Weather-ReadyToGo\Test-Data\ECMWF_2024-08-fc.grib"
ds = xr.open_dataset(file_path,engine='cfgrib')
msl_ecmwf=ds.msl.loc['2024-08-01':'2024-08-06',25.5:20,109.5:117.5]
msl_ecmwf=np.array(np.mean(msl_ecmwf,axis=(1,2))/100)
# 设置文件夹路径
folder_path = r"D:\Pangu-Weather-ReadyToGo\Test-Data\output_nc\24h"  # 替换为你的文件夹路径

# 获取以 "mslp-10mU-10mV-2mT" 开头的所有 .nc 文件
file_list = glob.glob(os.path.join(folder_path, 'mslp-10mU-10mV-2mT*.nc'))#z-q-t-u-v*.nc
# 按照文件名中的日期部分排序（假设日期部分为文件名的第20-29位）
file_list.sort(key=lambda x: x.split('-')[-1].split('.')[0])
print(file_list)
combined_data=[]
# 逐个读取文件并提取数据
for file in file_list:
    nc_file=xr.open_dataset(file)
    #v_component_of_wind_10m u_component_of_wind_10m mean_sea_level_pressure
    t2m = nc_file.mean_sea_level_pressure.loc[25.5:20,109.5:117.5]
    t2m=np.mean(t2m, axis=(0, 1))
    combined_data.append(t2m)
combined_data=np.array(combined_data)/100
# 计算误差的绝对值
# 计算误差的绝对值
error1 = np.abs(var_mean - combined_data)
error2 = np.abs(var_mean - msl_ecmwf)
plt.figure(figsize=[8,4])
# 设置主图风格为Nature风格
plt.rcParams.update({
    'font.size': 10,
    'axes.linewidth': 1.0,
    'axes.spines.top': False,
    'axes.spines.right': False,
    'xtick.major.width': 1.0,
    'ytick.major.width': 1.0,
    'legend.frameon': False,
    'legend.fontsize': 10
})
plt.title('Sea Level Pressure [hpa]', fontsize=12)
plt.ylim(1006, 1010)
plt.plot(time, var_mean, label='ERA5', c='k', lw=1.5)
plt.plot(time, combined_data, label='PanGu', c='red', lw=1.5)
plt.plot(time,msl_ecmwf,label='EC', c='grey', lw=1.5)
plt.legend(loc='lower right')
# 创建嵌入的误差柱状图
inset_ax = plt.gca().inset_axes([0.12, 0.15, 0.25, 0.3])  # [x, y, width, height]
inset_ax.bar(np.arange(1,7,1),error1, color='red', alpha=1)
inset_ax.bar(np.arange(1,7,1),error2, color='k', alpha=0.5)
# # 对误差进行线性回归拟合
# slope, intercept, r_value, p_value, std_err = stats.linregress(np.arange(1,7,1), error)
# fit_line = slope * np.arange(1,7,1) + intercept
# inset_ax.plot(np.arange(1,7,1),fit_line,color='blue', lw=1.5)
# 设置误差图的标题和样式
inset_ax.set_title('Absolute Error', fontsize=6)
inset_ax.tick_params(axis='both', which='major', labelsize=8)
inset_ax.spines['top'].set_visible(False)
inset_ax.spines['right'].set_visible(False)
inset_ax.spines['bottom'].set_linewidth(1.0)
inset_ax.spines['left'].set_linewidth(1.0)
plt.show()

结果如下：