Publish PBL-Emulator engineering reproduction

.gitattributes

@@ -1,35 +1,2 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
 *.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,152 @@
+---
+license: apache-2.0
+language:
+- en
+tags:
+- OneScience
+- Earth Science
+- Planetary Boundary Layer
+- Domain-aware Neural Network
+frameworks: PyTorch
+---
+
+<p align="center">
+  <strong><span style="font-size: 30px;">PBL-Emulator</span></strong>
+</p>
+
+# 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](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)
+
+## 2. Download and Installation
+
+```bash
+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**
+
+```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
+pip install numpy pyyaml matplotlib
+```
+
+**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
+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.
+
+```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
+```
+
+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:
+
+```text
+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
+
+```bash
+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:
+
+```text
+result/output/predictions.npz
+```
+
+### Evaluation and Visualization
+
+```bash
+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:
+
+```text
+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.

README_zh.md

@@ -0,0 +1,153 @@
+---
+license: apache-2.0
+language:
+- en
+- zh
+tags:
+- OneScience
+- 地球科学
+- 行星边界层
+- Domain-aware Neural Network
+frameworks: PyTorch
+---
+
+<p align="center">
+  <strong><span style="font-size: 30px;">PBL-Emulator</span></strong>
+</p>
+
+# 模型介绍
+
+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 编程](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)
+
+## 2.下载安装
+
+```bash
+modelscope download --model OneScience/PBL-Emulator --local_dir ./PBL-Emulator
+cd PBL-Emulator
+```
+
+### 环境依赖
+
+**硬件要求**
+
+- 推荐使用 GPU 或 DCU 运行。
+- CPU 可用于默认小样本配置的流程验证。
+- DCU 用户需预先安装 DTK，建议使用 DTK 25.04.2 以上版本或与当前集群匹配的 OneScience 推荐版本。
+
+**DCU环境**
+
+```bash
+# 请首先激活 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环境**
+
+```bash
+# 请首先激活 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 的数据分布、数据规模或论文性能。
+
+```bash
+python scripts/fake_data.py
+```
+
+### 训练
+
+单卡训练可使用：
+
+```bash
+python scripts/train.py
+```
+
+多卡训练可使用：
+
+```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
+```
+
+默认配置只将训练轮数从论文协议的 1000 epochs 缩小为 6 epochs，不缩小输入或输出维度。正式实验应使用 1984-2005 年 WRF/NCEP-R2/YSU 数据和完整训练轮数，训练产物保存到：
+
+```text
+result/checkpoints/pbl_emulator.pt
+result/training/metrics.json
+```
+
+### 训练权重
+
+本仓库不在 `weight/` 中内置论文官方权重，也未发现论文作者发布的可直接加载预训练 checkpoint。本地训练生成的 `result/checkpoints/pbl_emulator.pt` 是当前数据对应的工程 checkpoint，不得描述为官方预训练权重。
+
+### 推理
+
+```bash
+python scripts/inference.py
+```
+
+推理加载训练 checkpoint，为测试样本生成同一时刻的 PBL 垂直剖面诊断结果。完整数值结果保存到：
+
+```text
+result/output/predictions.npz
+```
+
+### 评估和可视化
+
+```bash
+python scripts/result.py
+```
+
+评估提供分变量误差和相关性，并绘制目标与预测剖面图。虚拟数据结果仅用于工程验证，不代表论文性能，结果保存到：
+
+```text
+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。
+
+本仓库代码、官方模型权重和数据的使用仍应以各自项目中的许可证及使用条款为准。

conf/config.yaml

