GraphCast: Learning skillful medium-range global weather forecasting

This package contains example code to run and train GraphCast. It also provides three pretrained models:

GraphCast, the high-resolution model used in the GraphCast paper (0.25 degree resolution, 37 pressure levels), trained on ERA5 data from 1979 to 2017,
GraphCast_small, a smaller, low-resolution version of GraphCast (1 degree resolution, 13 pressure levels, and a smaller mesh), trained on ERA5 data from 1979 to 2015, useful to run a model with lower memory and compute constraints,
GraphCast_operational, a high-resolution model (0.25 degree resolution, 13 pressure levels) pre-trained on ERA5 data from 1979 to 2017 and fine-tuned on HRES data from 2016 to 2021. This model can be initialized from HRES data (does not require precipitation inputs).

The model weights, normalization statistics, and example inputs are available on Google Cloud Bucket.

Full model training requires downloading the ERA5 dataset, available from ECMWF. This can best be accessed as Zarr from Weatherbench2's ERA5 data (see the 6h downsampled versions).

Overview of files

The best starting point is to open graphcast_demo.ipynb in Colaboratory, which gives an example of loading data, generating random weights or load a pre-trained snapshot, generating predictions, computing the loss and computing gradients. The one-step implementation of GraphCast architecture, is provided in graphcast.py.

Brief description of library files:

autoregressive.py: Wrapper used to run (and train) the one-step GraphCast to produce a sequence of predictions by auto-regressively feeding the outputs back as inputs at each step, in JAX a differentiable way.
casting.py: Wrapper used around GraphCast to make it work using BFloat16 precision.
checkpoint.py: Utils to serialize and deserialize trees.
data_utils.py: Utils for data preprocessing.
deep_typed_graph_net.py: General purpose deep graph neural network (GNN) that operates on TypedGraph's where both inputs and outputs are flat vectors of features for each of the nodes and edges. graphcast.py uses three of these for the Grid2Mesh GNN, the Multi-mesh GNN and the Mesh2Grid GNN, respectively.
graphcast.py: The main GraphCast model architecture for one-step of predictions.
grid_mesh_connectivity.py: Tools for converting between regular grids on a sphere and triangular meshes.
icosahedral_mesh.py: Definition of an icosahedral multi-mesh.
losses.py: Loss computations, including latitude-weighting.
model_utils.py: Utilities to produce flat node and edge vector features from input grid data, and to manipulate the node output vectors back into a multilevel grid data.
normalization.py: Wrapper for the one-step GraphCast used to normalize inputs according to historical values, and targets according to historical time differences.
predictor_base.py: Defines the interface of the predictor, which GraphCast and all of the wrappers implement.
rollout.py: Similar to autoregressive.py but used only at inference time using a python loop to produce longer, but non-differentiable trajectories.
solar_radiation.py: Computes Top-Of-the-Atmosphere (TOA) incident solar radiation compatible with ERA5. This is used as a forcing variable and thus needs to be computed for target lead times in an operational setting.
typed_graph.py: Definition of TypedGraph's.
typed_graph_net.py: Implementation of simple graph neural network building blocks defined over TypedGraph's that can be combined to build deeper models.
xarray_jax.py: A wrapper to let JAX work with xarrays.
xarray_tree.py: An implementation of tree.map_structure that works with xarrays.

Dependencies.

Chex, Dask, Haiku, JAX, JAXline, Jraph, Numpy, Pandas, Python, SciPy, Tree, Trimesh and XArray.

License and attribution

The Colab notebook and the associated code are licensed under the Apache License, Version 2.0. You may obtain a copy of the License at: https://www.apache.org/licenses/LICENSE-2.0.

The model weights are made available for use under the terms of the Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International (CC BY-NC-SA 4.0). You may obtain a copy of the License at: https://creativecommons.org/licenses/by-nc-sa/4.0/.

The weights were trained on ECMWF's ERA5 and HRES data. The colab includes a few examples of ERA5 and HRES data that can be used as inputs to the models. ECMWF data product are subject to the following terms:

Source www.ecmwf.int
Licence Statement: ECMWF data is published under a Creative Commons Attribution 4.0 International (CC BY 4.0). https://creativecommons.org/licenses/by/4.0/
Disclaimer: ECMWF does not accept any liability whatsoever for any error or omission in the data, their availability, or for any loss or damage arising from their use.

Disclaimer

This is not an officially supported Google product.

Citation

If you use this work, consider citing our paper (blog post, Science, arXiv):

@article{lam2023learning,
  title={Learning skillful medium-range global weather forecasting},
  author={Lam, Remi and Sanchez-Gonzalez, Alvaro and Willson, Matthew and Wirnsberger, Peter and Fortunato, Meire and Alet, Ferran and Ravuri, Suman and Ewalds, Timo and Eaton-Rosen, Zach and Hu, Weihua and others},
  journal={Science},
  volume={382},
  number={6677},
  pages={1416--1421},
  year={2023},
  publisher={American Association for the Advancement of Science}
}