neuralgcm-torch

<p align="center"> <img src="https://huggingface.co/it4lia/neuralgcm-torch/resolve/main/neuralgcm-torch.png" width="296" alt="neuralgcm-torch"> </p> <p align="center"> <a href="https://github.com/DSIP-FBK/neuralgcm-torch"><img src="https://img.shields.io/badge/code-DSIP--FBK%2Fneuralgcm--torch-181717?logo=github" alt="GitHub"></a> <a href="https://pypi.org/project/neuralgcm-torch/"><img src="https://img.shields.io/pypi/v/neuralgcm-torch?logo=pypi&logoColor=white" alt="PyPI"></a> <a href="https://dsip-fbk.github.io/neuralgcm-torch/"><img src="https://img.shields.io/badge/docs-Jupyter%20Book-F37726?logo=jupyter&logoColor=white" alt="Docs"></a> <img src="https://img.shields.io/badge/weights-CC%20BY--SA%204.0-lightgrey" alt="Weights: CC BY-SA 4.0"> </p> <p align="center"> <img src="https://huggingface.co/it4lia/neuralgcm-torch/resolve/main/forecast_0_7_deg.gif" width="360" alt="NeuralGCM-0.7° forecast — 850 hPa specific humidity on a rotating globe"> </p>

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

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.

<p align="center"> <table align="center" style="border: none; border-collapse: collapse;"> <tr style="border: none;"> <td align="center" style="border: none; padding: 10px;"><a href="https://www.fbk.eu"><img src="https://webvalley.fbk.eu/static/img/logos/fbk-logo-blue.png" width="105" alt="Fondazione Bruno Kessler"></a></td> <td align="center" style="border: none; padding: 10px;"><a href="https://it4lia-aifactory.eu/"><img src="https://it4lia-aifactory.eu/wp-content/uploads/2025/05/logo-IT4LIA-AI-factory.svg" width="125" alt="IT4LIA AI Factory"></a></td> </tr> </table> </p>