# GraphCast: Learning skillful medium-range global weather forecasting

This package contains example code to run and train [GraphCast](https://www.science.org/doi/10.1126/science.adi2336).
It also provides three pretrained models:

1.  `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,

2.  `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,

3.  `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](https://console.cloud.google.com/storage/browser/dm_graphcast).

Full model training requires downloading the
[ERA5](https://www.ecmwf.int/en/forecasts/datasets/reanalysis-datasets/era5)
dataset, available from [ECMWF](https://www.ecmwf.int/). This can best be
accessed as Zarr from [Weatherbench2's ERA5 data](https://weatherbench2.readthedocs.io/en/latest/data-guide.html#era5) (see the 6h downsampled versions).

## Overview of files

The best starting point is to open `graphcast_demo.ipynb` in [Colaboratory](https://colab.research.google.com/github/deepmind/graphcast/blob/master/graphcast_demo.ipynb), 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 `xarray`s.
*   `xarray_tree.py`: An implementation of tree.map_structure that works with
    `xarray`s.


### Dependencies.

[Chex](https://github.com/deepmind/chex),
[Dask](https://github.com/dask/dask),
[Haiku](https://github.com/deepmind/dm-haiku),
[JAX](https://github.com/google/jax),
[JAXline](https://github.com/deepmind/jaxline),
[Jraph](https://github.com/deepmind/jraph),
[Numpy](https://numpy.org/),
[Pandas](https://pandas.pydata.org/),
[Python](https://www.python.org/),
[SciPy](https://scipy.org/),
[Tree](https://github.com/deepmind/tree),
[Trimesh](https://github.com/mikedh/trimesh) and
[XArray](https://github.com/pydata/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:

1. Copyright statement: Copyright "© 2023 European Centre for Medium-Range Weather Forecasts (ECMWF)".
2. Source www.ecmwf.int
3. 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/
4. 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.

Copyright 2023 DeepMind Technologies Limited.

### Citation

If you use this work, consider citing our paper ([blog post](https://deepmind.google/discover/blog/graphcast-ai-model-for-faster-and-more-accurate-global-weather-forecasting/), [Science](https://www.science.org/doi/10.1126/science.adi2336), [arXiv](https://arxiv.org/abs/2212.12794)):

```latex
@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}
}
```