it4lia/neuralgcm-torch

NeuralGCM checkpoints — PyTorch

NeuralGCM is a hybrid ML + physics global circulation model: a differentiable spectral dynamical core coupled to learned physics, for medium-range weather forecasting and decade-long climate simulation originally developed in JAX. This repository hosts the published NeuralGCM v1 checkpoints converted to PyTorch, ready to load with neuralgcm-torch — real nn.Modules with registered parameters, no JAX, gin or haiku at runtime and no conversion step at load time.

Install

pip install 'neuralgcm-torch[hub,notebooks]'

Usage

import neuralgcm_torch as neuralgcm
from neuralgcm_torch import pretrained

# downloads from this repo on first use (cached)
path = pretrained.fetch_checkpoint('deterministic_2_8_deg')
model = neuralgcm.PressureLevelModel.from_checkpoint(path, device='cuda')

The default Hub repo is it4lia/neuralgcm-torch — override per call with repo_id= or globally with NEURALGCM_TORCH_HF_REPO.

Checkpoints

file	resolution	parameters	kind
deterministic_0_7_deg.pt	0.7° (TL255)	31.1M	deterministic
deterministic_1_4_deg.pt	1.4° (TL127)	18.3M	deterministic
deterministic_2_8_deg.pt	2.8° (TL63)	14.5M	deterministic
stochastic_1_4_deg.pt	1.4° (TL127)	11.5M	stochastic (NeuralGCM-ENS)
stochastic_precip_2_8_deg.pt	2.8° (TL63)	11.1M	stochastic, precipitation
stochastic_evap_2_8_deg.pt	2.8° (TL63)	11.1M	stochastic, evaporation
tl63_stochastic_mini.pt	TL63 toy	0.19M	stochastic toy / test fixture

Rules of thumb: the 1.4° stochastic model is the best general-purpose weather/ensemble model, the 2.8° models are stable for multi-decade climate runs, and the 0.7° model is the most accurate at short lead times. Each was converted once, offline from the original JAX pickle (gin config + haiku parameter tree) into a plain torch.save dictionary — structured config, auxiliary arrays, and parameter tensors.

Examples & documentation

Runnable example notebooks — forecasting at every resolution, batched ensembles, precipitation/evaporation, multi-decade climate stability, and the model internals — are rendered as a Jupyter Book at dsip-fbk.github.io/neuralgcm-torch. Source code, the full API and design notes live in the GitHub repository.

License & attribution

These weights are derivative works of the NeuralGCM v1 checkpoints published by Google at gs://neuralgcm/models/ under CC BY-SA 4.0. As required by ShareAlike, they are redistributed under the same license (CC BY-SA 4.0).

• Original checkpoints © Google — NeuralGCM.
• Deterministic / stochastic models: Kochkov et al., Nature 2024.
• Precipitation / evaporation models: Yuval et al., 2024.

This is an independent reimplementation, not affiliated with the original NeuralGCM authors or Google. The neuralgcm-torch model code is separately licensed under Apache 2.0.

Acknowledgements

The PyTorch port and checkpoint conversion were developed at Fondazione Bruno Kessler (FBK) and are hosted on the IT4LIA AI Factory.