OneScience-Group/NeuralGCM

NeuralGCM

Model Introduction

NeuralGCM (Neural General Circulation Models) is an open-source hybrid machine-learning and physics-based atmospheric model developed by Google Research for weather forecasting and climate simulation.

Paper: Neural General Circulation Models for Weather and Climate

https://arxiv.org/abs/2311.07222

Model Description

NeuralGCM is built around a differentiable atmospheric dynamical core. Neural networks represent unresolved physical processes, the encoder, and the decoder, improving forecast efficiency while retaining physical constraints.

Profile	Resolution	Type	Bundled official checkpoint
weather_forecast	0.7 degrees (512 x 256)	Deterministic weather forecasting for approximately 2 to 15 days	weight/models_v1_deterministic_0_7_deg.pkl
climate_scale	1.4 degrees (256 x 128)	Deterministic climate-scale simulation	weight/models_v1_deterministic_1_4_deg.pkl
forecast_2_8_deg	2.8 degrees (128 x 64)	Deterministic weather forecasting	weight/models_v1_deterministic_2_8_deg.pkl
stochastic_1_4_deg	1.4 degrees (256 x 128)	Stochastic weather forecasting	weight/models_v1_stochastic_1_4_deg.pkl

Use Cases

Scenario	Description
Global weather forecasting	Train the 0.7-degree model on ERA5 data for short- to medium-range weather forecasting.
Climate-scale simulation	Train the 1.4-degree model on ERA5 data for longer atmospheric simulations.
Low-resolution experiments	Use the 2.8-degree data profile for lower-cost weather forecasting experiments.
Local quick validation	Generate HDF5 data with the required channel protocol using scripts/fake_data.py and validate the data, model, and checkpoint workflows.
ModelScope / OneCode execution	Download the standalone model package, install the OneScience and JAX dependencies, and run the scripts directly.
Multi-device training	Run synchronous data-parallel training on multiple local accelerators.

Usage Guide

1. OneCode Usage

Experience intelligent one-click AI4S programming through the OneCode online environment:

Click to Experience Intelligent One-Click AI4S Programming

2. Manual Installation and Usage

Hardware Requirements

• A GPU or DCU is recommended.
• A CPU can be used for import checks and small-scale connectivity validation, but full training and inference will be slow.
• DCU users must install DTK in advance. DTK 25.04.2 or later, or the OneScience-recommended version compatible with the current cluster, is recommended.

Download the Model Package

hf download OneScience-Group/NeuralGCM --local-dir ./NeuralGCM
cd NeuralGCM

Install the Runtime Environment

DCU Environment

# Activate DTK and conda first.
conda create -n onescience311 python=3.11 -y
conda activate onescience311
# Installation with uv is also supported.
pip install onescience[earth-dcu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai

GPU Environment

# Activate conda first.
conda create -n onescience311 python=3.11 -y libstdcxx-ng=12 libgcc-ng=12 gcc_linux-64=12 gxx_linux-64=12
conda activate onescience311
# Installation with uv is also supported.
pip install onescience[earth-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai

Training Data

The OneScience community provides an ERA5 data slice for training. Download it with the following command and confirm that the data path in conf/config.yaml is correct:

hf download --repo-type dataset OneScience-Group/ERA5 --local-dir ./data

Generate Synthetic Data

python scripts/fake_data.py

The script creates yearly HDF5 files under data/data/, writes synthetic static fields to data/static.nc, and saves channel, time-window, and grid metadata to data/metadata/dataset_card.json. The synthetic fields use approximate physical units but are intended only for shape, loading, regridding, and numerical-stability checks.

Training

Single device:

# 0.7-degree deterministic short- to medium-range weather forecasting
python scripts/train_weather_forecast.py
# 1.4-degree deterministic climate-scale simulation
python scripts/train_climate_scale.py
# 2.8-degree deterministic low-resolution weather forecasting
python scripts/train_forecast_2_8_deg.py
# 1.4-degree stochastic weather forecasting
python scripts/train_stochastic_1_4_deg.py

Multiple devices:

# 0.7-degree deterministic short- to medium-range weather forecasting
python scripts/train_weather_forecast.py --devices 8
# 1.4-degree deterministic climate-scale simulation
python scripts/train_climate_scale.py --devices 8
# 2.8-degree deterministic low-resolution weather forecasting
python scripts/train_forecast_2_8_deg.py --devices 8
# 1.4-degree stochastic weather forecasting
python scripts/train_stochastic_1_4_deg.py --devices 8

Fine-tuning

Fine-tuning can start from either a checkpoint produced by local training or the bundled official checkpoint for the selected profile.

# Use the bundled official checkpoint for each profile.
python scripts/train_weather_forecast.py --finetune weight/models_v1_deterministic_0_7_deg.pkl
python scripts/train_climate_scale.py --finetune weight/models_v1_deterministic_1_4_deg.pkl
python scripts/train_forecast_2_8_deg.py --finetune weight/models_v1_deterministic_2_8_deg.pkl
python scripts/train_stochastic_1_4_deg.py --finetune weight/models_v1_stochastic_1_4_deg.pkl

# Alternatively, provide a local checkpoint explicitly.
python scripts/train_weather_forecast.py --finetune ./data/checkpoint/model_bak.pkl

For multi-device fine-tuning, add --devices to the corresponding command.

Pre-trained Weights

This project includes the following official pre-trained checkpoints:

Local file	Official release path
weight/models_v1_deterministic_0_7_deg.pkl	gs://neuralgcm/models/v1/deterministic_0_7_deg.pkl
weight/models_v1_deterministic_1_4_deg.pkl	gs://neuralgcm/models/v1/deterministic_1_4_deg.pkl
weight/models_v1_deterministic_2_8_deg.pkl	gs://neuralgcm/models/v1/deterministic_2_8_deg.pkl
weight/models_v1_stochastic_1_4_deg.pkl	gs://neuralgcm/models/v1/stochastic_1_4_deg.pkl

Inference

# 0.7-degree deterministic short- to medium-range weather forecasting
python scripts/inference.py --mode weather_forecast --checkpoint weight/models_v1_deterministic_0_7_deg.pkl
# 1.4-degree deterministic climate-scale simulation
python scripts/inference.py --mode climate_scale --checkpoint weight/models_v1_deterministic_1_4_deg.pkl
# 2.8-degree deterministic low-resolution weather forecasting
python scripts/inference.py --mode forecast_2_8_deg --checkpoint weight/models_v1_deterministic_2_8_deg.pkl
# 1.4-degree stochastic weather forecasting
python scripts/inference.py --mode stochastic_1_4_deg --checkpoint weight/models_v1_stochastic_1_4_deg.pkl

Without an explicit --checkpoint, inference first checks ./data/checkpoint/model_bak.pkl. The default output is results/predictions.nc, containing pressure-level variables with their official names and rollout time coordinates.

Evaluation and Visualization

python scripts/result.py

Official OneScience Resources

Platform	OneScience Main Repository	Skills Repository
Gitee	https://gitee.com/onescience-ai/onescience	https://gitee.com/onescience-ai/oneskills
GitHub	https://github.com/onescience-ai/OneScience	https://github.com/onescience-ai/oneskills

Citation and License

• This repository is a reproduction of the original NeuralGCM paper.
• The repository code is provided under the Apache License 2.0.
• The trained model weights released by Google, including the four checkpoints in this directory, are licensed under the Creative Commons Attribution-ShareAlike 4.0 International license (CC BY-SA 4.0). Redistribution or adaptation of the weights must preserve attribution and use the same license as required by those terms.