OneScience-Group/StableNN-Phys

StableNN-Phys

Model Introduction

StableNN-Phys is an engineering reproduction of a neural-network unified physics parameterization for an atmospheric single-column model, continuously predicting thermodynamic-state evolution at three-hour intervals from column states and surface fluxes.

Paper: Prognostic Validation of a Neural Network Unified Physics Parameterization
https://doi.org/10.1029/2018GL078510

Model Description

The method reproduced by StableNN-Phys was proposed by researchers from the Department of Atmospheric Sciences at the University of Washington. The model takes 34-level liquid-water static energy and total-water states with surface sensible heat, latent heat, and incoming solar radiation, maps 71 input features to 68 physics tendencies at each step, and integrates multiple steps with advective forcing and no teacher forcing. The model is suitable for neural physics-parameterization training, three-hourly single-column prediction, 64-step long-term stability validation, and column-water-budget diagnosis.

Use Cases

Use Case	Description
Single-column prediction	Continuously predict three-hourly state evolution from 34-level thermodynamic states and surface fluxes.
Physics-parameterization training	Learn the mapping from 71 input features to 68 physics tendencies with T=20 multistep windows.
Long-term stability validation	Run a fixed 64-step, eight-day single-column integration without teacher forcing.
Water-budget evaluation	Diagnose precipitation from column-water storage, latent heat flux, and advective moisture convergence.
Hugging Face/OneCode execution	Validate training, inference, evaluation, visualization, and checkpoint workflows in Hugging Face or OneCode environments.

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/StableNN-Phys --local-dir ./StableNN-Phys
cd StableNN-Phys

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 states and external physical forcing. Inputs contain 34-level static energy, 34-level total water, surface fluxes, and insolation for 71 values. Targets are 34-level temperature and moisture physics tendencies for the same column, giving 68 values. Samples form continuous T=20 windows at three-hour intervals. This repository uses a small amount of synthetic data to validate training, inference, and evaluation and does not represent the paper data distribution, training scale, or formal 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 uses Adam to fit T=20 multistep state sequences, with layer-mass-weighted MAD as the default paper loss. The recorded paper configuration uses a learning rate of 0.01, batch size 200, five epochs, and hidden width 128, while the default engineering configuration reduces the hidden width and sample count for quick workflow validation; results are saved to:

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

Trained Weights

This repository does not include pretrained weights under weight/. Checkpoints generated with synthetic data validate the engineering workflow only and are not official paper weights.

Inference

python scripts/inference.py

Inference performs a fixed 64-step, eight-day single-column rollout without teacher forcing and saves the complete state, tendency, and forcing sequence at every three-hour time. Complete numerical results are saved to result/output/rollout.npz.

Evaluation and Visualization

python scripts/result.py

Evaluation reports length-65 per-lead mass-weighted MAD and bias and computes R2 by data source. Precipitation is diagnosed from column-water storage, latent heat flux, and advective moisture convergence while the budget residual is checked; synthetic-data results validate the engineering workflow only and do not represent formal paper performance. Structured results and the figure are saved to:

result/evaluation/metrics.json
result/evaluation/state_precipitation_timeseries.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 StableNN-Phys specifications.

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