地球科学AI极端降水预报之baseline（Datawhale AI 夏令营）

        在本次博客中，我们将详细介绍如何利用伏羲气象大模型和ERA5数据，构建一个AI模型，用于预测未来72小时的极端降水情况。本文不仅涵盖了整个比赛的数据处理和模型构建流程，还提供了详细的代码说明，帮助参赛者理解和实现预报模型。

比赛背景

        随着全球气候变化，极端天气事件的频率和强度不断增加，准确预测这些事件对于防灾减灾具有重要意义。传统的数值天气预报方法存在一定局限，而AI技术可以通过学习大量历史数据，提高预报的精度和效率。本次比赛旨在通过AI技术，实现更精确的极端降水预报。

数据介绍

        比赛数据分为两部分：

1. ERA5 数据：

   • 来源：欧洲中期天气预报中心(ECMWF)的再分析数据。
   • 时间间隔：1小时，空间分辨率0.25度。
   • 目标变量：累积降水（tp）。
   • 覆盖区域：N18 ~ N32, E103 ~ E123。

2. FuXi 数据：

   • 来源：伏羲气象大模型。
   • 时间间隔：1小时，空间分辨率0.25度。
   • 输入变量包括高空和地面的多个气象要素，例如温度、风速、云量等。

数据处理与模型构建

        接下来，我们将逐步介绍数据处理、模型构建和训练的具体步骤。

第一步：安装必要的库并解压数据集

        首先，安装所需的Python库，并解压下载的数据集。

!pip install xarray[complete]
!unzip -q -n weather.round1.train.gt.2019-2021.zip -d groundtruth
!unzip -q -n weather.round1.test.zip -d test
print("数据集解压完成！")

第二步：导入必要的库函数

import os
import pandas as pd
import xarray as xr
from torch.utils.data import Dataset, DataLoader
import torch.nn as nn
import torch

第三步：数据集路径配置

        配置数据集路径，定义所使用的年份和预测时间步长。

feature_path = 'feature'
gt_path = 'groundtruth'
years = ['2021']
fcst_steps = list(range(1, 73, 1))

第四步：定义数据集类

        定义Feature类和GroundTruth类，用于读取和处理数据。

class Feature:
    def __init__(self):
        self.path = feature_path
        self.years = years
        self.fcst_steps = fcst_steps
        self.features_paths_dict = self.get_features_paths()

    def get_features_paths(self):
        init_time_path_dict = {}
        for year in self.years:
            init_time_dir_year = os.listdir(os.path.join(self.path, year))
            for init_time in sorted(init_time_dir_year):
                init_time_path_dict[pd.to_datetime(init_time)] = os.path.join(self.path, year, init_time)
        return init_time_path_dict

    def get_fts(self, init_time):
        return xr.open_mfdataset(self.features_paths_dict.get(init_time) + '/*').sel(lead_time=self.fcst_steps).isel(time=0)

class GT:
    def __init__(self):
        self.path = gt_path
        self.years = years
        self.fcst_steps = fcst_steps
        self.gt_paths = [os.path.join(self.path, f'{year}.nc') for year in self.years]
        self.gts = xr.open_mfdataset(self.gt_paths)

    def parser_gt_timestamps(self, init_time):
        return [init_time + pd.Timedelta(f'{fcst_step}h') for fcst_step in self.fcst_steps]

    def get_gts(self, init_time):
        return self.gts.sel(time=self.parser_gt_timestamps(init_time))

第五步：定义自定义数据集类

        定义mydataset类，将特征和对应的真值整合。

class mydataset(Dataset):
    def __init__(self):
        self.ft = Feature()
        self.gt = GT()
        self.features_paths_dict = self.ft.features_paths_dict
        self.init_times = list(self.features_paths_dict.keys())

    def __getitem__(self, index):
        init_time = self.init_times[index]
        ft_item = self.ft.get_fts(init_time).to_array().isel(variable=0).values
        gt_item = self.gt.get_gts(init_time).to_array().isel(variable=0).values
        return ft_item, gt_item

    def __len__(self):
        return len(list(self.init_times))

