neuralgcm-torch

NeuralGCM is a hybrid ML + physics global circulation model that pairs a differentiable spectral dynamical core with learned physics to forecast weather and run climate-scale simulations, originally written in JAX. neuralgcm-torch brings it to PyTorch — load the published NeuralGCM checkpoints (converted to torch) and forecast in a few lines, with no JAX, gin or haiku at runtime.

A 12-day NeuralGCM-0.7° forecast — 850 hPa specific humidity — on a slowly rotating globe.

These pages are the project’s example notebooks, rendered from their committed outputs. For the full API and design notes, see the README on GitHub.

Not affiliated with NeuralGCM or Google

This is an independent PyTorch reimplementation built on the original team’s published research and open-source weights. All credit for the models and science goes to them — see the NeuralGCM repository and please cite Kochkov et al., “Neural general circulation models for weather and climate”, Nature 632 (2024).

Running these notebooks yourself

Download any notebook with the button at the top-right of its page, then run it against an environment with the package installed — see Installation for the PyPI and clone-and-run setups. A CUDA GPU and network access are recommended; the higher-resolution and climate-stability runs assume a GPU.

The notebooks

• Forecasting — forecast_quickstart (the 2.8° ERA5 walkthrough) and the 1.4° / 0.7° / ensemble / precipitation / evaporation variants.

• Climate — climate_stability drives the stable models far past weather lead times (months to years) with seasonal ERA5 forcing.

• Data & model internals — regridding, model internals and autograd, and editing the converted config.

The original NeuralGCM team’s pre-computed simulation outputs (the gs://neuralgcm/… Zarr stores) are catalogued in the upstream NeuralGCM simulation datasets documentation.

Use the sidebar to open any notebook.