README.md

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

<p align="center">
  <strong><span style="font-size: 30px;">StableNN-Phys</span></strong>
</p>

# 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.