PBL-Emulator / README_zh.md

zhangrenchao

Publish PBL-Emulator engineering reproduction

702a3cf verified 2 days ago

preview code

Raw

History Blame Contribute Delete

5.2 kB

metadata

license: apache-2.0
language:
  - en
  - zh
tags:
  - OneScience
  - 地球科学
  - 行星边界层
  - Domain-aware Neural Network
frameworks: PyTorch

PBL-Emulator

模型介绍

PBL-Emulator 使用领域感知神经网络，根据近地面状态与强迫变量离线诊断同一时刻的行星边界层风、温度和水汽垂直剖面。

论文：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

模型描述

该方法由 Argonne National Laboratory 环境科学部、数学与计算机科学部的研究团队提出。论文使用 WRF v3.3.1 在 NCEP-R2 驱动下采用 YSU 行星边界层方案生成的 1984-2005 年数据。模型适用于离线诊断同一时刻的 PBL 垂直剖面，不用于未来时刻预报。

适用场景

场景	说明
离线 PBL 剖面诊断	根据已有近地面状态与强迫变量诊断同一时刻的风、温度和水汽剖面。
HPC/HAC 垂向依赖研究	验证 HPC 的相邻低层条件依赖和 HAC 的全部低层条件依赖。
工程验证	使用结构化虚拟数据验证数据生成、训练、推理、评估和可视化流程。
ModelScope/OneCode 运行	在 ModelScope 或 OneCode 环境中验证结构化数据、训练、推理、PBL 剖面指标和可视化流程。
多卡训练	通过 `torchrun` 验证分布式训练和 checkpoint 流程。

使用说明

1.OneCode

可通过 OneCode 在线环境体验智能化一键式 AI4S 编程：

点击体验智能化一键式 AI4S 编程

2.下载安装

modelscope download --model OneScience/PBL-Emulator --local_dir ./PBL-Emulator
cd PBL-Emulator

环境依赖

硬件要求

推荐使用 GPU 或 DCU 运行。
CPU 可用于默认小样本配置的流程验证。
DCU 用户需预先安装 DTK，建议使用 DTK 25.04.2 以上版本或与当前集群匹配的 OneScience 推荐版本。

DCU环境

# 请首先激活 DTK 及 Conda
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环境

# 请首先激活 Conda
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

训练数据介绍

每个样本由 16 个近地面状态和强迫变量输入映射到同一时刻 17 个垂直位置上的 5 个变量输出，数据形状为 [N,16] 到 [N,17,5]。结构化虚拟数据包含昼夜和季节变化、垂向结构以及热量、湿度和风场等物理关联。这些数据仅用于工程验证，不代表 WRF 的数据分布、数据规模或论文性能。

python scripts/fake_data.py

训练

单卡训练可使用：

python scripts/train.py

多卡训练可使用：

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

默认配置只将训练轮数从论文协议的 1000 epochs 缩小为 6 epochs，不缩小输入或输出维度。正式实验应使用 1984-2005 年 WRF/NCEP-R2/YSU 数据和完整训练轮数，训练产物保存到：

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

训练权重

本仓库不在 weight/ 中内置论文官方权重，也未发现论文作者发布的可直接加载预训练 checkpoint。本地训练生成的 result/checkpoints/pbl_emulator.pt 是当前数据对应的工程 checkpoint，不得描述为官方预训练权重。

推理

python scripts/inference.py

推理加载训练 checkpoint，为测试样本生成同一时刻的 PBL 垂直剖面诊断结果。完整数值结果保存到：

result/output/predictions.npz

评估和可视化

python scripts/result.py

评估提供分变量误差和相关性，并绘制目标与预测剖面图。虚拟数据结果仅用于工程验证，不代表论文性能，结果保存到：

result/evaluation/metrics.json
result/evaluation/profiles.png

OneScience官方信息

平台	OneScience 主仓库	Skills 仓库
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

引用与许可证

本仓库为 PBL-Emulator 公开规格的独立工程复现版本；论文文本采用 CC BY 4.0，论文官方代码采用 BSD-3-Clause。

本仓库代码、官方模型权重和数据的使用仍应以各自项目中的许可证及使用条款为准。