@@ -0,0 +1,53 @@
+project:
+  name: PBL-Emulator
+  format_version: "1.0"
+  seed: 19
+  independent_code_license: Apache-2.0
+
+paper_specification:
+  title: Fast domain-aware neural network emulation of a planetary boundary layer parameterization in a numerical weather forecast model
+  doi: 10.5194/gmd-12-4261-2019
+  dimensions: {input: 16, levels: 17, variables_per_level: 5}
+  output_variables: [U, V, W, tk, QVAPOR]
+  architectures: {FFN_hidden_dense: 34, hierarchical_hidden_dense_per_level: 3, width: 16}
+  training: {optimizer: Adam, learning_rate: 0.001, epochs: 1000, batch_size: 64, early_stopping_patience: 10}
+  temporal_protocol: "3-hourly input and same-timestamp diagnostic output"
+
+engineering_assumptions:
+  default_architecture: HAC
+  synthetic_data: "structured engineering data, not WRF observations or paper data"
+  split: "chronological 70/15/15 percent"
+  vertical_order: bottom_to_top
+
+paths:
+  data: data/pbl_emulator_synthetic.npz
+  checkpoint: result/checkpoints/pbl_emulator.pt
+  training_metrics: result/training/metrics.json
+  predictions: result/output/predictions.npz
+  evaluation_metrics: result/evaluation/metrics.json
+  figure: result/evaluation/profiles.png
+
+data:
+  samples: 768
+  seed: 7
+  interval_hours: 3
+
+model:
+  architecture: HAC
+  width: 16
+  levels: 17
+  output_variables: 5
+
+training:
+  epochs: 6
+  batch_size: 64
+  learning_rate: 0.001
+  early_stopping_patience: 10
+  seed: 19
+  resume: false
+  num_workers: 0
+
+paper_model:
+  epochs: 1000
+  batch_size: 64
+  width: 16

config.json

@@ -0,0 +1,43 @@
+{
+  "model_name": "PBL-Emulator",
+  "model_type": "pbl_emulator",
+  "architectures": ["PBLEmulator"],
+  "framework": "PyTorch",
+  "domain": "earth-science",
+  "task": "same-time-pbl-profile-diagnosis",
+  "implementation": {
+    "entry_point": "model/pbl_emulator.py",
+    "scope": "independent paper-method engineering reproduction",
+    "train_script": "scripts/train.py",
+    "inference_script": "scripts/inference.py",
+    "evaluation_script": "scripts/result.py",
+    "synthetic_data_script": "scripts/fake_data.py"
+  },
+  "architecture": {
+    "input_shape": ["B", 16],
+    "output_shape": ["B", 17, 5],
+    "default": "HAC",
+    "width": 16,
+    "parameter_counts": {"FFN": 10693, "HPC": 16597, "HAC": 26197}
+  },
+  "data": {
+    "format": "NPZ",
+    "protocol": "pbl_emulator_synthetic_v1",
+    "split": "chronological 70/15/15 percent",
+    "interval_hours": 3,
+    "input_shape": ["N", 16],
+    "target_shape": ["N", 17, 5],
+    "input_variables": ["Q2", "T2", "U10", "V10", "GRDFLX", "SWDOWN", "GLW", "LH", "HFX", "PBLH", "UST", "TSK", "TSLB", "SMOIS", "Ug", "Vg"],
+    "output_variables": ["U", "V", "W", "tk", "QVAPOR"],
+    "metadata": ["timestamps_hours", "virtual_heights_m", "synthetic_pblh_m"],
+    "synthetic": true
+  },
+  "configuration_sources": [
+    "conf/config.yaml",
+    "model/pbl_emulator.py",
+    "scripts/fake_data.py",
+    "scripts/train.py",
+    "scripts/inference.py",
+    "scripts/result.py"
+  ]
+}

configuration.json

@@ -0,0 +1,12 @@
+{
+  "framework": "PyTorch",
+  "task": "Offline emulation of WRF YSU planetary boundary layer diagnostics",
+  "model": "PBL-Emulator",
+  "input_format": "N16 near-surface and forcing variables",
+  "output_format": "N17x5 [U,V,W,tk,QVAPOR] at the same 3-hour timestamp",
+  "default_config": "conf/config.yaml",
+  "training": "scripts/train.py",
+  "inference": "scripts/inference.py",
+  "evaluation": "scripts/result.py",
+  "visualization": "scripts/result.py"
+}

model/pbl_emulator.py

