zhangrenchao commited on
Commit
fa2b79f
·
verified ·
1 Parent(s): 9e0f589

Publish WoFS-StormCal engineering reproduction

Browse files
.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
- *.zst filter=lfs diff=lfs merge=lfs -text
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