OneScience-Group/GCRM-ConvectionNN

GCRM-ConvectionNN

Model Introduction

GCRM-ConvectionNN learns convective heating and moistening tendencies in atmospheric columns and diagnoses propagating instabilities through linear response functions and gravity-wave spectra, supporting research on GCRM convection parameterization and online stability.

Paper: Interpreting and Stabilizing Machine-Learning Parametrizations of Convection
https://doi.org/10.1175/JAS-D-20-0082.1

Model Description

The method was proposed by research teams at Vulcan Inc., the University of California Irvine, Columbia University, and the University of Washington. The paper trains and analyzes neural networks with coarse-grained near-global SAM cloud-resolving simulations and SPCAM superparameterized climate simulations. It is intended for atmospheric convection parameterization, linear stability diagnosis, and stabilization of machine-learning physics schemes.

Use Cases

Use Case	Description
GCRM convection parameterization	Predict heating and moistening tendencies from 34-level thermodynamic columns.
Stability diagnosis	Compute automatic-differentiation LRFs, upper-input ablations, and two-dimensional wave spectra.
ModelScope/OneCode execution	Validate structured data, training, inference, spectral metrics, and visualization in ModelScope or OneCode environments.
Multi-GPU training	Validate distributed training and checkpointing through torchrun.

Usage Instructions

1.OneCode

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

Try intelligent, one-click AI4S programming

2. Download and Installation

hf download OneScience-Group/GCRM-ConvectionNN --local-dir ./GCRM-ConvectionNN
cd GCRM-ConvectionNN

Environment Dependencies

Hardware Requirements

• A GPU or DCU is recommended.
• A CPU can be used for connectivity validation with the default small-sample configuration.
• DCU users must install DTK first. DTK 25.04.2 or later, or the OneScience-recommended version matching the cluster, is recommended.

DCU Environment

# Activate DTK and Conda first
conda create -n onescience311 python=3.11 -y
conda activate onescience311
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
pip install onescience[earth-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai

Training Data

Training data consist of atmospheric-column thermodynamic samples with 34-level total-water and liquid-ice-static-energy profiles. Inputs also include sea-surface temperature and top-of-atmosphere insolation, giving 70 values. Targets are 34-level heating and moistening tendencies for the same column, giving 68 values at three-hour intervals. This repository uses only a small number of synthetic samples to validate training, stability diagnosis, inference, and evaluation; they do not represent the official data distribution, training scale, or formal paper performance.

python scripts/fake_data.py

Training

For single-device training, use:

python scripts/train.py

For multi-GPU training, use:

torchrun --nproc_per_node=8 --nnodes=1 --rdzv_id=1000 --rdzv_backend=c10d --max_restarts=0 --master_addr="localhost" --master_port=29500 scripts/train.py

Training applies mass-weighted MSE and the configurable 20-member balance regularizer and validates data versions and the 70/68 dimensions. Training outputs are saved to:

result/checkpoints/gcrm_convectionnn.pt
result/training/metrics.json

Trained Weights

No pretrained weights are bundled under weight/. The engineering checkpoint generated from synthetic data is not an official paper checkpoint.

Inference

python scripts/inference.py

Inference loads the trained checkpoint, generates Q1/Q2 convective tendencies from 34-level thermodynamic columns, and computes linear-response and wave-spectrum results for full and upper-input-ablated configurations. Complete numerical results are saved to:

result/output/predictions.npz

Evaluation and Visualization

python scripts/result.py

Evaluation computes mass-weighted MSE, dangerous propagating-mode counts, maximum growth, and maximum propagation speed for full and ablated configurations; this is neither a multi-step prediction nor a classification task, so no per-lead or per-class metrics are saved. It also produces Q1/Q2 target-prediction profiles and spectral comparison plots. Synthetic-data results validate the engineering workflow only and do not represent formal paper performance.

result/evaluation/metrics.json
result/evaluation/profiles_spectrum.png

Official OneScience Information

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 an independent engineering reproduction of the public GCRM-ConvectionNN specifications.

Use of this repository's code, official model weights, and data remains subject to the licenses and terms of their respective projects.