@@ -0,0 +1,338 @@
+from __future__ import annotations
+
+import json
+import math
+import os
+import random
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+from torch import nn
+from torch.nn.parallel import DistributedDataParallel
+from torch.utils.data import DataLoader, TensorDataset
+
+
+INPUT_NAMES = ["Q2", "T2", "U10", "V10", "GRDFLX", "SWDOWN", "GLW", "LH", "HFX", "PBLH", "UST", "TSK", "TSLB", "SMOIS", "Ug", "Vg"]
+OUTPUT_NAMES = ["U", "V", "W", "tk", "QVAPOR"]
+PARAMETER_COUNTS = {"FFN": 10693, "HPC": 16597, "HAC": 26197}
+CHECKPOINT_FORMAT_VERSION = "1.0"
+
+
+def load_yaml(path: str | Path) -> dict[str, Any]:
+    try:
+        import yaml
+    except ImportError as exc:
+        raise RuntimeError("PyYAML is required to read conf/config.yaml") from exc
+    with open(path, "r", encoding="utf-8") as handle:
+        return yaml.safe_load(handle)
+
+
+class FFN(nn.Module):
+    def __init__(self, width: int = 16, levels: int = 17, variables: int = 5):
+        super().__init__()
+        layers: list[nn.Module] = []
+        in_features = 16
+        for _ in range(34):
+            layers.extend((nn.Linear(in_features, width), nn.ReLU()))
+            in_features = width
+        self.hidden = nn.Sequential(*layers)
+        self.output = nn.Linear(width, levels * variables)
+        self.levels, self.variables = levels, variables
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.output(self.hidden(x)).reshape(-1, self.levels, self.variables)
+
+
+class HierarchicalNetwork(nn.Module):
+    def __init__(self, mode: str, width: int = 16, levels: int = 17, variables: int = 5):
+        super().__init__()
+        self.mode, self.levels, self.variables = mode, levels, variables
+        blocks = []
+        for level in range(levels):
+            conditioned_outputs = variables * (level if mode == "HAC" else min(level, 1))
+            blocks.append(nn.Sequential(
+                nn.Linear(16 + conditioned_outputs, width), nn.ReLU(),
+                nn.Linear(width, width), nn.ReLU(),
+                nn.Linear(width, width), nn.ReLU(),
+                nn.Linear(width, variables),
+            ))
+        self.blocks = nn.ModuleList(blocks)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        outputs = []
+        for block in self.blocks:
+            if not outputs:
+                conditioned = x
+            elif self.mode == "HPC":
+                conditioned = torch.cat((x, outputs[-1]), dim=-1)
+            else:
+                conditioned = torch.cat((x, *outputs), dim=-1)
+            outputs.append(block(conditioned))
+        return torch.stack(outputs, dim=1)
+
+
+class HAC(HierarchicalNetwork):
+    def __init__(self, width: int = 16, levels: int = 17, variables: int = 5):
+        super().__init__("HAC", width, levels, variables)
+
+
+class PBLEmulator(nn.Module):
+    def __init__(self, architecture: str = "HAC", width: int = 16, levels: int = 17, variables: int = 5):
+        super().__init__()
+        self.architecture = architecture.upper()
+        if self.architecture == "FFN":
+            self.model = FFN(width, levels, variables)
+        elif self.architecture == "HPC":
+            self.model = HierarchicalNetwork(self.architecture, width, levels, variables)
+        elif self.architecture == "HAC":
+            self.model = HAC(width, levels, variables)
+        else:
+            raise ValueError(f"Unknown architecture: {architecture}")
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.model(x)
+
+
+def build_model(architecture: str = "HAC", width: int = 16, levels: int = 17, variables: int = 5) -> PBLEmulator:
+    model = PBLEmulator(architecture, width, levels, variables)
+    count = sum(parameter.numel() for parameter in model.parameters())
+    if width == 16 and levels == 17 and variables == 5:
+        assert count == PARAMETER_COUNTS[model.architecture], f"{model.architecture}: expected {PARAMETER_COUNTS[model.architecture]}, got {count}"
+    return model
+
+
+@dataclass
+class ColumnScaler:
+    mean: np.ndarray
+    scale: np.ndarray
+    minimum: np.ndarray
+    span: np.ndarray
+
+    @classmethod
+    def fit(cls, values: np.ndarray) -> "ColumnScaler":
+        flat = np.asarray(values, dtype=np.float64).reshape(len(values), -1)
+        mean = flat.mean(axis=0)
+        scale = flat.std(axis=0)
+        scale[scale < 1e-12] = 1.0
+        standardized = (flat - mean) / scale
+        minimum = standardized.min(axis=0)
+        span = standardized.max(axis=0) - minimum
+        span[span < 1e-12] = 1.0
+        return cls(mean, scale, minimum, span)
+
+    def transform(self, values: np.ndarray) -> np.ndarray:
+        shape = values.shape
+        flat = np.asarray(values, dtype=np.float64).reshape(len(values), -1)
+        return (((flat - self.mean) / self.scale - self.minimum) / self.span).reshape(shape).astype(np.float32)
+
+    def inverse_transform(self, values: np.ndarray) -> np.ndarray:
+        shape = values.shape
+        flat = np.asarray(values, dtype=np.float64).reshape(len(values), -1)
+        return ((flat * self.span + self.minimum) * self.scale + self.mean).reshape(shape).astype(np.float32)
+
+    def state_dict(self) -> dict[str, np.ndarray]:
+        return {"mean": self.mean, "scale": self.scale, "minimum": self.minimum, "span": self.span}
+
+    @classmethod
+    def from_state_dict(cls, state: dict[str, Any]) -> "ColumnScaler":
+        return cls(*(np.asarray(state[key]) for key in ("mean", "scale", "minimum", "span")))
+
+
+def _distributed() -> tuple[bool, int, int, int]:
+    world_size = int(os.environ.get("WORLD_SIZE", "1"))
+    rank = int(os.environ.get("RANK", "0"))
+    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
+    if world_size > 1 and not torch.distributed.is_initialized():
+        torch.distributed.init_process_group("nccl" if torch.cuda.is_available() else "gloo")
+    return world_size > 1, rank, local_rank, world_size
+
+
+def _reduce_loss(total: float, count: int, device: torch.device) -> float:
+    pair = torch.tensor([total, count], dtype=torch.float64, device=device)
+    if torch.distributed.is_initialized():
+        torch.distributed.all_reduce(pair, op=torch.distributed.ReduceOp.SUM)
+    return float(pair[0] / pair[1].clamp_min(1))
+
+
+def _epoch(model: nn.Module, loader: DataLoader, device: torch.device, optimizer: torch.optim.Optimizer | None) -> float:
+    model.train(optimizer is not None)
+    total, count = 0.0, 0
+    context = torch.enable_grad() if optimizer is not None else torch.no_grad()
+    with context:
+        for x_batch, y_batch in loader:
+            x_batch, y_batch = x_batch.to(device), y_batch.to(device)
+            if optimizer is not None:
+                optimizer.zero_grad(set_to_none=True)
+            loss = torch.mean((model(x_batch) - y_batch) ** 2)
+            if optimizer is not None:
+                loss.backward()
+                optimizer.step()
+            total += float(loss.detach()) * len(x_batch)
+            count += len(x_batch)
+    return _reduce_loss(total, count, device)
+
+
+def train_model(data_path: str | Path, checkpoint_path: str | Path, metrics_path: str | Path, settings: dict[str, Any]) -> dict[str, Any]:
+    distributed, rank, local_rank, world_size = _distributed()
+    seed = int(settings.get("seed", 19))
+    random.seed(seed + rank); np.random.seed(seed + rank); torch.manual_seed(seed + rank)
+    if torch.cuda.is_available():
+        torch.cuda.set_device(local_rank)
+        device = torch.device("cuda", local_rank)
+    else:
+        device = torch.device("cpu")
+    raw = np.load(data_path)
+    x_train, y_train = raw["x_train"], raw["y_train"]
+    x_val, y_val = raw["x_val"], raw["y_val"]
+    assert x_train.shape[1:] == (16,) and y_train.shape[1:] == (17, 5)
+    model_config = {
+        "architecture": settings.get("architecture", "HAC"),
+        "width": int(settings.get("width", 16)),
+        "levels": int(settings.get("levels", 17)),
+        "output_variables": int(settings.get("output_variables", 5)),
+    }
+    start_epoch, history, best_loss = 0, [], math.inf
+    checkpoint_path = Path(checkpoint_path)
+    saved = None
+    if settings.get("resume") and checkpoint_path.exists():
+        saved = torch.load(checkpoint_path, map_location=device, weights_only=False)
+        if saved.get("format_version") != CHECKPOINT_FORMAT_VERSION:
+            raise ValueError(f"Unsupported checkpoint format_version: {saved.get('format_version')!r}; expected {CHECKPOINT_FORMAT_VERSION!r}")
+        model_config = saved["model_config"]
+        x_scaler = ColumnScaler.from_state_dict(saved["x_scaler"])
+        y_scaler = ColumnScaler.from_state_dict(saved["y_scaler"])
+    else:
+        x_scaler, y_scaler = ColumnScaler.fit(x_train), ColumnScaler.fit(y_train)
+    x_train, y_train = x_scaler.transform(x_train), y_scaler.transform(y_train)
+    x_val, y_val = x_scaler.transform(x_val), y_scaler.transform(y_val)
+    model = build_model(model_config["architecture"], model_config["width"], model_config["levels"], model_config["output_variables"]).to(device)
+    optimizer = torch.optim.Adam(model.parameters(), lr=float(settings.get("learning_rate", 0.001)))
+    best_state = None
+    if saved is not None:
+        model.load_state_dict(saved.get("last_model", saved["model"]))
+        optimizer.load_state_dict(saved["optimizer_state"])
+        start_epoch, history, best_loss = saved["epoch"] + 1, saved["history"], saved["best_val_loss"]
+        best_state = {key: value.detach().cpu().clone() for key, value in saved["model"].items()}
+    sampler = torch.utils.data.distributed.DistributedSampler(TensorDataset(torch.from_numpy(x_train), torch.from_numpy(y_train)), shuffle=True) if distributed else None
+    train_set = sampler.dataset if sampler else TensorDataset(torch.from_numpy(x_train), torch.from_numpy(y_train))
+    train_loader = DataLoader(train_set, batch_size=int(settings.get("batch_size", 64)), sampler=sampler, shuffle=sampler is None, num_workers=int(settings.get("num_workers", 0)))
+    val_loader = DataLoader(TensorDataset(torch.from_numpy(x_val), torch.from_numpy(y_val)), batch_size=int(settings.get("batch_size", 64)), shuffle=False)
+    if distributed:
+        model = DistributedDataParallel(model, device_ids=[local_rank] if device.type == "cuda" else None)
+    patience, stale = int(settings.get("early_stopping_patience", 10)), 0
+    for epoch in range(start_epoch, int(settings.get("epochs", 6))):
+        if sampler is not None:
+            sampler.set_epoch(epoch)
+        train_loss = _epoch(model, train_loader, device, optimizer)
+        val_loss = _epoch(model, val_loader, device, None)
+        history.append({"epoch": epoch, "train_mse": train_loss, "val_mse": val_loss})
+        if val_loss < best_loss:
+            best_loss, stale = val_loss, 0
+            best_state = {key: value.detach().cpu().clone() for key, value in (model.module if distributed else model).state_dict().items()}
+        else:
+            stale += 1
+        if stale >= patience:
+            break
+    if rank == 0:
+        checkpoint_path.parent.mkdir(parents=True, exist_ok=True)
+        final_model = model.module if distributed else model
+        last_state = {key: value.detach().cpu().clone() for key, value in final_model.state_dict().items()}
+        if best_state is not None:
+            final_model.load_state_dict(best_state)
+        checkpoint = {
+            "format_version": CHECKPOINT_FORMAT_VERSION,
+            "model_config": model_config,
+            "model": best_state or last_state,
+            "last_model": last_state,
+            "optimizer_state": optimizer.state_dict(), "epoch": history[-1]["epoch"],
+            "best_val_loss": best_loss, "history": history, "settings": settings, "x_scaler": x_scaler.state_dict(), "y_scaler": y_scaler.state_dict(),
+            "random_state": {"python": random.getstate(), "numpy": np.random.get_state(), "torch": torch.get_rng_state()},
+            "world_size": world_size,
+        }
+        torch.save(checkpoint, checkpoint_path)
+        metrics_path = Path(metrics_path); metrics_path.parent.mkdir(parents=True, exist_ok=True)
+        metrics_path.write_text(json.dumps({"architecture": model_config["architecture"], "parameters": sum(p.numel() for p in final_model.parameters()), "best_val_mse_normalized": best_loss, "epochs_completed": len(history), "history": history, "world_size": world_size}, indent=2), encoding="utf-8")
+    if distributed:
+        torch.distributed.barrier(); torch.distributed.destroy_process_group()
+    return {"best_val_mse_normalized": best_loss, "epochs_completed": len(history)}
+
+
+def run_inference(data_path: str | Path, checkpoint_path: str | Path, output_path: str | Path) -> dict[str, Any]:
+    data = np.load(data_path)
+    checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
+    if checkpoint.get("format_version") != CHECKPOINT_FORMAT_VERSION:
+        raise ValueError(f"Unsupported checkpoint format_version: {checkpoint.get('format_version')!r}; expected {CHECKPOINT_FORMAT_VERSION!r}")
+    settings = checkpoint["model_config"]
+    model = build_model(settings["architecture"], int(settings["width"]), int(settings["levels"]), int(settings["output_variables"]))
+    model.load_state_dict(checkpoint["model"]); model.eval()
+    x_scaler, y_scaler = ColumnScaler.from_state_dict(checkpoint["x_scaler"]), ColumnScaler.from_state_dict(checkpoint["y_scaler"])
+    with torch.no_grad():
+        prediction_scaled = model(torch.from_numpy(x_scaler.transform(data["x_test"]))).numpy()
+    prediction = y_scaler.inverse_transform(prediction_scaled)
+    output_path = Path(output_path); output_path.parent.mkdir(parents=True, exist_ok=True)
+    np.savez_compressed(output_path, inputs=data["x_test"], targets=data["y_test"], predictions=prediction,
+                        timestamps=data["time_test"], heights_m=data["heights_m"], pblh_m=data["pblh_test"],
+                        input_names=np.asarray(INPUT_NAMES), output_names=np.asarray(OUTPUT_NAMES), architecture=np.asarray(settings["architecture"]))
+    return {"samples": len(prediction), "shape": list(prediction.shape), "architecture": settings["architecture"]}
+
+
+def _pearson(a: np.ndarray, b: np.ndarray) -> float:
+    a, b = a.ravel(), b.ravel()
+    if len(a) < 2 or np.std(a) < 1e-12 or np.std(b) < 1e-12:
+        return 0.0
+    return float(np.corrcoef(a, b)[0, 1])
+
+
+def _scores(target: np.ndarray, prediction: np.ndarray) -> dict[str, float]:
+    return {"rmse": float(np.sqrt(np.mean((prediction - target) ** 2))), "pearson": _pearson(target, prediction)}
+
+
+def evaluate(predictions_path: str | Path, metrics_path: str | Path, figure_path: str | Path) -> dict[str, Any]:
+    data = np.load(predictions_path)
+    target, prediction = data["targets"], data["predictions"]
+    heights, pblh = data["heights_m"], data["pblh_m"]
+    mask = heights[None, :] <= pblh[:, None]
+    metrics: dict[str, Any] = {
+        "physical_scale": True,
+        "primary_protocol": "synthetic virtual-level height <= synthetic PBLH mask",
+        "primary": {"standardized_rmse_by_variable": {}},
+        "full_17_level_diagnostics": {"by_variable": {}, "by_level_and_variable": {}},
+    }
+    for index, name in enumerate(OUTPUT_NAMES):
+        masked_target, masked_prediction = target[:, :, index][mask], prediction[:, :, index][mask]
+        scale = max(float(np.std(masked_target)), 1e-12)
+        metrics["primary"]["standardized_rmse_by_variable"][name] = float(np.sqrt(np.mean(((masked_prediction - masked_target) / scale) ** 2)))
+        metrics["full_17_level_diagnostics"]["by_variable"][name] = _scores(target[:, :, index], prediction[:, :, index])
+    primary_values = metrics["primary"]["standardized_rmse_by_variable"].values()
+    metrics["primary"]["macro_mean_standardized_rmse"] = float(np.mean(list(primary_values)))
+    metrics["primary"]["definition"] = "Unweighted mean of per-variable RMSE divided by that variable's target standard deviation within the synthetic PBLH mask; no physical units are mixed."
+    for level, height in enumerate(heights):
+        metrics["full_17_level_diagnostics"]["by_level_and_variable"][str(level)] = {
+            "height_m": float(height),
+            "by_variable": {name: _scores(target[:, level, index], prediction[:, level, index]) for index, name in enumerate(OUTPUT_NAMES)},
+        }
+    speed_true = np.hypot(target[:, :, 0], target[:, :, 1]); speed_pred = np.hypot(prediction[:, :, 0], prediction[:, :, 1])
+    direction_true = np.mod(1.5 * np.pi - np.arctan2(target[:, :, 1], target[:, :, 0]), 2 * np.pi)
+    direction_pred = np.mod(1.5 * np.pi - np.arctan2(prediction[:, :, 1], prediction[:, :, 0]), 2 * np.pi)
+    delta = np.arctan2(np.sin(direction_pred - direction_true), np.cos(direction_pred - direction_true))
+    metrics["wind_speed"] = _scores(speed_true, speed_pred)
+    metrics["wind_direction"] = {"convention": "meteorological direction from: 0 degrees from north, increasing clockwise", "circular_rmse_degrees": float(np.degrees(np.sqrt(np.mean(delta ** 2)))), "mean_absolute_circular_error_degrees": float(np.degrees(np.mean(np.abs(delta)))), "circular_correlation_cosine": float(np.mean(np.cos(delta)))}
+    metrics["synthetic_pblh_mask"] = {"synthetic": True, "definition": "virtual level height <= synthetic PBLH; not a paper or real-WRF PBL mask", "sample_level_pairs": int(mask.sum())}
+    metrics_path = Path(metrics_path); metrics_path.parent.mkdir(parents=True, exist_ok=True)
+    metrics_path.write_text(json.dumps(metrics, indent=2), encoding="utf-8")
+    try:
+        import matplotlib.pyplot as plt
+        figure_path = Path(figure_path); figure_path.parent.mkdir(parents=True, exist_ok=True)
+        fig, axes = plt.subplots(1, 5, figsize=(15, 4), sharey=True)
+        for index, (axis, name) in enumerate(zip(axes, OUTPUT_NAMES)):
+            axis.plot(target[:, :, index].mean(0), heights, label="target")
+            axis.plot(prediction[:, :, index].mean(0), heights, "--", label="prediction")
+            axis.set_title(name); axis.grid(alpha=0.25)
+        axes[0].set_ylabel("synthetic height (m)"); axes[-1].legend()
+        fig.tight_layout(); fig.savefig(figure_path, dpi=150); plt.close(fig)
+    except ImportError:
+        metrics["figure_note"] = "matplotlib unavailable; numerical evaluation completed"
+    return metrics

