NeuralGCM documentation

Overview

The NeuralGCM codebase consists of a handful of different components, suitable for reproducing and extending results from our paper, Neural General Circulation Models for Weather and Climate:

  1. Dynamics: The atmospheric dynamical core is distributed in the separate Dinosaur package.

  2. ML modules: Haiku modules for defining neural network layers.

  3. ML training: Pseudo-code for training NeuralGCM models can be found in the reference_code subdirectory.

  4. ML inference: Code for running forecasts with pre-trained models, encapsulated in the PressureLevelModel class.

  5. Evaluation: Code for evaluating NeuralGCM weather forecasts can be found in the WeatherBench2 project.

The documentation here focuses mostly on our API for inference (i.e., running trained NeuralGCM atmospheric models), which we believe is the most immediately useful part of the NeuralGCM code for third parties. It is also a part of our code that we can commit to supporting in roughly its current form.

We would love to support training, modifying and fine-tuning NeuralGCM models, but with the present codebase based on Haiku and Gin this is much trickier than it needs to be. We are currently (in May 2024) refactoring the modeling code to improve usability.

Tip

To stay up to date on NeuralGCM, subscribe to our mailing list.

Questions?

The best place to ask for help using NeuralGCM models or datasets is on GitHub.

You can also reach the NeuralGCM team directly at neuralgcm@google.com.

Contents