A Differentiable General Circulation Model

JAX-GCM (JCM) is a groundbreaking fully differentiable General Circulation Model (GCM) designed to bridge traditional climate modeling with modern machine learning techniques. Written entirely in JAX, it represents a new paradigm in climate simulation that enables capabilities previously impossible with conventional models.

Key Features

Fully Differentiable Every component of the model supports automatic differentiation, enabling gradient-based optimization, data assimilation, and inverse modeling—capabilities that traditional climate models cannot provide.

GPU/TPU Accelerated Leverages JAX’s just-in-time (JIT) compilation and hardware acceleration to run efficiently on GPUs and TPUs, dramatically speeding up simulations and enabling rapid iteration.

Modular Physics Includes comprehensive atmospheric physics packages:

• Radiation schemes (shortwave and longwave)
• Convection parameterizations
• Cloud physics
• Surface processes and fluxes
• Vertical diffusion

Flexible Resolution Supports spectral resolutions from T31 (coarse, ~4°) to T425 (high-resolution, ~0.28°), allowing users to balance computational cost with simulation fidelity.

ML-Ready Architecture Designed from the ground up for hybrid physics-ML workflows, parameter optimization, and emulator development. Perfect for neural network integration and data-driven climate science.

Technical Foundation

JAX-GCM combines:

• Dinosaur Dynamical Core: Google Research’s spectral dynamical core handles atmospheric dynamics
• SPEEDY Physics: Adapted from F. Molteni’s simplified parameterization scheme, implemented fully in JAX
• Modern Python Stack: Built on Python ≥3.11, XArray, and the scientific Python ecosystem

Use Cases

The differentiability of JAX-GCM opens new research frontiers:

• Parameter Optimization: Gradient descent through the entire model for optimal parameter estimation
• Data Assimilation: Efficient incorporation of observational data into model states
• Hybrid ML Models: Seamless integration of neural networks into physical parameterizations
• Climate Sensitivity Studies: Rapid exploration of model responses to perturbations
• Emulator Development: Training fast surrogate models with full-physics guidance

Get Started

Visit the JAX-GCM GitHub repository to explore the code, documentation, and examples. The project welcomes contributions from the climate science and machine learning communities.

Learn more Code