OneScience-Group/StormCast

StormCast

Model Introduction

StormCast is a generative regional weather forecasting model proposed by NVIDIA, targeting high-resolution nowcasting of mesoscale convective weather.

Paper: StormCast: A Machine Learning Method for Meso-β-Scale Convection-resolving Weather Forecasting

https://arxiv.org/abs/2408.10958

Model Description

StormCast constrains the evolution of regional states with large-scale weather backgrounds, and uses a generative diffusion approach to supplement the fine-scale structures that deterministic forecasts struggle to represent.

Use Cases

Scenario	Description
Two-Stage Weather Forecast Training	Train a deterministic regression model and a conditional residual diffusion model in sequence.
Local Quick Validation	Use synthetic data to verify data loading, model training, inference, and inference result visualization.
ModelScope / OneCode Execution	Download as a standalone model package, install dependencies, and run scripts directly.
Multi-GPU Training	Launch multi-process training via torchrun.

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

• Training and inference require a GPU or DCU recognized by PyTorch; CPU can be used to generate synthetic data and verify configuration, but cannot run the current training and inference scripts.
• Multi-GPU training uses the NCCL backend. Please make sure the device driver, communication libraries, and PyTorch version are compatible.
• DCU users must install DTK in advance. DTK 25.04.2 or above, or the OneScience recommended version matching your cluster, is recommended.

Download the Model Package

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

Install the Runtime Environment

DCU Environment

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

GPU Environment

# Please 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
# uv installation is supported
pip install onescience[earth-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/  --trusted-host mirrors.onescience.ai

Training Data Introduction

The OneScience community provides ERA5 data for training (due to file size limits, the current repository contains a slice of the full dataset). Users can download it with the command below and confirm that the data path in conf/config.yaml is set correctly:

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

Generate Synthetic Data for Pipeline Validation

python scripts/fake_data.py

Synthetic data is only used to verify the data protocol and program flow; it does not represent the model's scientific forecasting capability.

Training

Single GPU:

# Train the deterministic regression model; weights are saved to data/checkpoint/regression/model_bak.pt by default
python scripts/train.py --stage regression
# Train the residual diffusion model; weights are saved to data/checkpoint/diffusion/model_bak.pt by default
python scripts/train.py --stage diffusion

Multi-GPU

# Train the deterministic regression model
torchrun --nproc_per_node=2 scripts/train.py --stage regression
# Train the residual diffusion model
torchrun --nproc_per_node=2 scripts/train.py --stage diffusion

Training Weights

This repository provides weights trained on ERA5 reanalysis data in the weights/ folder. The weight files will be uploaded soon and are expected to be available in the near future.

Inference

Run autoregressive prediction with the default configuration and the two sets of weights saved during training:

python scripts/inference.py

Evaluation and Visualization

python scripts/result.py

Plots are saved to outputs/inference/plots/ by default.

OneScience Official 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 & License

• This repository is a reproduction of the original StormCast paper.