OneScience-Group/PBL-Emulator

PBL-Emulator

Model Introduction

PBL-Emulator uses domain-aware neural networks to diagnose same-timestamp vertical profiles of planetary-boundary-layer winds, temperature, and water vapor offline from near-surface states and forcing variables.

Paper: Fast domain-aware neural network emulation of a planetary boundary layer parameterization in a numerical weather forecast model
https://doi.org/10.5194/gmd-12-4261-2019

Model Description

The method was proposed by research teams in the Environmental Science Division and Mathematics and Computer Science Division at Argonne National Laboratory. The paper uses data from 1984-2005 generated by WRF v3.3.1 with NCEP-R2 forcing and the YSU planetary boundary layer scheme. The model is intended for offline diagnosis of PBL vertical profiles at the same timestamp, not forecasting future states.

Applicable Scenarios

Scenario	Description
Offline PBL profile diagnosis	Diagnose same-timestamp wind, temperature, and water-vapor profiles from existing near-surface states and forcing variables.
HPC/HAC vertical-dependency studies	Validate adjacent-lower-level conditioning in HPC and all-lower-level conditioning in HAC.
Engineering validation	Validate data generation, training, inference, evaluation, and visualization with structured synthetic data.
ModelScope/OneCode execution	Validate structured data, training, inference, PBL-profile metrics, and visualization in ModelScope or OneCode environments.
Multi-GPU training	Validate distributed training and the checkpoint workflow 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/PBL-Emulator --local-dir ./PBL-Emulator
cd PBL-Emulator

Environment Dependencies

Hardware Requirements

• A GPU or DCU is recommended.
• A CPU can be used to validate the workflow with the default small-sample configuration.
• DCU users must install DTK in advance. DTK 25.04.2 or later, or the OneScience-recommended version matching the current 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
pip install numpy pyyaml matplotlib

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
pip install numpy pyyaml matplotlib

Training Data

Each sample maps 16 near-surface state and forcing inputs to five outputs at 17 vertical positions at the same timestamp, from [N,16] to [N,17,5]. The structured synthetic data include diurnal and seasonal cycles, vertical structure, and physical relationships among heat, moisture, and wind. These data are for engineering validation only and do not represent WRF data distributions, dataset scale, or the paper's 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

The default configuration reduces only the training duration from the paper protocol's 1,000 epochs to six epochs; it does not reduce the input or output dimensions. Formal experiments require the 1984-2005 WRF/NCEP-R2/YSU data and the complete training duration, and training artifacts are saved to:

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

Trained Weights

No official paper weights are bundled under weight/, and no directly loadable pretrained checkpoint released by the paper authors was identified. The locally generated result/checkpoints/pbl_emulator.pt is an engineering checkpoint for the current data and must not be represented as an official pretrained weight.

Inference

python scripts/inference.py

Inference loads the training checkpoint and generates same-timestamp PBL vertical-profile diagnoses for the test samples. Complete numerical results are saved to:

result/output/predictions.npz

Evaluation and Visualization

python scripts/result.py

Evaluation reports per-variable errors and correlations and plots target and predicted profiles. Results on synthetic data are for engineering validation only and do not represent the paper's performance; outputs are saved to:

result/evaluation/metrics.json
result/evaluation/profiles.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 PBL-Emulator specifications; the paper text is licensed under CC BY 4.0, while the official paper code uses BSD-3-Clause.

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