Skip to content

Commit 9b687ac

Browse files
liaoxin2HydrogenSulfate
liaoxin2
and
HydrogenSulfate
authored
kmcast example (#1247)
* Modify a bug in the gencast model * add kmcast model * add kmcast code * Modify log-related code * add kmcast example * Update examples/kmcast/kmcast.py * Update examples/kmcast/kmcast.py * Update examples/kmcast/kmcast.py * Update examples/kmcast/kmcast.py --------- Co-authored-by: HydrogenSulfate <490868991@qq.com>
1 parent 6e71cc2 commit 9b687ac

File tree

15 files changed

+1508
-0
lines changed

15 files changed

+1508
-0
lines changed

‎README.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -133,6 +133,7 @@ PaddleScience 是一个基于深度学习框架 PaddlePaddle 开发的科学计
133133

134134
| 问题类型 | 案例名称 | 优化算法 | 模型类型 | 训练方式 | 数据集 | 参考资料 |
135135
|-----|---------|-----|---------|----|---------|---------|
136+
| 气象降尺度 | [KMCast](https://paddlescience-docs.readthedocs.io/zh-cn/latest/zh/examples/kmcast) | 数据驱动 | Diffusion | 监督学习 | [GFS](https://rda.ucar.edu/datasets/d084006/) | - |
136137
| 天气预报 | [Extformer-MoE 气象预报](https://paddlescience-docs.readthedocs.io/zh-cn/latest/zh/examples/extformer_moe) | 数据驱动 | Transformer | 监督学习 | [enso](https://tianchi.aliyun.com/dataset/98942) | - |
137138
| 天气预报 | [FourCastNet 气象预报](https://paddlescience-docs.readthedocs.io/zh-cn/latest/zh/examples/fourcastnet) | 数据驱动 | AFNO | 监督学习 | [ERA5](https://app.globus.org/file-manager?origin_id=945b3c9e-0f8c-11ed-8daf-9f359c660fbd&origin_path=%2F~%2Fdata%2F) | [Paper](https://arxiv.org/pdf/2202.11214.pdf) |
138139
| 天气预报 | [NowCastNet 气象预报](https://paddlescience-docs.readthedocs.io/zh-cn/latest/zh/examples/nowcastnet) | 数据驱动 | GAN | 监督学习 | [MRMS](https://app.globus.org/file-manager?origin_id=945b3c9e-0f8c-11ed-8daf-9f359c660fbd&origin_path=%2F~%2Fdata%2F) | [Paper](https://www.nature.com/articles/s41586-023-06184-4) |

‎docs/index.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -150,6 +150,7 @@
150150

151151
| 问题类型 | 案例名称 | 优化算法 | 模型类型 | 训练方式 | 数据集 | 参考资料 |
152152
|-----|---------|-----|---------|----|---------|---------|
153+
| 气象降尺度 | [KMCast](./zh/examples/kmcast.md) | 数据驱动 | Diffusion | 监督学习 | [GFS](https://rda.ucar.edu/datasets/d084006/) | - |
153154
| 天气预报 | [Extformer-MoE 气象预报](./zh/examples/extformer_moe.md) | 数据驱动 | Transformer | 监督学习 | [enso](https://tianchi.aliyun.com/dataset/98942) | - |
154155
| 天气预报 | [FourCastNet 气象预报](./zh/examples/fourcastnet.md) | 数据驱动 | AFNO | 监督学习 | [ERA5](https://app.globus.org/file-manager?origin_id=945b3c9e-0f8c-11ed-8daf-9f359c660fbd&origin_path=%2F~%2Fdata%2F) | [Paper](https://arxiv.org/pdf/2202.11214.pdf) |
155156
| 天气预报 | [NowCastNet 气象预报](./zh/examples/nowcastnet.md) | 数据驱动 | GAN | 监督学习 | [MRMS](https://app.globus.org/file-manager?origin_id=945b3c9e-0f8c-11ed-8daf-9f359c660fbd&origin_path=%2F~%2Fdata%2F) | [Paper](https://www.nature.com/articles/s41586-023-06184-4) |

‎docs/zh/examples/kmcast.md‎

Lines changed: 166 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,166 @@
1+
# KMCast
2+
3+
<a href="https://aistudio.baidu.com/projectdetail/9801067" class="md-button md-button--primary" style>AI Studio快速体验</a>
4+
5+
开始训练、评估前，请下载数据集:
6+
7+
``` sh
8+
# GFS 数据
9+
wget -nc https://paddle-org.bj.bcebos.com/paddlescience/datasets/kmcast/GFS_all_spinup.nc -P ./dataset/
10+
# WRF 数据
11+
wget -nc https://paddle-org.bj.bcebos.com/paddlescience/datasets/kmcast/WRF_all_spinup.nc -P ./dataset/
12+
```
13+
14+
=== "模型训练命令"
15+
16+
``` sh
17+
python kmcast.py
18+
```
19+
20+
=== "模型评估命令"
21+
22+
``` sh
23+
# 下载模型权重文件
24+
wget -nc https://paddle-org.bj.bcebos.com/paddlescience/models/kmcast/I597500_E184_gen.pdparams
25+
wget -nc https://paddle-org.bj.bcebos.com/paddlescience/models/kmcast/I597500_E184_opt.pdparams
26+
# 运行评估
27+
python kmcast.py mode=eval eval.pretrained_model_path=./I597500_E184
28+
```
29+
30+
## 1 背景简介
31+
32+
随着全球风电装机容量的迅速扩张，风电从资源评估、机组控制到电网调度等环节，都在逐渐依赖高时空分辨率的气象信息。风速在海岸、岛屿及复杂地形区域往往受到海陆风环流、局地急流、山谷风、边界层结构等多种中小尺度过程的影响，这些过程决定着机组来流风的细节，也决定着风电场的功率稳定性和极端风风险。如果风速场过于平滑或偏差过大，就会直接导致发电量预测错误、调度策略不合理，或在长期规划中形成不正确的投资判断。因此，一个能够刻画公里级空间结构的风场对于风电行业至关重要。
33+
34+
然而，目前最常使用的天气和气候模式在分辨率上仍受到严重限制。全球预报模式（如 GFS）以及多项 AI 天气模式通常在 25 公里量级，这种尺度无法解析海岸线细节、地形突变、小型岛屿或海湾，因此不能表现真实存在的局地加速带、遮蔽尾迹和风速突变。气候模式的分辨率更为粗糙，往往在百公里量级，使得整个风电场甚至多个沿海站点在模式上表现得几乎一致，与真实世界的差异非常明显。模式无法显式模拟对流和湍流过程，而这些过程的参数化正是风速偏差与“场过于光滑”的重要来源。
35+
36+
为了获得更接近真实的风场结构，许多研究依赖区域数值模式（例如 WRF）的动力降尺度，通过更细的网格和显式对流来重建局地风场。这样的模拟确实能够显著改善风场结构和极端统计，但其计算成本极高。即便仅限于单一地区、每天模拟一次，其运算仍可能需要大量计算资源，这使得在实际风电运营中每日运行高分辨率模拟几乎不现实；在气候研究中对长时间序列进行公里级动力降尺度更是成本难以承受。
37+
38+
近年来，深度学习成为一种潜在的替代方案，通过学习高分辨率模拟的统计结构，将粗分辨率模式的输出“补充”成更精细的风场图像。但是现有的机器学习降尺度研究主要聚焦于降水、温度等变量，对于近地面风场尤其是海岸和复杂地形区域的风速精细结构关注较少。此外，大多数方法仍然将天气预报和气候模拟视为两个截然独立的任务，缺乏一种能够同时处理短期天气预测与长期气候分析的统一框架，也缺乏对风电行业最关心的极端风速、局地梯度和不确定性的系统刻画。
39+
40+
因此，本研究背景可归结为：在风电行业对公里级气象信息需求持续增加的形势下，传统预报模式分辨率不足、动力降尺度成本过高，而现有深度学习方法覆盖有限、缺乏统一的天气–气候框架。在这样的背景下，开发一种能够以极低成本生成可信的公里级风场、既服务短时预报又服务长期风气候评估的新方法变得十分迫切。KMCast 正是在这一需求缺口下提出的，通过生成式扩散模型学习高分辨率模拟的空间细节，以弥补粗分辨率模式的不足，同时建立一种贯通天气预报和气候模拟的统一方法体系。
41+
42+
## 2 模型原理
43+
44+
### 2.1 扩散模型
45+
46+
KMCast主要采用条件扩散模型，采用逐步去噪的方式，在低分辨率输入的条件下生成高分辨率图像。在前向扩散过程中，逐步向高分辨率图像中加入高斯噪声，经过多个时间步，图像逐渐转变为各向同性的噪声分布。通过以下随机过程，我们将目标分布 $p(x_0)$ 转变为标准高斯噪声：
47+
48+
$$
49+
q(x_t|x_{t-1}) = \mathcal{N}(\sqrt{1 - \beta_t}x_{t-1}, \beta_t I)
50+
$$
51+
52+
其中， $x_t$ 是随时间变化的潜在变量，时间索引 $t$ 取值范围为 $[0, T]$ ，超参数 $\beta_t$ 随着时间步 $t$ 递增，呈单调增长的趋势。
53+
54+
在逆向阶段，模型以低分辨率输入 $y$ 为条件，利用U-Net架构逐步预测并去除噪声，通过梯度下降优化实现图像复原。$\mu_\theta$ 和 $\Sigma_\theta$ 分别为模型中可学习的均值向量和协方差矩阵，这些参数通过最大化变分下界（变分推断）对负对数似然的优化进行训练。逆向的条件概率表示为：
55+
56+
$$
57+
p_\theta(x_{t-1}|x_t, y) = \mathcal{N}(\mu_\theta(x_t, y, t), \Sigma_\theta(x_t, y, t))
58+
$$
59+
60+
模型的损失函数定义为：
61+
62+
$$
63+
L:=E_q\left[-\log p\left(x_T\right)-\sum_{t \geq 1} \log \frac{p_\theta\left(x_{t-1} \mid x_t, y\right)}{q\left(x_t \mid x_{t-1}\right)}\right] \geq E\left[-\log p_\theta\left(x_0\right)\right]
64+
$$
65+
66+
这一双向机制使得模型能够从纯噪声中逐步重建出高分辨率图像，同时保持图像的结构完整性。
67+
68+
69+
### 2.2 网络架构与训练细节
70+
71+
为了在多尺度生成过程中保持模型的兼容性和稳定性，我们采用了UNet架构作为噪声预测器。该模型经过优化，增加了6个编码层和6个解码层。基本通道数设置为32，通道倍增方案为[1, 2, 4, 8, 8]。每一层包含两个残差块。整个UNet模型拥有约2300万参数。
72+
73+
输入数据在扩散训练过程中由14个粗分辨率的GFS通道组成。为了降低计算负担，同时保留关键气象特征，我们选择了从原始WRF输出数据中提取的代表性子区域。北哥伦比亚区域定义为经度77.5°W至69.5°W，纬度8°N至13.5°N；长江三角洲区域定义为经度117°E至123.5°E，纬度29°N至34.5°N。两个区域的空间尺寸均为192×256像素。模型输出的高分辨率风场包括东西向和南北向分量，空间尺寸与输入保持一致。
74+
75+
在训练过程中，采用Adam优化器，学习率设为 $1\times 10^{-4}$ ， $\beta_1=0.9,\beta_2=0.999$ 。KMCast条件扩散模型在其网络架构中使用SiLU激活函数。为了防止过拟合，训练中采用了0.2的dropout率，并对训练数据进行随机打乱。此外， $\beta_1$ 参数遵循线性调度，从 $1\times 10^{-6}$ 开始，经过2000个时间步逐步增加至 $1\times 10^{-2}$ 。模型在一台A100 GPU上训练了58万步，总批次大小为4，整个训练时间约为2天。
76+
77+
为气候预测目的，我们还训练了一个每日尺度的模型，用于适应气候模型的输出。该模型的架构与小时尺度模型相同，唯一的区别在于时间分辨率和输入仅包括东西向和南北向风分量，适用于5天预报。
78+
79+
### 2.3 评价指标
80+
81+
为了全面衡量模型预测的准确性和不确定性，我们采用了两项关键指标：平均绝对误差（MAE）和均方根误差（RMSE）。具体定义如下：
82+
83+
$$
84+
\text{MAE} = \frac{1}{n} \sum_{i=1}^{n} |x_i - \hat{x}_i|
85+
$$
86+
87+
$$
88+
\text{RMSE} = \sqrt{\frac{1}{n} \sum_{i=1}^{n} (x_i - \hat{x}_i)^2}
89+
$$
90+
91+
其中， $\hat{x}_i$ 代表模型的预测值， $x_i$ 为真实值， $n$ 为样本数量。 KMCast模型整体架构如图所示：
92+
93+
<figure markdown>
94+
![kmcast.png](https://paddle-org.bj.bcebos.com/paddlescience/docs/kmcast/kmcast.png){ loading=lazy }
95+
</figure>
96+
97+
## 3 数据说明
98+
99+
### 3.1 观测数据集
100+
101+
在哥伦比亚地区，收集了来自四个气象站的观测数据：
102+
103+
1. 阿尔米兰特·帕迪利亚（Almirante Padilla，东经-72.926°，北纬11.526°）
104+
2. 西蒙·玻利瓦尔（Simon Bolivar，东经-74.231°，北纬11.120°）
105+
3. 拉斐尔·努涅斯（Rafael Nunez，东经-75.513°，北纬10.442°）
106+
4. 欧内斯特·科尔蒂索斯（Ernesto Cortissoz，东经-74.781°，北纬10.890°）
107+
108+
所有站点提供2022年全年每小时的观测数据。
109+
110+
为了更好地评估长期模拟的性能，还在中国进行了补充验证。中国拥有更多的气象站点和更长的观测时间段。我们从中国日报气象要素站观测数据中获取了2008年至2018年的10米东西向（U分量）和南北向（V分量）风场数据，该数据由全国超过4000个气象站每日计算值组成。
111+
112+
用于哥伦比亚北部地区的GFS（全球预报系统）数据涵盖874天，每次初始化后可进行10天预报，时间分辨率为3小时（共81个预报步骤），每个预报中选择15个时间点，共计13,110个时刻。对于WRF（微尺度气象模型）数据，我们提取了相同的13,110个时间点，以确保与GFS数据的时间对齐。在模型训练和测试策略方面，最新的2,000个时间点用于测试，其余数据用于模型训练。
113+
114+
中国区域的WRF数据覆盖了334天，每个预报选择15个时间点，共计5,110个时间点，全部用于气候模型的训练。
115+
116+
输入数据集来自GFS，经过双线性插值处理，空间分辨率调整至3公里。我们采用每日0000 UTC启动的10天预报，空间分辨率为25公里，时间分辨率为3小时，覆盖2020年至2022年的全年数据。
117+
118+
用于千米尺度气候模拟的输入为FGOALS-f3-H的历史数据，该数据经过分位数映射偏差校正，基于GFS数据进行校正。校正后数据保持GFS的空间分辨率，但时间分辨率为每日，覆盖2012年至2014年的全年数据。
119+
120+
为了训练KMCast模型以预测哥伦比亚北部地区的风，选取了六个与风预报最相关的气象变量：位势高度、相对湿度、温度、大气压力，以及风矢量的东西向和南北向分量。这些变量在不同的垂直层面上组成共计14个输入通道。具体细节见下表。
121+
122+
| Variable | Selected Layers |
123+
| --- | --- |
124+
| Zonal/Meridional Wind | 10m above ground, 500hPa, 200hPa |
125+
| Temperature | 2m above ground, 925hPa, 850hPa, 700hPa, 500hPa |
126+
| Geopotential Height | 850hPa |
127+
| Relative Humidity | 2m above ground |
128+
| Pressure | Surface |
129+
130+
NCEP全球预报系统（GFS）分析资料和预报数据可在 [https://rda.ucar.edu/datasets/d084006/](https://rda.ucar.edu/datasets/d084006/) 获取。
131+
132+
### 3.2 WRF配置
133+
134+
我们选择哥伦比亚北部地区和中国大陆作为目标区域，利用WRF模型对粗分辨率的GFS数据进行动态降尺度，生成高分辨率数据集。该数据集具有15分钟的时间分辨率和3公里的空间分辨率，采用WRF-ARW V4.5模型系统进行模拟。
135+
哥伦比亚地区的地理范围为北纬8°至14°，西经79.5°至72.5°；中国地区则定义在北纬27°至34°，东经117°至125°。为了实现3公里的模拟精度，采用两层嵌套域结构。两个域都采用兰伯特等角投影。
136+
在中国区域，父域分辨率为9公里，网格点为309×269；嵌套子域的分辨率为3公里，网格点为787×697。哥伦比亚区域的父域网格为173×130，嵌套区域为319×238。
137+
在垂直方向上，两个域均包含34个sigma层，模型顶部设置在50hPa。哥伦比亚北部区域的WRF模拟每天在UTC 0000时启动，覆盖2020年1月至6月以及2021年和2022年的全年。中国大陆区域的模拟也在每天UTC 0000时启动，覆盖2022年全年度。
138+
139+
140+
## 4 模型代码说明
141+
142+
- **`conf/kmcast.yaml`**：配置文件，定义模型运行的参数和设置，用于控制模型的行为和参数配置。
143+
- **`core/metrics.py`**：性能指标计算模块，用于评估模型的预测效果。
144+
- **`data/LRHR_dataset.py`**：定义数据集加载与预处理流程的脚本。
145+
- **`model/sr3_modules`**：包含模型中的核心模块和子模型。
146+
- **`diffusion.py`**：扩散模型相关代码。
147+
- **`unet.py`**：U-Net结构实现，用于图像处理。
148+
- **`model/base_model.py`**：定义模型的基础结构和框架。
149+
- **`model/model.py`**：模型的主定义文件。
150+
- **`model/netsworks.py`**：网络结构定义。
151+
- **`kmcast.py`**：主执行脚本，整合模型运行、训练和预测流程。
152+
153+
## 5 结果展示
154+
155+
如图所示，左侧为KMCast模型预测结果，中间为观测到的真实风场数据，右侧为输入的低分辨率GFS数据。
156+
157+
<figure markdown>
158+
![result.png](https://paddle-org.bj.bcebos.com/paddlescience/docs/kmcast/result.png
159+
){ loading=lazy }
160+
</figure>
161+
162+
如图所示，KMCast模型在哥伦比亚北部预测结果与观测到的风场数据高度一致，验证了其在不同地理区域和气候条件下的泛化能力。
163+
164+
## 6 参考资料
165+
166+
- Fast, High Resolution Wind Information for Operations and Planning via Generative Downscaling

‎examples/kmcast/conf/kmcast.yaml‎

Lines changed: 95 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,95 @@
1+
defaults:
2+
- ppsci_default
3+
- TRAIN: train_default
4+
- TRAIN/ema: ema_default
5+
- TRAIN/swa: swa_default
6+
- EVAL: eval_default
7+
- INFER: infer_default
8+
- hydra/job/config/override_dirname/exclude_keys: exclude_keys_default
9+
- _self_
10+
11+
hydra:
12+
run:
13+
# dynamic output directory according to running time and override name
14+
dir: kmcast/${now:%Y-%m-%d}/${now:%H-%M-%S}/${hydra.job.override_dirname}
15+
job:
16+
name: ${mode} # name of logfile
17+
chdir: false # keep current working directory unchanged
18+
callbacks:
19+
init_callback:
20+
_target_: ppsci.utils.callbacks.InitCallback
21+
sweep:
22+
# output directory for multirun
23+
dir: ${hydra.run.dir}
24+
subdir: ./
25+
26+
# general settings
27+
mode: train # running mode: train/eval
28+
seed: 2024
29+
log_freq: 20
30+
output_dir: ${hydra:run.dir}
31+
gpu_ids: [0]
32+
distributed: False
33+
34+
path:
35+
log: ${hydra:run.dir}/logs
36+
tb_logger: ${hydra:run.dir}/tb_logger
37+
results: ${hydra:run.dir}/results
38+
checkpoint: ${hydra:run.dir}/checkpoint
39+
datasets:
40+
val_samples: 2000
41+
dataroot: dataset
42+
ds_gfs_path: ./dataset/GFS_all_spinup.nc
43+
ds_wrf_path: ./dataset/WRF_all_spinup.nc
44+
train:
45+
batch_size: 4
46+
num_workers: 1
47+
use_shuffle: true
48+
data_len: -1 # -1 represents all data used in train
49+
eval:
50+
batch_size: 32
51+
data_len: 50 # data length in validation
52+
53+
model:
54+
which_model_G: sr3 # use the ddpm or sr3 network structure
55+
finetune_norm: false
56+
unet:
57+
in_channel: 3
58+
out_channel: 1
59+
inner_channel: 32
60+
channel_multiplier: [1,2,4,8,8]
61+
res_blocks: 2
62+
dropout: 0.2
63+
norm_groups: 32
64+
beta_schedule:
65+
train:
66+
schedule: linear
67+
n_timestep: 2000
68+
linear_start: 1e-6
69+
linear_end: 1e-2
70+
eval:
71+
schedule: linear
72+
n_timestep: 2000
73+
linear_start: 1e-6
74+
linear_end: 1e-2
75+
diffusion:
76+
image_H: 196
77+
image_W: 256
78+
channels: 1 # sample channel
79+
conditional: true # unconditional generation or super-resolution
80+
81+
train:
82+
n_iter: 597500
83+
val_freq: 2500
84+
save_checkpoint_freq: 2500
85+
print_freq: 50
86+
optimizer:
87+
type: adam
88+
lr: 0.0001
89+
90+
eval:
91+
pretrained_model_path: null
92+
hr_min: -33.122757
93+
hr_max: 57.7706
94+
lr_min: -34.006493
95+
lr_max: 34.886944

‎examples/kmcast/core/metrics.py‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
import paddle
2+
3+
4+
def inv_transform_minus_one_to_one(data, min_val, max_val):
5+
return data * (max_val - min_val) + min_val
6+
7+
8+
def tensor2rawdata(tensor, min_val, max_val, norm_min_max=(0, 1)):
9+
"""Not Support auto convert *.clamp_, please judge whether it is Pytorch API and convert by yourself"""
10+
tensor = tensor.squeeze().clip_(*norm_min_max)
11+
tensor = inv_transform_minus_one_to_one(tensor, min_val, max_val)
12+
return tensor
13+
14+
15+
def save_img(img_turple, latlon_turple, img_path):
16+
sr_img, hr_img, lr_img = img_turple
17+
lat, lon = latlon_turple
18+
import matplotlib
19+
20+
matplotlib.use("Agg")
21+
import matplotlib.pyplot as plt
22+
23+
fig, axes = plt.subplots(3, 3, figsize=(15, 15))
24+
axes[0, 0].set_title("SR (Model Output)")
25+
axes[0, 1].set_title("HR (Ground Truth)")
26+
axes[0, 2].set_title("LR (Bilinear)")
27+
vmin = min(sr_img.min(), hr_img.min(), lr_img.min())
28+
vmax = max(sr_img.max(), hr_img.max(), lr_img.max())
29+
cmap = plt.cm.RdYlBu_r
30+
for i in range(3):
31+
axes[i, 0].pcolormesh(lon, lat, sr_img[i], vmin=vmin, vmax=vmax, cmap=cmap)
32+
axes[i, 0].set_ylabel(f"Sample {i + 1}")
33+
axes[i, 1].pcolormesh(lon, lat, hr_img[i], vmin=vmin, vmax=vmax, cmap=cmap)
34+
axes[i, 2].pcolormesh(lon, lat, lr_img[i, 1], vmin=vmin, vmax=vmax, cmap=cmap)
35+
plt.tight_layout()
36+
plt.savefig(img_path, bbox_inches="tight", dpi=300)
37+
plt.close()
38+
39+
40+
def calculate_rmse_sum(img1, img2):
41+
mse = paddle.mean(x=(img1 - img2) ** 2, axis=(-2, -1))
42+
return paddle.sqrt(x=mse).sum()

0 commit comments

Comments
 (0)