--- license: apache-2.0 language: - en tags: - OneScience - Earth Science - Physics Parameterization frameworks: PyTorch ---

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](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home) ## 2. Download and Installation ```bash 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** ```bash # 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** ```bash # 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. ```bash python scripts/fake_data.py ``` ### Training For single-device training, use: ```bash python scripts/train.py ``` For multi-GPU training, use: ```bash 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: ```text 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 ```bash 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 ```bash 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: ```text 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.