gencast (#1130)

* Modify the document of heat_pinn

* add gencast code

* add gencast code

* add gencast code

* add gencast code

* add gencast code

* modify gencast code

Author: liaoxin2

README.md

@@ -117,6 +117,7 @@ PaddleScience 是一个基于深度学习框架 PaddlePaddle 开发的科学计
 | 天气预报 | [FourCastNet 气象预报](https://paddlescience-docs.readthedocs.io/zh-cn/latest/zh/examples/fourcastnet) | 数据驱动 | FourCastNet | 监督学习 | [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) |
 | 天气预报 | [NowCastNet 气象预报](https://paddlescience-docs.readthedocs.io/zh-cn/latest/zh/examples/nowcastnet) | 数据驱动 | NowCastNet | 监督学习 | [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) |
 | 天气预报 | [GraphCast 气象预报](https://paddlescience-docs.readthedocs.io/zh-cn/latest/zh/examples/graphcast) | 数据驱动 | GraphCastNet | 监督学习 | - | [Paper](https://arxiv.org/abs/2212.12794) |
+| 天气预报 | [GenCast 气象预报](https://paddlescience-docs.readthedocs.io/zh-cn/latest/zh/examples/gencast) | 数据驱动 | Diffusion | 监督学习 | [Gencast](https://console.cloud.google.com/storage/browser/dm_graphcast) | [Paper](https://arxiv.org/abs/2312.15796) |
 | 天气预报 | [FengWu 气象预报](https://paddlescience-docs.readthedocs.io/zh-cn/latest/zh/examples/fengwu) | 数据驱动 | Transformer | 监督学习 | - | [Paper](https://arxiv.org/pdf/2304.02948) |
 | 天气预报 | [Pangu-Weather 气象预报](https://paddlescience-docs.readthedocs.io/zh-cn/latest/zh/examples/pangu_weather) | 数据驱动 | Transformer | 监督学习 | - | [Paper](https://arxiv.org/pdf/2211.02556) |
 | 大气污染物 | [UNet 污染物扩散](https://aistudio.baidu.com/projectdetail/5663515?channel=0&channelType=0&sUid=438690&shared=1&ts=1698221963752) | 数据驱动 | UNet | 监督学习 | [Data](https://aistudio.baidu.com/datasetdetail/198102) | - |

docs/index.md

@@ -150,6 +150,7 @@
 | 天气预报 | [FourCastNet 气象预报](./zh/examples/fourcastnet.md) | 数据驱动 | FourCastNet | 监督学习 | [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) |
 | 天气预报 | [NowCastNet 气象预报](./zh/examples/nowcastnet.md) | 数据驱动 | NowCastNet | 监督学习 | [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) |
 | 天气预报 | [GraphCast 气象预报](./zh/examples/graphcast.md) | 数据驱动 | GraphCastNet | 监督学习 | - | [Paper](https://arxiv.org/abs/2212.12794) |
+| 天气预报 | [GenCast 气象预报](./zh/examples/gencast.md) | 数据驱动 | Diffusion | 监督学习 | [Gencast](https://console.cloud.google.com/storage/browser/dm_graphcast) | [Paper](https://arxiv.org/abs/2312.15796) |
 | 天气预报 | [FengWu 气象预报](./zh/examples/fengwu.md) | 数据驱动 | Transformer | 监督学习 | - | [Paper](https://arxiv.org/pdf/2304.02948) |
 | 天气预报 | [Pangu-Weather 气象预报](./zh/examples/pangu_weather.md) | 数据驱动 | Transformer | 监督学习 | - | [Paper](https://arxiv.org/pdf/2211.02556) |
 | 大气污染物 | [UNet 污染物扩散](https://aistudio.baidu.com/projectdetail/5663515?channel=0&channelType=0&sUid=438690&shared=1&ts=1698221963752) | 数据驱动 | UNet | 监督学习 | [Data](https://aistudio.baidu.com/datasetdetail/198102) | - |

docs/zh/examples/gencast.md

@@ -0,0 +1,99 @@
+# GenCast
+
+开始评估前，请在 [Google Cloud Bucket](https://console.cloud.google.com/storage/browser/dm_graphcast) 上获取相关数据，并将之放到`gencast.yaml`文件中数据配置的路径下。
+
+- 下载目录`dm_graphcast/gencast/stats`下的所有文件放入`./data/stats/`目录下。
+- 下载目录`dm_graphcast/gencast/dataset`下的任意或所有文件（例如：source-era5_date-2019-03-29_res-1.0_levels-13_steps-12.nc）放入`./data/dataset/`目录下。
+
+=== "模型评估命令"
+
+    ``` sh
+    # 设置路径到 PaddleScience/jointContribution 文件夹
+    cd PaddleScience/jointContribution
+    export PYTHONPATH=$PWD:$PYTHONPAT
+    # 下载模型参数
+    cd gencast/
+    wget -nc https://paddle-org.bj.bcebos.com/paddlescience/models/gencast/gencast_params_GenCast-1p0deg-Mini-_2019.pdparams -P ./data/params/
+    # 运行评估脚本
+    python run_gencast.py
+    ```
+
+## 1. 背景简介
+
+天气预报本质上存在不确定性，因此预测可能天气情景的范围对于许多重要决策至关重要，从警告公众危险天气到规划可再生能源的使用。在此，我们介绍了 GenCast，这是一种概率性天气模型，其技能和速度优于世界顶级的中期天气预报——欧洲中期天气预报中心（ECMWF）的集合预报 ENS。与基于数值天气预报（NWP）的传统方法不同，GenCast 是一种机器学习天气预报（MLWP）方法，基于数十年的再分析数据进行训练。GenCast 能够在 8 分钟内生成一个随机的 15 天全球预报集合，以 12 小时为步长，0.25 度的纬度-经度分辨率，覆盖 80 多个地表和大气变量。在我们评估的 1320 个目标中，GenCast 在 97.4% 上表现优于 ENS，并能更好地预测极端天气、热带气旋和风力发电。该工作帮助开启了操作性天气预报的下一个篇章，使依赖天气的重要决策能够以更高的准确性和效率做出。
+
+## 2. 模型原理
+
+在这里，我们介绍了一种概率性天气模型——GenCast，它以0.25°的分辨率生成全球15天的集合预报，首次实现了比顶级操作性集合系统ECMWF的ENS更高的准确性。在云TPUv5设备上生成一个单一的15天GenCast预报大约需要8分钟，可以并行生成多个预报集合。
+
+GenCast 模型化了未来天气状态 $X^{t+1}$ 的条件概率分布 $p(X^{t+1} | X^t, X^{t-1})$，这个分布是基于当前和之前的天气状态的条件来进行的。长度为 $T$ 的预报轨迹 $X^{1:T}$ 是通过对初始和之前状态 $(X^0, X^{-1})$ 进行条件化来建模的，并对连续状态的联合分布进行分解：
+
+$$
+p(X^{1:T} | X^0, X^{-1}) = \prod_{t=0}^{T-1} p(X^{t+1} | X^t, X^{t-1})
+$$
+
+每个状态都是通过自回归采样得出的。
+
+全球天气状态 $X$ 的表示包括6个地表变量和13个垂直压力层上的6个大气变量，分布在0.25°的纬度-经度网格上（详见表B1）。预报时长为15天，连续步骤 $t$ 和 $t+1$ 之间的间隔为12小时，因此 $T = 30$。
+
+GenCast 实现为一个条件扩散模型，这是一种生成式机器学习模型，用于从给定数据分布生成新样本，这为自然图像、声音和视频建模的许多最新进展提供了支持，被称为“生成式 AI”。扩散模型通过迭代细化的过程运行。未来的大气状态 $X^{t+1}$ 是通过迭代细化候选状态 $Z_0^{t+1}$ 产生的，该状态纯粹从噪声初始化，并以之前的两个大气状态 $(X^t, X^{t-1})$ 为条件。图中的蓝色框显示了第一个预报步骤如何从初始条件生成，以及整个轨迹 $X^{1:T}$ 如何通过自回归生成。由于预报中的每个时间步都是用噪声（即 $Z_0^{t+1}$）初始化的，因此可以用不同的噪声样本重复该过程，以生成轨迹集合。
+
+<figure markdown>
+  ![gencast.png](https://paddle-org.bj.bcebos.com/paddlescience/docs/gencast/gencast.png){ loading=lazy }
+</figure>
+
+在迭代细化过程的每个阶段，GenCast 应用一个由编码器、处理器和解码器组成的神经网络架构。编码器组件将输入 $Z_n^{t+1}$ 和条件 $(X^t, X^{t-1})$ 从原始的纬度-经度网格映射到六次细化的二十面体网格上的内部学习表示。处理器组件是一个Graph Transformer，其中每个节点关注其在内部网格上的k跳邻居。解码器组件将内部网格表示映射回 $Z_{n+1}^{t+1}$，其定义在纬度-经度网格上。
+
+GenCast 在40年的ERA5再分析数据上进行训练，时间范围从1979年到2018年，使用标准的扩散模型去噪目标。重要的是，尽管只在单步预测任务上直接训练GenCast，但它可以通过自回归展开来生成15天的集合预报。
+
+## 3. 模型构建
+
+### 3.1 环境依赖
+
+* paddlepaddle
+* matpoltlib （用于图像绘制）
+* pickle （用于存储和加载图模板）
+* xarray （用于加载.nc数据）
+* trimesh （用于制作mesh数据）
+* scipy （用于球谐变换过程中的稀疏矩阵操作）
+* math （用于球谐变换过程中的数学计算）
+
+### 3.2 模型相关文件说明
+
+- **xarray_tree.py**: 一种适用于 xarray 的 tree.map_structure 实现。
+
+- **denoiser.py**: 用于一步预测的 GenCast 去噪器。
+
+- **dpm_solver_plus_plus_2s.py**: 使用 [1] 中的 DPM-Solver++ 2S 的采样器。
+
+- **gencast.py**: 将 GenCast 模型架构与采样器结合，作为去噪器封装以生成预测。
+
+- **samplers_base.py**: 定义采样器的接口。
+
+- **samplers_utils.py**: 采样器的实用方法。
+
+- **sparse_transformer.py**: 通用稀疏变压器，作用于 TypedGraph，其中输入和输出都是每个节点和边的特征平坦向量。`predictor.py` 使用其中一个用于网格图神经网络（GNN）。
+
+- **spherical_harmonic.py**: 球面谐波基础评估和微分算子。
+
+- **main.py**: 评估和可视化脚本。
+
+[1] DPM-Solver++: Fast Solver for Guided Sampling of Diffusion Probabilistic Models, https://arxiv.org/abs/2211.01095
+
+## 4. 结果展示
+
+下图展示了2米温度的真值结果、预测结果和误差。
+
+<figure markdown>
+  ![gencast_2m_t.png](https://paddle-org.bj.bcebos.com/paddlescience/docs/gencast/gencast_2m_t.png){ loading=lazy style="margin:0 auto;"}
+  <figcaption>真值结果（"targets"）、预测结果（"prediction"）和误差（"diff"）</figcaption>
+</figure>
+
+可以看到模型预测结果与真实结果基本一致。
+
+## 4. 参考资料
+
+* [GenCast: Diffusion-based ensemble forecasting for medium-range weather](https://arxiv.org/abs/2312.15796)
+* [GraphCast: Learning skillful medium-range global weather forecasting](https://arxiv.org/abs/2212.12794)
+* [GenCast Github地址](https://github.com/deepmind/graphcast)
+* [dinosaur Github地址](https://github.com/neuralgcm/dinosaur)

jointContribution/gencast/conf/gencast.yaml

@@ -0,0 +1,99 @@
+
+hydra:
+  run:
+    # dynamic output directory according to running time and override name
+    dir: gencast/${now:%Y-%m-%d}/${now:%H-%M-%S}/${hydra.job.override_dirname}
+  job:
+    name: ${mode} # name of logfile
+    chdir: false # keep current working direcotry unchaned
+  sweep:
+    # output directory for multirun
+    dir: ${hydra.run.dir}
+    subdir: ./
+
+# general settings
+mode: eval # running mode: train/eval
+seed: 2024
+output_dir: ${hydra:run.dir}
+log_freq: 20
+num_ensemble_members: 8
+input_duration: "24h"
+target_lead_times: "12h"
+
+type: gencast
+data_path: data/dataset/source-era5_date-2019-03-29_res-1.0_levels-13_steps-12.nc
+stddev_diffs_path: data/stats/gencast_stats_diffs_stddev_by_level.nc
+stddev_path: data/stats/gencast_stats_stddev_by_level.nc
+mean_path: data/stats/gencast_stats_mean_by_level.nc
+min_path: data/stats/gencast_stats_min_by_level.nc
+param_path: data/params/gencast_params_GenCast-1p0deg-Mini-_2019.pdparams
+
+sampler_config:
+  max_noise_level: 80.0
+  min_noise_level: 0.03
+  num_noise_levels: 20
+  rho: 7.0
+  stochastic_churn_rate: 2.5
+  churn_min_noise_level: 0.75
+  churn_max_noise_level: inf
+  noise_level_inflation_factor: 1.05
+
+noise_config:
+  training_noise_level_rho: 7.0
+  training_max_noise_level: 88.0
+  training_min_noise_level: 0.02
+
+noise_encoder_config:
+  apply_log_first: true
+  base_period: 16.0
+  num_frequencies: 32
+  output_sizes: [32, 16]
+
+denoiser_architecture_config:
+  sparse_transformer_config:
+    attention_k_hop: 16
+    d_model: 512
+    num_layers: 16
+    num_heads: 4
+    attention_type: triblockdiag_mha
+    mask_type: lazy
+    block_q: 1024
+    block_kv: 512
+    block_kv_compute: 256
+    block_q_dkv: 512
+    block_kv_dkv: 1024
+    block_kv_dkv_compute: 1024
+    ffw_winit_final_mult: 0.0
+    attn_winit_final_mult: 0.0
+    ffw_hidden: 2048
+    mesh_node_dim: 186
+    mesh_node_emb_dim: 512
+    ffw_winit_mult: 2.0
+    value_size: 128
+    key_size: 128
+    norm_conditioning_feat: 16
+    activation: gelu
+  mesh_size: 4
+  latent_size: 512
+  hidden_layers: 1
+  radius_query_fraction_edge_length: 0.6
+  norm_conditioning_features: ['noise_level_encodings']
+  grid2mesh_aggregate_normalization: null
+  node_output_size: 84
+  grid_node_dim: 267
+  grid_node_emb_dim: 512
+  mesh_node_dim: 267
+  mesh_node_emb_dim: 512
+  mesh_edge_emb_dim: 512
+  mesh_edge_dim: 4
+  grid2mesh_edge_dim: 4
+  grid2mesh_edge_emb_dim: 512
+  mesh2grid_edge_dim: 4
+  mesh2grid_edge_emb_dim: 512
+  gnn_msg_steps: 16
+  node_output_dim: 84
+  norm_conditioning_feat: 16
+  mesh_node_num: 2562
+  grid_node_num: 65160
+  resolution: 1.0
+  name: gencast

jointContribution/gencast/data/dataset/.gitkeep

@@ -0,0 +1,171 @@
+# Copyright 2024 DeepMind Technologies Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS-IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Support for wrapping a general Predictor to act as a Denoiser."""
+
+import copy
+import os
+import pickle
+from typing import Optional
+from typing import Sequence
+
+import numpy as np
+import paddle
+import paddle.nn as nn
+import xarray as xr
+from graphcast import datasets
+from graphcast import graphcast
+from graphcast import graphtype
+from graphcast import utils
+
+
+class FourierFeaturesMLP(nn.Layer):
+    """A simple MLP applied to Fourier features of values or their logarithms."""
+
+    def __init__(
+        self,
+        base_period: float,
+        num_frequencies: int,
+        output_sizes: Sequence[int],
+        apply_log_first: bool = False,
+        w_init: Optional[nn.initializer.Initializer] = None,
+        activation: Optional[nn.Layer] = nn.GELU(),
+        **mlp_kwargs,
+    ):
+        """Initializes the module.
+
+        Args:
+        base_period:
+            See model_utils.fourier_features. Note this would apply to log inputs if
+            apply_log_first is used.
+        num_frequencies:
+            See model_utils.fourier_features.
+        output_sizes:
+            Layer sizes for the MLP.
+        apply_log_first:
+            Whether to take the log of the inputs before computing Fourier features.
+        w_init:
+            Weights initializer for the MLP, default setting aims to produce
+            approx unit-variance outputs given the input sin/cos features.
+        activation:
+        **mlp_kwargs:
+            Further settings for the MLP.
+        """
+        super(FourierFeaturesMLP, self).__init__()
+        self._base_period = base_period
+        self._num_frequencies = num_frequencies
+        self._apply_log_first = apply_log_first
+
+        # 创建 MLP
+        layers = []
+        input_size = 2 * num_frequencies
+        num_layers = len(output_sizes)
+        for i, output_size in enumerate(output_sizes):
+            linear_layer = nn.Linear(input_size, output_size)
+            layers.append(linear_layer)
+            if i < num_layers - 1:
+                layers.append(activation)
+            input_size = output_size
+
+        self._mlp = nn.Sequential(*layers)
+
+    def forward(self, values: paddle.Tensor) -> paddle.Tensor:
+        if self._apply_log_first:
+            values = paddle.log(values)
+        features = utils.fourier_features(
+            values, self._base_period, self._num_frequencies
+        )
+
+        return self._mlp(features)
+
+
+class Denoiser(nn.Layer):
+    """Wraps a general deterministic Predictor to act as a Denoiser.
+
+    This passes an encoding of the noise level as an additional input to the
+    Predictor as an additional input 'noise_level_encodings' with shape
+    ('batch', 'noise_level_encoding_channels'). It passes the noisy_targets as
+    additional forcings (since they are also per-target-timestep data that the
+    predictor needs to condition on) with the same names as the original target
+    variables.
+    """
+
+    def __init__(
+        self,
+        cfg,
+    ):
+        super(Denoiser, self).__init__()
+        self.cfg = cfg
+        self._predictor = graphcast.GraphCastNet(
+            config=cfg.denoiser_architecture_config,
+        )
+
+        self._noise_level_encoder = FourierFeaturesMLP(**cfg.noise_encoder_config)
+
+    def forward(
+        self,
+        inputs: xr.Dataset,
+        noisy_targets: xr.Dataset,
+        noise_levels: xr.DataArray,
+        forcings: Optional[xr.Dataset] = None,
+        **kwargs,
+    ) -> xr.Dataset:
+
+        if forcings is None:
+            forcings = xr.Dataset()
+        forcings = forcings.assign(**noisy_targets)
+
+        if noise_levels.dims != ("batch",):
+            raise ValueError("noise_levels expected to be shape (batch,).")
+
+        noise_level_encodings = self._noise_level_encoder(
+            paddle.to_tensor(noise_levels.values)
+        )
+
+        stacked_inputs = datasets.dataset_to_stacked(inputs)
+
+        stacked_forcings = datasets.dataset_to_stacked(forcings)
+        stacked_inputs = xr.concat([stacked_inputs, stacked_forcings], dim="channels")
+
+        stacked_inputs = stacked_inputs.transpose("lat", "lon", ...)
+        lat_dim, lon_dim, batch_dim, feat_dim = stacked_inputs.shape
+        stacked_inputs = stacked_inputs.data.reshape(lat_dim * lon_dim, batch_dim, -1)
+
+        graph_template_path = os.path.join(
+            "data", "template_graph", f"{self.cfg.type}.pkl"
+        )
+        if os.path.exists(graph_template_path):
+            graph_template = pickle.load(open(graph_template_path, "rb"))
+        else:
+            graph_template = graphtype.GraphGridMesh(
+                self.cfg.denoiser_architecture_config
+            )
+        graph = copy.deepcopy(graph_template)
+
+        graph.grid_node_feat = np.concatenate(
+            [stacked_inputs, graph.grid_node_feat], axis=-1
+        )
+        mesh_node_feat = np.zeros([graph.mesh_num_nodes, batch_dim, feat_dim])
+        graph.mesh_node_feat = np.concatenate(
+            [mesh_node_feat, graph.mesh_node_feat], axis=-1
+        )
+        graph.global_norm_conditioning = noise_level_encodings
+
+        predictor = self._predictor(graph=graphtype.convert_np_to_tensor(graph))
+
+        grid_node_outputs = predictor.grid_node_feat
+        raw_predictions = predictor.grid_node_outputs_to_prediction(
+            grid_node_outputs, noisy_targets
+        )
+
+        return raw_predictions