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

PyTorch implementations of NeuralGCM and its spectral dynamical core, Dinosaur — idiomatic nn.Module models that run the published NeuralGCM checkpoints with no JAX, gin or haiku at runtime.

This repository contains two packages:

• neuralgcm-torch — the NeuralGCM model (hybrid ML + physics atmospheric model): xarray-in/out inference, encode / advance / decode, ensembles, training, multi-GPU.
• dinosaur-torch — the standalone spectral dynamical core (spherical-harmonic transforms, sigma coordinates, primitive-equation IMEX time stepping).

Install

pip install neuralgcm-torch        # pulls in dinosaur-torch
# checkpoint downloads also need: pip install "neuralgcm-torch[hub]"

Quick start

import neuralgcm_torch as neuralgcm
from neuralgcm_torch import pretrained

path = pretrained.fetch_checkpoint('deterministic_2_8_deg')  # cached Hub download
model = neuralgcm.PressureLevelModel.from_checkpoint(path, device='cuda')

See neuralgcm-torch/README.md and dinosaur-torch/README.md for details, and neuralgcm-torch/notebooks/ for end-to-end forecast, ensemble and climate-stability examples.

Documentation

The example notebooks — forecasting at every resolution, batched ensembles, precipitation / evaporation, multi-decade climate stability, and the model internals — are rendered with their outputs as a Jupyter Book at dsip-fbk.github.io/neuralgcm-torch.

Development

uv sync                  # both packages editable + dev tools
pre-commit install
DINOSAUR_TORCH_TEST_DEVICE=cpu uv run pytest dinosaur-torch

License

Apache 2.0 (code) — see LICENSE and NOTICE. These packages are a PyTorch port of Google's NeuralGCM and Dinosaur. The converted model checkpoints, hosted on the Hugging Face Hub, are © Google LLC and licensed under CC BY-SA 4.0 (see the model card).