JAX-GCM (JCM)

JAX-GCM is a differentiable atmospheric general circulation model written entirely in JAX. A pluggable dynamical-core interface couples the Dinosaur spectral backend to modular SPEEDY, Held-Suarez, and ECHAM-style physics packages.

• Flexible. Switch from lightweight SPEEDY physics to full-fat ECHAM with online aerosol by changing one flag — and compose anything in between, term by term, through one composable-physics API.
• Differentiable. Any output is differentiable with respect to the physics parameters, the initial conditions, and the boundary conditions, so gradient calibration, data assimilation, and hybrid physics-ML experiments come for free from JAX's jit, grad, and vmap.
• Fast, on whatever you have. The same configuration runs on CPU, GPU, or TPU. ECHAM T63L47 climate with RRTMGP radiation costs roughly 22.7 s per simulated day on a single GPU — about 4x faster again with the neural-emulated radiation backend — and SPEEDY is cheaper still.

pip install jcm

JCM requires Python 3.11 or newer. See the getting-started guide for the development install and the full dependency set.

Upgrading from v2? v3.0 changes unit contracts, public APIs, configuration group names and checkpoint compatibility — see the v2-to-v3 migration guide.

Run a short SPEEDY aquaplanet integration from Python:

from jcm.model import Model
from jcm.physics.speedy.speedy_coords import get_speedy_coords

# T31 spectral resolution, 8 vertical levels. Omitting time_step lets the
# Model resolve a numerically stable default from the physics and grid
# (SPEEDY picks 30 min here; the no-limit default for ECHAM/Held-Suarez is
# 12 min).
coords = get_speedy_coords(layers=8, spectral_truncation=31)
model = Model(coords=coords)

predictions = model.run(save_interval=10.0, total_time=120.0)  # days
ds = predictions.to_xarray()
print(ds)

• Getting started — the Python quick start: coords, terrain, physics, running, and analysing output.
• Running at scale — the python -m jcm.main Hydra CLI, validated configurations, chunked/resumable runs, Docker, and GPU/batch patterns.
• ECHAM physics and SPEEDY physics — scheme notes and references.
• JAM optics backend seam — how an out-of-tree aerosol-optics pathway (e.g. a neural Mie emulator) plugs into the JAM optics as a JamOpticsTerm subclass.
• Example notebooks live in notebooks/.

Full documentation is hosted at jax-gcm.readthedocs.io.

If you use JAX-GCM in your research, please cite:

@article{jcm_gmd_2026,
  title   = {{JCM} v1.1: a differentiable, intermediate-complexity atmospheric model},
  author  = {Davenport, Ellen H. and Madan, J. Varan and Gjini, Rebecca and
             Brzenski, Jared and Ho, Nick and Hsu, Tien-Yiao and Liang, Yueshan and
             Liu, Zhixing and Manivannan, Veeramakali and Pham, Eric and
             Vutukuru, Rohith and Williams, Andrew I. L. and Yang, Zhiqi and
             Yu, Rose and Lutsko, Nicholas J. and Hoyer, Stephan and
             Watson-Parris, Duncan},
  journal = {Geoscientific Model Development},
  volume  = {19},
  pages   = {6451--6466},
  year    = {2026},
  doi     = {10.5194/gmd-19-6451-2026},
  url     = {https://gmd.copernicus.org/articles/19/6451/2026/}
}

JAX-GCM is licensed under Apache 2.0. See LICENSE.