GraphCast：学习熟练的中程全球天气预报

本软件包包含运行和训练 GraphCast 的示例代码。它还提供了三个预训练模型：

GraphCast，GraphCast 论文中使用的高分辨率模型（0.25 度分辨率，37 个压力层），在 1979 年至 2017 年的 ERA5 数据上训练。
GraphCast_small，GraphCast 的小型低分辨率版本（1 度分辨率，13 个压力层，以及更小的网格），在 1979 年至 2015 年的 ERA5 数据上训练，适用于在较低内存和计算约束下运行模型。
GraphCast_operational，一个高分辨率模型（0.25 度分辨率，13 个压力层），在 1979 年至 2017 年的 ERA5 数据上预训练，并在 2016 年至 2021 年的 HRES 数据上微调。该模型可以从 HRES 数据初始化（不需要降水输入）。

模型权重、归一化统计数据和示例输入可在 Google Cloud Bucket 上获取。

完整的模型训练需要下载 ERA5 数据集，可从 ECMWF 获得。最好以 Zarr 格式从 Weatherbench2 的 ERA5 数据（参见 6 小时降采样版本）访问。

文件概览

最佳起点是在 Colaboratory 中打开 graphcast_demo.ipynb，其中提供了加载数据、生成随机权重或加载预训练快照、生成预测、计算损失和计算梯度的示例。GraphCast 架构的单步实现在 graphcast.py 中提供。

库文件简述：

autoregressive.py：用于以 JAX 可微分方式运行（和训练）单步 GraphCast 的包装器，通过自回归方式将输出作为每一步的输入来生成预测序列。
casting.py：围绕 GraphCast 的包装器，用于使用 BFloat16 精度。
checkpoint.py：用于序列化和反序列化树的工具。
data_utils.py：数据预处理工具。
deep_typed_graph_net.py：通用深度图神经网络 (GNN)，操作 TypedGraph，其中输入和输出都是每个节点和边的特征的平面向量。graphcast.py 分别使用其中三个用于 Grid2Mesh GNN、Multi-mesh GNN 和 Mesh2Grid GNN。
graphcast.py：用于单步预测的主要 GraphCast 模型架构。
grid_mesh_connectivity.py：用于在球面上的规则网格和三角形网格之间转换的工具。
icosahedral_mesh.py：二十面体多重网格的定义。
losses.py：损失计算，包括纬度加权。
model_utils.py：用于从输入网格数据生成平面节点和边向量特征，以及将节点输出向量操作回多层网格数据的工具。
normalization.py：单步 GraphCast 的包装器，用于根据历史值归一化输入，并根据历史时间差异归一化目标。
predictor_base.py：定义了预测器的接口，GraphCast 和所有包装器都实现了这个接口。
rollout.py：类似于 autoregressive.py，但仅在推理时使用 Python 循环生成更长但不可微分的轨迹。
solar_radiation.py：计算与 ERA5 兼容的大气顶 (TOA) 入射太阳辐射。这用作强制变量，因此需要在操作设置中计算目标提前时间。
typed_graph.py：TypedGraph 的定义。
typed_graph_net.py：在 TypedGraph 上定义的简单图神经网络构建块的实现，可以组合成更深的模型。
xarray_jax.py：让 JAX 与 xarray 配合使用的包装器。
xarray_tree.py：适用于 xarray 的 tree.map_structure 实现。

依赖项

Chex、 Dask、 Haiku、 JAX、 JAXline、 Jraph、 Numpy、 Pandas、 Python、 SciPy、 Tree、 Trimesh 和 XArray。

许可和归属

Colab 笔记本和相关代码根据 Apache License, Version 2.0 授权。您可以在以下地址获取许可证副本： https://www.apache.org/licenses/LICENSE-2.0。

模型权重根据知识共享署名-非商业性使用-相同方式共享 4.0 国际许可协议（CC BY-NC-SA 4.0）提供使用。您可以在以下地址获取许可证副本： https://creativecommons.org/licenses/by-nc-sa/4.0/。

这些权重是在 ECMWF 的 ERA5 和 HRES 数据上训练的。Colab 包含一些 ERA5 和 HRES 数据的示例，可用作模型的输入。 ECMWF 数据产品受以下条款约束：

来源 www.ecmwf.int
许可声明：ECMWF 数据根据知识共享署名 4.0 国际许可协议（CC BY 4.0）发布。https://creativecommons.org/licenses/by/4.0/
免责声明：ECMWF 不对数据中的任何错误或遗漏、其可用性或因使用而造成的任何损失或损害承担任何责任。

免责声明

这不是 Google 官方支持的产品。

引用

如果您使用本工作，请考虑引用我们的论文：

@article{lam2022graphcast,
      title={GraphCast: Learning skillful medium-range global weather forecasting},
      author={Remi Lam and Alvaro Sanchez-Gonzalez and Matthew Willson and Peter Wirnsberger and Meire Fortunato and Alexander Pritzel and Suman Ravuri and Timo Ewalds and Ferran Alet and Zach Eaton-Rosen and Weihua Hu and Alexander Merose and Stephan Hoyer and George Holland and Jacklynn Stott and Oriol Vinyals and Shakir Mohamed and Peter Battaglia},
      year={2022},
      eprint={2212.12794},
      archivePrefix={arXiv},
      primaryClass={cs.LG}
}