第六步：实例化数据集并查看数据

        通过实例化mydataset类，查看数据数量和加载数据。

my_data = mydataset()
print('样本数量:', mydataset().__len__())
train_loader = DataLoader(my_data, batch_size=1, shuffle=True)

第七步：构建模型

        构建一个简单的卷积神经网络。

class Model(nn.Module):
    def __init__(self, num_in_ch, num_out_ch):
        super(Model, self).__init__()
        self.conv1 = nn.Conv2d(num_in_ch, num_out_ch, 3, 1, 1)

    def forward(self, x):
        B, S, C, W, H = tuple(x.shape)
        x = x.reshape(B, -1, W, H)
        out = self.conv1(x)
        out = out.reshape(B, S, W, H)
        return out

in_varibales = 24
in_times = len(fcst_steps)
out_varibales = 1
out_times = len(fcst_steps)
input_size = in_times * in_varibales
output_size = out_times * out_varibales
model = Model(input_size, output_size).cuda()

第八步：定义损失函数

        定义模型的损失函数，用于训练时的反向传播。

loss_func = nn.MSELoss()

第九步：模型训练

        训练模型并保存训练好的权重。

num_epochs = 1
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

for epoch in range(num_epochs):
    for index, (ft_item, gt_item) in enumerate(train_loader):
        ft_item = ft_item.cuda().float()
        gt_item = gt_item.cuda().float()
        
        # Forward pass
        output_item = model(ft_item)
        loss = loss_func(output_item, gt_item)
        
        # Backward and optimize
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        if (index+1) % 10 == 0:
            print(f"Epoch [{epoch+1}/{num_epochs}], Step [{index+1}/{len(train_loader)}], Loss: {loss.item():.4f}")

torch.save(model.state_dict(), 'model_weights.pth')

第十步：模型推理

        加载训练好的模型，并使用测试数据进行预测。

model.load_state_dict(torch.load('model_weights.pth'))
model.eval()

test_data_path = "test/weather.round1.test"
os.makedirs("./output", exist_ok=True)

for index, test_data_file in enumerate(os.listdir(test_data_path)):
    test_data = torch.load(os.path.join(test_data_path, test_data_file))
    test_data = test_data.cuda().float()
    
    # Forward pass
    output_item = model(test_data)
    
    # Print the output shape
    print(f"Output shape for sample {test_data_file.split('.')[0]}: {output_item.shape}")
    
    # Save the output
    output_path = f"output/{test_data_file}"
    torch.save(output_item.cpu(), output_path)

!zip -r output.zip output

print("模型推理完成，结果已保存为output.zip")

代码实现逻辑

        本次实现的AI极端降水预报模型主要包括数据准备、模型构建、模型训练和模型推理四个部分：首先，配置环境并解压和读取数据，然后定义数据处理类以便加载特征和真值数据；接着，构建一个简单的卷积神经网络模型，并定义损失函数；在训练过程中，通过前向传播、计算损失和反向传播来更新模型参数，最后使用训练好的模型对测试数据进行推理，并保存预测结果。

结语

        通过本文，我们详细介绍了AI极端降水预报模型的构建过程，包括数据处理、模型构建、训练和推理等多个环节。希望这篇博客能帮助大家更好地理解和实现AI极端降水预报模型，提升预报的精度和效率。在实际应用中，该模型有助于更准确地预测极端降水事件，为防灾减灾提供重要参考。祝各位参赛者在比赛中取得优异成绩，也希望本次的实现能够为未来的气象预报研究提供一些新的思路和方法。

如果你觉得这篇博文对你有帮助，请点赞、收藏、关注我，并且可以打赏支持我！

欢迎关注我的后续博文，我将分享更多关于人工智能、自然语言处理和计算机视觉的精彩内容。

谢谢大家的支持！