Publish WoFS-StormCal engineering reproduction
Browse files- .gitattributes +1 -24
- README.md +137 -0
- README_zh.md +138 -0
- conf/config.yaml +37 -0
- config.json +46 -0
- configuration.json +13 -0
- model/wofsstormcal.py +132 -0
- scripts/fake_data.py +86 -0
- scripts/inference.py +49 -0
- scripts/result.py +116 -0
- scripts/train.py +129 -0
- weight/.gitkeep +0 -0
.gitattributes
CHANGED
|
@@ -1,35 +1,12 @@
|
|
| 1 |
-
*.7z filter=lfs diff=lfs merge=lfs -text
|
| 2 |
-
*.arrow filter=lfs diff=lfs merge=lfs -text
|
| 3 |
*.bin filter=lfs diff=lfs merge=lfs -text
|
| 4 |
-
*.bz2 filter=lfs diff=lfs merge=lfs -text
|
| 5 |
*.ckpt filter=lfs diff=lfs merge=lfs -text
|
| 6 |
-
*.ftz filter=lfs diff=lfs merge=lfs -text
|
| 7 |
-
*.gz filter=lfs diff=lfs merge=lfs -text
|
| 8 |
-
*.h5 filter=lfs diff=lfs merge=lfs -text
|
| 9 |
*.joblib filter=lfs diff=lfs merge=lfs -text
|
| 10 |
-
*.lfs.* filter=lfs diff=lfs merge=lfs -text
|
| 11 |
-
*.mlmodel filter=lfs diff=lfs merge=lfs -text
|
| 12 |
*.model filter=lfs diff=lfs merge=lfs -text
|
| 13 |
-
*.msgpack filter=lfs diff=lfs merge=lfs -text
|
| 14 |
*.npy filter=lfs diff=lfs merge=lfs -text
|
| 15 |
*.npz filter=lfs diff=lfs merge=lfs -text
|
| 16 |
-
*.onnx filter=lfs diff=lfs merge=lfs -text
|
| 17 |
-
*.ot filter=lfs diff=lfs merge=lfs -text
|
| 18 |
-
*.parquet filter=lfs diff=lfs merge=lfs -text
|
| 19 |
-
*.pb filter=lfs diff=lfs merge=lfs -text
|
| 20 |
-
*.pickle filter=lfs diff=lfs merge=lfs -text
|
| 21 |
*.pkl filter=lfs diff=lfs merge=lfs -text
|
| 22 |
*.pt filter=lfs diff=lfs merge=lfs -text
|
| 23 |
*.pth filter=lfs diff=lfs merge=lfs -text
|
| 24 |
-
*.rar filter=lfs diff=lfs merge=lfs -text
|
| 25 |
*.safetensors filter=lfs diff=lfs merge=lfs -text
|
| 26 |
-
saved_model/**/* filter=lfs diff=lfs merge=lfs -text
|
| 27 |
-
*.tar.* filter=lfs diff=lfs merge=lfs -text
|
| 28 |
-
*.tar filter=lfs diff=lfs merge=lfs -text
|
| 29 |
-
*.tflite filter=lfs diff=lfs merge=lfs -text
|
| 30 |
-
*.tgz filter=lfs diff=lfs merge=lfs -text
|
| 31 |
-
*.wasm filter=lfs diff=lfs merge=lfs -text
|
| 32 |
-
*.xz filter=lfs diff=lfs merge=lfs -text
|
| 33 |
*.zip filter=lfs diff=lfs merge=lfs -text
|
| 34 |
-
*
|
| 35 |
-
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
|
|
|
|
|
|
|
|
|
| 1 |
*.bin filter=lfs diff=lfs merge=lfs -text
|
|
|
|
| 2 |
*.ckpt filter=lfs diff=lfs merge=lfs -text
|
|
|
|
|
|
|
|
|
|
| 3 |
*.joblib filter=lfs diff=lfs merge=lfs -text
|
|
|
|
|
|
|
| 4 |
*.model filter=lfs diff=lfs merge=lfs -text
|
|
|
|
| 5 |
*.npy filter=lfs diff=lfs merge=lfs -text
|
| 6 |
*.npz filter=lfs diff=lfs merge=lfs -text
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 7 |
*.pkl filter=lfs diff=lfs merge=lfs -text
|
| 8 |
*.pt filter=lfs diff=lfs merge=lfs -text
|
| 9 |
*.pth filter=lfs diff=lfs merge=lfs -text
|
|
|
|
| 10 |
*.safetensors filter=lfs diff=lfs merge=lfs -text
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 11 |
*.zip filter=lfs diff=lfs merge=lfs -text
|
| 12 |
+
weight/* filter=lfs diff=lfs merge=lfs -text
|
|
|
README.md
ADDED
|
@@ -0,0 +1,137 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
---
|
| 2 |
+
license: apache-2.0
|
| 3 |
+
language:
|
| 4 |
+
- en
|
| 5 |
+
tags:
|
| 6 |
+
- OneScience
|
| 7 |
+
- Earth Science
|
| 8 |
+
- Severe Weather
|
| 9 |
+
- Probability Calibration
|
| 10 |
+
frameworks: PyTorch
|
| 11 |
+
---
|
| 12 |
+
|
| 13 |
+
<p align="center"><strong><span style="font-size: 30px;">WoFS-StormCal</span></strong></p>
|
| 14 |
+
|
| 15 |
+
# Model Introduction
|
| 16 |
+
|
| 17 |
+
WoFS-StormCal addresses calibration of short-term, storm-scale severe-weather probabilities by combining storm state, surrounding environment, and object morphology from ensemble storm tracks to estimate the likelihood of tornadoes, severe hail, and severe wind. It is primarily intended to improve the reliability of ensemble guidance and support nowcasting, severe-weather risk assessment, and forecaster decision-making.
|
| 18 |
+
|
| 19 |
+
Paper: Using Machine Learning to Calibrate Storm-Scale Probabilistic Guidance of Severe Weather Hazards in the Warn-on-Forecast System
|
| 20 |
+
https://arxiv.org/abs/2012.00679
|
| 21 |
+
|
| 22 |
+
# Model Description
|
| 23 |
+
|
| 24 |
+
WoFS-StormCal was proposed by research teams from the University of Oklahoma, the Cooperative Institute for Mesoscale Meteorological Studies, and the NOAA National Severe Storms Laboratory. The paper trained and validated the models with WoFS ensemble forecasts and local storm reports from the 2017-2019 NOAA Hazardous Weather Testbed Spring Forecasting Experiments. The model supports short-term storm-scale probabilistic forecasting and calibration for tornadoes, severe hail, and severe wind.
|
| 25 |
+
|
| 26 |
+
# Use Cases
|
| 27 |
+
|
| 28 |
+
| Use Case | Description |
|
| 29 |
+
| :---: | :--- |
|
| 30 |
+
| Storm-scale probabilistic guidance | Predict tornado, severe-hail, and severe-wind probabilities for ensemble storm-track objects. |
|
| 31 |
+
| Probability calibration | Fit monotone isotonic mappings from cross-validated probabilities. |
|
| 32 |
+
| Lead-group modeling | Process first-hour and second-hour 30-minute storm-track windows separately. |
|
| 33 |
+
| ModelScope/OneCode execution | Validate training, inference, probabilistic evaluation, visualization, and checkpoint workflows. |
|
| 34 |
+
| Multi-GPU training | Launch distributed data-parallel training with `torchrun`. |
|
| 35 |
+
|
| 36 |
+
# Usage Instructions
|
| 37 |
+
|
| 38 |
+
## 1.OneCode
|
| 39 |
+
|
| 40 |
+
Experience intelligent, one-click AI4S programming through the OneCode online environment:
|
| 41 |
+
|
| 42 |
+
[Try intelligent, one-click AI4S programming](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)
|
| 43 |
+
|
| 44 |
+
## 2. Download and Installation
|
| 45 |
+
|
| 46 |
+
```bash
|
| 47 |
+
hf download OneScience-Group/WoFS-StormCal --local-dir ./WoFS-StormCal
|
| 48 |
+
cd WoFS-StormCal
|
| 49 |
+
```
|
| 50 |
+
|
| 51 |
+
### Environment Dependencies
|
| 52 |
+
|
| 53 |
+
**Hardware Requirements**
|
| 54 |
+
|
| 55 |
+
- A GPU or DCU is recommended.
|
| 56 |
+
- A CPU can be used for connectivity validation with the default small-sample configuration.
|
| 57 |
+
- DCU users must install DTK in advance. DTK 25.04.2 or later, or the OneScience-recommended version matching the current cluster, is recommended.
|
| 58 |
+
|
| 59 |
+
**DCU Environment**
|
| 60 |
+
|
| 61 |
+
```bash
|
| 62 |
+
# Activate DTK and Conda first
|
| 63 |
+
conda create -n onescience311 python=3.11 -y
|
| 64 |
+
conda activate onescience311
|
| 65 |
+
pip install onescience[earth-dcu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
|
| 66 |
+
```
|
| 67 |
+
|
| 68 |
+
**GPU Environment**
|
| 69 |
+
|
| 70 |
+
```bash
|
| 71 |
+
# Activate Conda first
|
| 72 |
+
conda create -n onescience311 python=3.11 -y libstdcxx-ng=12 libgcc-ng=12 gcc_linux-64=12 gxx_linux-64=12
|
| 73 |
+
conda activate onescience311
|
| 74 |
+
pip install onescience[earth-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
|
| 75 |
+
```
|
| 76 |
+
|
| 77 |
+
### Training Data
|
| 78 |
+
|
| 79 |
+
Synthetic features use storm organization, rotation, instability, shear, cold-pool, and ensemble-spread latent drivers to create within-group and cross-group correlation while preserving all 113 dimensions. The data validates feature loading, probabilistic training, calibration, inference, and evaluation only; it does not represent the real severe-weather sample distribution or paper performance.
|
| 80 |
+
|
| 81 |
+
```bash
|
| 82 |
+
python scripts/fake_data.py
|
| 83 |
+
```
|
| 84 |
+
|
| 85 |
+
### Training
|
| 86 |
+
|
| 87 |
+
For single-device training, use:
|
| 88 |
+
|
| 89 |
+
```bash
|
| 90 |
+
python scripts/train.py
|
| 91 |
+
```
|
| 92 |
+
|
| 93 |
+
For multi-GPU training, use:
|
| 94 |
+
|
| 95 |
+
```bash
|
| 96 |
+
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
|
| 97 |
+
```
|
| 98 |
+
|
| 99 |
+
Training outputs are saved to:
|
| 100 |
+
|
| 101 |
+
```text
|
| 102 |
+
result/checkpoints/wofsstormcal.pt
|
| 103 |
+
result/training/metrics.json
|
| 104 |
+
```
|
| 105 |
+
|
| 106 |
+
### Trained Weights
|
| 107 |
+
|
| 108 |
+
No weights are bundled under `weight/`. The paper provides no confirmed official checkpoint, and the engineering checkpoint is not claimed to be compatible with external weights.
|
| 109 |
+
|
| 110 |
+
### Inference
|
| 111 |
+
|
| 112 |
+
```bash
|
| 113 |
+
python scripts/inference.py
|
| 114 |
+
```
|
| 115 |
+
|
| 116 |
+
Inference loads the training checkpoint, produces calibrated probabilities for three hazards from 113-dimensional object features, and saves them to `result/output/predictions.npz`.
|
| 117 |
+
|
| 118 |
+
### Evaluation and Visualization
|
| 119 |
+
|
| 120 |
+
```bash
|
| 121 |
+
python scripts/result.py
|
| 122 |
+
```
|
| 123 |
+
|
| 124 |
+
Evaluation reports probabilistic skill, classification skill, and reliability metrics for each hazard and lead group in `result/evaluation/metrics.json`. It also generates performance and reliability diagrams. Synthetic-data results are only for engineering-workflow validation and do not represent paper performance on the real test set.
|
| 125 |
+
|
| 126 |
+
# Official OneScience Information
|
| 127 |
+
|
| 128 |
+
| Platform | OneScience Main Repository | Skills Repository |
|
| 129 |
+
| --- | --- | --- |
|
| 130 |
+
| Gitee | https://gitee.com/onescience-ai/onescience | https://gitee.com/onescience-ai/oneskills |
|
| 131 |
+
| GitHub | https://github.com/onescience-ai/OneScience | https://github.com/onescience-ai/oneskills |
|
| 132 |
+
|
| 133 |
+
# Citation and License
|
| 134 |
+
|
| 135 |
+
This repository is an independent engineering reproduction of the public WoFS-StormCal paper specifications.
|
| 136 |
+
|
| 137 |
+
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
ADDED
|
@@ -0,0 +1,138 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
---
|
| 2 |
+
license: Apache-2.0
|
| 3 |
+
language:
|
| 4 |
+
- en
|
| 5 |
+
- zh
|
| 6 |
+
tags:
|
| 7 |
+
- OneScience
|
| 8 |
+
- 地球科学
|
| 9 |
+
- 强天气预报
|
| 10 |
+
- 概率校准
|
| 11 |
+
frameworks: PyTorch
|
| 12 |
+
---
|
| 13 |
+
|
| 14 |
+
<p align="center"><strong><span style="font-size: 30px;">WoFS-StormCal</span></strong></p>
|
| 15 |
+
|
| 16 |
+
# 模型介绍
|
| 17 |
+
|
| 18 |
+
WoFS-StormCal 用于解决短时风暴尺度强天气概率预报的校准问题,综合集合风暴轨迹中的风暴状态、周边环境和对象形态信息,判断风暴产生龙卷风、严重冰雹或严重大风的可能性。模型主要用于提高集合预报概率的可靠性,为临近预报、强天气风险研判和预报员决策提供概率指导。
|
| 19 |
+
|
| 20 |
+
论文:Using Machine Learning to Calibrate Storm-Scale Probabilistic Guidance of Severe Weather Hazards in the Warn-on-Forecast System
|
| 21 |
+
https://arxiv.org/abs/2012.00679
|
| 22 |
+
|
| 23 |
+
# 模型描述
|
| 24 |
+
|
| 25 |
+
WoFS-StormCal 由 University of Oklahoma、Cooperative Institute for Mesoscale Meteorological Studies 和 NOAA National Severe Storms Laboratory 的研究团队提出。论文使用 2017 至 2019 年 NOAA Hazardous Weather Testbed Spring Forecasting Experiments 的 WoFS 集合预报和本地风暴报告训练与验证。模型适用于龙卷风、严重冰雹和严重大风的短时风暴尺度概率预报与校准。
|
| 26 |
+
|
| 27 |
+
# 适用场景
|
| 28 |
+
|
| 29 |
+
| 场景 | 说明 |
|
| 30 |
+
| :---: | :--- |
|
| 31 |
+
| 风暴尺度概率预报 | 根据集合风暴轨迹对象预测龙卷风、严重冰雹和严重大风概率。 |
|
| 32 |
+
| 概率校准 | 使用交叉验证概率拟合单调 isotonic 映射。 |
|
| 33 |
+
| 分时效建模 | 分别处理 first hour 和 second hour 的 30 分钟风暴轨迹窗口。 |
|
| 34 |
+
| ModelScope/OneCode 运行 | 验证训练、推理、概率评估、可视化和 checkpoint 流程。 |
|
| 35 |
+
| 多卡训练 | 通过 `torchrun` 启动分布式数据并行训练。 |
|
| 36 |
+
|
| 37 |
+
# 使用说明
|
| 38 |
+
|
| 39 |
+
## 1.OneCode
|
| 40 |
+
|
| 41 |
+
可通过 OneCode 在线环境体验智能化一键式 AI4S 编程:
|
| 42 |
+
|
| 43 |
+
[点击体验智能化一键式 AI4S 编程](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)
|
| 44 |
+
|
| 45 |
+
## 2.下载安装
|
| 46 |
+
|
| 47 |
+
```bash
|
| 48 |
+
modelscope download --model OneScience/WoFS-StormCal --local_dir ./WoFS-StormCal
|
| 49 |
+
cd WoFS-StormCal
|
| 50 |
+
```
|
| 51 |
+
|
| 52 |
+
### 环境依赖
|
| 53 |
+
|
| 54 |
+
**硬件要求**
|
| 55 |
+
|
| 56 |
+
- 推荐使用 GPU 或 DCU 运行。
|
| 57 |
+
- CPU 可用于默认小样本配置的连通性验证。
|
| 58 |
+
- DCU 用户需预先安装 DTK,建议使用 DTK 25.04.2 以上版本或与当前集群匹配的 OneScience 推荐版本。
|
| 59 |
+
|
| 60 |
+
**DCU环境**
|
| 61 |
+
|
| 62 |
+
```bash
|
| 63 |
+
# 请首先激活DTK及CONDA
|
| 64 |
+
conda create -n onescience311 python=3.11 -y
|
| 65 |
+
conda activate onescience311
|
| 66 |
+
pip install onescience[earth-dcu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
|
| 67 |
+
```
|
| 68 |
+
|
| 69 |
+
**GPU环境**
|
| 70 |
+
|
| 71 |
+
```bash
|
| 72 |
+
# 请首先激活CONDA
|
| 73 |
+
conda create -n onescience311 python=3.11 -y libstdcxx-ng=12 libgcc-ng=12 gcc_linux-64=12 gxx_linux-64=12
|
| 74 |
+
conda activate onescience311
|
| 75 |
+
pip install onescience[earth-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
|
| 76 |
+
```
|
| 77 |
+
|
| 78 |
+
### 训练数据介绍
|
| 79 |
+
|
| 80 |
+
虚拟数据使用组织度、旋转、热力不稳定度、风切变、冷池和集合离散度等潜变量构造组内及跨组相关结构,并保持真实 113 维不缩减。该数据仅用于验证特征加载、概率训练、校准、推理和评估流程,不代表真实强天气样本分布与论文性能。
|
| 81 |
+
|
| 82 |
+
```bash
|
| 83 |
+
python scripts/fake_data.py
|
| 84 |
+
```
|
| 85 |
+
|
| 86 |
+
### 训练
|
| 87 |
+
|
| 88 |
+
单卡训练可使用:
|
| 89 |
+
|
| 90 |
+
```bash
|
| 91 |
+
python scripts/train.py
|
| 92 |
+
```
|
| 93 |
+
|
| 94 |
+
多卡训练可使用:
|
| 95 |
+
|
| 96 |
+
```bash
|
| 97 |
+
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
|
| 98 |
+
```
|
| 99 |
+
|
| 100 |
+
训练输出保存到:
|
| 101 |
+
|
| 102 |
+
```text
|
| 103 |
+
result/checkpoints/wofsstormcal.pt
|
| 104 |
+
result/training/metrics.json
|
| 105 |
+
```
|
| 106 |
+
|
| 107 |
+
### 训练权重
|
| 108 |
+
|
| 109 |
+
本仓库不在 `weight/` 中内置权重。论文未提供可确认的官方 checkpoint,当前工程 checkpoint 不声明兼容外部权重。
|
| 110 |
+
|
| 111 |
+
### 推理
|
| 112 |
+
|
| 113 |
+
```bash
|
| 114 |
+
python scripts/inference.py
|
| 115 |
+
```
|
| 116 |
+
|
| 117 |
+
推理加载训练 checkpoint,根据 113 维对象特征输出三个灾种的校准概率,并保存到 `result/output/predictions.npz`。
|
| 118 |
+
|
| 119 |
+
### 评估和可视化
|
| 120 |
+
|
| 121 |
+
```bash
|
| 122 |
+
python scripts/result.py
|
| 123 |
+
```
|
| 124 |
+
|
| 125 |
+
评估结果包含各灾种和时效组的概率技巧、分类技巧与可靠性指标,并保存到 `result/evaluation/metrics.json`。脚本同时生成性能图和可靠性图。虚拟数据结果仅用于验证工程流程,不代表论文真实测试集性能。
|
| 126 |
+
|
| 127 |
+
# OneScience官方信息
|
| 128 |
+
|
| 129 |
+
| 平台 | OneScience 主仓库 | Skills 仓库 |
|
| 130 |
+
| --- | --- | --- |
|
| 131 |
+
| Gitee | https://gitee.com/onescience-ai/onescience | https://gitee.com/onescience-ai/oneskills |
|
| 132 |
+
| GitHub | https://github.com/onescience-ai/OneScience | https://github.com/onescience-ai/oneskills |
|
| 133 |
+
|
| 134 |
+
# 引用与许可证
|
| 135 |
+
|
| 136 |
+
本仓库为 WoFS-StormCal 论文公开规格的独立工程复现版本。
|
| 137 |
+
|
| 138 |
+
本仓库代码、官方模型权重和数据的使用仍应以各自项目中的许可证及使用条款为准。
|
conf/config.yaml
ADDED
|
@@ -0,0 +1,37 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
seed: 42
|
| 2 |
+
data:
|
| 3 |
+
root: data
|
| 4 |
+
format_version: wofs_storm_object_113_v1
|
| 5 |
+
train_samples: 1200
|
| 6 |
+
test_samples: 500
|
| 7 |
+
input_features: 113
|
| 8 |
+
amplitude_features: 30
|
| 9 |
+
spatial_features: 76
|
| 10 |
+
object_properties: 7
|
| 11 |
+
hazards: [tornado, hail, wind]
|
| 12 |
+
lead_groups: [first_hour, second_hour]
|
| 13 |
+
ensemble_members: 18
|
| 14 |
+
grid_spacing_km: 3
|
| 15 |
+
forecast_window_minutes: 30
|
| 16 |
+
forecast_interval_minutes: 5
|
| 17 |
+
model:
|
| 18 |
+
calibration_points: 256
|
| 19 |
+
train:
|
| 20 |
+
epochs: 30
|
| 21 |
+
batch_size: 128
|
| 22 |
+
learning_rate: 0.03
|
| 23 |
+
l1_strength: 0.002
|
| 24 |
+
l2_strength: 0.01
|
| 25 |
+
calibration_fraction: 0.2
|
| 26 |
+
gradient_clip_norm: 5.0
|
| 27 |
+
num_workers: 0
|
| 28 |
+
runtime:
|
| 29 |
+
device: auto
|
| 30 |
+
paths:
|
| 31 |
+
checkpoint: result/checkpoints/wofsstormcal.pt
|
| 32 |
+
training_metrics: result/training/metrics.json
|
| 33 |
+
inference_dir: result/output
|
| 34 |
+
evaluation_dir: result/evaluation
|
| 35 |
+
evaluation:
|
| 36 |
+
probability_bins: 10
|
| 37 |
+
threshold_count: 201
|
config.json
ADDED
|
@@ -0,0 +1,46 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"model_name": "WoFS-StormCal",
|
| 3 |
+
"model_type": "wofsstormcal",
|
| 4 |
+
"architectures": ["WoFSStormCal"],
|
| 5 |
+
"framework": "PyTorch",
|
| 6 |
+
"domain": "meteorology",
|
| 7 |
+
"task": "storm-object-multilabel-severe-hazard-probability",
|
| 8 |
+
"license": "Apache-2.0",
|
| 9 |
+
"implementation": {
|
| 10 |
+
"entry_point": "model/wofsstormcal.py",
|
| 11 |
+
"train_script": "scripts/train.py",
|
| 12 |
+
"inference_script": "scripts/inference.py",
|
| 13 |
+
"evaluation_script": "scripts/result.py",
|
| 14 |
+
"synthetic_data_script": "scripts/fake_data.py"
|
| 15 |
+
},
|
| 16 |
+
"architecture": {
|
| 17 |
+
"family": "elastic-net logistic regression with isotonic calibration",
|
| 18 |
+
"input_shape": ["N", 113],
|
| 19 |
+
"output_shape": ["N", 3],
|
| 20 |
+
"feature_groups": {"amplitude": 30, "spatial": 76, "object_property": 7},
|
| 21 |
+
"hazards": ["tornado", "hail", "wind"],
|
| 22 |
+
"lead_groups": ["first_hour", "second_hour"]
|
| 23 |
+
},
|
| 24 |
+
"paper_model": {
|
| 25 |
+
"families": ["random_forest", "xgboost", "elastic_net_logistic_regression"],
|
| 26 |
+
"probability_calibration": "isotonic regression",
|
| 27 |
+
"cross_validation_folds": 5
|
| 28 |
+
},
|
| 29 |
+
"forecast_metadata": {
|
| 30 |
+
"ensemble_members": 18,
|
| 31 |
+
"grid_spacing_km": 3,
|
| 32 |
+
"forecast_window_minutes": 30,
|
| 33 |
+
"forecast_interval_minutes": 5,
|
| 34 |
+
"first_hour_start_minutes": [0, 60],
|
| 35 |
+
"second_hour_start_minutes": [65, 120]
|
| 36 |
+
},
|
| 37 |
+
"data": {"protocol": "wofs_storm_object_113_v1", "format": "NPZ"},
|
| 38 |
+
"configuration_sources": [
|
| 39 |
+
"conf/config.yaml",
|
| 40 |
+
"model/wofsstormcal.py",
|
| 41 |
+
"scripts/fake_data.py",
|
| 42 |
+
"scripts/train.py",
|
| 43 |
+
"scripts/inference.py",
|
| 44 |
+
"scripts/result.py"
|
| 45 |
+
]
|
| 46 |
+
}
|
configuration.json
ADDED
|
@@ -0,0 +1,13 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"framework": "PyTorch",
|
| 3 |
+
"task": "storm-object-severe-hazard-probability-calibration",
|
| 4 |
+
"model": "WoFS-StormCal",
|
| 5 |
+
"input_format": "N113",
|
| 6 |
+
"output_format": "N3",
|
| 7 |
+
"protocol": "113-feature storm-object elastic-net probability calibration for three hazards and two lead groups",
|
| 8 |
+
"default_config": "conf/config.yaml",
|
| 9 |
+
"training": "scripts/train.py",
|
| 10 |
+
"inference": "scripts/inference.py",
|
| 11 |
+
"evaluation": "scripts/result.py",
|
| 12 |
+
"visualization": "scripts/result.py"
|
| 13 |
+
}
|
model/wofsstormcal.py
ADDED
|
@@ -0,0 +1,132 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Elastic-net logistic calibration of WoFS ensemble storm-track hazards."""
|
| 2 |
+
|
| 3 |
+
from __future__ import annotations
|
| 4 |
+
|
| 5 |
+
import numpy as np
|
| 6 |
+
import torch
|
| 7 |
+
from torch import nn
|
| 8 |
+
|
| 9 |
+
|
| 10 |
+
HAZARDS = ("tornado", "hail", "wind")
|
| 11 |
+
LEAD_GROUPS = ("first_hour", "second_hour")
|
| 12 |
+
|
| 13 |
+
|
| 14 |
+
def _pav(values: np.ndarray, targets: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
|
| 15 |
+
"""Fit an isotonic map with the pool-adjacent-violators algorithm."""
|
| 16 |
+
order = np.argsort(values, kind="stable")
|
| 17 |
+
x, y = values[order], targets[order].astype(np.float64)
|
| 18 |
+
starts, ends, sums, counts = [], [], [], []
|
| 19 |
+
for index, target in enumerate(y):
|
| 20 |
+
starts.append(index); ends.append(index); sums.append(float(target)); counts.append(1)
|
| 21 |
+
while len(sums) > 1 and sums[-2] / counts[-2] > sums[-1] / counts[-1]:
|
| 22 |
+
ends[-2] = ends[-1]
|
| 23 |
+
sums[-2] += sums[-1]
|
| 24 |
+
counts[-2] += counts[-1]
|
| 25 |
+
starts.pop(); ends.pop(); sums.pop(); counts.pop()
|
| 26 |
+
xp, yp = [], []
|
| 27 |
+
for start, end, total, count in zip(starts, ends, sums, counts):
|
| 28 |
+
level = total / count
|
| 29 |
+
xp.extend((float(x[start]), float(x[end])))
|
| 30 |
+
yp.extend((level, level))
|
| 31 |
+
xp = np.maximum.accumulate(np.asarray(xp, dtype=np.float32))
|
| 32 |
+
return xp, np.asarray(yp, dtype=np.float32)
|
| 33 |
+
|
| 34 |
+
|
| 35 |
+
class WoFSStormCal(nn.Module):
|
| 36 |
+
"""Two lead-group linear classifiers with portable isotonic calibration."""
|
| 37 |
+
|
| 38 |
+
input_dim = 113
|
| 39 |
+
output_dim = 3
|
| 40 |
+
hazards = HAZARDS
|
| 41 |
+
lead_groups = LEAD_GROUPS
|
| 42 |
+
ensemble_members = 18
|
| 43 |
+
grid_spacing_km = 3
|
| 44 |
+
forecast_window_minutes = 30
|
| 45 |
+
forecast_interval_minutes = 5
|
| 46 |
+
|
| 47 |
+
def __init__(self, calibration_points: int = 256):
|
| 48 |
+
super().__init__()
|
| 49 |
+
self.calibration_points = int(calibration_points)
|
| 50 |
+
self.weight = nn.Parameter(torch.empty(2, 3, self.input_dim))
|
| 51 |
+
self.bias = nn.Parameter(torch.zeros(2, 3))
|
| 52 |
+
nn.init.normal_(self.weight, std=0.01)
|
| 53 |
+
self.register_buffer("feature_mean", torch.zeros(2, self.input_dim))
|
| 54 |
+
self.register_buffer("feature_scale", torch.ones(2, self.input_dim))
|
| 55 |
+
grid = torch.linspace(0, 1, self.calibration_points)
|
| 56 |
+
self.register_buffer("calibration_x", grid.expand(2, 3, -1).clone())
|
| 57 |
+
self.register_buffer("calibration_y", grid.expand(2, 3, -1).clone())
|
| 58 |
+
self.register_buffer("calibration_length", torch.full((2, 3), self.calibration_points, dtype=torch.long))
|
| 59 |
+
|
| 60 |
+
@staticmethod
|
| 61 |
+
def validate_features(features: torch.Tensor) -> None:
|
| 62 |
+
if features.ndim != 2 or features.shape[1] != 113:
|
| 63 |
+
raise ValueError(f"features must have shape [N,113], got {tuple(features.shape)}")
|
| 64 |
+
if not torch.isfinite(features).all():
|
| 65 |
+
raise ValueError("features contain NaN or Inf")
|
| 66 |
+
|
| 67 |
+
@staticmethod
|
| 68 |
+
def validate_lead_group(lead_group: torch.Tensor, samples: int) -> None:
|
| 69 |
+
if lead_group.ndim != 1 or len(lead_group) != samples:
|
| 70 |
+
raise ValueError(f"lead_group must have shape [N], got {tuple(lead_group.shape)}")
|
| 71 |
+
if bool(((lead_group < 0) | (lead_group > 1)).any()):
|
| 72 |
+
raise ValueError("lead_group values must be 0 (first hour) or 1 (second hour)")
|
| 73 |
+
|
| 74 |
+
def set_normalization(self, mean: torch.Tensor, scale: torch.Tensor) -> None:
|
| 75 |
+
if mean.shape != (2, 113) or scale.shape != (2, 113):
|
| 76 |
+
raise ValueError("normalization statistics must both have shape [2,113]")
|
| 77 |
+
self.feature_mean.copy_(mean)
|
| 78 |
+
self.feature_scale.copy_(scale.clamp_min(1e-6))
|
| 79 |
+
|
| 80 |
+
def logits(self, features: torch.Tensor, lead_group: torch.Tensor) -> torch.Tensor:
|
| 81 |
+
self.validate_features(features)
|
| 82 |
+
lead_group = lead_group.to(device=features.device, dtype=torch.long)
|
| 83 |
+
self.validate_lead_group(lead_group, len(features))
|
| 84 |
+
normalized = (features - self.feature_mean[lead_group]) / self.feature_scale[lead_group]
|
| 85 |
+
return torch.einsum("ni,noi->no", normalized, self.weight[lead_group]) + self.bias[lead_group]
|
| 86 |
+
|
| 87 |
+
def _calibrate(self, probabilities: torch.Tensor, lead_group: torch.Tensor) -> torch.Tensor:
|
| 88 |
+
result = torch.empty_like(probabilities)
|
| 89 |
+
for group in range(2):
|
| 90 |
+
mask = lead_group == group
|
| 91 |
+
if not bool(mask.any()):
|
| 92 |
+
continue
|
| 93 |
+
for hazard in range(3):
|
| 94 |
+
length = int(self.calibration_length[group, hazard])
|
| 95 |
+
xp = self.calibration_x[group, hazard, :length]
|
| 96 |
+
yp = self.calibration_y[group, hazard, :length]
|
| 97 |
+
value = probabilities[mask, hazard].clamp(xp[0], xp[-1])
|
| 98 |
+
upper = torch.searchsorted(xp.contiguous(), value.contiguous()).clamp(1, length - 1)
|
| 99 |
+
lower = upper - 1
|
| 100 |
+
fraction = (value - xp[lower]) / (xp[upper] - xp[lower]).clamp_min(1e-7)
|
| 101 |
+
result[mask, hazard] = yp[lower] + fraction * (yp[upper] - yp[lower])
|
| 102 |
+
return result.clamp(0, 1)
|
| 103 |
+
|
| 104 |
+
def forward(self, features: torch.Tensor, lead_group: torch.Tensor, calibrated: bool = True) -> torch.Tensor:
|
| 105 |
+
probabilities = torch.sigmoid(self.logits(features, lead_group))
|
| 106 |
+
return self._calibrate(probabilities, lead_group.to(probabilities.device)) if calibrated else probabilities
|
| 107 |
+
|
| 108 |
+
@torch.no_grad()
|
| 109 |
+
def fit_calibration(self, features: torch.Tensor, targets: torch.Tensor, lead_group: torch.Tensor) -> None:
|
| 110 |
+
if targets.shape != (len(features), 3):
|
| 111 |
+
raise ValueError(f"targets must have shape [N,3], got {tuple(targets.shape)}")
|
| 112 |
+
probabilities = torch.sigmoid(self.logits(features, lead_group)).cpu().numpy()
|
| 113 |
+
target_array, groups = targets.cpu().numpy(), lead_group.cpu().numpy()
|
| 114 |
+
for group in range(2):
|
| 115 |
+
for hazard in range(3):
|
| 116 |
+
mask = groups == group
|
| 117 |
+
xp, yp = _pav(probabilities[mask, hazard], target_array[mask, hazard])
|
| 118 |
+
if len(xp) > self.calibration_points:
|
| 119 |
+
selected = np.linspace(0, len(xp) - 1, self.calibration_points).round().astype(int)
|
| 120 |
+
xp, yp = xp[selected], yp[selected]
|
| 121 |
+
if len(xp) == 1:
|
| 122 |
+
xp, yp = np.repeat(xp, 2), np.repeat(yp, 2)
|
| 123 |
+
length = len(xp)
|
| 124 |
+
self.calibration_x[group, hazard, :length] = torch.from_numpy(xp).to(self.calibration_x)
|
| 125 |
+
self.calibration_y[group, hazard, :length] = torch.from_numpy(yp).to(self.calibration_y)
|
| 126 |
+
self.calibration_length[group, hazard] = length
|
| 127 |
+
|
| 128 |
+
def elastic_net_loss(self, logits: torch.Tensor, targets: torch.Tensor, l1: float, l2: float) -> torch.Tensor:
|
| 129 |
+
if logits.shape != targets.shape or logits.ndim != 2 or logits.shape[1] != 3:
|
| 130 |
+
raise ValueError("logits and targets must both have shape [N,3]")
|
| 131 |
+
bce = nn.functional.binary_cross_entropy_with_logits(logits, targets)
|
| 132 |
+
return bce + float(l1) * self.weight.abs().mean() + 0.5 * float(l2) * self.weight.square().mean()
|
scripts/fake_data.py
ADDED
|
@@ -0,0 +1,86 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Generate correlated 113-feature WoFS ensemble storm-track examples."""
|
| 2 |
+
|
| 3 |
+
import argparse
|
| 4 |
+
from pathlib import Path
|
| 5 |
+
|
| 6 |
+
import numpy as np
|
| 7 |
+
import yaml
|
| 8 |
+
|
| 9 |
+
|
| 10 |
+
ROOT = Path(__file__).resolve().parents[1]
|
| 11 |
+
|
| 12 |
+
|
| 13 |
+
def correlated_block(rng, latent, count, noise, offset):
|
| 14 |
+
projection_rng = np.random.default_rng(1000 + int(offset * 100))
|
| 15 |
+
weights = projection_rng.normal(0, 0.35, (latent.shape[1], count))
|
| 16 |
+
for index in range(min(latent.shape[1], count)):
|
| 17 |
+
weights[index, index::latent.shape[1]] += 0.8
|
| 18 |
+
values = latent @ weights + rng.normal(0, noise, (len(latent), count))
|
| 19 |
+
values += 0.08 * np.sin(np.arange(count)[None, :] * 0.21 + offset)
|
| 20 |
+
return values.astype(np.float32)
|
| 21 |
+
|
| 22 |
+
|
| 23 |
+
def make_split(path, samples, config, seed):
|
| 24 |
+
rng = np.random.default_rng(seed)
|
| 25 |
+
lead_start = rng.choice(np.arange(0, 125, 5), samples)
|
| 26 |
+
lead_group = (lead_start > 60).astype(np.int64)
|
| 27 |
+
organization = rng.normal(size=samples)
|
| 28 |
+
rotation = 0.65 * organization + rng.normal(0, 0.75, samples)
|
| 29 |
+
instability = rng.normal(size=samples)
|
| 30 |
+
shear = 0.35 * instability + rng.normal(0, 0.9, samples)
|
| 31 |
+
cold_pool = 0.45 * organization + rng.normal(0, 0.85, samples)
|
| 32 |
+
spread = np.maximum(0.15, rng.lognormal(-0.2 + 0.2 * lead_group, 0.35, samples))
|
| 33 |
+
latent = np.column_stack((organization, rotation, instability, shear, cold_pool, spread))
|
| 34 |
+
amplitude = correlated_block(rng, latent, 30, 0.32, 0.0)
|
| 35 |
+
spatial = correlated_block(rng, latent, 76, 0.48 + 0.08 * lead_group[:, None], 0.7)
|
| 36 |
+
area = np.exp(0.45 * organization + rng.normal(5.8, 0.35, samples))
|
| 37 |
+
eccentricity = 1 / (1 + np.exp(-(0.5 * shear + rng.normal(0, 0.6, samples))))
|
| 38 |
+
major = np.sqrt(area) * (1.2 + eccentricity)
|
| 39 |
+
minor = area / np.maximum(major, 1)
|
| 40 |
+
orientation = np.arctan2(spatial[:, 4], spatial[:, 3]) / np.pi
|
| 41 |
+
extent = np.clip(0.72 - 0.15 * spread + rng.normal(0, 0.08, samples), 0.15, 1)
|
| 42 |
+
initialization_time = rng.uniform(0, 1, samples)
|
| 43 |
+
object_properties = np.column_stack((np.log1p(area), eccentricity, orientation, major / 100,
|
| 44 |
+
minor / 100, extent, initialization_time)).astype(np.float32)
|
| 45 |
+
features = np.concatenate((amplitude, spatial, object_properties), axis=1).astype(np.float32)
|
| 46 |
+
risk = np.column_stack((1.25 * rotation + 0.70 * shear + 0.35 * organization,
|
| 47 |
+
1.15 * instability + 0.55 * organization + 0.45 * amplitude[:, 5],
|
| 48 |
+
1.00 * cold_pool + 0.65 * shear + 0.35 * spatial[:, 12]))
|
| 49 |
+
risk -= lead_group[:, None] * np.array((0.35, 0.22, 0.18))
|
| 50 |
+
# Elevated synthetic rates make the tiny engineering dataset trainable; paper rates were about 1.2/2.5/4%.
|
| 51 |
+
intercept = np.array((-2.35, -1.75, -1.45))
|
| 52 |
+
probabilities = 1 / (1 + np.exp(-(risk + intercept)))
|
| 53 |
+
targets = (rng.random((samples, 3)) < probabilities).astype(np.float32)
|
| 54 |
+
np.savez_compressed(path, features=features, targets=targets, lead_group=lead_group,
|
| 55 |
+
lead_start_minutes=lead_start.astype(np.int16),
|
| 56 |
+
lead_end_minutes=(lead_start + 30).astype(np.int16),
|
| 57 |
+
format_version=np.asarray(config["data"]["format_version"]),
|
| 58 |
+
feature_group_sizes=np.asarray((30, 76, 7), dtype=np.int16),
|
| 59 |
+
hazards=np.asarray(config["data"]["hazards"]),
|
| 60 |
+
lead_group_names=np.asarray(config["data"]["lead_groups"]),
|
| 61 |
+
ensemble_members=np.asarray(18), grid_spacing_km=np.asarray(3),
|
| 62 |
+
forecast_window_minutes=np.asarray(30), forecast_interval_minutes=np.asarray(5),
|
| 63 |
+
data_source=np.asarray("structured_synthetic_elevated_event_rates"),
|
| 64 |
+
paper_event_rates=np.asarray((0.012, 0.025, 0.040), dtype=np.float32))
|
| 65 |
+
return targets.mean(0)
|
| 66 |
+
|
| 67 |
+
|
| 68 |
+
def main():
|
| 69 |
+
parser = argparse.ArgumentParser()
|
| 70 |
+
parser.add_argument("--force", action="store_true")
|
| 71 |
+
args = parser.parse_args()
|
| 72 |
+
config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
|
| 73 |
+
output = ROOT / config["data"]["root"]
|
| 74 |
+
output.mkdir(parents=True, exist_ok=True)
|
| 75 |
+
for offset, (name, count) in enumerate((("train.npz", config["data"]["train_samples"]),
|
| 76 |
+
("test.npz", config["data"]["test_samples"]))):
|
| 77 |
+
path = output / name
|
| 78 |
+
if args.force or not path.exists():
|
| 79 |
+
rates = make_split(path, int(count), config, int(config["seed"]) + offset)
|
| 80 |
+
else:
|
| 81 |
+
rates = np.load(path)["targets"].mean(0)
|
| 82 |
+
print(f"generated={path.relative_to(ROOT)} shape=({count},113) event_rates={rates.round(3).tolist()}")
|
| 83 |
+
|
| 84 |
+
|
| 85 |
+
if __name__ == "__main__":
|
| 86 |
+
main()
|
scripts/inference.py
ADDED
|
@@ -0,0 +1,49 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Restore a checkpoint and infer calibrated tornado, hail, and wind probabilities."""
|
| 2 |
+
|
| 3 |
+
import sys
|
| 4 |
+
from pathlib import Path
|
| 5 |
+
|
| 6 |
+
import numpy as np
|
| 7 |
+
import torch
|
| 8 |
+
import yaml
|
| 9 |
+
from torch.utils.data import DataLoader
|
| 10 |
+
|
| 11 |
+
|
| 12 |
+
ROOT = Path(__file__).resolve().parents[1]
|
| 13 |
+
sys.path.insert(0, str(ROOT))
|
| 14 |
+
from model.wofsstormcal import WoFSStormCal
|
| 15 |
+
from train import HazardDataset, device_from_config
|
| 16 |
+
|
| 17 |
+
|
| 18 |
+
def main():
|
| 19 |
+
config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
|
| 20 |
+
device = device_from_config(config)
|
| 21 |
+
checkpoint = torch.load(ROOT / config["paths"]["checkpoint"], map_location=device, weights_only=False)
|
| 22 |
+
if checkpoint["format_version"] != config["data"]["format_version"]:
|
| 23 |
+
raise ValueError("checkpoint and data format versions differ")
|
| 24 |
+
model = WoFSStormCal(int(config["model"]["calibration_points"])).to(device)
|
| 25 |
+
model.load_state_dict(checkpoint["model"]); model.eval()
|
| 26 |
+
dataset = HazardDataset(ROOT / config["data"]["root"] / "test.npz", config)
|
| 27 |
+
loader = DataLoader(dataset, batch_size=int(config["train"]["batch_size"]), shuffle=False)
|
| 28 |
+
predictions = []
|
| 29 |
+
with torch.no_grad():
|
| 30 |
+
for features, _, lead_group in loader:
|
| 31 |
+
prediction = model(features.to(device), lead_group.to(device))
|
| 32 |
+
if prediction.shape != (len(features), 3):
|
| 33 |
+
raise RuntimeError("model output must have shape [N,3]")
|
| 34 |
+
predictions.append(prediction.cpu().numpy())
|
| 35 |
+
predictions = np.concatenate(predictions)
|
| 36 |
+
if predictions.shape != (len(dataset), 3) or not np.isfinite(predictions).all():
|
| 37 |
+
raise FloatingPointError("inference output is invalid")
|
| 38 |
+
source = dataset.data
|
| 39 |
+
output = ROOT / config["paths"]["inference_dir"] / "predictions.npz"
|
| 40 |
+
output.parent.mkdir(parents=True, exist_ok=True)
|
| 41 |
+
np.savez_compressed(output, probabilities=predictions, targets=source["targets"],
|
| 42 |
+
lead_group=source["lead_group"], lead_start_minutes=source["lead_start_minutes"],
|
| 43 |
+
lead_end_minutes=source["lead_end_minutes"], hazards=source["hazards"],
|
| 44 |
+
lead_group_names=source["lead_group_names"], format_version=source["format_version"])
|
| 45 |
+
print(f"predictions={output.relative_to(ROOT)} shape={predictions.shape} range=({predictions.min():.3f},{predictions.max():.3f})")
|
| 46 |
+
|
| 47 |
+
|
| 48 |
+
if __name__ == "__main__":
|
| 49 |
+
main()
|
scripts/result.py
ADDED
|
@@ -0,0 +1,116 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Compute paper verification metrics and plot performance and reliability diagrams."""
|
| 2 |
+
|
| 3 |
+
import json
|
| 4 |
+
from pathlib import Path
|
| 5 |
+
|
| 6 |
+
import matplotlib
|
| 7 |
+
matplotlib.use("Agg")
|
| 8 |
+
import matplotlib.pyplot as plt
|
| 9 |
+
import numpy as np
|
| 10 |
+
import yaml
|
| 11 |
+
from sklearn.metrics import average_precision_score, roc_auc_score
|
| 12 |
+
|
| 13 |
+
|
| 14 |
+
ROOT = Path(__file__).resolve().parents[1]
|
| 15 |
+
|
| 16 |
+
|
| 17 |
+
def curves(probability, target, threshold_count):
|
| 18 |
+
thresholds = np.linspace(0, 1, threshold_count)
|
| 19 |
+
pod, pofd, sr, far, csi, bias = [], [], [], [], [], []
|
| 20 |
+
for threshold in thresholds:
|
| 21 |
+
forecast = probability >= threshold
|
| 22 |
+
event = target == 1
|
| 23 |
+
hits = np.sum(forecast & event); false_alarms = np.sum(forecast & ~event)
|
| 24 |
+
misses = np.sum(~forecast & event); negatives = np.sum(~forecast & ~event)
|
| 25 |
+
pod.append(hits / max(hits + misses, 1)); pofd.append(false_alarms / max(false_alarms + negatives, 1))
|
| 26 |
+
sr.append(hits / max(hits + false_alarms, 1)); far.append(false_alarms / max(hits + false_alarms, 1))
|
| 27 |
+
csi.append(hits / max(hits + false_alarms + misses, 1)); bias.append((hits + false_alarms) / max(hits + misses, 1))
|
| 28 |
+
return {key: np.asarray(value) for key, value in (("threshold", thresholds), ("pod", pod), ("pofd", pofd),
|
| 29 |
+
("sr", sr), ("far", far), ("csi", csi), ("bias", bias))}
|
| 30 |
+
|
| 31 |
+
|
| 32 |
+
def reliability(probability, target, bins):
|
| 33 |
+
edges = np.linspace(0, 1, bins + 1)
|
| 34 |
+
index = np.minimum(np.digitize(probability, edges[1:-1]), bins - 1)
|
| 35 |
+
records, component = [], 0.0
|
| 36 |
+
for bin_index in range(bins):
|
| 37 |
+
mask = index == bin_index
|
| 38 |
+
if not mask.any():
|
| 39 |
+
continue
|
| 40 |
+
forecast_mean, observed_frequency = float(probability[mask].mean()), float(target[mask].mean())
|
| 41 |
+
component += mask.mean() * (forecast_mean - observed_frequency) ** 2
|
| 42 |
+
records.append({"count": int(mask.sum()), "forecast_probability": forecast_mean,
|
| 43 |
+
"observed_frequency": observed_frequency})
|
| 44 |
+
return records, float(component)
|
| 45 |
+
|
| 46 |
+
|
| 47 |
+
def metrics(probability, target, threshold_count, bins):
|
| 48 |
+
curve = curves(probability, target, threshold_count)
|
| 49 |
+
prevalence = float(target.mean())
|
| 50 |
+
auc = float(roc_auc_score(target, probability)) if 0 < target.sum() < len(target) else 0.5
|
| 51 |
+
aupdc = float(average_precision_score(target, probability)) if target.sum() else 0.0
|
| 52 |
+
positives, negatives = int(target.sum()), int(len(target) - target.sum())
|
| 53 |
+
minimum_aupdc = float(np.mean(np.arange(1, positives + 1) / (np.arange(1, positives + 1) + negatives))) if positives else 0.0
|
| 54 |
+
best = int(np.argmax(curve["csi"])); max_csi = float(curve["csi"][best])
|
| 55 |
+
brier = float(np.mean((probability - target) ** 2))
|
| 56 |
+
reference = prevalence * (1 - prevalence)
|
| 57 |
+
records, reliability_component = reliability(probability, target, bins)
|
| 58 |
+
return {"samples": len(target), "event_rate": prevalence, "auc": auc, "aupdc": aupdc,
|
| 59 |
+
"minimum_aupdc": minimum_aupdc, "naupdc": (aupdc - minimum_aupdc) / max(1 - minimum_aupdc, 1e-12),
|
| 60 |
+
"max_csi": max_csi, "ncsi": (max_csi - prevalence) / max(1 - prevalence, 1e-12),
|
| 61 |
+
"max_csi_threshold": float(curve["threshold"][best]), "pod": float(curve["pod"][best]),
|
| 62 |
+
"pofd": float(curve["pofd"][best]), "sr": float(curve["sr"][best]), "far": float(curve["far"][best]),
|
| 63 |
+
"csi": max_csi, "bias": float(curve["bias"][best]), "brier_score": brier,
|
| 64 |
+
"brier_skill_score": 1 - brier / reference if reference > 0 else 0.0,
|
| 65 |
+
"reliability_component": reliability_component, "reliability_bins": records}, curve
|
| 66 |
+
|
| 67 |
+
|
| 68 |
+
def main():
|
| 69 |
+
config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
|
| 70 |
+
data = np.load(ROOT / config["paths"]["inference_dir"] / "predictions.npz")
|
| 71 |
+
if str(data["format_version"]) != config["data"]["format_version"]:
|
| 72 |
+
raise ValueError("incompatible prediction format")
|
| 73 |
+
probabilities, targets, groups = data["probabilities"], data["targets"], data["lead_group"]
|
| 74 |
+
if probabilities.shape != targets.shape or probabilities.shape[1:] != (3,):
|
| 75 |
+
raise ValueError("probabilities and targets must have shape [N,3]")
|
| 76 |
+
hazards, group_names = data["hazards"].tolist(), data["lead_group_names"].tolist()
|
| 77 |
+
report, all_curves = {}, {}
|
| 78 |
+
for group_index, group_name in enumerate(group_names):
|
| 79 |
+
report[group_name] = {}
|
| 80 |
+
mask = groups == group_index
|
| 81 |
+
for hazard_index, hazard in enumerate(hazards):
|
| 82 |
+
result, curve = metrics(probabilities[mask, hazard_index], targets[mask, hazard_index],
|
| 83 |
+
int(config["evaluation"]["threshold_count"]), int(config["evaluation"]["probability_bins"]))
|
| 84 |
+
report[group_name][hazard] = result; all_curves[(group_name, hazard)] = curve
|
| 85 |
+
numeric = [value for group in report.values() for hazard in group.values() for value in hazard.values()
|
| 86 |
+
if isinstance(value, (int, float))]
|
| 87 |
+
if not np.isfinite(numeric).all():
|
| 88 |
+
raise FloatingPointError("evaluation contains NaN or Inf")
|
| 89 |
+
output = ROOT / config["paths"]["evaluation_dir"]
|
| 90 |
+
output.mkdir(parents=True, exist_ok=True)
|
| 91 |
+
(output / "metrics.json").write_text(json.dumps({"by_lead_group_and_hazard": report,
|
| 92 |
+
"metric_protocol": "Flora et al. AUC, weighted-average-precision AUPDC, minimum-AUPDC NAUPDC, climatology-normalized max CSI, and Brier verification"}, indent=2) + "\n")
|
| 93 |
+
colors = {"tornado": "#b3261e", "hail": "#2e7d32", "wind": "#1565c0"}
|
| 94 |
+
figure, axes = plt.subplots(1, 2, figsize=(10, 4.5), sharex=True, sharey=True)
|
| 95 |
+
for axis, group_name in zip(axes, group_names):
|
| 96 |
+
for hazard in hazards:
|
| 97 |
+
curve = all_curves[(group_name, hazard)]
|
| 98 |
+
axis.plot(curve["sr"], curve["pod"], color=colors[hazard], label=hazard)
|
| 99 |
+
axis.set(xlim=(0, 1), ylim=(0, 1), xlabel="Success ratio", ylabel="Probability of detection", title=group_name.replace("_", " ").title())
|
| 100 |
+
axis.grid(alpha=0.25); axis.legend()
|
| 101 |
+
figure.tight_layout(); figure.savefig(output / "performance.png", dpi=160); plt.close(figure)
|
| 102 |
+
figure, axes = plt.subplots(1, 2, figsize=(10, 4.5), sharex=True, sharey=True)
|
| 103 |
+
for axis, group_name in zip(axes, group_names):
|
| 104 |
+
axis.plot((0, 1), (0, 1), "k--", linewidth=1, label="perfect")
|
| 105 |
+
for hazard in hazards:
|
| 106 |
+
bins = report[group_name][hazard]["reliability_bins"]
|
| 107 |
+
axis.plot([item["forecast_probability"] for item in bins], [item["observed_frequency"] for item in bins],
|
| 108 |
+
marker="o", color=colors[hazard], label=hazard)
|
| 109 |
+
axis.set(xlim=(0, 1), ylim=(0, 1), xlabel="Forecast probability", ylabel="Observed frequency", title=group_name.replace("_", " ").title())
|
| 110 |
+
axis.grid(alpha=0.25); axis.legend()
|
| 111 |
+
figure.tight_layout(); figure.savefig(output / "reliability.png", dpi=160); plt.close(figure)
|
| 112 |
+
print(f"evaluation={output.relative_to(ROOT)} groups={len(group_names)} hazards={len(hazards)}")
|
| 113 |
+
|
| 114 |
+
|
| 115 |
+
if __name__ == "__main__":
|
| 116 |
+
main()
|
scripts/train.py
ADDED
|
@@ -0,0 +1,129 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Train, calibrate, and checkpoint WoFS elastic-net logistic models."""
|
| 2 |
+
|
| 3 |
+
import argparse
|
| 4 |
+
import json
|
| 5 |
+
import os
|
| 6 |
+
import sys
|
| 7 |
+
from pathlib import Path
|
| 8 |
+
|
| 9 |
+
import numpy as np
|
| 10 |
+
import torch
|
| 11 |
+
import yaml
|
| 12 |
+
from torch.nn.parallel import DistributedDataParallel
|
| 13 |
+
from torch.utils.data import DataLoader, Dataset, DistributedSampler, Subset
|
| 14 |
+
|
| 15 |
+
|
| 16 |
+
ROOT = Path(__file__).resolve().parents[1]
|
| 17 |
+
sys.path.insert(0, str(ROOT))
|
| 18 |
+
from model.wofsstormcal import WoFSStormCal
|
| 19 |
+
|
| 20 |
+
|
| 21 |
+
class HazardDataset(Dataset):
|
| 22 |
+
def __init__(self, path, config):
|
| 23 |
+
self.data = np.load(path)
|
| 24 |
+
if str(self.data["format_version"]) != config["data"]["format_version"]:
|
| 25 |
+
raise ValueError("incompatible WoFS storm-object data format")
|
| 26 |
+
count = len(self.data["features"])
|
| 27 |
+
if self.data["features"].shape != (count, 113):
|
| 28 |
+
raise ValueError("features must have shape [N,113]")
|
| 29 |
+
if self.data["targets"].shape != (count, 3):
|
| 30 |
+
raise ValueError("targets must have shape [N,3]")
|
| 31 |
+
if self.data["lead_group"].shape != (count,):
|
| 32 |
+
raise ValueError("lead_group must have shape [N]")
|
| 33 |
+
if not np.isfinite(self.data["features"]).all() or not np.isfinite(self.data["targets"]).all():
|
| 34 |
+
raise ValueError("data contain NaN or Inf")
|
| 35 |
+
|
| 36 |
+
def __len__(self):
|
| 37 |
+
return len(self.data["features"])
|
| 38 |
+
|
| 39 |
+
def __getitem__(self, index):
|
| 40 |
+
return (torch.from_numpy(self.data["features"][index]).float(),
|
| 41 |
+
torch.from_numpy(self.data["targets"][index]).float(),
|
| 42 |
+
torch.as_tensor(self.data["lead_group"][index], dtype=torch.long))
|
| 43 |
+
|
| 44 |
+
|
| 45 |
+
def device_from_config(config, rank=0):
|
| 46 |
+
requested = config["runtime"]["device"]
|
| 47 |
+
if requested == "auto":
|
| 48 |
+
return torch.device("cuda", rank) if torch.cuda.is_available() else torch.device("cpu")
|
| 49 |
+
return torch.device(requested)
|
| 50 |
+
|
| 51 |
+
|
| 52 |
+
def main():
|
| 53 |
+
parser = argparse.ArgumentParser()
|
| 54 |
+
parser.add_argument("--resume", action="store_true", help="restore model and optimizer state before training")
|
| 55 |
+
args = parser.parse_args()
|
| 56 |
+
config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
|
| 57 |
+
seed = int(config["seed"])
|
| 58 |
+
np.random.seed(seed); torch.manual_seed(seed)
|
| 59 |
+
distributed = int(os.environ.get("WORLD_SIZE", "1")) > 1
|
| 60 |
+
local_rank = int(os.environ.get("LOCAL_RANK", "0"))
|
| 61 |
+
if distributed:
|
| 62 |
+
torch.distributed.init_process_group("nccl" if torch.cuda.is_available() else "gloo")
|
| 63 |
+
rank = torch.distributed.get_rank() if distributed else 0
|
| 64 |
+
device = device_from_config(config, local_rank)
|
| 65 |
+
if device.type == "cuda":
|
| 66 |
+
torch.cuda.set_device(device)
|
| 67 |
+
dataset = HazardDataset(ROOT / config["data"]["root"] / "train.npz", config)
|
| 68 |
+
split = int(len(dataset) * (1 - float(config["train"]["calibration_fraction"])))
|
| 69 |
+
fit_set = Subset(dataset, range(split))
|
| 70 |
+
sampler = DistributedSampler(fit_set, shuffle=True) if distributed else None
|
| 71 |
+
loader = DataLoader(fit_set, batch_size=int(config["train"]["batch_size"]), sampler=sampler,
|
| 72 |
+
shuffle=sampler is None, num_workers=int(config["train"]["num_workers"]))
|
| 73 |
+
model = WoFSStormCal(int(config["model"]["calibration_points"])).to(device)
|
| 74 |
+
features = dataset.data["features"][:split]
|
| 75 |
+
groups = dataset.data["lead_group"][:split]
|
| 76 |
+
means, scales = [], []
|
| 77 |
+
for group in range(2):
|
| 78 |
+
group_features = features[groups == group]
|
| 79 |
+
means.append(group_features.mean(0)); scales.append(group_features.std(0).clip(1e-6))
|
| 80 |
+
model.set_normalization(torch.from_numpy(np.stack(means)).to(device), torch.from_numpy(np.stack(scales)).to(device))
|
| 81 |
+
wrapped = DistributedDataParallel(model, device_ids=[local_rank] if device.type == "cuda" else None) if distributed else model
|
| 82 |
+
optimizer = torch.optim.Adam(wrapped.parameters(), lr=float(config["train"]["learning_rate"]))
|
| 83 |
+
checkpoint_path = ROOT / config["paths"]["checkpoint"]
|
| 84 |
+
start_epoch, history = 0, []
|
| 85 |
+
if args.resume:
|
| 86 |
+
checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
|
| 87 |
+
model.load_state_dict(checkpoint["model"])
|
| 88 |
+
optimizer.load_state_dict(checkpoint["optimizer"])
|
| 89 |
+
start_epoch = int(checkpoint["epoch"])
|
| 90 |
+
history = checkpoint.get("history", [])
|
| 91 |
+
for epoch in range(start_epoch, start_epoch + int(config["train"]["epochs"])):
|
| 92 |
+
if sampler is not None:
|
| 93 |
+
sampler.set_epoch(epoch)
|
| 94 |
+
total, steps = 0.0, 0
|
| 95 |
+
for batch_features, targets, lead_group in loader:
|
| 96 |
+
batch_features, targets, lead_group = batch_features.to(device), targets.to(device), lead_group.to(device)
|
| 97 |
+
active_model = wrapped.module if distributed else wrapped
|
| 98 |
+
logits = wrapped(batch_features, lead_group, False)
|
| 99 |
+
logits = torch.logit(logits.clamp(1e-6, 1 - 1e-6))
|
| 100 |
+
loss = active_model.elastic_net_loss(logits, targets, config["train"]["l1_strength"], config["train"]["l2_strength"])
|
| 101 |
+
optimizer.zero_grad(set_to_none=True); loss.backward()
|
| 102 |
+
torch.nn.utils.clip_grad_norm_(wrapped.parameters(), float(config["train"]["gradient_clip_norm"]))
|
| 103 |
+
optimizer.step(); total += float(loss.detach()); steps += 1
|
| 104 |
+
if rank == 0:
|
| 105 |
+
history.append({"epoch": epoch + 1, "elastic_net_loss": total / max(steps, 1)})
|
| 106 |
+
model = wrapped.module if distributed else wrapped
|
| 107 |
+
if rank == 0:
|
| 108 |
+
calibration_features = torch.from_numpy(dataset.data["features"][split:]).float().to(device)
|
| 109 |
+
calibration_targets = torch.from_numpy(dataset.data["targets"][split:]).float().to(device)
|
| 110 |
+
calibration_groups = torch.from_numpy(dataset.data["lead_group"][split:]).long().to(device)
|
| 111 |
+
model.fit_calibration(calibration_features, calibration_targets, calibration_groups)
|
| 112 |
+
checkpoint_path.parent.mkdir(parents=True, exist_ok=True)
|
| 113 |
+
payload = {"model": model.state_dict(), "optimizer": optimizer.state_dict(), "epoch": start_epoch + int(config["train"]["epochs"]),
|
| 114 |
+
"history": history, "format_version": config["data"]["format_version"],
|
| 115 |
+
"model_metadata": {"input_shape": ["N", 113], "output_shape": ["N", 3],
|
| 116 |
+
"hazards": model.hazards, "lead_groups": model.lead_groups,
|
| 117 |
+
"ensemble_members": 18, "grid_spacing_km": 3,
|
| 118 |
+
"forecast_window_minutes": 30, "forecast_interval_minutes": 5}}
|
| 119 |
+
torch.save(payload, checkpoint_path)
|
| 120 |
+
metrics = ROOT / config["paths"]["training_metrics"]
|
| 121 |
+
metrics.parent.mkdir(parents=True, exist_ok=True)
|
| 122 |
+
metrics.write_text(json.dumps({"history": history, "calibration_samples": len(dataset) - split}, indent=2) + "\n")
|
| 123 |
+
print(f"checkpoint={checkpoint_path.relative_to(ROOT)} output_shape=(N,3) epoch={payload['epoch']}")
|
| 124 |
+
if distributed:
|
| 125 |
+
torch.distributed.destroy_process_group()
|
| 126 |
+
|
| 127 |
+
|
| 128 |
+
if __name__ == "__main__":
|
| 129 |
+
main()
|
weight/.gitkeep
ADDED
|
File without changes
|