scripts/fake_data.py

@@ -0,0 +1,55 @@
+from pathlib import Path
+import sys
+
+import numpy as np
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+
+from model.pbl_emulator import INPUT_NAMES, OUTPUT_NAMES, load_yaml
+
+
+def main() -> None:
+    config = load_yaml(ROOT / "conf/config.yaml")
+    count, seed = int(config["data"]["samples"]), int(config["data"]["seed"])
+    interval_hours = int(config["data"]["interval_hours"])
+    rng = np.random.default_rng(seed)
+    hours = (np.arange(count) * interval_hours) % 24
+    days = (np.arange(count) * interval_hours / 24) % 365
+    daylight = np.maximum(0.0, np.sin(np.pi * (hours - 6) / 12))
+    season = np.sin(2 * np.pi * (days - 172) / 365)
+    swdown = (760 + 110 * season) * daylight
+    t2 = 284 + 10 * season + 5 * np.sin(2 * np.pi * (hours - 14) / 24) + rng.normal(0, 0.8, count)
+    q2 = np.clip(0.0065 + 0.004 * (season + 1) / 2 + 0.0012 * daylight + rng.normal(0, 0.0003, count), 0.001, None)
+    ug = 5.5 + 1.8 * np.sin(2 * np.pi * days / 9) + rng.normal(0, 0.5, count)
+    vg = 2.0 + 1.5 * np.cos(2 * np.pi * days / 13) + rng.normal(0, 0.5, count)
+    u10 = 0.62 * ug + rng.normal(0, 0.45, count); v10 = 0.62 * vg + rng.normal(0, 0.45, count)
+    hfx = 12 + 150 * daylight * (0.75 + 0.25 * season) + rng.normal(0, 5, count)
+    lh = 25 + 110 * daylight * (0.8 + 0.2 * season) + rng.normal(0, 5, count)
+    pblh = np.clip(170 + 1200 * daylight * (0.75 + 0.25 * season) + rng.normal(0, 45, count), 80, 1800)
+    ust = np.clip(0.12 + 0.0014 * hfx + 0.022 * np.hypot(u10, v10), 0.05, 1.2)
+    tsk = t2 + 1.5 + 4.0 * daylight; tslb = 283 + 7.0 * season
+    smois = np.clip(0.25 + 0.04 * np.sin(2 * np.pi * days / 40) - 0.025 * season + rng.normal(0, 0.008, count), 0.08, 0.45)
+    grdflx = 0.12 * hfx - 8 * (1 - daylight); glw = 285 + 2.2 * (t2 - 273.15) + 450 * q2
+    x = np.column_stack((q2, t2, u10, v10, grdflx, swdown, glw, lh, hfx, pblh, ust, tsk, tslb, smois, ug, vg)).astype(np.float32)
+    heights = np.array([10, 35, 70, 110, 160, 220, 290, 370, 460, 560, 680, 820, 980, 1160, 1370, 1600, 1880], dtype=np.float32)
+    z = heights[None, :]; mix = np.exp(-z / pblh[:, None])
+    u = ug[:, None] + (u10 - ug)[:, None] * mix + 0.18 * np.sin(z / 260)
+    v = vg[:, None] + (v10 - vg)[:, None] * mix + 0.14 * np.cos(z / 300)
+    w = 0.20 * daylight[:, None] * np.sin(np.pi * np.minimum(z / pblh[:, None], 1)) - 0.015 * (1 - daylight[:, None]) + rng.normal(0, 0.008, (count, 17))
+    tk = t2[:, None] - 0.0062 * z + 0.9 * mix * daylight[:, None]
+    qv = np.clip(q2[:, None] * np.exp(-z / 2100) * (0.88 + 0.12 * mix) + 1e-4 * smois[:, None], 1e-5, None)
+    y = np.stack((u, v, w, tk, qv), axis=-1).astype(np.float32)
+    y += rng.normal(size=y.shape).astype(np.float32) * np.array([0.035, 0.035, 0.003, 0.04, 0.000015], dtype=np.float32)
+    n_train, n_val = int(0.70 * count), int(0.15 * count)
+    slices = {"train": slice(0, n_train), "val": slice(n_train, n_train + n_val), "test": slice(n_train + n_val, count)}
+    output = ROOT / config["paths"]["data"]; output.parent.mkdir(parents=True, exist_ok=True)
+    payload = {"heights_m": heights, "input_names": np.asarray(INPUT_NAMES), "output_names": np.asarray(OUTPUT_NAMES)}
+    for name, section in slices.items():
+        payload.update({f"x_{name}": x[section], f"y_{name}": y[section], f"time_{name}": (np.arange(count)[section] * interval_hours).astype(np.int64), f"pblh_{name}": pblh[section].astype(np.float32)})
+    np.savez_compressed(output, **payload)
+    print(f"saved {output}: x={x.shape}, y={y.shape}, split={[len(range(*s.indices(count))) for s in slices.values()]}")
+
+
+if __name__ == "__main__":
+    main()

scripts/inference.py

@@ -0,0 +1,12 @@
+from pathlib import Path
+import sys
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+
+from model.pbl_emulator import load_yaml, run_inference
+
+if __name__ == "__main__":
+    config = load_yaml(ROOT / "conf/config.yaml")
+    summary = run_inference(ROOT / config["paths"]["data"], ROOT / config["paths"]["checkpoint"], ROOT / config["paths"]["predictions"])
+    print(summary)

scripts/result.py

@@ -0,0 +1,12 @@
+from pathlib import Path
+import sys
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+
+from model.pbl_emulator import evaluate, load_yaml
+
+if __name__ == "__main__":
+    config = load_yaml(ROOT / "conf/config.yaml")
+    summary = evaluate(ROOT / config["paths"]["predictions"], ROOT / config["paths"]["evaluation_metrics"], ROOT / config["paths"]["figure"])
+    print({"primary": summary["primary"], "wind_direction": summary["wind_direction"]})

scripts/train.py

@@ -0,0 +1,13 @@
+from pathlib import Path
+import sys
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+
+from model.pbl_emulator import load_yaml, train_model
+
+if __name__ == "__main__":
+    config = load_yaml(ROOT / "conf/config.yaml")
+    settings = {**config["model"], **config["training"], "format_version": config["project"]["format_version"]}
+    result = train_model(ROOT / config["paths"]["data"], ROOT / config["paths"]["checkpoint"], ROOT / config["paths"]["training_metrics"], settings)
+    print(result)