Open released TrackFormer 1.2 training source, exact corpus and optimizer checkpoints
Browse files- .gitattributes +7 -0
- README.md +23 -26
- docs/repository_layout.md +4 -0
- docs/trackformer_1_2_training_data.md +3 -1
- training/trackformer_1_2/LICENSE +21 -0
- training/trackformer_1_2/README.md +63 -0
- training/trackformer_1_2/SOURCE_MANIFEST.json +178 -0
- training/trackformer_1_2/THIRD_PARTY_NOTICES.md +9 -0
- training/trackformer_1_2/assets.json +60 -0
- training/trackformer_1_2/checkpoint-history.json +26 -0
- training/trackformer_1_2/checkpoints/v168_epoch_003.pt +3 -0
- training/trackformer_1_2/checkpoints/v173_epoch_004.pt +3 -0
- training/trackformer_1_2/config.json +29 -0
- training/trackformer_1_2/data/basin_all_int8.npz +3 -0
- training/trackformer_1_2/data/basin_slp_atlas_float16.npy +3 -0
- training/trackformer_1_2/data/geography.npz +3 -0
- training/trackformer_1_2/data/manifest.json +1 -0
- training/trackformer_1_2/data/plan.json +0 -0
- training/trackformer_1_2/data/pressure_hpa.npy +3 -0
- training/trackformer_1_2/data/track_windows_v13.npz +3 -0
- training/trackformer_1_2/download_assets.py +77 -0
- training/trackformer_1_2/requirements.txt +3 -0
- training/trackformer_1_2/source/data_preparation/requirements.txt +7 -0
- training/trackformer_1_2/source/data_preparation/v164_dataset_contract.md +41 -0
- training/trackformer_1_2/source/data_preparation/v164_download_pressure_geography.py +174 -0
- training/trackformer_1_2/source/data_preparation/v164_global_regional_train.py +845 -0
- training/trackformer_1_2/source/data_preparation/v164_prepare_data.py +348 -0
- training/trackformer_1_2/source/data_preparation/v164_prepare_mac.py +36 -0
- training/trackformer_1_2/source/data_preparation/v164_reuse_impl.py +283 -0
- training/trackformer_1_2/source/data_preparation/v164_structure_research.md +142 -0
- training/trackformer_1_2/source/v173/DESIGN.md +52 -0
- training/trackformer_1_2/source/v173/baseline_model.py +158 -0
- training/trackformer_1_2/source/v173/model.py +74 -0
- training/trackformer_1_2/source/v173/test_model.py +90 -0
- training/trackformer_1_2/source/v173/v165_base.py +657 -0
- training/trackformer_1_2/source/v173/v173_pressure_train.py +176 -0
- training/trackformer_1_2/train.py +33 -0
.gitattributes
CHANGED
|
@@ -209,3 +209,10 @@ docs/report_figures/mangkhut_observed_r34.png filter=lfs diff=lfs merge=lfs -tex
|
|
| 209 |
docs/report_figures/jebi_pressure_comparison.png filter=lfs diff=lfs merge=lfs -text
|
| 210 |
docs/report_figures/jebi_track_pressure.png filter=lfs diff=lfs merge=lfs -text
|
| 211 |
docs/report_figures/jebi_observed_r34.png filter=lfs diff=lfs merge=lfs -text
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 209 |
docs/report_figures/jebi_pressure_comparison.png filter=lfs diff=lfs merge=lfs -text
|
| 210 |
docs/report_figures/jebi_track_pressure.png filter=lfs diff=lfs merge=lfs -text
|
| 211 |
docs/report_figures/jebi_observed_r34.png filter=lfs diff=lfs merge=lfs -text
|
| 212 |
+
training/trackformer_1_2/data/track_windows_v13.npz filter=lfs diff=lfs merge=lfs -text
|
| 213 |
+
training/trackformer_1_2/data/basin_all_int8.npz filter=lfs diff=lfs merge=lfs -text
|
| 214 |
+
training/trackformer_1_2/data/basin_slp_atlas_float16.npy filter=lfs diff=lfs merge=lfs -text
|
| 215 |
+
training/trackformer_1_2/data/pressure_hpa.npy filter=lfs diff=lfs merge=lfs -text
|
| 216 |
+
training/trackformer_1_2/data/geography.npz filter=lfs diff=lfs merge=lfs -text
|
| 217 |
+
training/trackformer_1_2/checkpoints/v168_epoch_003.pt filter=lfs diff=lfs merge=lfs -text
|
| 218 |
+
training/trackformer_1_2/checkpoints/v173_epoch_004.pt filter=lfs diff=lfs merge=lfs -text
|
README.md
CHANGED
|
@@ -1,19 +1,10 @@
|
|
| 1 |
-
---
|
| 2 |
-
license: mit
|
| 3 |
-
tags:
|
| 4 |
-
- tropical-cyclone
|
| 5 |
-
- weather-forecasting
|
| 6 |
-
- sea-level-pressure
|
| 7 |
-
- research
|
| 8 |
-
---
|
| 9 |
-
|
| 10 |
# TrackFormer 1.2 — pressure maps and five-day typhoon forecasts
|
| 11 |
|
| 12 |
An AI weather-research model designed to be trained and run on an ordinary home computer.
|
| 13 |
|
| 14 |
TrackFormer 1.2 predicts **Western Pacific pressure maps**, then obtains the typhoon's route and central pressure from the predicted pressure field. It updates the surrounding weather and a moving storm core together, every six hours from **+6 to +120 hours**. Horizontal winds at three pressure levels are updated within the model, providing a basis for studying northeasterly and southwesterly flow; this report's scored results concern track and pressure, not demonstrated monsoon-wind skill.
|
| 15 |
|
| 16 |
-
[Get the model](https://github.com/yu314-coder/typhoon-predict/releases/tag/trackformer-1.2) · [Hugging Face](https://huggingface.co/euler314/typhoon-predict) · [中文研究說明(與報告對照)](
|
| 17 |
|
| 18 |
This page follows the Chinese research report **「訓練AI模型探討西北太平洋氣壓系統影響臺灣附近颱風的預測評估」** and its approved figure/text updates. It describes the released **1.2** model; later rainfall research is not a capability of these weights. The existing technical paper and older evaluation records remain separate, versioned references.
|
| 19 |
|
|
@@ -30,7 +21,7 @@ This is a **research release**, not an operational warning service. Follow offic
|
|
| 30 |
|
| 31 |
## See the forecast: Mangkhut (2018)
|
| 32 |
|
| 33 |
-
|
| 34 |
|
| 35 |
[Watch / download MP4](https://huggingface.co/euler314/typhoon-predict/resolve/87a6e366b42bb4cc0edc95d2c50c55fca21a2c93/docs/trackformer_1_2_mangkhut.mp4) · **11 September 2018, 00 UTC** · **50-member mean** · **20 seconds**
|
| 36 |
|
|
@@ -40,15 +31,15 @@ The animation shows twenty actual six-hour forecasts through +120 h. The Western
|
|
| 40 |
|
| 41 |
This selected example illustrates the model, not typical skill. Observations are verification only. The 50 members are distinct seeded perturbations of **past and issue-time inputs**, not separately trained networks. Member fields are registered geographically before averaging; the mean centre need not locate the minimum of the mean map. Display interpolation does not add resolution to the approximately 20 km core.
|
| 42 |
|
| 43 |
-
[Pressure arrays](
|
| 44 |
|
| 45 |
## Results in the research report
|
| 46 |
|
| 47 |
The completed comparison uses **1,473 daily starts / 270 storms**, with all twenty +6 to +120 h leads. The pressure comparison uses a **subset of those same 270 storms**, not a different 270-storm dataset.
|
| 48 |
|
| 49 |
-
**DeepMind checkpoint: `WeatherNextCyclones_Mini_<2024` · trained through 2023 · native 1° · one member · WeatherNext software v0.3.0.** This is the official **Mini, not the full-sized WeatherNext model**. All reference forecasts completed on an **NVIDIA GeForce RTX 3070 using CUDA**. [Model identity and protocol](
|
| 50 |
|
| 51 |
-
, mean error is **258.7 km** for Mini versus **469.4 km** for 1.2, on 61 starts / 18 storms. Calendar 2024 is excluded from that subset. The combined cohort contains 230 storms from 1980–1999 and 40 recent storms; historical dates overlap Mini's training years. Different inputs and member counts make this a development comparison, **not a controlled architecture ablation, equal-compute test or certified untouched holdout**.
|
| 67 |
|
| 68 |
-
[Full saved scores and common support](
|
| 69 |
|
| 70 |
## Three report examples
|
| 71 |
|
|
@@ -73,13 +64,13 @@ The report compares predicted and reference pressure at **+48 h**, and routes an
|
|
| 73 |
|
| 74 |
| Example | Forecast issue (UTC) | Report figures |
|
| 75 |
| --- | --- | --- |
|
| 76 |
-
| Fung-wong / 鳳凰 (2025) | 2025-11-07 00:00 | [Pressure](
|
| 77 |
-
| Mangkhut / 山竹 (2018) | 2018-09-11 00:00 | [Pressure](
|
| 78 |
-
| Jebi / 燕子 (2018) | 2018-08-31 00:00 | [Pressure](
|
| 79 |
|
| 80 |
All three use the same released checkpoint and 50-input-member policy. Fung-wong and Mangkhut use saved Metal forecasts; Jebi uses CPU forecasts from the same weights. Mangkhut and Jebi lie within the fitting years and are retrospective demonstrations, not unseen-case accuracy evidence. **The R34 panels are IBTrACS observed 34-knot quadrant-radius records, not model radius predictions or full wind maps.** Missing radii are left missing. There is **no native wind-radius forecast head** in 1.2.
|
| 81 |
|
| 82 |
-
[All report panels with captions](
|
| 83 |
|
| 84 |
## Training data and year cutoffs
|
| 85 |
|
|
@@ -104,11 +95,11 @@ The eastward/northward wind pairs **u/v carry both direction and speed** at 850,
|
|
| 104 |
|
| 105 |
The checksum-pinned source caches extend further: the derived typhoon archive reaches **13 July 2026**, and the NCEP basin archive reaches **17 March 2026, 18 UTC**. Those endpoints describe archive availability, **not the fitting cutoff**. Later live weather uses separately documented experimental GFS transfer; it does not update the released weights. This release has no satellite cloud-image or humidity input channel. Previously inspected test/showcase storms remain development evidence, not a newly untouched holdout.
|
| 106 |
|
| 107 |
-
[Training-source details and exact date ranges](
|
| 108 |
|
| 109 |
## How Trackformer 1.2 works
|
| 110 |
|
| 111 |
-
, or use this repository with the [Hugging Face weights](https://huggingface.co/euler314/typhoon-predict/resolve/main/models/trackformer_1_2_field/weights.pt). Put `weights.pt` in `models/trackformer_1_2_field/`. The inference weights and learned model equations are unchanged; the package includes the moving-core output and map renderer.
|
| 143 |
|
| 144 |
-
Prepare a normalized causal issue packet following the [input schema](
|
| 145 |
|
| 146 |
```bash
|
| 147 |
python models/trackformer_1_2_field/predict.py causal_issue_packet.npz forecast.npz --device mps
|
|
@@ -155,7 +152,7 @@ For an immediate whole-WP and detailed-core PNG, add `--pressure-map pressure_12
|
|
| 155 |
python models/trackformer_1_2_field/plot_pressure.py forecast.npz pressure_24h.png --lead 24 --interval 2
|
| 156 |
```
|
| 157 |
|
| 158 |
-
Blue denotes low pressure and red high pressure. The renderer uses the actual saved moving field, coordinates and masks; it never draws a replacement vortex around a scalar pressure or observed track. See the [pressure-output schema](
|
| 159 |
|
| 160 |
The inference weight SHA-256 is `db49f36e85a3766defc4c172746897a1f783705d1ce8e6f9dfb8e87ae1d902cb`.
|
| 161 |
|
|
@@ -171,10 +168,10 @@ The archive covers a pinned Western Pacific plan: six-hour starts from 1996 onwa
|
|
| 171 |
|
| 172 |
The auxiliary wind's averaging period and surface-wind skill are not validated. Optional pressure-derived wind/radius estimates are experimental diagnostics, not agency-equivalent forecasts. Pressure-map interpolation is not evidence of finer effective resolution. Retrospective analyses and experimental GFS live transfer must not be presented as operational validation.
|
| 173 |
|
| 174 |
-
See [wind-estimation assumptions](
|
| 175 |
|
| 176 |
## Ongoing work: Taiwan rainfall
|
| 177 |
|
| 178 |
Ongoing research combines moisture, wind direction, weather outside Taiwan and terrain to study ordinary and heavy rainfall over Taiwan, with a separate observation-based evaluation. **Released 1.2 has no rainfall forecast head or humidity input.** These are research directions, not demonstrated rainfall improvements from the pressure model.
|
| 179 |
|
| 180 |
-
Model code is in [`models/`](
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
# TrackFormer 1.2 — pressure maps and five-day typhoon forecasts
|
| 2 |
|
| 3 |
An AI weather-research model designed to be trained and run on an ordinary home computer.
|
| 4 |
|
| 5 |
TrackFormer 1.2 predicts **Western Pacific pressure maps**, then obtains the typhoon's route and central pressure from the predicted pressure field. It updates the surrounding weather and a moving storm core together, every six hours from **+6 to +120 hours**. Horizontal winds at three pressure levels are updated within the model, providing a basis for studying northeasterly and southwesterly flow; this report's scored results concern track and pressure, not demonstrated monsoon-wind skill.
|
| 6 |
|
| 7 |
+
[Get the model](https://github.com/yu314-coder/typhoon-predict/releases/tag/trackformer-1.2) · [Hugging Face](https://huggingface.co/euler314/typhoon-predict) · [中文研究說明(與報告對照)](docs/research_report_zh.md) · [Live forecast](https://trackformer-weatherlab.rudin-euler-8253.chatgpt.site/live) · [Technical paper](paper/trackformer.pdf)
|
| 8 |
|
| 9 |
This page follows the Chinese research report **「訓練AI模型探討西北太平洋氣壓系統影響臺灣附近颱風的預測評估」** and its approved figure/text updates. It describes the released **1.2** model; later rainfall research is not a capability of these weights. The existing technical paper and older evaluation records remain separate, versioned references.
|
| 10 |
|
|
|
|
| 21 |
|
| 22 |
## See the forecast: Mangkhut (2018)
|
| 23 |
|
| 24 |
+
[](https://huggingface.co/euler314/typhoon-predict/resolve/87a6e366b42bb4cc0edc95d2c50c55fca21a2c93/docs/trackformer_1_2_mangkhut.mp4)
|
| 25 |
|
| 26 |
[Watch / download MP4](https://huggingface.co/euler314/typhoon-predict/resolve/87a6e366b42bb4cc0edc95d2c50c55fca21a2c93/docs/trackformer_1_2_mangkhut.mp4) · **11 September 2018, 00 UTC** · **50-member mean** · **20 seconds**
|
| 27 |
|
|
|
|
| 31 |
|
| 32 |
This selected example illustrates the model, not typical skill. Observations are verification only. The 50 members are distinct seeded perturbations of **past and issue-time inputs**, not separately trained networks. Member fields are registered geographically before averaging; the mean centre need not locate the minimum of the mean map. Display interpolation does not add resolution to the approximately 20 km core.
|
| 33 |
|
| 34 |
+
[Pressure arrays](evaluation/release_data/pressure_core/mangkhut/common-pressure.npz) · [Field verification](evaluation/release_data/pressure_core/mangkhut/verification.json) · [Playback verification](evaluation/release_data/mangkhut_playback_verification.json) · [More examples and diagnostics](docs/showcase_archive.md)
|
| 35 |
|
| 36 |
## Results in the research report
|
| 37 |
|
| 38 |
The completed comparison uses **1,473 daily starts / 270 storms**, with all twenty +6 to +120 h leads. The pressure comparison uses a **subset of those same 270 storms**, not a different 270-storm dataset.
|
| 39 |
|
| 40 |
+
**DeepMind checkpoint: `WeatherNextCyclones_Mini_<2024` · trained through 2023 · native 1° · one member · WeatherNext software v0.3.0.** This is the official **Mini, not the full-sized WeatherNext model**. All reference forecasts completed on an **NVIDIA GeForce RTX 3070 using CUDA**. [Model identity and protocol](docs/deepmind_daily_benchmark.md) · [Completion receipt](evaluation/deepmind_daily/verification.json).
|
| 41 |
|
| 42 |
+

|
| 43 |
|
| 44 |
The bars show exact **24, 48, 72, 96 and 120 h** leads, not averages over each day. Track uses 270 storms; pressure uses 39–40 storms per displayed lead, with the counts printed in the figure. The table below averages all available six-hour leads using the common scoring procedure.
|
| 45 |
|
|
|
|
| 56 |
|
| 57 |
On this mixed historical/recent development cohort, 1.2 has lower mean position and heading errors; Mini has lower central-pressure MAE. **The recent-only track results favour Mini**: for strict issue years after 2024 (2025–2026), mean error is **258.7 km** for Mini versus **469.4 km** for 1.2, on 61 starts / 18 storms. Calendar 2024 is excluded from that subset. The combined cohort contains 230 storms from 1980–1999 and 40 recent storms; historical dates overlap Mini's training years. Different inputs and member counts make this a development comparison, **not a controlled architecture ablation, equal-compute test or certified untouched holdout**.
|
| 58 |
|
| 59 |
+
[Full saved scores and common support](evaluation/deepmind_daily/benchmark.json) · [Period breakdowns](evaluation/deepmind_daily/period_comparison.json) · [Independent publication audit](evaluation/deepmind_daily/publication_audit.json) · [Daily scoring protocol](docs/daily_storm_benchmark.md) · [Evaluation details and older comparisons](docs/trackformer_1_2_evaluation.md)
|
| 60 |
|
| 61 |
## Three report examples
|
| 62 |
|
|
|
|
| 64 |
|
| 65 |
| Example | Forecast issue (UTC) | Report figures |
|
| 66 |
| --- | --- | --- |
|
| 67 |
+
| Fung-wong / 鳳凰 (2025) | 2025-11-07 00:00 | [Pressure](docs/report_figures/fung_wong_pressure_comparison.png) · [Track / central pressure](docs/report_figures/fung_wong_track_pressure.png) · [Observed R34](docs/report_figures/fung_wong_observed_r34.png) |
|
| 68 |
+
| Mangkhut / 山竹 (2018) | 2018-09-11 00:00 | [Pressure](docs/report_figures/mangkhut_pressure_comparison.png) · [Track / central pressure](docs/report_figures/mangkhut_track_pressure.png) · [Observed R34](docs/report_figures/mangkhut_observed_r34.png) |
|
| 69 |
+
| Jebi / 燕子 (2018) | 2018-08-31 00:00 | [Pressure](docs/report_figures/jebi_pressure_comparison.png) · [Track / central pressure](docs/report_figures/jebi_track_pressure.png) · [Observed R34](docs/report_figures/jebi_observed_r34.png) |
|
| 70 |
|
| 71 |
All three use the same released checkpoint and 50-input-member policy. Fung-wong and Mangkhut use saved Metal forecasts; Jebi uses CPU forecasts from the same weights. Mangkhut and Jebi lie within the fitting years and are retrospective demonstrations, not unseen-case accuracy evidence. **The R34 panels are IBTrACS observed 34-knot quadrant-radius records, not model radius predictions or full wind maps.** Missing radii are left missing. There is **no native wind-radius forecast head** in 1.2.
|
| 72 |
|
| 73 |
+
[All report panels with captions](docs/research_report_zh.md#三個颱風展示) · [Figure hashes and sources](docs/report_figures/provenance.json) · [Earlier films and diagnostic archive](docs/showcase_archive.md)
|
| 74 |
|
| 75 |
## Training data and year cutoffs
|
| 76 |
|
|
|
|
| 95 |
|
| 96 |
The checksum-pinned source caches extend further: the derived typhoon archive reaches **13 July 2026**, and the NCEP basin archive reaches **17 March 2026, 18 UTC**. Those endpoints describe archive availability, **not the fitting cutoff**. Later live weather uses separately documented experimental GFS transfer; it does not update the released weights. This release has no satellite cloud-image or humidity input channel. Previously inspected test/showcase storms remain development evidence, not a newly untouched holdout.
|
| 97 |
|
| 98 |
+
[Training-source details and exact date ranges](docs/trackformer_1_2_training_data.md) · [Recomputed split counts and verified dataset hashes](evaluation/training_data/trackformer_1_2_provenance.json)
|
| 99 |
|
| 100 |
## How Trackformer 1.2 works
|
| 101 |
|
| 102 |
+

|
| 103 |
|
| 104 |
1. **Initialize from the past.** Nine six-hour analyses provide MSLP, 500 hPa height and winds at 850/500/200 hPa, ending at issue time. Static geography, the current observed centre, recent motion, current intensity and validity masks accompany the history. Native regional pressure detail is used when available; missing detail stays flagged. Current observations initialize the core, not future labels.
|
| 105 |
2. **Update the surrounding weather.** The model estimates changes in pressure and horizontal winds and carries information forward to the next six-hour step.
|
|
|
|
| 114 |
| Moving core | 65 × 65; approximately 20 km spacing |
|
| 115 |
| Forecast steps | 20 × 6 hours |
|
| 116 |
|
| 117 |
+
Training increases the prediction horizon from one to two, three and five days, checking the full five-day forecast each epoch. The report uses the selected **epoch-4** checkpoint. [Model source and input contract](models/trackformer_1_2_field/README.md) · [Exact architecture and provenance](models/trackformer_1_2_field/manifest.json).
|
| 118 |
|
| 119 |
### What “50 forecasts averaged” means
|
| 120 |
|
|
|
|
| 128 |
- **NVIDIA GeForce RTX 3070 / CUDA:** verified execution of the completed DeepMind reference benchmark. This is hardware, not a software package.
|
| 129 |
- **Python / PyTorch**, NumPy, SciPy, xarray, netCDF4 and Matplotlib support modelling, data handling and figures. The Chinese report uses XeLaTeX / xeCJK.
|
| 130 |
|
| 131 |
+
## Train the released 1.2 architecture
|
| 132 |
+
|
| 133 |
+
The [open training package](training/trackformer_1_2/README.md) contains the original trainer, architecture, losses, sampler, resume logic, tests, data-preparation utilities and explicit launch configuration. The **exact six derived training-array files**, original manifest, **v168 epoch-3 warm start** and selected **v173 epoch-4 full checkpoint with optimizer/RNG state** are public on [Hugging Face](https://huggingface.co/euler314/typhoon-predict/tree/main/training/trackformer_1_2). A checksum-verifying downloader works without the old Drive setup. Source and instructions are mirrored on both repositories; large binary assets are served by Hugging Face and linked from GitHub.
|
| 134 |
+
|
| 135 |
+
This opens the actual released training path; it does not claim a newly reproduced training run or an untouched benchmark. The archived raw IBTrACS snapshot ID was not recorded. The report's selected epoch 4 is **not** the minimum aggregate validation-loss checkpoint: the saved history's minimum is epoch 2. See the training package for historical metadata and remaining reproducibility limits.
|
| 136 |
+
|
| 137 |
## Get started
|
| 138 |
|
| 139 |
[Download weights and the complete pressure-map exporter](https://github.com/yu314-coder/typhoon-predict/releases/download/trackformer-1.2/trackformer_1_2_field_pressure_export_v2.tar.gz), or use this repository with the [Hugging Face weights](https://huggingface.co/euler314/typhoon-predict/resolve/main/models/trackformer_1_2_field/weights.pt). Put `weights.pt` in `models/trackformer_1_2_field/`. The inference weights and learned model equations are unchanged; the package includes the moving-core output and map renderer.
|
| 140 |
|
| 141 |
+
Prepare a normalized causal issue packet following the [input schema](models/trackformer_1_2_field/README.md), then run with PyTorch and NumPy:
|
| 142 |
|
| 143 |
```bash
|
| 144 |
python models/trackformer_1_2_field/predict.py causal_issue_packet.npz forecast.npz --device mps
|
|
|
|
| 152 |
python models/trackformer_1_2_field/plot_pressure.py forecast.npz pressure_24h.png --lead 24 --interval 2
|
| 153 |
```
|
| 154 |
|
| 155 |
+
Blue denotes low pressure and red high pressure. The renderer uses the actual saved moving field, coordinates and masks; it never draws a replacement vortex around a scalar pressure or observed track. See the [pressure-output schema](models/trackformer_1_2_field/README.md#render-a-pressure-map-directly).
|
| 156 |
|
| 157 |
The inference weight SHA-256 is `db49f36e85a3766defc4c172746897a1f783705d1ce8e6f9dfb8e87ae1d902cb`.
|
| 158 |
|
|
|
|
| 168 |
|
| 169 |
The auxiliary wind's averaging period and surface-wind skill are not validated. Optional pressure-derived wind/radius estimates are experimental diagnostics, not agency-equivalent forecasts. Pressure-map interpolation is not evidence of finer effective resolution. Retrospective analyses and experimental GFS live transfer must not be presented as operational validation.
|
| 170 |
|
| 171 |
+
See [wind-estimation assumptions](models/trackformer_1_2_field/WIND_ESTIMATION.md), [full intensity evaluation](docs/intensity_benchmark.md), [evaluation limits](docs/trackformer_1_2_evaluation.md), and the [example archive](docs/showcase_archive.md). Source, normalization and checkpoint identities are recorded in the [manifest](models/trackformer_1_2_field/manifest.json). Older release comparisons remain in the linked technical records, not the report's two-model headline table.
|
| 172 |
|
| 173 |
## Ongoing work: Taiwan rainfall
|
| 174 |
|
| 175 |
Ongoing research combines moisture, wind direction, weather outside Taiwan and terrain to study ordinary and heavy rainfall over Taiwan, with a separate observation-based evaluation. **Released 1.2 has no rainfall forecast head or humidity input.** These are research directions, not demonstrated rainfall improvements from the pressure model.
|
| 176 |
|
| 177 |
+
Model code is in [`models/`](models/), methods in [`paper/`](paper/), protocols in [`docs/`](docs/), results in [`evaluation/`](evaluation/) and reproducible utilities in [`release_tools/`](release_tools/).
|
docs/repository_layout.md
CHANGED
|
@@ -17,3 +17,7 @@ The 1.1-versus-1.2 comparison arrays, TIP diagnostic, current paper, weights, mo
|
|
| 17 |
## Recovery
|
| 18 |
|
| 19 |
The complete pre-cleanup repository remains in [commit 459a3d014aa4794652ceba3ecdffcde0beada2cd](https://github.com/yu314-coder/typhoon-predict/tree/459a3d014aa4794652ceba3ecdffcde0beada2cd). The [Trackformer 1.1 release](https://github.com/yu314-coder/typhoon-predict/releases/tag/trackformer-1.1) remains available. This cleanup does not rewrite Git history.
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 17 |
## Recovery
|
| 18 |
|
| 19 |
The complete pre-cleanup repository remains in [commit 459a3d014aa4794652ceba3ecdffcde0beada2cd](https://github.com/yu314-coder/typhoon-predict/tree/459a3d014aa4794652ceba3ecdffcde0beada2cd). The [Trackformer 1.1 release](https://github.com/yu314-coder/typhoon-predict/releases/tag/trackformer-1.1) remains available. This cleanup does not rewrite Git history.
|
| 20 |
+
|
| 21 |
+
## Released 1.2 training
|
| 22 |
+
|
| 23 |
+
`training/trackformer_1_2/` contains the original training source, launch wrapper, requirements, historical preparation tools, configuration, hashes and public asset downloader. Large arrays and optimizer checkpoints live at the identical Hugging Face path and are referenced from GitHub.
|
docs/trackformer_1_2_training_data.md
CHANGED
|
@@ -57,4 +57,6 @@ With the original frozen dataset available, reproduce the small receipt without
|
|
| 57 |
python release_tools/audit_training_data.py --dataset-root /Volumes/D/typhoon_predict/data/v164_reuse --output /Volumes/D/typhoon_predict/remote_typhoon_predict_repo/evaluation/training_data/trackformer_1_2_provenance.json
|
| 58 |
```
|
| 59 |
|
| 60 |
-
The original full training
|
|
|
|
|
|
|
|
|
| 57 |
python release_tools/audit_training_data.py --dataset-root /Volumes/D/typhoon_predict/data/v164_reuse --output /Volumes/D/typhoon_predict/remote_typhoon_predict_repo/evaluation/training_data/trackformer_1_2_provenance.json
|
| 58 |
```
|
| 59 |
|
| 60 |
+
The exact derived arrays, original manifest, warm start and selected full training checkpoint are now publicly available in [the open training package](../training/trackformer_1_2/README.md). Original source bytes and frozen model weights are unchanged. The additive publication does not run training or modify any active trainer.
|
| 61 |
+
|
| 62 |
+
The source archive includes historical ERA5 preparation utilities; its small full-field pilot is not the released six-file corpus. Use the published exact corpus to reproduce the released data boundary. No missing raw IBTrACS provenance has been invented.
|
training/trackformer_1_2/LICENSE
ADDED
|
@@ -0,0 +1,21 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
MIT License
|
| 2 |
+
|
| 3 |
+
Copyright (c) 2026 Yu Yao-Hsing (yu314-coder)
|
| 4 |
+
|
| 5 |
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
| 6 |
+
of this software and associated documentation files (the "Software"), to deal
|
| 7 |
+
in the Software without restriction, including without limitation the rights
|
| 8 |
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
| 9 |
+
copies of the Software, and to permit persons to whom the Software is
|
| 10 |
+
furnished to do so, subject to the following conditions:
|
| 11 |
+
|
| 12 |
+
The above copyright notice and this permission notice shall be included in all
|
| 13 |
+
copies or substantial portions of the Software.
|
| 14 |
+
|
| 15 |
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
| 16 |
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
| 17 |
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
| 18 |
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
| 19 |
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
| 20 |
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
| 21 |
+
SOFTWARE.
|
training/trackformer_1_2/README.md
ADDED
|
@@ -0,0 +1,63 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# TrackFormer 1.2: open training package
|
| 2 |
+
|
| 3 |
+
This is the actual released pressure-field model: development **1.2.73 (v173), epoch 4**, 21,452,595 parameters. It is separate from older route-only 1.2 experiments and later rainfall models.
|
| 4 |
+
|
| 5 |
+
## Published contents
|
| 6 |
+
|
| 7 |
+
- `source/v173/`: original, byte-identical model, baseline architecture, archive loader, objective, curriculum, sampler, trainer, resume logic, design and original Metal tests.
|
| 8 |
+
- `source/data_preparation/`: original historical ERA5 download/plan utilities and dataset contracts. The small full-field pilot described there is **not** the six-file training corpus.
|
| 9 |
+
- `assets.json`: exact sizes and SHA-256 values for all six original derived archives, the original manifest, warm-start checkpoint and selected full training checkpoint.
|
| 10 |
+
- `data/`: those exact public arrays, including original split plan, labels/masks and static geography. Binary files are hosted on Hugging Face at the identical path, not committed into GitHub source history.
|
| 11 |
+
- `checkpoints/v168_epoch_003.pt`: the original, mandatory warm start, previously split across Drive files.
|
| 12 |
+
- `checkpoints/v173_epoch_004.pt`: the selected full training checkpoint, with model, optimizer, normalization, four-epoch history and Python/NumPy/Torch CPU RNG states.
|
| 13 |
+
- `config.json`, `checkpoint-history.json`, `SOURCE_MANIFEST.json`: archived recipe, exact saved loss history and source hashes.
|
| 14 |
+
- `train.py`: a path-portability wrapper; it calls the unmodified original trainer with explicit paths and settings. No module import starts training.
|
| 15 |
+
- Model export, physical pressure-map rendering and benchmark tools remain in the repository's `release_tools/` and `models/trackformer_1_2_field/`.
|
| 16 |
+
|
| 17 |
+
Both [GitHub](https://github.com/yu314-coder/typhoon-predict/tree/main/training/trackformer_1_2) and [Hugging Face](https://huggingface.co/euler314/typhoon-predict/tree/main/training/trackformer_1_2) carry the source and instructions. Hugging Face serves the large assets (2,644,261,044 bytes in total); no Drive permission is required.
|
| 18 |
+
|
| 19 |
+
## Download and verify
|
| 20 |
+
|
| 21 |
+
From the repository root, after obtaining this source folder:
|
| 22 |
+
|
| 23 |
+
```bash
|
| 24 |
+
python training/trackformer_1_2/download_assets.py
|
| 25 |
+
python training/trackformer_1_2/download_assets.py --verify-only
|
| 26 |
+
python -m pip install -r training/trackformer_1_2/requirements.txt
|
| 27 |
+
```
|
| 28 |
+
|
| 29 |
+
The downloader uses only Python's standard library. It publishes each local file only after exact size/SHA verification and does not overwrite a differing existing file. Use `--revision <immutable-HF-commit>` from `PUBLICATION.md` to pin the upload; even the default `main` URL must match the fixed hashes. For a writable data/output directory on an external disk, copy the package there or set `--destination`; then pass its `data` and `checkpoints` paths to the launcher.
|
| 30 |
+
|
| 31 |
+
## Preflight, then explicitly start training
|
| 32 |
+
|
| 33 |
+
```bash
|
| 34 |
+
python training/trackformer_1_2/train.py --device mps --output runs/retrain-1.2 --preflight
|
| 35 |
+
python training/trackformer_1_2/train.py --device mps --output runs/retrain-1.2
|
| 36 |
+
```
|
| 37 |
+
|
| 38 |
+
Use `--device cuda` for NVIDIA or `--device cpu` for CPU. This wrapper explicitly sets the archived notebook/design recipe: effective batch 10, microbatch 4, workers 3, four epochs; pass flags to change execution settings deliberately. It uses FP32 even though the historical notebook's compatibility AMP positional flag was 1. The unchanged original entry point is also available:
|
| 39 |
+
|
| 40 |
+
```bash
|
| 41 |
+
python training/trackformer_1_2/source/v173/v173_pressure_train.py 10 3 0 \
|
| 42 |
+
--microbatch 4 --epochs 4 --device mps \
|
| 43 |
+
--manifest training/trackformer_1_2/data/manifest.json \
|
| 44 |
+
--warm-start training/trackformer_1_2/checkpoints/v168_epoch_003.pt \
|
| 45 |
+
--output runs/retrain-1.2
|
| 46 |
+
```
|
| 47 |
+
|
| 48 |
+
The original loader hashes all arrays, reconstructs a writable basin mmap beside the dataset, recomputes train-only normalization and enforces counts **13,949 / 1,041 / 1,195** with native-pressure subsets **800 / 100 / 100**. Allow space for that extra approximately 390 MB mmap and approximately 260 MB checkpoint files. Do not point preflight/training at an active trainer's data/output directories. The trainer uses a per-output lock; the same output resumes its `last.pt`. SIGINT/SIGTERM request a checkpoint at an optimizer boundary. Training horizons are 4/8/12/20 six-hour steps; validation always uses 20 steps. Four-step truncated backpropagation, gradient clipping, AdamW settings and all losses are in the original source.
|
| 49 |
+
|
| 50 |
+
## Provenance and honest reproduction limits
|
| 51 |
+
|
| 52 |
+
- Dataset manifest: `04fd0d30620137aad20173dcd59e22da2a6c20fae2f17de593e91a59517459c1`.
|
| 53 |
+
- Warm start: `99ed776097e05ad6df03375465d9618741bed63a4ee57724d3ba98e295039d30`.
|
| 54 |
+
- Selected full checkpoint: `f194a23d3f91ea76ad776dfad942fabd669eeae8b3fd665815463095367e9ee0`.
|
| 55 |
+
- Inference-only weights: `db49f36e85a3766defc4c172746897a1f783705d1ce8e6f9dfb8e87ae1d902cb`.
|
| 56 |
+
|
| 57 |
+
The selected full checkpoint is epoch **4**, used by the report; **it is not the smallest saved aggregate validation loss** (epoch 2 is 0.6939101; epoch 4 is 0.7283264). A fresh run's `best.pt` follows the original minimum-validation criterion and need not be the report's epoch-4 checkpoint. Do not conflate these choices. The supplied warm start makes the actual v173 training path runnable without reconstructing every earlier experimental model. Historical v168 optimization logs and complete original runtime package versions/accelerator RNG states are not archived here. The recipe is preserved from the original notebook/design, not a newly recovered complete original process argv. Bitwise end-to-end retraining across devices has **not** been demonstrated.
|
| 58 |
+
|
| 59 |
+
The exact derived data are published; the raw IBTrACS release/download ID was not recorded in the frozen manifest. Current upstream downloads may differ and cannot be claimed byte-identical. Fitting uses 2000–2021, validation 2022–2023, original test 2024–2025, with complete storms separated. Later cache records are not fitting years. All future fields/positions are labels only. Repeatedly inspected tests/showcases are development evidence, not a fresh unused holdout. This publication does not rerun training, change released weights or modify active rainfall training.
|
| 60 |
+
|
| 61 |
+
## Licenses and attribution
|
| 62 |
+
|
| 63 |
+
Project source and project checkpoints retain the repository's MIT license. Upstream weather/track data retain their own terms and attributions, not the MIT software license; see `THIRD_PARTY_NOTICES.md`. Downloaded full checkpoints use Python pickle containers: load these identified, checksum-verified project artifacts only.
|
training/trackformer_1_2/SOURCE_MANIFEST.json
ADDED
|
@@ -0,0 +1,178 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"schema": "trackformer-1.2-open-training-source-v1",
|
| 3 |
+
"unchanged_original_sources": true,
|
| 4 |
+
"model_weights_modified": false,
|
| 5 |
+
"files": [
|
| 6 |
+
{
|
| 7 |
+
"path": "training/trackformer_1_2/LICENSE",
|
| 8 |
+
"sha256": "090a97ebbccef9de77cb3898c473f5c9ada7e140777957f1195fc6ab89b9a313",
|
| 9 |
+
"bytes": 1083
|
| 10 |
+
},
|
| 11 |
+
{
|
| 12 |
+
"path": "training/trackformer_1_2/requirements.txt",
|
| 13 |
+
"sha256": "37e15a665eff6c8ee94c47ac2da0f7c7e2d699b7e21623b0289da4cf5d3995d0",
|
| 14 |
+
"bytes": 45
|
| 15 |
+
},
|
| 16 |
+
{
|
| 17 |
+
"path": "training/trackformer_1_2/config.json",
|
| 18 |
+
"sha256": "813f5d05a3dec5ec160b4e91eb80970fb11d88416fc7709ffc9d147fc880316b",
|
| 19 |
+
"bytes": 820
|
| 20 |
+
},
|
| 21 |
+
{
|
| 22 |
+
"path": "training/trackformer_1_2/THIRD_PARTY_NOTICES.md",
|
| 23 |
+
"sha256": "f764dd203eb124fa6192edd3b7b49a7a546cd5d69d3a8a250589a14ce43e2115",
|
| 24 |
+
"bytes": 1697
|
| 25 |
+
},
|
| 26 |
+
{
|
| 27 |
+
"path": "training/trackformer_1_2/checkpoint-history.json",
|
| 28 |
+
"sha256": "8b9593a7ed58c1bd28dd27c8049c89a0c5c5644a30bfe5203a4b54c445ff2fb4",
|
| 29 |
+
"bytes": 459
|
| 30 |
+
},
|
| 31 |
+
{
|
| 32 |
+
"path": "training/trackformer_1_2/README.md",
|
| 33 |
+
"sha256": "342c43d94710355101ddc5bcad44923f47707f473eb803398f44fe2d77e8ace9",
|
| 34 |
+
"bytes": 6560
|
| 35 |
+
},
|
| 36 |
+
{
|
| 37 |
+
"path": "training/trackformer_1_2/train.py",
|
| 38 |
+
"sha256": "4502a87bfafb11667015070c6fad37b723b4a53723789febda075fba1c7770e8",
|
| 39 |
+
"bytes": 1436
|
| 40 |
+
},
|
| 41 |
+
{
|
| 42 |
+
"path": "training/trackformer_1_2/download_assets.py",
|
| 43 |
+
"sha256": "01d21d09c001394f49d807314e38b22a4a8c179783a3f22ddaa4c27e8c3f2011",
|
| 44 |
+
"bytes": 3049
|
| 45 |
+
},
|
| 46 |
+
{
|
| 47 |
+
"path": "training/trackformer_1_2/assets.json",
|
| 48 |
+
"sha256": "a59148e9ba2617599614784b47789c3654df1a7809ee039e78bcd00fe982b211",
|
| 49 |
+
"bytes": 1993
|
| 50 |
+
},
|
| 51 |
+
{
|
| 52 |
+
"path": "training/trackformer_1_2/source/v173/test_model.py",
|
| 53 |
+
"sha256": "493e0db88a098c89bb4db5e6c577fa24b4b7bcb83ce547dceac35237dbce7206",
|
| 54 |
+
"bytes": 5579
|
| 55 |
+
},
|
| 56 |
+
{
|
| 57 |
+
"path": "training/trackformer_1_2/source/v173/v165_base.py",
|
| 58 |
+
"sha256": "5c9473993a19d0fed00055494b5c1f50c1006aaf46551ff50145f67048c3374a",
|
| 59 |
+
"bytes": 39040
|
| 60 |
+
},
|
| 61 |
+
{
|
| 62 |
+
"path": "training/trackformer_1_2/source/v173/v173_pressure_train.py",
|
| 63 |
+
"sha256": "f40691cec3e787f4ac3c189f647efb0953c1eaa7e03fc86edaf4fee06990b2be",
|
| 64 |
+
"bytes": 11200
|
| 65 |
+
},
|
| 66 |
+
{
|
| 67 |
+
"path": "training/trackformer_1_2/source/v173/model.py",
|
| 68 |
+
"sha256": "5dab06c20224cf48c51302122455ee7e16948f977bf4b760fd8aff5d443ab0bc",
|
| 69 |
+
"bytes": 3428
|
| 70 |
+
},
|
| 71 |
+
{
|
| 72 |
+
"path": "training/trackformer_1_2/source/v173/DESIGN.md",
|
| 73 |
+
"sha256": "163ae39ff572f23c519d65e27ee6a8db23b47df9461670709229ccf95130049a",
|
| 74 |
+
"bytes": 3278
|
| 75 |
+
},
|
| 76 |
+
{
|
| 77 |
+
"path": "training/trackformer_1_2/source/v173/baseline_model.py",
|
| 78 |
+
"sha256": "6c05941fd93baaca04508877edc69f8f0394b70e0a64e734ef0d2c1e8832246a",
|
| 79 |
+
"bytes": 10040
|
| 80 |
+
},
|
| 81 |
+
{
|
| 82 |
+
"path": "training/trackformer_1_2/source/data_preparation/v164_prepare_data.py",
|
| 83 |
+
"sha256": "c1746e0c0f985890a0f02061a2d84b521cc7b29042c92972aa1ebab3354e4c95",
|
| 84 |
+
"bytes": 19312
|
| 85 |
+
},
|
| 86 |
+
{
|
| 87 |
+
"path": "training/trackformer_1_2/source/data_preparation/requirements.txt",
|
| 88 |
+
"sha256": "54918b8f29330299e18b8924e9cd93d05807e7f8cd99b2efd4805e9be94003c1",
|
| 89 |
+
"bytes": 55
|
| 90 |
+
},
|
| 91 |
+
{
|
| 92 |
+
"path": "training/trackformer_1_2/source/data_preparation/v164_structure_research.md",
|
| 93 |
+
"sha256": "a4b40e50b7d0d1433a162438f14ed110541702c4cec2a386959857293ea78141",
|
| 94 |
+
"bytes": 17119
|
| 95 |
+
},
|
| 96 |
+
{
|
| 97 |
+
"path": "training/trackformer_1_2/source/data_preparation/v164_download_pressure_geography.py",
|
| 98 |
+
"sha256": "b8602cd2fba131a287dc9e25781c8f037d3325fec8a23203576916c2f7d2189c",
|
| 99 |
+
"bytes": 9256
|
| 100 |
+
},
|
| 101 |
+
{
|
| 102 |
+
"path": "training/trackformer_1_2/source/data_preparation/v164_prepare_mac.py",
|
| 103 |
+
"sha256": "7e7313f865369d42f67645196b4af0e1b2b84216303f7b2b61e577d3734df79a",
|
| 104 |
+
"bytes": 1309
|
| 105 |
+
},
|
| 106 |
+
{
|
| 107 |
+
"path": "training/trackformer_1_2/source/data_preparation/v164_dataset_contract.md",
|
| 108 |
+
"sha256": "cfd838651e632aba0d2439ba202a59895bd4df22137185974df1d97f0e2525d4",
|
| 109 |
+
"bytes": 4664
|
| 110 |
+
},
|
| 111 |
+
{
|
| 112 |
+
"path": "training/trackformer_1_2/source/data_preparation/v164_global_regional_train.py",
|
| 113 |
+
"sha256": "c0808a1467695d2bc2ad8483aa4f4c6de1463628221024ac559f6204effe7978",
|
| 114 |
+
"bytes": 52581
|
| 115 |
+
},
|
| 116 |
+
{
|
| 117 |
+
"path": "training/trackformer_1_2/source/data_preparation/v164_reuse_impl.py",
|
| 118 |
+
"sha256": "2da0078d4a74faa93f8cf874e42707b01f286406f0ceaa4db1814aab0eceb7d1",
|
| 119 |
+
"bytes": 18414
|
| 120 |
+
}
|
| 121 |
+
],
|
| 122 |
+
"assets": [
|
| 123 |
+
{
|
| 124 |
+
"path": "training/trackformer_1_2/data/manifest.json",
|
| 125 |
+
"sha256": "04fd0d30620137aad20173dcd59e22da2a6c20fae2f17de593e91a59517459c1",
|
| 126 |
+
"bytes": 814,
|
| 127 |
+
"group": "data"
|
| 128 |
+
},
|
| 129 |
+
{
|
| 130 |
+
"path": "training/trackformer_1_2/data/track_windows_v13.npz",
|
| 131 |
+
"sha256": "cb8d7291aa8140124ca522feb448a1d44124b018efc4ba1d54824b60a4f73451",
|
| 132 |
+
"bytes": 42755049,
|
| 133 |
+
"group": "data"
|
| 134 |
+
},
|
| 135 |
+
{
|
| 136 |
+
"path": "training/trackformer_1_2/data/basin_all_int8.npz",
|
| 137 |
+
"sha256": "71e5f1bb46c44e5936c506a3f8af021c1fe94c238ed030c05c66e8f342374b6d",
|
| 138 |
+
"bytes": 274115577,
|
| 139 |
+
"group": "data"
|
| 140 |
+
},
|
| 141 |
+
{
|
| 142 |
+
"path": "training/trackformer_1_2/data/basin_slp_atlas_float16.npy",
|
| 143 |
+
"sha256": "9f6a4ec7b2f6beadfa1075af9176d2a653e39a824d504d989e5e56e19c571521",
|
| 144 |
+
"bytes": 111394928,
|
| 145 |
+
"group": "data"
|
| 146 |
+
},
|
| 147 |
+
{
|
| 148 |
+
"path": "training/trackformer_1_2/data/pressure_hpa.npy",
|
| 149 |
+
"sha256": "16cc1b153aeee5eff4b67337b599892dcd86a5716554596b8031048035d247ba",
|
| 150 |
+
"bytes": 1698356128,
|
| 151 |
+
"group": "data"
|
| 152 |
+
},
|
| 153 |
+
{
|
| 154 |
+
"path": "training/trackformer_1_2/data/plan.json",
|
| 155 |
+
"sha256": "117eb291e385a326191be2f73182ed6dcbf0e9f99745780ac9e49a29def3d8da",
|
| 156 |
+
"bytes": 295005,
|
| 157 |
+
"group": "data"
|
| 158 |
+
},
|
| 159 |
+
{
|
| 160 |
+
"path": "training/trackformer_1_2/data/geography.npz",
|
| 161 |
+
"sha256": "f5471471da4a1e15f85fbfb2529065806281357c8e50cb272fcd2b02e8553b96",
|
| 162 |
+
"bytes": 2547017,
|
| 163 |
+
"group": "data"
|
| 164 |
+
},
|
| 165 |
+
{
|
| 166 |
+
"path": "training/trackformer_1_2/checkpoints/v168_epoch_003.pt",
|
| 167 |
+
"sha256": "99ed776097e05ad6df03375465d9618741bed63a4ee57724d3ba98e295039d30",
|
| 168 |
+
"bytes": 256826555,
|
| 169 |
+
"group": "checkpoints"
|
| 170 |
+
},
|
| 171 |
+
{
|
| 172 |
+
"path": "training/trackformer_1_2/checkpoints/v173_epoch_004.pt",
|
| 173 |
+
"sha256": "f194a23d3f91ea76ad776dfad942fabd669eeae8b3fd665815463095367e9ee0",
|
| 174 |
+
"bytes": 257969971,
|
| 175 |
+
"group": "checkpoints"
|
| 176 |
+
}
|
| 177 |
+
]
|
| 178 |
+
}
|
training/trackformer_1_2/THIRD_PARTY_NOTICES.md
ADDED
|
@@ -0,0 +1,9 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Data sources and attribution
|
| 2 |
+
|
| 3 |
+
The MIT license covers project code and project model artifacts; it does not replace upstream data terms.
|
| 4 |
+
|
| 5 |
+
- **NOAA PSL NCEP/NCAR Reanalysis 1**: basin mean sea-level pressure, 500 hPa height and winds at 850/500/200 hPa. Data provided by NOAA PSL, Boulder, Colorado, USA. https://psl.noaa.gov/data/gridded/data.ncep.reanalysis.html
|
| 6 |
+
- **NOAA/NCEI IBTrACS**: derived storm observations, labels and masks; credit the original contributing meteorological agencies and NOAA/NCEI. Knapp et al. (2010), doi:10.1175/2009BAMS2755.1. https://www.ncei.noaa.gov/products/international-best-track-archive . Exact derived archive SHA is in assets.json; its original raw snapshot ID was not recorded.
|
| 7 |
+
- **Copernicus Climate Change Service (C3S), ERA5, accessed through Google's ARCO-ERA5 mirror**: native mean-sea-level-pressure patches and static land fraction/elevation. Contains modified Copernicus Climate Change Service information (2026). ERA5 hourly single-level data, doi:10.24381/cds.adbb2d47, https://cds.climate.copernicus.eu/datasets/reanalysis-era5-single-levels . The current dataset lists a CC-BY licence; attribution and upstream terms apply. ECMWF or the European Commission does not endorse this model. ARCO mirror and recipes: https://github.com/google-research/arco-era5 . The mirror's code licence is separate from the underlying meteorological data licence.
|
| 8 |
+
|
| 9 |
+
The published arrays are transformed subsets: quantized/coarse basin grids, float16 pressure, issue-centred patches, normalized derived track features and masks. These transformations do not create new native observational resolution. No claim of endorsement by NOAA, ECMWF, C3S or Google is made.
|
training/trackformer_1_2/assets.json
ADDED
|
@@ -0,0 +1,60 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"repository": "euler314/typhoon-predict",
|
| 3 |
+
"repo_type": "model",
|
| 4 |
+
"files": [
|
| 5 |
+
{
|
| 6 |
+
"path": "training/trackformer_1_2/data/manifest.json",
|
| 7 |
+
"sha256": "04fd0d30620137aad20173dcd59e22da2a6c20fae2f17de593e91a59517459c1",
|
| 8 |
+
"bytes": 814,
|
| 9 |
+
"group": "data"
|
| 10 |
+
},
|
| 11 |
+
{
|
| 12 |
+
"path": "training/trackformer_1_2/data/track_windows_v13.npz",
|
| 13 |
+
"sha256": "cb8d7291aa8140124ca522feb448a1d44124b018efc4ba1d54824b60a4f73451",
|
| 14 |
+
"bytes": 42755049,
|
| 15 |
+
"group": "data"
|
| 16 |
+
},
|
| 17 |
+
{
|
| 18 |
+
"path": "training/trackformer_1_2/data/basin_all_int8.npz",
|
| 19 |
+
"sha256": "71e5f1bb46c44e5936c506a3f8af021c1fe94c238ed030c05c66e8f342374b6d",
|
| 20 |
+
"bytes": 274115577,
|
| 21 |
+
"group": "data"
|
| 22 |
+
},
|
| 23 |
+
{
|
| 24 |
+
"path": "training/trackformer_1_2/data/basin_slp_atlas_float16.npy",
|
| 25 |
+
"sha256": "9f6a4ec7b2f6beadfa1075af9176d2a653e39a824d504d989e5e56e19c571521",
|
| 26 |
+
"bytes": 111394928,
|
| 27 |
+
"group": "data"
|
| 28 |
+
},
|
| 29 |
+
{
|
| 30 |
+
"path": "training/trackformer_1_2/data/pressure_hpa.npy",
|
| 31 |
+
"sha256": "16cc1b153aeee5eff4b67337b599892dcd86a5716554596b8031048035d247ba",
|
| 32 |
+
"bytes": 1698356128,
|
| 33 |
+
"group": "data"
|
| 34 |
+
},
|
| 35 |
+
{
|
| 36 |
+
"path": "training/trackformer_1_2/data/plan.json",
|
| 37 |
+
"sha256": "117eb291e385a326191be2f73182ed6dcbf0e9f99745780ac9e49a29def3d8da",
|
| 38 |
+
"bytes": 295005,
|
| 39 |
+
"group": "data"
|
| 40 |
+
},
|
| 41 |
+
{
|
| 42 |
+
"path": "training/trackformer_1_2/data/geography.npz",
|
| 43 |
+
"sha256": "f5471471da4a1e15f85fbfb2529065806281357c8e50cb272fcd2b02e8553b96",
|
| 44 |
+
"bytes": 2547017,
|
| 45 |
+
"group": "data"
|
| 46 |
+
},
|
| 47 |
+
{
|
| 48 |
+
"path": "training/trackformer_1_2/checkpoints/v168_epoch_003.pt",
|
| 49 |
+
"sha256": "99ed776097e05ad6df03375465d9618741bed63a4ee57724d3ba98e295039d30",
|
| 50 |
+
"bytes": 256826555,
|
| 51 |
+
"group": "checkpoints"
|
| 52 |
+
},
|
| 53 |
+
{
|
| 54 |
+
"path": "training/trackformer_1_2/checkpoints/v173_epoch_004.pt",
|
| 55 |
+
"sha256": "f194a23d3f91ea76ad776dfad942fabd669eeae8b3fd665815463095367e9ee0",
|
| 56 |
+
"bytes": 257969971,
|
| 57 |
+
"group": "checkpoints"
|
| 58 |
+
}
|
| 59 |
+
]
|
| 60 |
+
}
|
training/trackformer_1_2/checkpoint-history.json
ADDED
|
@@ -0,0 +1,26 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
[
|
| 2 |
+
{
|
| 3 |
+
"epoch": 1,
|
| 4 |
+
"horizon": 4,
|
| 5 |
+
"train": 0.1411035786937901,
|
| 6 |
+
"validation": 0.7017406511544373
|
| 7 |
+
},
|
| 8 |
+
{
|
| 9 |
+
"epoch": 2,
|
| 10 |
+
"horizon": 8,
|
| 11 |
+
"train": 0.18676476989154306,
|
| 12 |
+
"validation": 0.6939100875934489
|
| 13 |
+
},
|
| 14 |
+
{
|
| 15 |
+
"epoch": 3,
|
| 16 |
+
"horizon": 12,
|
| 17 |
+
"train": 0.23248198316925273,
|
| 18 |
+
"validation": 0.7200751373592149
|
| 19 |
+
},
|
| 20 |
+
{
|
| 21 |
+
"epoch": 4,
|
| 22 |
+
"horizon": 20,
|
| 23 |
+
"train": 0.3557772162892601,
|
| 24 |
+
"validation": 0.7283263609942383
|
| 25 |
+
}
|
| 26 |
+
]
|
training/trackformer_1_2/checkpoints/v168_epoch_003.pt
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:99ed776097e05ad6df03375465d9618741bed63a4ee57724d3ba98e295039d30
|
| 3 |
+
size 256826555
|
training/trackformer_1_2/checkpoints/v173_epoch_004.pt
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:f194a23d3f91ea76ad776dfad942fabd669eeae8b3fd665815463095367e9ee0
|
| 3 |
+
size 257969971
|
training/trackformer_1_2/config.json
ADDED
|
@@ -0,0 +1,29 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"public_version": "1.2",
|
| 3 |
+
"development_version": "1.2.73",
|
| 4 |
+
"seed": 168,
|
| 5 |
+
"optimizer": {
|
| 6 |
+
"name": "AdamW",
|
| 7 |
+
"lr": 3e-05,
|
| 8 |
+
"weight_decay": 0.01
|
| 9 |
+
},
|
| 10 |
+
"gradient_clip_norm": 1.0,
|
| 11 |
+
"epochs": 4,
|
| 12 |
+
"training_horizon_steps": [
|
| 13 |
+
4,
|
| 14 |
+
8,
|
| 15 |
+
12,
|
| 16 |
+
20
|
| 17 |
+
],
|
| 18 |
+
"validation_horizon_steps": 20,
|
| 19 |
+
"precision": "float32",
|
| 20 |
+
"effective_batch": 10,
|
| 21 |
+
"microbatch": 4,
|
| 22 |
+
"workers": 3,
|
| 23 |
+
"checkpoint_every_windows": 50,
|
| 24 |
+
"launch_provenance": "Archived notebook/design command; complete original machine/runtime/argv not recorded in selected checkpoint",
|
| 25 |
+
"selected_epoch": 4,
|
| 26 |
+
"minimum_saved_validation_loss_epoch": 2,
|
| 27 |
+
"warm_start_sha256": "99ed776097e05ad6df03375465d9618741bed63a4ee57724d3ba98e295039d30",
|
| 28 |
+
"selected_checkpoint_sha256": "f194a23d3f91ea76ad776dfad942fabd669eeae8b3fd665815463095367e9ee0"
|
| 29 |
+
}
|
training/trackformer_1_2/data/basin_all_int8.npz
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:71e5f1bb46c44e5936c506a3f8af021c1fe94c238ed030c05c66e8f342374b6d
|
| 3 |
+
size 274115577
|
training/trackformer_1_2/data/basin_slp_atlas_float16.npy
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:9f6a4ec7b2f6beadfa1075af9176d2a653e39a824d504d989e5e56e19c571521
|
| 3 |
+
size 111394928
|
training/trackformer_1_2/data/geography.npz
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:f5471471da4a1e15f85fbfb2529065806281357c8e50cb272fcd2b02e8553b96
|
| 3 |
+
size 2547017
|
training/trackformer_1_2/data/manifest.json
ADDED
|
@@ -0,0 +1 @@
|
|
|
|
|
|
|
| 1 |
+
{"schema":"v164-reuse-pressure-v1","pilot":false,"future_weather_as_inputs":false,"official_forecasts_as_inputs":false,"files":{"track":{"path":"track_windows_v13.npz","sha256":"cb8d7291aa8140124ca522feb448a1d44124b018efc4ba1d54824b60a4f73451"},"basin":{"path":"basin_all_int8.npz","sha256":"71e5f1bb46c44e5936c506a3f8af021c1fe94c238ed030c05c66e8f342374b6d"},"slp":{"path":"basin_slp_atlas_float16.npy","sha256":"9f6a4ec7b2f6beadfa1075af9176d2a653e39a824d504d989e5e56e19c571521"},"pressure":{"path":"pressure_hpa.npy","sha256":"16cc1b153aeee5eff4b67337b599892dcd86a5716554596b8031048035d247ba"},"plan":{"path":"plan.json","sha256":"117eb291e385a326191be2f73182ed6dcbf0e9f99745780ac9e49a29def3d8da"},"geography":{"path":"geography.npz","sha256":"f5471471da4a1e15f85fbfb2529065806281357c8e50cb272fcd2b02e8553b96"}}}
|
training/trackformer_1_2/data/plan.json
ADDED
|
The diff for this file is too large to render.
See raw diff
|
|
|
training/trackformer_1_2/data/pressure_hpa.npy
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:16cc1b153aeee5eff4b67337b599892dcd86a5716554596b8031048035d247ba
|
| 3 |
+
size 1698356128
|
training/trackformer_1_2/data/track_windows_v13.npz
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:cb8d7291aa8140124ca522feb448a1d44124b018efc4ba1d54824b60a4f73451
|
| 3 |
+
size 42755049
|
training/trackformer_1_2/download_assets.py
ADDED
|
@@ -0,0 +1,77 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Fetch exact public training inputs/checkpoints, verifying every SHA-256.
|
| 2 |
+
|
| 3 |
+
Uses Python's standard library. Does not train, install packages, or use Drive.
|
| 4 |
+
"""
|
| 5 |
+
import argparse
|
| 6 |
+
import hashlib
|
| 7 |
+
import json
|
| 8 |
+
from pathlib import Path
|
| 9 |
+
import os
|
| 10 |
+
import urllib.parse
|
| 11 |
+
import urllib.request
|
| 12 |
+
|
| 13 |
+
|
| 14 |
+
def sha(path):
|
| 15 |
+
h = hashlib.sha256()
|
| 16 |
+
with path.open('rb') as f:
|
| 17 |
+
for block in iter(lambda: f.read(8 * 1024**2), b''):
|
| 18 |
+
h.update(block)
|
| 19 |
+
return h.hexdigest()
|
| 20 |
+
|
| 21 |
+
|
| 22 |
+
def download(entry, root, base, prefix):
|
| 23 |
+
relative = Path(entry['path']).relative_to(prefix)
|
| 24 |
+
path = root / relative
|
| 25 |
+
if path.exists():
|
| 26 |
+
if path.stat().st_size == entry['bytes'] and sha(path) == entry['sha256']:
|
| 27 |
+
print('Verified existing', relative, flush=True)
|
| 28 |
+
return
|
| 29 |
+
raise ValueError('Existing file differs; preserve/move it before retrying: ' + str(path))
|
| 30 |
+
path.parent.mkdir(parents=True, exist_ok=True)
|
| 31 |
+
temp = path.with_name(path.name + '.downloading')
|
| 32 |
+
url = base + urllib.parse.quote(entry['path'], safe='/')
|
| 33 |
+
print('Downloading', relative, entry['bytes'], 'bytes', flush=True)
|
| 34 |
+
try:
|
| 35 |
+
h = hashlib.sha256()
|
| 36 |
+
size = 0
|
| 37 |
+
with urllib.request.urlopen(url, timeout=120) as response, temp.open('wb') as out:
|
| 38 |
+
for block in iter(lambda: response.read(8 * 1024**2), b''):
|
| 39 |
+
out.write(block)
|
| 40 |
+
h.update(block)
|
| 41 |
+
size += len(block)
|
| 42 |
+
out.flush()
|
| 43 |
+
os.fsync(out.fileno())
|
| 44 |
+
if size != entry['bytes'] or h.hexdigest() != entry['sha256']:
|
| 45 |
+
raise ValueError('Downloaded content differs: ' + str(relative))
|
| 46 |
+
temp.replace(path)
|
| 47 |
+
except Exception:
|
| 48 |
+
if temp.is_file():
|
| 49 |
+
temp.unlink()
|
| 50 |
+
raise
|
| 51 |
+
print('Verified', relative, flush=True)
|
| 52 |
+
|
| 53 |
+
|
| 54 |
+
def main():
|
| 55 |
+
here = Path(__file__).resolve().parent
|
| 56 |
+
p = argparse.ArgumentParser(description=__doc__)
|
| 57 |
+
p.add_argument('--destination', type=Path, default=here)
|
| 58 |
+
p.add_argument('--revision', default='main', help='Use the immutable HF revision in PUBLICATION.md for exact publication')
|
| 59 |
+
p.add_argument('--group', choices=['all', 'data', 'checkpoints'], default='all')
|
| 60 |
+
p.add_argument('--verify-only', action='store_true')
|
| 61 |
+
a = p.parse_args()
|
| 62 |
+
manifest = json.loads((here / 'assets.json').read_text())
|
| 63 |
+
base = 'https://huggingface.co/' + manifest['repository'] + '/resolve/' + urllib.parse.quote(a.revision, safe='') + '/'
|
| 64 |
+
for entry in manifest['files']:
|
| 65 |
+
if a.group != 'all' and entry['group'] != a.group:
|
| 66 |
+
continue
|
| 67 |
+
if a.verify_only:
|
| 68 |
+
path = a.destination / Path(entry['path']).relative_to('training/trackformer_1_2')
|
| 69 |
+
if not path.is_file() or path.stat().st_size != entry['bytes'] or sha(path) != entry['sha256']:
|
| 70 |
+
raise ValueError('Missing or different asset: ' + str(path))
|
| 71 |
+
print('Verified', path, flush=True)
|
| 72 |
+
else:
|
| 73 |
+
download(entry, a.destination, base, 'training/trackformer_1_2')
|
| 74 |
+
|
| 75 |
+
|
| 76 |
+
if __name__ == '__main__':
|
| 77 |
+
main()
|
training/trackformer_1_2/requirements.txt
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
numpy>=1.26
|
| 2 |
+
torch>=2.1
|
| 3 |
+
huggingface_hub>=0.24
|
training/trackformer_1_2/source/data_preparation/requirements.txt
ADDED
|
@@ -0,0 +1,7 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
numpy>=1.26
|
| 2 |
+
torch>=2.1
|
| 3 |
+
xarray
|
| 4 |
+
zarr
|
| 5 |
+
gcsfs
|
| 6 |
+
netCDF4
|
| 7 |
+
scipy
|
training/trackformer_1_2/source/data_preparation/v164_dataset_contract.md
ADDED
|
@@ -0,0 +1,41 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# v1.2.64 Combined-Archive Training
|
| 2 |
+
|
| 3 |
+
## Active Workflow
|
| 4 |
+
|
| 5 |
+
The same `ai_train.ipynb` has two cells: version information, then fetching and preflight. It uses Google Drive only, installs no packages, and does not start training. The trainer remains `v164_global_regional_train.py` for command compatibility. Its active architecture is `wp-coarse8-native-pressure1-v164`, with 53,435,692 parameters. The old 31-channel global experiment and six-window pilot are not used.
|
| 6 |
+
|
| 7 |
+
Run in the Colab terminal after setup reports READY:
|
| 8 |
+
|
| 9 |
+
```sh
|
| 10 |
+
cd /content/v164_colab_bundle
|
| 11 |
+
python v164_global_regional_train.py 20 3 1
|
| 12 |
+
```
|
| 13 |
+
|
| 14 |
+
Numbers mean effective batch, workers, and CUDA AMP. Physical microbatch is one; workers are capped at two. Twelve epochs use rollout horizons 2, 4, 8, 12, then 20 steps. First two epochs train fields; later epochs add route and intensity objectives. Best selection starts at epoch five with the stable full-horizon objective. Validation is never training data; test is excluded from selection.
|
| 15 |
+
|
| 16 |
+
## Data Coverage
|
| 17 |
+
|
| 18 |
+
- Old archives supply 13,949 train, 1,041 validation, and 1,195 test windows with complete six-hour context/target weather coverage inside the Western Pacific domain.
|
| 19 |
+
- Train years 2000-2021; validation 2022-2023; test 2024-2025. Whole storms and window boundaries cannot cross splits. Tip 1979 is excluded from all three.
|
| 20 |
+
- Eight 2.5-degree context variables: mean sea-level pressure, 500-hPa height, and u/v winds at 850/500/200 hPa. Domain: 0-60 N, 100-180 E. These are old NCEP reanalysis fields, not high-resolution ERA5.
|
| 21 |
+
- New native 0.25-degree ERA5 pressure patches provide 800/100/100 detailed windows within those splits. Each patch is 121x121 over 30 degrees, fixed at its observed issue-time center, with nine history and twenty target frames. These 1,000 windows are additional fine supervision, NOT the entire old corpus.
|
| 22 |
+
- Static geographic channels: latitude, longitude, land fraction and elevation. Fine pressure is never synthesized by enlarging the old coarse map.
|
| 23 |
+
- SST, ONI and fine upper-air variables are not consumed by this configuration; they are not silently substituted or claimed present.
|
| 24 |
+
|
| 25 |
+
## Causality And Outputs
|
| 26 |
+
|
| 27 |
+
Inputs are analyses at -48,-42,...,0 hours, static geography, observed issue-time center/intensity and past six-hour motion. Future analysis fields and best-track positions/intensities are labels only. Issue-time retrospective best-track estimates are a historical hindcast assumption, not proof of real-time operational availability. Normalization uses only the training years. No official forecasts or future weather are model inputs.
|
| 28 |
+
|
| 29 |
+
The model autoregressively predicts eight basin fields and a native pressure patch when native history is available. No future recentering or route-painted vortex is used. Its track head uses generated basin pressure/winds/features. Fine data outside the parent domain receives an explicit missing-coverage mask, not extrapolated fields. Central pressure is a sampled generated field plus a separately learned subgrid correction; maximum wind is a scalar head, not a predicted 10m-wind map. No exact physical conservation is claimed.
|
| 30 |
+
|
| 31 |
+
The archive-wide training processes every eligible coarse window. Only the detailed subset trains the fine branch. A case lacking native pressure history can run the coarse branch but is not represented as a genuine high-resolution pressure forecast.
|
| 32 |
+
|
| 33 |
+
## Storage And Resume
|
| 34 |
+
|
| 35 |
+
Drive folder: `typhoon_predict/data/v164_reuse` contains the manifest and bundle descriptor referencing existing archived data plus `data/v164_pressure_geography` parts. No large archives are duplicated in Drive. Total payload about 2.13 GB. Four independent Drive HTTP clients fetch parts concurrently. Existing verified files are reused; the manifest is published after complete-file checksums pass.
|
| 36 |
+
|
| 37 |
+
Checkpoints are written directly to `MyDrive/typhoon_predict/training/v164_reuse_pressure/last.pt` and `best.pt`, with history and run manifest. Resume restores the last completed epoch, optimizer, scaler, and RNG; an interrupted epoch restarts. Architecture and dataset hashes must match. A stale lock after runtime death requires checking that no other training process is active before removing it. Pilot checkpoints are incompatible and are not auto-loaded.
|
| 38 |
+
|
| 39 |
+
## Verification Limits
|
| 40 |
+
|
| 41 |
+
Mac checks cover all archive hashes and split counts, real coarse and fine samples, tiny-model +120h finite rollouts, backward gradients, completed-epoch checkpoint resume, and full-width two-step MPS gradients. These are software checks, not model-skill validation. CUDA/T4 execution and a complete training run have not been tested; Colab was not accessed.
|
training/trackformer_1_2/source/data_preparation/v164_download_pressure_geography.py
ADDED
|
@@ -0,0 +1,174 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Mac-only, resumable MSLP patches and geography. Not a full trainer dataset."""
|
| 2 |
+
import argparse
|
| 3 |
+
from collections import defaultdict
|
| 4 |
+
from concurrent.futures import ThreadPoolExecutor
|
| 5 |
+
import hashlib
|
| 6 |
+
import json
|
| 7 |
+
import os
|
| 8 |
+
from pathlib import Path
|
| 9 |
+
import sqlite3
|
| 10 |
+
import time
|
| 11 |
+
|
| 12 |
+
import numpy as np
|
| 13 |
+
import v164_prepare_data as prep
|
| 14 |
+
|
| 15 |
+
|
| 16 |
+
def make_plan(track):
|
| 17 |
+
if prep.trainer.digest(track) != prep.TRACK_SHA:
|
| 18 |
+
raise ValueError('Unexpected track archive')
|
| 19 |
+
rng = np.random.default_rng(164)
|
| 20 |
+
with np.load(track, allow_pickle=False) as z:
|
| 21 |
+
times, ids = z['base_time'], z['storm_id'].astype(str)
|
| 22 |
+
lat, lon = z['base_lat'], z['base_lon'] % 360
|
| 23 |
+
first = {}
|
| 24 |
+
for sid, ns in zip(ids, times):
|
| 25 |
+
first[sid] = min(first.get(sid, 9999), prep.year(ns))
|
| 26 |
+
rows = []
|
| 27 |
+
for split, lo, hi, count in [('train', 2000, 2021, 800), ('validation', 2022, 2023, 100), ('test', 2024, 2025, 100)]:
|
| 28 |
+
groups = defaultdict(list)
|
| 29 |
+
for i, (sid, ns) in enumerate(zip(ids, times)):
|
| 30 |
+
ns = int(ns)
|
| 31 |
+
if not (lo <= first[sid] <= hi and lo <= prep.year(ns-8*prep.STEP)
|
| 32 |
+
and prep.year(ns+20*prep.STEP) <= hi and ns % prep.STEP == 0
|
| 33 |
+
and np.isfinite(lat[i]) and np.isfinite(lon[i])
|
| 34 |
+
and 0 < lat[i] < 60 and 100 < lon[i] < 180):
|
| 35 |
+
continue
|
| 36 |
+
groups[sid].append(i)
|
| 37 |
+
for group in groups.values():
|
| 38 |
+
rng.shuffle(group)
|
| 39 |
+
storms = sorted(groups)
|
| 40 |
+
rng.shuffle(storms)
|
| 41 |
+
chosen = []
|
| 42 |
+
while len(chosen) < count:
|
| 43 |
+
before = len(chosen)
|
| 44 |
+
for sid in storms:
|
| 45 |
+
if groups[sid]:
|
| 46 |
+
chosen.append(groups[sid].pop())
|
| 47 |
+
if len(chosen) == count:
|
| 48 |
+
break
|
| 49 |
+
if len(chosen) == before:
|
| 50 |
+
raise ValueError(f'Insufficient eligible {split} windows')
|
| 51 |
+
for i in chosen:
|
| 52 |
+
rows.append({'storm_id': str(ids[i]), 'issue_ns': int(times[i]), 'split': split,
|
| 53 |
+
'center_lat': float(np.round(lat[i]*4)/4), 'center_lon': float(np.round(lon[i]*4)/4),
|
| 54 |
+
'center_source': 'retrospective issue-time best-track observation; no future positions',
|
| 55 |
+
'track_archive_row': int(i)})
|
| 56 |
+
plan = {'schema': 'v164-pressure-geography-v1', 'dataset_role': 'partial_inputs_not_training_ready',
|
| 57 |
+
'source': prep.SOURCE, 'track_sha256': prep.TRACK_SHA, 'rows': rows,
|
| 58 |
+
'lead_hours': list(range(-48, 121, 6)), 'shape': [len(rows), 29, 121, 121],
|
| 59 |
+
'units': 'hPa', 'resolution_degrees': .25, 'fixed_issue_time_crops': True,
|
| 60 |
+
'forecast_inputs': 'lead_hours <= 0 only', 'targets': 'lead_hours > 0 only',
|
| 61 |
+
'split_policy': 'existing 2000-2021/2022-2023/2024-2025 policy, pending separately requested changes',
|
| 62 |
+
'hindcast_not_operational_availability_claim': True}
|
| 63 |
+
plan['sha256'] = hashlib.sha256(json.dumps(plan, sort_keys=True).encode()).hexdigest()
|
| 64 |
+
return plan
|
| 65 |
+
|
| 66 |
+
|
| 67 |
+
def patch_indices(row):
|
| 68 |
+
lat = row['center_lat'] + 15 - np.arange(121)*.25
|
| 69 |
+
lon = row['center_lon'] - 15 + np.arange(121)*.25
|
| 70 |
+
return np.rint((90-lat)*4).astype(int), np.rint(lon*4).astype(int) % 1440
|
| 71 |
+
|
| 72 |
+
|
| 73 |
+
def patch_hash(a):
|
| 74 |
+
return hashlib.sha256(np.ascontiguousarray(a).tobytes()).hexdigest()
|
| 75 |
+
|
| 76 |
+
|
| 77 |
+
def run(root, track, workers, max_frames=None):
|
| 78 |
+
root.mkdir(parents=True, exist_ok=True)
|
| 79 |
+
fd = os.open(root/'download.lock', os.O_CREAT | os.O_EXCL | os.O_WRONLY)
|
| 80 |
+
os.write(fd, str(os.getpid()).encode())
|
| 81 |
+
os.close(fd)
|
| 82 |
+
db = None
|
| 83 |
+
try:
|
| 84 |
+
plan = make_plan(track)
|
| 85 |
+
path = root/'plan.json'
|
| 86 |
+
if path.exists() and json.loads(path.read_text()) != plan:
|
| 87 |
+
raise ValueError('Existing download plan differs; refusing overwrite')
|
| 88 |
+
prep.trainer.atomic_json(path, plan)
|
| 89 |
+
shape = tuple(plan['shape'])
|
| 90 |
+
output = root/'pressure_hpa.npy'
|
| 91 |
+
a = np.lib.format.open_memmap(output, mode='r+' if output.exists() else 'w+', dtype='float32', shape=shape)
|
| 92 |
+
if a.shape != shape or a.dtype != np.float32:
|
| 93 |
+
raise ValueError('Existing array shape/dtype differs')
|
| 94 |
+
indices = [patch_indices(r) for r in plan['rows']]
|
| 95 |
+
requests = defaultdict(list)
|
| 96 |
+
for i, row in enumerate(plan['rows']):
|
| 97 |
+
for j in range(29):
|
| 98 |
+
requests[row['issue_ns'] + (j-8)*prep.STEP].append((i, j))
|
| 99 |
+
db = sqlite3.connect(root/'progress.sqlite')
|
| 100 |
+
db.execute('CREATE TABLE IF NOT EXISTS patches (i INTEGER, j INTEGER, sha TEXT, PRIMARY KEY(i,j))')
|
| 101 |
+
verified = set()
|
| 102 |
+
for i, j, sha in db.execute('SELECT i,j,sha FROM patches'):
|
| 103 |
+
if patch_hash(a[i,j]) == sha:
|
| 104 |
+
verified.add((i,j))
|
| 105 |
+
pending = [(ns, slots) for ns, slots in sorted(requests.items()) if any(s not in verified for s in slots)]
|
| 106 |
+
print(json.dumps({'windows': len(plan['rows']), 'pressure_gib': a.nbytes/1024**3,
|
| 107 |
+
'unique_times': len(requests), 'remaining_times': len(pending),
|
| 108 |
+
'reused_verified_patches': len(verified), 'workers': workers}), flush=True)
|
| 109 |
+
source = prep.Era5Source()
|
| 110 |
+
geography = root/'geography.npz'
|
| 111 |
+
if not geography.exists():
|
| 112 |
+
ns = plan['rows'][0]['issue_ns']
|
| 113 |
+
land = source.read('land_sea_mask', ns)
|
| 114 |
+
height = source.read('geopotential_at_surface', ns)/9.80665
|
| 115 |
+
if not np.isfinite(land).all() or not np.isfinite(height).all():
|
| 116 |
+
raise ValueError('Invalid geography')
|
| 117 |
+
with (root/'geography.pending').open('wb') as f:
|
| 118 |
+
np.savez_compressed(f, latitude=source.lat, longitude=source.lon,
|
| 119 |
+
land_fraction=land, elevation_m=height)
|
| 120 |
+
(root/'geography.pending').replace(geography)
|
| 121 |
+
with np.load(geography, allow_pickle=False) as z:
|
| 122 |
+
if z['land_fraction'].shape != (721,1440) or z['elevation_m'].shape != (721,1440):
|
| 123 |
+
raise ValueError('Geography grid mismatch')
|
| 124 |
+
prep.trainer.atomic_json(root/'geography_manifest.json', {'sha256': prep.trainer.digest(geography),
|
| 125 |
+
'source': prep.SOURCE, 'size': geography.stat().st_size, 'units': {'elevation_m': 'm', 'land_fraction': '1'}})
|
| 126 |
+
def fetch(item):
|
| 127 |
+
ns, slots = item
|
| 128 |
+
pressure = source.read('mean_sea_level_pressure', ns)/100
|
| 129 |
+
if not np.isfinite(pressure).all() or pressure.min() < 800 or pressure.max() > 1120:
|
| 130 |
+
raise ValueError(f'Invalid MSLP at {ns}')
|
| 131 |
+
return ns, [(i,j,pressure[indices[i][0][:,None],indices[i][1][None,:]]) for i,j in slots if (i,j) not in verified]
|
| 132 |
+
started = time.monotonic()
|
| 133 |
+
initial = len(requests)-len(pending)
|
| 134 |
+
selected = pending if max_frames is None else pending[:max_frames]
|
| 135 |
+
# Process bounded groups: completed patches are committed before requesting more.
|
| 136 |
+
done = 0
|
| 137 |
+
with ThreadPoolExecutor(max_workers=workers) as pool:
|
| 138 |
+
for offset in range(0, len(selected), workers):
|
| 139 |
+
for ns, patches in pool.map(fetch, selected[offset:offset+workers]):
|
| 140 |
+
for i,j,values in patches:
|
| 141 |
+
a[i,j] = values
|
| 142 |
+
a.flush()
|
| 143 |
+
for i,j,_ in patches:
|
| 144 |
+
db.execute('INSERT OR REPLACE INTO patches VALUES (?,?,?)', (i,j,patch_hash(a[i,j])))
|
| 145 |
+
verified.add((i,j))
|
| 146 |
+
db.commit()
|
| 147 |
+
done += 1
|
| 148 |
+
elapsed = time.monotonic()-started
|
| 149 |
+
status = {'completed_times': initial+done, 'total_times': len(requests),
|
| 150 |
+
'verified_patches': len(verified), 'total_patches': shape[0]*shape[1],
|
| 151 |
+
'elapsed_s_this_run': round(elapsed,1),
|
| 152 |
+
'eta_s': round((len(pending)-done)*elapsed/done),
|
| 153 |
+
'last_time': str(np.datetime64(ns,'ns')), 'status': 'downloading'}
|
| 154 |
+
prep.trainer.atomic_json(root/'progress.json', status)
|
| 155 |
+
print(json.dumps(status), flush=True)
|
| 156 |
+
if len(verified) == shape[0]*shape[1]:
|
| 157 |
+
prep.trainer.atomic_json(root/'completion.json', {'status': 'pressure_and_geography_complete',
|
| 158 |
+
'training_ready': False, 'plan_sha256': plan['sha256'],
|
| 159 |
+
'pressure_sha256': prep.trainer.digest(output), 'geography_sha256': prep.trainer.digest(geography)})
|
| 160 |
+
print('Pressure and geography verified complete. Other training inputs still required.', flush=True)
|
| 161 |
+
finally:
|
| 162 |
+
if db is not None:
|
| 163 |
+
db.close()
|
| 164 |
+
(root/'download.lock').unlink(missing_ok=True)
|
| 165 |
+
|
| 166 |
+
|
| 167 |
+
if __name__ == '__main__':
|
| 168 |
+
p = argparse.ArgumentParser(description=__doc__)
|
| 169 |
+
p.add_argument('--output', type=Path, default=Path('/Volumes/D/typhoon_predict/data/v164_pressure_geography'))
|
| 170 |
+
p.add_argument('--track', type=Path, default=Path('/Volumes/D/typhoon_predict/track_build/track_windows_v13.npz'))
|
| 171 |
+
p.add_argument('--workers', type=int, choices=range(1,9), default=4)
|
| 172 |
+
p.add_argument('--max-frames', type=int)
|
| 173 |
+
args = p.parse_args()
|
| 174 |
+
run(args.output, args.track, args.workers, args.max_frames)
|
training/trackformer_1_2/source/data_preparation/v164_global_regional_train.py
ADDED
|
@@ -0,0 +1,845 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#!/usr/bin/env python3
|
| 2 |
+
"""v1.2.64 Western Pacific coarse-context/native-pressure forecaster.
|
| 3 |
+
|
| 4 |
+
Run: python v164_global_regional_train.py 10 3 1
|
| 5 |
+
Uses verified old coarse archives plus new native 0.25-degree pressure patches.
|
| 6 |
+
Predicts eight coarse fields and one fine pressure field, not 31 fine variables.
|
| 7 |
+
The legacy 31-channel definitions below are not the active CLI architecture.
|
| 8 |
+
No route-centered vortex is painted. This is not a numerical conservation solver.
|
| 9 |
+
"""
|
| 10 |
+
from __future__ import annotations
|
| 11 |
+
|
| 12 |
+
import argparse
|
| 13 |
+
import hashlib
|
| 14 |
+
import json
|
| 15 |
+
import math
|
| 16 |
+
import os
|
| 17 |
+
from pathlib import Path
|
| 18 |
+
import random
|
| 19 |
+
import socket
|
| 20 |
+
import time
|
| 21 |
+
from contextlib import nullcontext
|
| 22 |
+
|
| 23 |
+
import numpy as np
|
| 24 |
+
import torch
|
| 25 |
+
from torch import nn
|
| 26 |
+
from torch.nn import functional as F
|
| 27 |
+
from torch.utils.checkpoint import checkpoint
|
| 28 |
+
from torch.utils.data import Dataset, DataLoader
|
| 29 |
+
|
| 30 |
+
VERSION = "1.2.64"
|
| 31 |
+
ARCHITECTURE = "global-regional-field-first-v1"
|
| 32 |
+
ROOT = Path(__file__).resolve().parent
|
| 33 |
+
LEVELS = [850, 700, 500, 300, 200]
|
| 34 |
+
CHANNELS = ["mslp", "surface_pressure", "u10", "v10", "t2m", "sst"] + [
|
| 35 |
+
f"{v}{p}" for p in LEVELS for v in ("u", "v", "temperature", "humidity", "geopotential")]
|
| 36 |
+
UNITS = ["hPa", "hPa", "m/s", "m/s", "K", "K"] + [
|
| 37 |
+
unit for _ in LEVELS for unit in ("m/s", "m/s", "K", "kg/kg", "m2/s2")]
|
| 38 |
+
HISTORY = 9
|
| 39 |
+
LEADS = 20
|
| 40 |
+
HOUR = 3_600_000_000_000
|
| 41 |
+
|
| 42 |
+
|
| 43 |
+
def digest(path):
|
| 44 |
+
h = hashlib.sha256()
|
| 45 |
+
with Path(path).open("rb") as f:
|
| 46 |
+
for block in iter(lambda: f.read(8 * 1024**2), b""):
|
| 47 |
+
h.update(block)
|
| 48 |
+
return h.hexdigest()
|
| 49 |
+
|
| 50 |
+
|
| 51 |
+
def atomic_json(path, data):
|
| 52 |
+
path = Path(path)
|
| 53 |
+
tmp = path.with_suffix(path.suffix + ".tmp")
|
| 54 |
+
tmp.write_text(json.dumps(data, indent=2, allow_nan=False) + "\n")
|
| 55 |
+
tmp.replace(path)
|
| 56 |
+
|
| 57 |
+
|
| 58 |
+
def atomic_checkpoint(path, data):
|
| 59 |
+
path = Path(path)
|
| 60 |
+
tmp = path.with_suffix(".tmp")
|
| 61 |
+
with tmp.open("wb") as f:
|
| 62 |
+
torch.save(data, f)
|
| 63 |
+
f.flush()
|
| 64 |
+
os.fsync(f.fileno())
|
| 65 |
+
tmp.replace(path)
|
| 66 |
+
|
| 67 |
+
|
| 68 |
+
def device_for(name):
|
| 69 |
+
if name == "auto":
|
| 70 |
+
name = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
|
| 71 |
+
if name == "cuda" and not torch.cuda.is_available():
|
| 72 |
+
raise RuntimeError("CUDA unavailable in this Python. Install a CUDA PyTorch wheel before training.")
|
| 73 |
+
if name == "mps" and not torch.backends.mps.is_available():
|
| 74 |
+
raise RuntimeError("MPS requested but unavailable")
|
| 75 |
+
return torch.device(name)
|
| 76 |
+
|
| 77 |
+
|
| 78 |
+
class Block(nn.Module):
|
| 79 |
+
def __init__(self, width, periodic=False):
|
| 80 |
+
super().__init__()
|
| 81 |
+
self.periodic = periodic
|
| 82 |
+
self.depthwise = nn.Conv2d(width, width, 5, groups=width)
|
| 83 |
+
self.net = nn.Sequential(nn.GroupNorm(8, width), nn.Conv2d(width, width * 4, 1),
|
| 84 |
+
nn.GELU(), nn.Conv2d(width * 4, width, 1))
|
| 85 |
+
self.scale = nn.Parameter(torch.full((1, width, 1, 1), 0.01))
|
| 86 |
+
|
| 87 |
+
def forward(self, x):
|
| 88 |
+
padded = F.pad(x, (2, 2, 0, 0), mode="circular" if self.periodic else "replicate")
|
| 89 |
+
padded = F.pad(padded, (0, 0, 2, 2), mode="replicate")
|
| 90 |
+
return x + self.scale * self.net(self.depthwise(padded))
|
| 91 |
+
|
| 92 |
+
|
| 93 |
+
class FieldOperator(nn.Module):
|
| 94 |
+
def __init__(self, channels, base, periodic, attention_layers=0, tiny=False):
|
| 95 |
+
super().__init__()
|
| 96 |
+
widths = [base, base * 2, base * 4, base * 6]
|
| 97 |
+
self.stem = nn.Conv2d(channels * 3 + 4, base, 3, padding=1)
|
| 98 |
+
depths = [1, 1, 1, 1] if tiny else [2, 3, 5, 6]
|
| 99 |
+
self.enc = nn.ModuleList([nn.Sequential(*[Block(w, periodic) for _ in range(n)]) for w, n in zip(widths, depths)])
|
| 100 |
+
self.down = nn.ModuleList([nn.Conv2d(a, b, 3, stride=2, padding=1) for a, b in zip(widths, widths[1:])])
|
| 101 |
+
self.attention = nn.ModuleList([nn.TransformerEncoderLayer(widths[-1], 8,
|
| 102 |
+
dim_feedforward=widths[-1] * 4, batch_first=True,
|
| 103 |
+
dropout=0.0, norm_first=True) for _ in range(attention_layers)])
|
| 104 |
+
self.up = nn.ModuleList([nn.Conv2d(widths[i + 1], widths[i], 1) for i in (2, 1, 0)])
|
| 105 |
+
self.dec = nn.ModuleList([nn.Sequential(Block(widths[i], periodic), Block(widths[i], periodic)) for i in (2, 1, 0)])
|
| 106 |
+
self.out = nn.Conv2d(base, channels, 1)
|
| 107 |
+
self.object_features = nn.Conv2d(base, 32, 1)
|
| 108 |
+
self.checkpointing = True
|
| 109 |
+
nn.init.normal_(self.out.weight, std=1e-3)
|
| 110 |
+
nn.init.zeros_(self.out.bias)
|
| 111 |
+
|
| 112 |
+
def run(self, module, value):
|
| 113 |
+
return checkpoint(module, value, use_reentrant=False) if self.training and self.checkpointing else module(value)
|
| 114 |
+
|
| 115 |
+
def forward(self, state, context, mask, static):
|
| 116 |
+
x = self.stem(torch.cat([state, context, mask, static], 1))
|
| 117 |
+
skips = []
|
| 118 |
+
for i in range(4):
|
| 119 |
+
x = self.run(self.enc[i], x)
|
| 120 |
+
skips.append(x)
|
| 121 |
+
if i < 3:
|
| 122 |
+
x = self.down[i](x)
|
| 123 |
+
if self.attention:
|
| 124 |
+
shape = x.shape
|
| 125 |
+
tokens = x.flatten(2).transpose(1, 2)
|
| 126 |
+
for layer in self.attention:
|
| 127 |
+
tokens = self.run(layer, tokens)
|
| 128 |
+
x = tokens.transpose(1, 2).reshape(shape)
|
| 129 |
+
for up, dec, skip in zip(self.up, self.dec, reversed(skips[:-1])):
|
| 130 |
+
x = self.run(dec, F.interpolate(up(x), size=skip.shape[-2:], mode="bilinear", align_corners=False) + skip)
|
| 131 |
+
# A learned field tendency, with no fixed per-cell hPa cap or route warp.
|
| 132 |
+
return state + self.out(x), self.object_features(x)
|
| 133 |
+
|
| 134 |
+
|
| 135 |
+
class GlobalRegionalForecaster(nn.Module):
|
| 136 |
+
def __init__(self, global_lat, global_lon, regional_lat, regional_lon, mean, std, tiny=False):
|
| 137 |
+
super().__init__()
|
| 138 |
+
c = len(CHANNELS)
|
| 139 |
+
self.register_buffer("mean", torch.tensor(mean, dtype=torch.float32).view(1, c, 1, 1))
|
| 140 |
+
self.register_buffer("std", torch.tensor(std, dtype=torch.float32).view(1, c, 1, 1))
|
| 141 |
+
self.register_buffer("regional_lat", torch.tensor(regional_lat, dtype=torch.float32))
|
| 142 |
+
self.register_buffer("regional_lon", torch.tensor(regional_lon, dtype=torch.float32))
|
| 143 |
+
self.global_time = nn.Conv2d(c * HISTORY, c, 1)
|
| 144 |
+
self.regional_time = nn.Conv2d(c * HISTORY, c, 1)
|
| 145 |
+
self.global_operator = FieldOperator(c, 8 if tiny else 80, True, 1 if tiny else 3, tiny)
|
| 146 |
+
self.regional_operator = FieldOperator(c, 8 if tiny else 88, False, 0, tiny)
|
| 147 |
+
self.parent_fusion = nn.Conv2d(c * 2, c, 1)
|
| 148 |
+
self.object_score = nn.Conv2d(32 + 3, 1, 1)
|
| 149 |
+
self.intensity = nn.Sequential(nn.Linear(32 + 3, 64), nn.GELU(), nn.Linear(64, 2))
|
| 150 |
+
y, x = np.meshgrid(regional_lat, regional_lon, indexing="ij")
|
| 151 |
+
grid = np.stack([2 * (x - global_lon[0]) / (global_lon[-1] - global_lon[0]) - 1,
|
| 152 |
+
2 * (y - global_lat[0]) / (global_lat[-1] - global_lat[0]) - 1], -1)
|
| 153 |
+
self.register_buffer("parent_grid", torch.tensor(grid, dtype=torch.float32)[None])
|
| 154 |
+
self.register_buffer("yy", torch.tensor(y, dtype=torch.float32)[None])
|
| 155 |
+
self.register_buffer("xx", torch.tensor(x, dtype=torch.float32)[None])
|
| 156 |
+
|
| 157 |
+
def initial(self, inputs):
|
| 158 |
+
expected = {"global_history", "regional_history", "global_mask", "regional_mask",
|
| 159 |
+
"global_static", "regional_static", "center", "motion"}
|
| 160 |
+
if set(inputs) != expected:
|
| 161 |
+
raise ValueError("Forward inputs must contain only the declared causal input keys")
|
| 162 |
+
gh, rh = inputs["global_history"], inputs["regional_history"]
|
| 163 |
+
return {"g": gh[:, -1], "r": rh[:, -1],
|
| 164 |
+
"gc": self.global_time(gh.flatten(1, 2)), "rc": self.regional_time(rh.flatten(1, 2)),
|
| 165 |
+
"gm": inputs["global_mask"], "rm": inputs["regional_mask"],
|
| 166 |
+
"gs": inputs["global_static"], "rs": inputs["regional_static"],
|
| 167 |
+
"center": inputs["center"], "motion": inputs["motion"]}
|
| 168 |
+
|
| 169 |
+
def step(self, state):
|
| 170 |
+
g, _ = self.global_operator(state["g"], state["gc"], state["gm"], state["gs"])
|
| 171 |
+
parent = F.grid_sample(g, self.parent_grid.expand(len(g), -1, -1, -1), align_corners=True)
|
| 172 |
+
context = self.parent_fusion(torch.cat([state["rc"], parent], 1))
|
| 173 |
+
r, features = self.regional_operator(state["r"], context, state["rm"], state["rs"])
|
| 174 |
+
physical = r.float() * self.std + self.mean
|
| 175 |
+
pressure = physical[:, 0]
|
| 176 |
+
anomaly = pressure - F.avg_pool2d(F.pad(pressure[:, None], (4, 4, 4, 4), mode="replicate"), 9, stride=1)[:, 0]
|
| 177 |
+
idx_u, idx_v = CHANNELS.index("u850"), CHANNELS.index("v850")
|
| 178 |
+
u, v = physical[:, idx_u], physical[:, idx_v]
|
| 179 |
+
evidence = torch.cat([features.float(), anomaly[:, None] / 10, u[:, None] / 30, v[:, None] / 30], 1)
|
| 180 |
+
score = self.object_score(evidence)[:, 0] - anomaly / 10
|
| 181 |
+
lat0, lon0 = state["center"][:, 0, None, None], state["center"][:, 1, None, None]
|
| 182 |
+
prior_lat = lat0 + state["motion"][:, 1, None, None] / 111.2
|
| 183 |
+
prior_lon = lon0 + state["motion"][:, 0, None, None] / (111.2 * lat0.deg2rad().cos().clamp_min(.2))
|
| 184 |
+
distance2 = ((self.yy - prior_lat) * 111.2).square() + ((self.xx - prior_lon) * 111.2 * lat0.deg2rad().cos()).square()
|
| 185 |
+
score = score - distance2 / (2 * 600**2)
|
| 186 |
+
# Select one spatial mode, then refine locally; never average distinct lows.
|
| 187 |
+
peak = score.flatten(1).argmax(1)
|
| 188 |
+
h, w = pressure.shape[-2:]
|
| 189 |
+
iy = torch.arange(h, device=r.device)[None, :, None]
|
| 190 |
+
ix = torch.arange(w, device=r.device)[None, None, :]
|
| 191 |
+
local = ((iy - (peak // w)[:, None, None]).abs() <= 2) & ((ix - (peak % w)[:, None, None]).abs() <= 2)
|
| 192 |
+
pdf = torch.softmax(score.masked_fill(~local, -1e9).flatten(1), 1).reshape_as(score)
|
| 193 |
+
center = torch.stack([(pdf * self.yy).sum((1, 2)), (pdf * self.xx).sum((1, 2))], 1)
|
| 194 |
+
pooled = (evidence * pdf[:, None]).sum((2, 3))
|
| 195 |
+
intensity = self.intensity(pooled)
|
| 196 |
+
grid_pressure = (pressure * pdf).sum((1, 2))
|
| 197 |
+
central_pressure = grid_pressure - F.softplus(intensity[:, 0]) * 5
|
| 198 |
+
surface_speed = torch.sqrt(physical[:, 2].square() + physical[:, 3].square() + 1e-6)
|
| 199 |
+
vmax = (surface_speed * pdf).sum((1, 2)) * 1.943844 + F.softplus(intensity[:, 1]) * 5
|
| 200 |
+
mean_lat = (center[:, 0] + state["center"][:, 0]) / 2
|
| 201 |
+
motion = torch.stack([(center[:, 1] - state["center"][:, 1]) * 111.2 * mean_lat.deg2rad().cos(),
|
| 202 |
+
(center[:, 0] - state["center"][:, 0]) * 111.2], 1)
|
| 203 |
+
next_state = dict(state, g=g, r=r, center=center, motion=motion)
|
| 204 |
+
return next_state, {"global": g, "regional": r, "center": center, "pressure": central_pressure,
|
| 205 |
+
"vmax": vmax, "grid_pressure": grid_pressure, "association": pdf}
|
| 206 |
+
|
| 207 |
+
def forward(self, inputs, steps=20):
|
| 208 |
+
if not 1 <= steps <= LEADS:
|
| 209 |
+
raise ValueError("steps must be 1..20")
|
| 210 |
+
state = self.initial(inputs)
|
| 211 |
+
outputs = []
|
| 212 |
+
for _ in range(steps):
|
| 213 |
+
state, out = self.step(state)
|
| 214 |
+
outputs.append(out)
|
| 215 |
+
return {k: torch.stack([o[k] for o in outputs], 1) for k in outputs[0]}
|
| 216 |
+
|
| 217 |
+
|
| 218 |
+
def require_full_dataset(m, times=None, rows=None):
|
| 219 |
+
if m.get('dataset_role') != 'full_training' or m.get('synthetic', False):
|
| 220 |
+
raise ValueError('FULL DATA REQUIRED: pilot, synthetic or unclassified datasets cannot be used by the training CLI. Prepare the full 2000-2025 dataset on the Mac first.')
|
| 221 |
+
if times is not None:
|
| 222 |
+
expected = np.arange(np.datetime64('2000-01-01', 'ns').astype('int64'),
|
| 223 |
+
np.datetime64('2026-01-01', 'ns').astype('int64'), 6*HOUR, dtype='int64')
|
| 224 |
+
if not np.array_equal(times, expected):
|
| 225 |
+
raise ValueError('Full dataset must contain every six-hour field from 2000 through 2025; partial coverage is not full training')
|
| 226 |
+
if rows is not None:
|
| 227 |
+
for split, years in [('train', range(2000, 2022)), ('validation', range(2022, 2024)), ('test', range(2024, 2026))]:
|
| 228 |
+
covered = {int(np.datetime64(int(times[r['issue_index']]), 'ns').astype('datetime64[Y]').astype(int))+1970
|
| 229 |
+
for r in rows if r['split'] == split}
|
| 230 |
+
if covered != set(years):
|
| 231 |
+
raise ValueError(f'Full dataset lacks sample coverage in {split}: {sorted(set(years)-covered)}')
|
| 232 |
+
|
| 233 |
+
|
| 234 |
+
def load_manifest(path, verify=True, smoke=False, require_full=False):
|
| 235 |
+
path = Path(path)
|
| 236 |
+
if not path.is_file():
|
| 237 |
+
raise FileNotFoundError(f"V164 DATA NOT READY: {path}. Requires native global <=1-degree and regional <=0.25-degree multilevel fields. The v163 2.5-degree archive cannot substitute. See v164_dataset_contract.md.")
|
| 238 |
+
m = json.loads(path.read_text())
|
| 239 |
+
if require_full:
|
| 240 |
+
require_full_dataset(m)
|
| 241 |
+
if m.get("schema") != "v164-atmosphere-v1" or m.get("channels") != CHANNELS or m.get("units") != UNITS:
|
| 242 |
+
raise ValueError("Dataset schema/channel order/physical units do not match v164")
|
| 243 |
+
if m.get("future_weather_as_inputs") is not False or m.get("official_forecasts_as_inputs") is not False:
|
| 244 |
+
raise ValueError("Dataset must explicitly disable future weather and official forecast inputs")
|
| 245 |
+
if not smoke and m.get("synthetic", False):
|
| 246 |
+
raise ValueError("Synthetic data is only permitted by the smoke-test API")
|
| 247 |
+
for key in ("global", "regional", "global_static", "regional_static", "time", "samples"):
|
| 248 |
+
item = m["files"][key]
|
| 249 |
+
p = (path.parent / item["path"]).resolve()
|
| 250 |
+
if not p.is_relative_to(path.parent.resolve()) or not p.is_file():
|
| 251 |
+
raise ValueError(f"Missing or outside-dataset file: {key}")
|
| 252 |
+
if verify and digest(p) != item["sha256"]:
|
| 253 |
+
raise ValueError(f"Dataset checksum mismatch: {key}")
|
| 254 |
+
item["resolved"] = str(p)
|
| 255 |
+
times = np.load(m["files"]["time"]["resolved"], mmap_mode="r")
|
| 256 |
+
if times.dtype != np.dtype("int64") or np.any(np.diff(times) <= 0) or np.any(times % (6 * HOUR)):
|
| 257 |
+
raise ValueError("Time must be increasing int64 UTC nanoseconds aligned to six hours")
|
| 258 |
+
if m["normalization"]["fit_split"] != "train" or m["normalization"]["end_year"] > 2021:
|
| 259 |
+
raise ValueError("Normalization must be fitted exclusively to training years <=2021")
|
| 260 |
+
mean, std = np.array(m["normalization"]["mean"]), np.array(m["normalization"]["std"])
|
| 261 |
+
if mean.shape != (len(CHANNELS),) or std.shape != mean.shape or not np.isfinite(mean).all() or not np.isfinite(std).all() or (std <= 0).any():
|
| 262 |
+
raise ValueError("Invalid training normalization")
|
| 263 |
+
for region, max_step in (("global", 1), ("regional", .25)):
|
| 264 |
+
lat, lon = np.array(m[region + "_lat"]), np.array(m[region + "_lon"])
|
| 265 |
+
if min(len(lat), len(lon)) < 2 or np.any(np.diff(lat) >= 0) or np.any(np.diff(lon) <= 0):
|
| 266 |
+
raise ValueError("Coordinates must be north-to-south, west-to-east")
|
| 267 |
+
if not np.allclose(np.diff(lat), lat[1]-lat[0], atol=1e-5) or not np.allclose(np.diff(lon), lon[1]-lon[0], atol=1e-5):
|
| 268 |
+
raise ValueError("This operator requires regular latitude/longitude grids")
|
| 269 |
+
if not smoke and (np.max(-np.diff(lat)) > max_step + 1e-5 or np.max(np.diff(lon)) > max_step + 1e-5):
|
| 270 |
+
raise ValueError(f"{region} resolution too coarse")
|
| 271 |
+
if region == "global" and not smoke and (lat[0] < 89 or lat[-1] > -89 or lon[0] != 0 or lon[-1] < 359):
|
| 272 |
+
raise ValueError("Global parent must cover the globe without duplicated longitude seam")
|
| 273 |
+
array = np.load(m["files"][region]["resolved"], mmap_mode="r")
|
| 274 |
+
static = np.load(m["files"][region + "_static"]["resolved"], mmap_mode="r")
|
| 275 |
+
if array.shape != (len(times), len(CHANNELS), len(lat), len(lon)) or static.shape != (4, len(lat), len(lon)):
|
| 276 |
+
raise ValueError(f"{region} field/static shape mismatch")
|
| 277 |
+
if not np.isfinite(static).all():
|
| 278 |
+
raise ValueError("Non-finite static geography")
|
| 279 |
+
if not (min(m["global_lat"]) <= min(m["regional_lat"]) <= max(m["regional_lat"]) <= max(m["global_lat"]) and
|
| 280 |
+
min(m["global_lon"]) <= min(m["regional_lon"]) <= max(m["regional_lon"]) <= max(m["global_lon"])):
|
| 281 |
+
raise ValueError("Regional grid lies outside global parent coverage")
|
| 282 |
+
rows = [json.loads(line) for line in Path(m["files"]["samples"]["resolved"]).read_text().splitlines() if line.strip()]
|
| 283 |
+
membership = {}
|
| 284 |
+
for row in rows:
|
| 285 |
+
i = row["issue_index"]
|
| 286 |
+
if i < HISTORY - 1 or i + LEADS >= len(times):
|
| 287 |
+
raise ValueError("Sample lacks full causal history or target coverage")
|
| 288 |
+
if np.any(np.diff(times[i-8:i+21]) != 6 * HOUR):
|
| 289 |
+
raise ValueError("Sample crosses a gap in causal history or future target coverage")
|
| 290 |
+
year = int(np.datetime64(int(times[i]), "ns").astype("datetime64[Y]").astype(int)) + 1970
|
| 291 |
+
valid_years = {"train": (2000, 2021), "validation": (2022, 2023), "test": (2024, 2025)}
|
| 292 |
+
low, high = valid_years[row["split"]]
|
| 293 |
+
first_year = int(np.datetime64(int(times[i-8]), "ns").astype("datetime64[Y]").astype(int)) + 1970
|
| 294 |
+
last_year = int(np.datetime64(int(times[i+20]), "ns").astype("datetime64[Y]").astype(int)) + 1970
|
| 295 |
+
if not smoke and not (low <= first_year <= year <= last_year <= high):
|
| 296 |
+
raise ValueError("Sample violates fixed chronological split")
|
| 297 |
+
storm = row.get("storm_id")
|
| 298 |
+
if storm:
|
| 299 |
+
if "1979" in storm or storm.upper() == "TIP":
|
| 300 |
+
raise ValueError("Tip is excluded from training, validation and test")
|
| 301 |
+
if storm in membership and membership[storm] != row["split"]:
|
| 302 |
+
raise ValueError("Storm crosses dataset splits")
|
| 303 |
+
membership[storm] = row["split"]
|
| 304 |
+
if "center" not in row or "target_centers" not in row:
|
| 305 |
+
raise ValueError("Storm samples require observed center and masked target centers")
|
| 306 |
+
centers = np.asarray(row["target_centers"])
|
| 307 |
+
route_mask = np.asarray(row.get("route_mask", [1] * LEADS))
|
| 308 |
+
if centers.shape != (LEADS, 2) or route_mask.shape != (LEADS,) or not np.isfinite(centers).all():
|
| 309 |
+
raise ValueError("Invalid storm target coordinates/mask")
|
| 310 |
+
for field in ("target_pressure_hpa", "target_vmax_kt", "pressure_mask", "vmax_mask"):
|
| 311 |
+
values = np.asarray(row.get(field, [0] * LEADS))
|
| 312 |
+
if values.shape != (LEADS,) or not np.isfinite(values).all():
|
| 313 |
+
raise ValueError(f"Invalid target field: {field}")
|
| 314 |
+
for field in ("route_mask", "pressure_mask", "vmax_mask"):
|
| 315 |
+
values = np.asarray(row.get(field, [0] * LEADS))
|
| 316 |
+
if not np.isin(values, [0, 1]).all():
|
| 317 |
+
raise ValueError("Target masks must be binary")
|
| 318 |
+
center = np.asarray(row.get("center", [20, 140]))
|
| 319 |
+
motion = np.asarray(row.get("motion_km_per_6h", [0, 0]))
|
| 320 |
+
if center.shape != (2,) or motion.shape != (2,) or not np.isfinite([center, motion]).all():
|
| 321 |
+
raise ValueError("Invalid causal center or motion")
|
| 322 |
+
if not (min(m["regional_lat"]) <= center[0] <= max(m["regional_lat"]) and min(m["regional_lon"]) <= center[1] <= max(m["regional_lon"])):
|
| 323 |
+
raise ValueError("Issue center outside regional domain")
|
| 324 |
+
if not smoke and not any(not r.get("storm_id") for r in rows if r["split"] == "train"):
|
| 325 |
+
raise ValueError("Field pretraining requires non-storm atmospheric samples")
|
| 326 |
+
if require_full:
|
| 327 |
+
require_full_dataset(m, times, rows)
|
| 328 |
+
m["rows"] = rows
|
| 329 |
+
m["identity"] = digest(path)
|
| 330 |
+
return m
|
| 331 |
+
|
| 332 |
+
|
| 333 |
+
class AtmosphereDataset(Dataset):
|
| 334 |
+
def __init__(self, manifest, split, training=False):
|
| 335 |
+
self.m, self.training = manifest, training
|
| 336 |
+
self.rows = [r for r in manifest["rows"] if r["split"] == split]
|
| 337 |
+
self.arrays = None
|
| 338 |
+
|
| 339 |
+
def __len__(self):
|
| 340 |
+
return len(self.rows)
|
| 341 |
+
|
| 342 |
+
def __getitem__(self, index):
|
| 343 |
+
if self.arrays is None:
|
| 344 |
+
self.arrays = {k: np.load(self.m["files"][k]["resolved"], mmap_mode="r") for k in ("global", "regional", "global_static", "regional_static")}
|
| 345 |
+
row = self.rows[index]
|
| 346 |
+
i = row["issue_index"]
|
| 347 |
+
mean = np.array(self.m["normalization"]["mean"], dtype="float32")[:, None, None]
|
| 348 |
+
std = np.array(self.m["normalization"]["std"], dtype="float32")[:, None, None]
|
| 349 |
+
inputs, targets = {}, {}
|
| 350 |
+
drop_sst = self.training and np.random.random() < .15
|
| 351 |
+
for name in ("global", "regional"):
|
| 352 |
+
history = np.array(self.arrays[name][i - 8:i + 1], dtype="float32")
|
| 353 |
+
available = np.isfinite(history)
|
| 354 |
+
# SST can be masked; all atmospheric predictor channels are required.
|
| 355 |
+
if not available[:, [k for k in range(len(CHANNELS)) if k != 5]].all():
|
| 356 |
+
raise ValueError(f"Missing required atmospheric inputs at sample {i}")
|
| 357 |
+
if drop_sst:
|
| 358 |
+
available[:, 5] = False
|
| 359 |
+
history = np.where(available, (history - mean) / std, 0)
|
| 360 |
+
inputs[name + "_history"] = torch.from_numpy(history)
|
| 361 |
+
inputs[name + "_mask"] = torch.from_numpy(available.all(0).astype("float32"))
|
| 362 |
+
inputs[name + "_static"] = torch.from_numpy(np.array(self.arrays[name + "_static"], dtype="float32"))
|
| 363 |
+
future = np.array(self.arrays[name][i + 1:i + 21], dtype="float32")
|
| 364 |
+
mask = np.isfinite(future)
|
| 365 |
+
if not mask[:, [k for k in range(len(CHANNELS)) if k != 5]].all():
|
| 366 |
+
raise ValueError(f"Missing required atmospheric verification fields at sample {i}")
|
| 367 |
+
targets[name] = torch.from_numpy(np.where(mask, (future - mean) / std, 0))
|
| 368 |
+
targets[name + "_mask"] = torch.from_numpy(mask.astype("float32"))
|
| 369 |
+
inputs["center"] = torch.tensor(row.get("center", [20, 140]), dtype=torch.float32)
|
| 370 |
+
inputs["motion"] = torch.tensor(row.get("motion_km_per_6h", [0, 0]), dtype=torch.float32)
|
| 371 |
+
storm = bool(row.get("storm_id"))
|
| 372 |
+
targets["center"] = torch.tensor(row.get("target_centers", [[0, 0]] * 20), dtype=torch.float32)
|
| 373 |
+
targets["pressure"] = torch.tensor(row.get("target_pressure_hpa", [0] * 20), dtype=torch.float32)
|
| 374 |
+
targets["vmax"] = torch.tensor(row.get("target_vmax_kt", [0] * 20), dtype=torch.float32)
|
| 375 |
+
targets["storm_mask"] = torch.tensor(row.get("route_mask", [float(storm)] * 20), dtype=torch.float32)
|
| 376 |
+
targets["pressure_mask"] = torch.tensor(row.get("pressure_mask", [0] * 20), dtype=torch.float32)
|
| 377 |
+
targets["vmax_mask"] = torch.tensor(row.get("vmax_mask", [0] * 20), dtype=torch.float32)
|
| 378 |
+
return inputs, targets
|
| 379 |
+
|
| 380 |
+
|
| 381 |
+
def masked_mean(value, mask):
|
| 382 |
+
return (value * mask).sum() / mask.expand_as(value).sum().clamp_min(1)
|
| 383 |
+
|
| 384 |
+
|
| 385 |
+
def objective(output, target, step, manifest, joint):
|
| 386 |
+
terms = {}
|
| 387 |
+
for region in ("global", "regional"):
|
| 388 |
+
p, t = output[region].float(), target[region][:, step]
|
| 389 |
+
mask = target[region + "_mask"][:, step]
|
| 390 |
+
lat = torch.tensor(manifest[region + "_lat"], device=p.device, dtype=torch.float32).deg2rad().cos()[None, None, :, None]
|
| 391 |
+
if region == "regional" and "regional_lat_weight" in target:
|
| 392 |
+
lat = target["regional_lat_weight"][:, step]
|
| 393 |
+
terms[region] = masked_mean(F.smooth_l1_loss(p, t, reduction="none"), mask * lat)
|
| 394 |
+
dx = (p[..., 1:] - p[..., :-1]) - (t[..., 1:] - t[..., :-1])
|
| 395 |
+
dy = (p[..., 1:, :] - p[..., :-1, :]) - (t[..., 1:, :] - t[..., :-1, :])
|
| 396 |
+
terms[region + "_gradient"] = .5 * (masked_mean(dx.abs(), mask[..., 1:] * mask[..., :-1]) + masked_mean(dy.abs(), mask[..., 1:, :] * mask[..., :-1, :]))
|
| 397 |
+
delta = output["center"] - target["center"][:, step]
|
| 398 |
+
scale = torch.stack([torch.ones_like(delta[:, 0]), target["center"][:, step, 0].deg2rad().cos()], 1) * 111.2 / 100
|
| 399 |
+
terms["route"] = masked_mean(F.smooth_l1_loss(delta * scale, torch.zeros_like(delta), reduction="none"), target["storm_mask"][:, step, None])
|
| 400 |
+
terms["pressure"] = masked_mean((output["pressure"] - target["pressure"][:, step]).abs() / 20, target["pressure_mask"][:, step])
|
| 401 |
+
terms["vmax"] = masked_mean((output["vmax"] - target["vmax"][:, step]).abs() / 30, target["vmax_mask"][:, step])
|
| 402 |
+
value = terms["global"] + terms["regional"] + .2 * (terms["global_gradient"] + terms["regional_gradient"])
|
| 403 |
+
if joint:
|
| 404 |
+
value = value + .2 * terms["route"] + .1 * (terms["pressure"] + terms["vmax"])
|
| 405 |
+
return value, terms
|
| 406 |
+
|
| 407 |
+
|
| 408 |
+
def make_model(m, tiny=False):
|
| 409 |
+
return GlobalRegionalForecaster(*(m[k] for k in ("global_lat", "global_lon", "regional_lat", "regional_lon")),
|
| 410 |
+
m["normalization"]["mean"], m["normalization"]["std"], tiny)
|
| 411 |
+
|
| 412 |
+
|
| 413 |
+
def main():
|
| 414 |
+
parser = argparse.ArgumentParser(description=__doc__)
|
| 415 |
+
parser.add_argument("batch", type=int, nargs="?", default=10, help="Effective examples per optimizer update; physical microbatch is 1")
|
| 416 |
+
parser.add_argument("workers", type=int, nargs="?", default=3)
|
| 417 |
+
parser.add_argument("amp", type=int, nargs="?", default=1, choices=[0, 1])
|
| 418 |
+
persistent = Path("/content/drive/MyDrive/typhoon_predict") if Path("/content").is_dir() else ROOT
|
| 419 |
+
default_manifest = persistent / "data/v164/manifest.json"
|
| 420 |
+
parser.add_argument("--manifest", type=Path, default=Path(os.getenv("V164_MANIFEST", default_manifest)))
|
| 421 |
+
parser.add_argument("--output", type=Path, default=os.getenv("V164_OUTPUT_DIR"))
|
| 422 |
+
parser.add_argument("--epochs", type=int, default=12)
|
| 423 |
+
parser.add_argument("--device", default="auto", choices=["auto", "cuda", "mps", "cpu"])
|
| 424 |
+
parser.add_argument("--preflight", action="store_true")
|
| 425 |
+
parser.add_argument("--model-count", action="store_true")
|
| 426 |
+
args = parser.parse_args()
|
| 427 |
+
if args.batch < 1 or args.workers < 0 or args.epochs < 1:
|
| 428 |
+
parser.error("batch/epochs must be positive; workers must be nonnegative")
|
| 429 |
+
if args.model_count:
|
| 430 |
+
m = {"global_lat": np.linspace(90, -90, 181), "global_lon": np.arange(360),
|
| 431 |
+
"regional_lat": np.linspace(60, 0, 241), "regional_lon": np.linspace(100, 180, 321),
|
| 432 |
+
"normalization": {"mean": [0] * len(CHANNELS), "std": [1] * len(CHANNELS)}}
|
| 433 |
+
model = make_model(m)
|
| 434 |
+
print(json.dumps({"version": VERSION, "architecture": ARCHITECTURE, "parameters": sum(p.numel() for p in model.parameters()), "channels": len(CHANNELS)}))
|
| 435 |
+
return
|
| 436 |
+
m = load_manifest(args.manifest, require_full=True)
|
| 437 |
+
if args.output is None:
|
| 438 |
+
args.output = persistent / "training/v164"
|
| 439 |
+
print(json.dumps({"version": VERSION, "dataset_role": m.get("dataset_role", "full"),
|
| 440 |
+
"pilot_only": False, "samples": len(m["rows"]), "output": str(args.output)}), flush=True)
|
| 441 |
+
if args.preflight:
|
| 442 |
+
for split in ("train", "validation", "test"):
|
| 443 |
+
dataset = AtmosphereDataset(m, split)
|
| 444 |
+
if not len(dataset):
|
| 445 |
+
raise ValueError(f"Empty {split} split")
|
| 446 |
+
dataset[0]
|
| 447 |
+
print("PASS dataset integrity and first-sample checks. No training started.")
|
| 448 |
+
return
|
| 449 |
+
if Path("/content").is_dir() and not args.output.resolve().is_relative_to(Path("/content/drive/MyDrive/typhoon_predict")):
|
| 450 |
+
raise ValueError("Colab outputs must persist in MyDrive/typhoon_predict")
|
| 451 |
+
if Path("/content").is_dir() and not Path("/content/drive/MyDrive").is_dir():
|
| 452 |
+
raise RuntimeError("Mount Drive before training; refusing temporary checkpoints")
|
| 453 |
+
device = device_for(args.device)
|
| 454 |
+
random.seed(164); np.random.seed(164); torch.manual_seed(164)
|
| 455 |
+
model = make_model(m).to(device)
|
| 456 |
+
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=.01)
|
| 457 |
+
scaler = torch.amp.GradScaler("cuda", enabled=device.type == "cuda" and bool(args.amp))
|
| 458 |
+
args.output.mkdir(parents=True, exist_ok=True)
|
| 459 |
+
lock = args.output / "training.lock"
|
| 460 |
+
try:
|
| 461 |
+
fd = os.open(lock, os.O_CREAT | os.O_EXCL | os.O_WRONLY)
|
| 462 |
+
except FileExistsError as exc:
|
| 463 |
+
raise RuntimeError(f"Output locked: {lock}. Verify no other runtime is training before removing a stale lock.") from exc
|
| 464 |
+
os.write(fd, json.dumps({"pid": os.getpid(), "host": socket.gethostname()}).encode()); os.close(fd)
|
| 465 |
+
try:
|
| 466 |
+
train(model, optimizer, scaler, m, args, device)
|
| 467 |
+
finally:
|
| 468 |
+
lock.unlink(missing_ok=True)
|
| 469 |
+
|
| 470 |
+
|
| 471 |
+
def train(model, optimizer, scaler, m, args, device):
|
| 472 |
+
last, best = args.output / "last.pt", args.output / "best.pt"
|
| 473 |
+
start, history, best_value = 1, [], float("inf")
|
| 474 |
+
if last.exists():
|
| 475 |
+
saved = torch.load(last, map_location="cpu", weights_only=False)
|
| 476 |
+
if saved["architecture"] != ARCHITECTURE or saved["dataset_sha256"] != m["identity"]:
|
| 477 |
+
raise ValueError("Resume architecture/dataset mismatch")
|
| 478 |
+
model.load_state_dict(saved["model"]); optimizer.load_state_dict(saved["optimizer"]); scaler.load_state_dict(saved["scaler"])
|
| 479 |
+
start, history, best_value = saved["epoch"] + 1, saved["history"], saved["best_value"]
|
| 480 |
+
random.setstate(saved["python_rng"]); np.random.set_state(saved["numpy_rng"]); torch.set_rng_state(saved["torch_rng"])
|
| 481 |
+
if device.type == "cuda" and saved.get("cuda_rng") is not None:
|
| 482 |
+
torch.cuda.set_rng_state_all(saved["cuda_rng"])
|
| 483 |
+
atomic_json(args.output / "manifest.json", {"version": VERSION, "architecture": ARCHITECTURE,
|
| 484 |
+
"parameters": sum(p.numel() for p in model.parameters()), "dataset_sha256": m["identity"],
|
| 485 |
+
"script_sha256": digest(__file__), "future_weather_as_inputs": False,
|
| 486 |
+
"official_forecasts_as_inputs": False, "ensemble": False,
|
| 487 |
+
"resume_policy": "last completed epoch with optimizer/scaler/RNG",
|
| 488 |
+
"physical_microbatch": 1, "effective_batch": args.batch})
|
| 489 |
+
workers = min(args.workers, max(0, (os.cpu_count() or 1) - 1), 2)
|
| 490 |
+
for epoch in range(start, args.epochs + 1):
|
| 491 |
+
horizon = [2, 4, 8, 12, 20][min(epoch - 1, 4)]
|
| 492 |
+
joint = epoch > 2
|
| 493 |
+
epoch_start = time.time()
|
| 494 |
+
metrics = {}
|
| 495 |
+
for split in ("train", "validation"):
|
| 496 |
+
training = split == "train"
|
| 497 |
+
model.train(training)
|
| 498 |
+
dataset = AtmosphereDataset(m, split, training)
|
| 499 |
+
if not len(dataset):
|
| 500 |
+
raise ValueError(f"Empty {split} split")
|
| 501 |
+
loader = DataLoader(dataset, batch_size=1, shuffle=training, num_workers=workers,
|
| 502 |
+
**({"prefetch_factor": 1} if workers else {}))
|
| 503 |
+
optimizer.zero_grad(set_to_none=True)
|
| 504 |
+
sums = []
|
| 505 |
+
for index, (inputs, targets) in enumerate(loader, 1):
|
| 506 |
+
inputs = {k: v.to(device) for k, v in inputs.items()}
|
| 507 |
+
steps = horizon if training else 20
|
| 508 |
+
loss_block, total = 0, 0.0
|
| 509 |
+
with torch.set_grad_enabled(training):
|
| 510 |
+
state = model.initial(inputs)
|
| 511 |
+
for step in range(steps):
|
| 512 |
+
with torch.autocast(device_type="cuda", dtype=torch.float16) if scaler.is_enabled() else nullcontext():
|
| 513 |
+
state, out = model.step(state)
|
| 514 |
+
# Transfer one target frame at a time to bound VRAM.
|
| 515 |
+
frame = {k: v[:, step:step+1].to(device) for k, v in targets.items()}
|
| 516 |
+
loss, _ = objective(out, frame, 0, m, joint)
|
| 517 |
+
if not torch.isfinite(loss):
|
| 518 |
+
raise RuntimeError(f"Non-finite loss epoch={epoch} sample={index} step={step}; checkpoint not advanced")
|
| 519 |
+
total += float(loss.detach()) / steps
|
| 520 |
+
loss_block = loss_block + loss / (steps * args.batch)
|
| 521 |
+
if training and ((step + 1) % 2 == 0 or step + 1 == steps):
|
| 522 |
+
scaler.scale(loss_block).backward()
|
| 523 |
+
state = {k: v.detach() for k, v in state.items()}
|
| 524 |
+
loss_block = 0
|
| 525 |
+
if training and (index % args.batch == 0 or index == len(loader)):
|
| 526 |
+
scaler.unscale_(optimizer)
|
| 527 |
+
remainder = index % args.batch
|
| 528 |
+
if index == len(loader) and remainder:
|
| 529 |
+
for p in model.parameters():
|
| 530 |
+
if p.grad is not None:
|
| 531 |
+
p.grad.mul_(args.batch / remainder)
|
| 532 |
+
norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
|
| 533 |
+
if not torch.isfinite(norm):
|
| 534 |
+
raise RuntimeError("Non-finite gradients; checkpoint not advanced. Retry last epoch with AMP=0.")
|
| 535 |
+
scaler.step(optimizer); scaler.update(); optimizer.zero_grad(set_to_none=True)
|
| 536 |
+
sums.append(total)
|
| 537 |
+
if index == 1 or index % 10 == 0:
|
| 538 |
+
print(f"{split} epoch={epoch}/{args.epochs} sample={index}/{len(loader)} horizon={steps} loss={total:.5f} elapsed_s={time.time()-epoch_start:.0f}", flush=True)
|
| 539 |
+
metrics[split] = float(np.mean(sums))
|
| 540 |
+
# Selection only after reaching the full horizon, on a stable joint objective.
|
| 541 |
+
promote = epoch >= 5 and metrics["validation"] < best_value
|
| 542 |
+
if promote:
|
| 543 |
+
best_value = metrics["validation"]
|
| 544 |
+
history.append({"epoch": epoch, "horizon": horizon, "joint": joint, **metrics})
|
| 545 |
+
payload = {"version": VERSION, "architecture": ARCHITECTURE, "dataset_sha256": m["identity"],
|
| 546 |
+
"data_contract": {k: m[k] for k in ("channels", "units", "global_lat", "global_lon", "regional_lat", "regional_lon", "normalization")},
|
| 547 |
+
"epoch": epoch, "model": model.state_dict(), "optimizer": optimizer.state_dict(),
|
| 548 |
+
"scaler": scaler.state_dict(), "history": history, "best_value": best_value,
|
| 549 |
+
"python_rng": random.getstate(), "numpy_rng": np.random.get_state(), "torch_rng": torch.get_rng_state(),
|
| 550 |
+
"cuda_rng": torch.cuda.get_rng_state_all() if device.type == "cuda" else None}
|
| 551 |
+
atomic_checkpoint(last, payload)
|
| 552 |
+
if promote:
|
| 553 |
+
atomic_checkpoint(best, payload)
|
| 554 |
+
atomic_json(args.output / "history.json", history)
|
| 555 |
+
print(f"Saved completed epoch {epoch}: {last}; promoted_best={promote}", flush=True)
|
| 556 |
+
print("Training completed. Checkpoints and history are persistent. Test set was not used for selection.", flush=True)
|
| 557 |
+
|
| 558 |
+
|
| 559 |
+
# BEGIN REUSE IMPLEMENTATION
|
| 560 |
+
"""Embedded in the standalone v164 trainer: coarse archive plus native MSLP patches."""
|
| 561 |
+
REUSE_ARCH = 'wp-coarse8-native-pressure1-v164'
|
| 562 |
+
REUSE_CHANNELS = ['mslp', 'hgt500', 'uwnd850', 'vwnd850', 'uwnd500', 'vwnd500', 'uwnd200', 'vwnd200']
|
| 563 |
+
|
| 564 |
+
|
| 565 |
+
def load_reuse(path):
|
| 566 |
+
path = Path(path)
|
| 567 |
+
if not path.is_file():
|
| 568 |
+
raise FileNotFoundError('Combined dataset missing. Run the updated Drive setup cell first: '+str(path))
|
| 569 |
+
m = json.loads(path.read_text())
|
| 570 |
+
if m.get('schema') != 'v164-reuse-pressure-v1' or m.get('pilot', True):
|
| 571 |
+
raise ValueError('Combined non-pilot archive required; no pilot fallback')
|
| 572 |
+
for name, entry in m['files'].items():
|
| 573 |
+
p = path.parent/entry['path']
|
| 574 |
+
print('Checking dataset file:', name, flush=True)
|
| 575 |
+
if digest(p) != entry['sha256']:
|
| 576 |
+
raise ValueError('Dataset checksum mismatch: '+name)
|
| 577 |
+
entry['resolved'] = str(p)
|
| 578 |
+
with np.load(m['files']['basin']['resolved'], allow_pickle=False) as z:
|
| 579 |
+
if z['channels'].tolist() != REUSE_CHANNELS[1:]:
|
| 580 |
+
raise ValueError('Coarse channel order mismatch')
|
| 581 |
+
times, lat, lon = z['time'], z['lat'], z['lon']
|
| 582 |
+
q = z['q']
|
| 583 |
+
scale, offset = z['scale'], z['offset']
|
| 584 |
+
if q.shape != (len(times),7,25,33) or not np.all(np.diff(times) == 6*HOUR):
|
| 585 |
+
raise ValueError('Coarse grid/time mismatch')
|
| 586 |
+
qpath = path.parent/'basin_q_verified.npy'
|
| 587 |
+
# Refresh the compact worker mmap from the checksum-verified source every preflight.
|
| 588 |
+
temporary = qpath.with_suffix('.pending.npy')
|
| 589 |
+
np.save(temporary, q)
|
| 590 |
+
os.replace(temporary, qpath)
|
| 591 |
+
slp = np.load(m['files']['slp']['resolved'], mmap_mode='r')
|
| 592 |
+
if slp.shape != (len(times),25,33) or not np.isfinite(slp).all():
|
| 593 |
+
raise ValueError('Coarse pressure coverage invalid')
|
| 594 |
+
with np.load(m['files']['track']['resolved'], allow_pickle=False) as z:
|
| 595 |
+
track = {k:z[k] for k in ['base_time','storm_id','base_lat','base_lon','target','target_mask']}
|
| 596 |
+
issue_intensity = z['track'][:,-1,4:6]*z['track_std'][4:6]+z['track_mean'][4:6]
|
| 597 |
+
plan = json.loads(Path(m['files']['plan']['resolved']).read_text())
|
| 598 |
+
if len(plan['rows']) != 1000 or plan['shape'] != [1000,29,121,121] or plan['lead_hours'] != list(range(-48,121,6)):
|
| 599 |
+
raise ValueError('Expected the verified 1000-window native pressure package, not the six-sample pilot')
|
| 600 |
+
pressure = np.load(m['files']['pressure']['resolved'], mmap_mode='r')
|
| 601 |
+
if pressure.shape != (1000,29,121,121):
|
| 602 |
+
raise ValueError('Native pressure shape mismatch')
|
| 603 |
+
patches = {}
|
| 604 |
+
for pi,row in enumerate(plan['rows']):
|
| 605 |
+
ti = row['track_archive_row']
|
| 606 |
+
if int(track['base_time'][ti]) != row['issue_ns'] or str(track['storm_id'][ti]) != row['storm_id']:
|
| 607 |
+
raise ValueError('Native patch/track identity mismatch')
|
| 608 |
+
patches[ti] = pi
|
| 609 |
+
first_year = {}
|
| 610 |
+
years = track['base_time'].astype('datetime64[ns]').astype('datetime64[Y]').astype(int)+1970
|
| 611 |
+
for sid,y in zip(track['storm_id'],years):
|
| 612 |
+
first_year[str(sid)] = min(first_year.get(str(sid),9999),int(y))
|
| 613 |
+
lookup = {(str(s),int(t)):i for i,(s,t) in enumerate(zip(track['storm_id'],track['base_time']))}
|
| 614 |
+
rows = []
|
| 615 |
+
for i,ns in enumerate(track['base_time']):
|
| 616 |
+
sid = str(track['storm_id'][i])
|
| 617 |
+
y = first_year[sid]
|
| 618 |
+
split = 'train' if 2000 <= y <= 2021 else 'validation' if 2022 <= y <= 2023 else 'test' if 2024 <= y <= 2025 else None
|
| 619 |
+
if split is None:
|
| 620 |
+
continue
|
| 621 |
+
low,high = {'train':(2000,2021),'validation':(2022,2023),'test':(2024,2025)}[split]
|
| 622 |
+
bounds = np.array([ns-48*HOUR,ns+120*HOUR],dtype='int64').astype('datetime64[ns]').astype('datetime64[Y]').astype(int)+1970
|
| 623 |
+
ai = int(np.searchsorted(times,ns))
|
| 624 |
+
center = [float(track['base_lat'][i]),float(track['base_lon'][i])%360]
|
| 625 |
+
if bounds[0] < low or bounds[1] > high or ai < 8 or ai+20 >= len(times) or times[ai] != ns:
|
| 626 |
+
continue
|
| 627 |
+
if not np.isfinite(center).all() or not (0 < center[0] < 60 and 100 < center[1] < 180):
|
| 628 |
+
continue
|
| 629 |
+
previous = lookup.get((sid,int(ns-6*HOUR)))
|
| 630 |
+
motion = [0.,0.]
|
| 631 |
+
if previous is not None:
|
| 632 |
+
old = [float(track['base_lat'][previous]),float(track['base_lon'][previous])%360]
|
| 633 |
+
if np.isfinite(old).all():
|
| 634 |
+
motion = [(center[1]-old[1])*111.2*np.cos(np.deg2rad(center[0])),(center[0]-old[0])*111.2]
|
| 635 |
+
future = [lookup.get((sid,int(ns+j*6*HOUR))) for j in range(1,21)]
|
| 636 |
+
target_centers = [[float(track['base_lat'][j]),float(track['base_lon'][j])%360] if j is not None else [0.,0.] for j in future]
|
| 637 |
+
rm = [float(j is not None and np.isfinite(target_centers[k]).all()) for k,j in enumerate(future)]
|
| 638 |
+
target_centers = np.nan_to_num(target_centers).tolist()
|
| 639 |
+
wind,p = track['target'][i,:,2],track['target'][i,:,3]
|
| 640 |
+
wm = track['target_mask'][i,:,2].astype(bool)&np.isfinite(wind)&(wind>=0)&(wind<250)
|
| 641 |
+
pm = track['target_mask'][i,:,3].astype(bool)&np.isfinite(p)&(p>800)&(p<1100)
|
| 642 |
+
pi = patches.get(i,-1)
|
| 643 |
+
if pi >= 0 and plan['rows'][pi]['split'] != split:
|
| 644 |
+
raise ValueError('Patch split mismatch')
|
| 645 |
+
issue = issue_intensity[i]
|
| 646 |
+
iv = [bool(np.isfinite(issue[0]) and 0<=issue[0]<250),bool(np.isfinite(issue[1]) and 800<issue[1]<1100)]
|
| 647 |
+
rows.append({'split':split,'storm_id':sid,'atlas':ai,'patch':pi,'center':center,'motion':motion,
|
| 648 |
+
'issue_intensity':[float(issue[k]) if iv[k] else 0. for k in range(2)],'issue_mask':iv,
|
| 649 |
+
'target_centers':target_centers,'route_mask':rm,
|
| 650 |
+
'pressure':np.where(pm,p,0).tolist(),'pressure_mask':pm.tolist(),
|
| 651 |
+
'vmax':np.where(wm,wind,0).tolist(),'vmax_mask':wm.tolist()})
|
| 652 |
+
counts = {s:sum(r['split']==s for r in rows) for s in ['train','validation','test']}
|
| 653 |
+
patch_counts = {s:sum(r['split']==s and r['patch']>=0 for r in rows) for s in counts}
|
| 654 |
+
if counts['train'] < 10000 or patch_counts != {'train':800,'validation':100,'test':100}:
|
| 655 |
+
raise ValueError(f'Incomplete combined coverage: {counts}, native patches {patch_counts}')
|
| 656 |
+
train_times = np.flatnonzero((times.astype('datetime64[ns]').astype('datetime64[Y]').astype(int)+1970>=2000)&
|
| 657 |
+
(times.astype('datetime64[ns]').astype('datetime64[Y]').astype(int)+1970<=2021))
|
| 658 |
+
sums,squares = np.zeros(8),np.zeros(8)
|
| 659 |
+
n = 0
|
| 660 |
+
for start in range(0,len(train_times),256):
|
| 661 |
+
ii = train_times[start:start+256]
|
| 662 |
+
b = np.concatenate([np.asarray(slp[ii],dtype='float32')[:,None],q[ii].astype('float32')*scale[None,:,None,None]+offset[None,:,None,None]],1).astype('float64')
|
| 663 |
+
if not np.isfinite(b).all():
|
| 664 |
+
raise ValueError('Nonfinite coarse training fields')
|
| 665 |
+
sums += b.sum((0,2,3)); squares += (b*b).sum((0,2,3)); n += len(ii)*25*33
|
| 666 |
+
mean = sums/n
|
| 667 |
+
std = np.maximum(np.sqrt(np.maximum(squares/n-mean*mean,0)),1e-4)
|
| 668 |
+
m.update(rows=rows, plan=plan, qpath=str(qpath), scale=scale.tolist(), offset=offset.tolist(),
|
| 669 |
+
channels=REUSE_CHANNELS,units=['hPa','m','m/s','m/s','m/s','m/s','m/s','m/s'],
|
| 670 |
+
global_lat=lat.tolist(),global_lon=lon.tolist(),regional_lat=list(np.linspace(15,-15,121)),
|
| 671 |
+
regional_lon=list(np.linspace(-15,15,121)),normalization={'mean':mean.tolist(),'std':std.tolist(),'fit_split':'train','end_year':2021},
|
| 672 |
+
identity=digest(path),counts=counts,patch_counts=patch_counts)
|
| 673 |
+
print(json.dumps({'all_archive_windows':counts,'native_pressure_windows':patch_counts,'architecture':REUSE_ARCH}),flush=True)
|
| 674 |
+
return m
|
| 675 |
+
|
| 676 |
+
|
| 677 |
+
class ReuseDataset(Dataset):
|
| 678 |
+
def __init__(self,m,split,training=False):
|
| 679 |
+
self.m = m
|
| 680 |
+
self.rows = [r for r in m['rows'] if r['split']==split]
|
| 681 |
+
self.arrays = None
|
| 682 |
+
|
| 683 |
+
def __len__(self):
|
| 684 |
+
return len(self.rows)
|
| 685 |
+
|
| 686 |
+
def __getitem__(self,i):
|
| 687 |
+
m,row = self.m,self.rows[i]
|
| 688 |
+
if self.arrays is None:
|
| 689 |
+
self.arrays = {k:np.load(m['files'][k]['resolved'],mmap_mode='r') for k in ['slp','pressure']}
|
| 690 |
+
self.arrays['q'] = np.load(m['qpath'],mmap_mode='r')
|
| 691 |
+
with np.load(m['files']['geography']['resolved'],allow_pickle=False) as geo:
|
| 692 |
+
self.arrays['land'] = geo['land_fraction']
|
| 693 |
+
self.arrays['height'] = geo['elevation_m']/8000
|
| 694 |
+
ai = row['atlas']
|
| 695 |
+
ii = slice(ai-8,ai+21)
|
| 696 |
+
coarse = np.concatenate([np.array(self.arrays['slp'][ii],dtype='float32')[:,None],
|
| 697 |
+
np.array(self.arrays['q'][ii],dtype='float32')*np.array(m['scale'],dtype='float32')[None,:,None,None]+np.array(m['offset'],dtype='float32')[None,:,None,None]],1)
|
| 698 |
+
mean,std = np.array(m['normalization']['mean'],dtype='float32'),np.array(m['normalization']['std'],dtype='float32')
|
| 699 |
+
coarse = (coarse-mean[None,:,None,None])/std[None,:,None,None]
|
| 700 |
+
pi = row['patch']
|
| 701 |
+
anchor = m['plan']['rows'][pi] if pi>=0 else {'center_lat':round(row['center'][0]*4)/4,'center_lon':round(row['center'][1]*4)/4}
|
| 702 |
+
lat = anchor['center_lat']+np.linspace(15,-15,121,dtype='float32')
|
| 703 |
+
lon = anchor['center_lon']+np.linspace(-15,15,121,dtype='float32')
|
| 704 |
+
yy,xx = np.meshgrid(lat,lon,indexing='ij')
|
| 705 |
+
y = np.rint((90-yy)*4).astype(int); x = np.rint(xx*4).astype(int)%1440
|
| 706 |
+
rs = np.stack([yy/90,xx/180-1,self.arrays['land'][y,x],self.arrays['height'][y,x]]).astype('float32')
|
| 707 |
+
cy,cx = np.meshgrid(m['global_lat'],m['global_lon'],indexing='ij')
|
| 708 |
+
iy = np.rint((90-cy)*4).astype(int); ix = np.rint(cx*4).astype(int)%1440
|
| 709 |
+
gs = np.stack([cy/90,cx/180-1,self.arrays['land'][iy,ix],self.arrays['height'][iy,ix]]).astype('float32')
|
| 710 |
+
detail = np.zeros((29,1,121,121),dtype='float32')
|
| 711 |
+
if pi>=0:
|
| 712 |
+
detail[:,0] = (np.array(self.arrays['pressure'][pi],dtype='float32')-mean[0])/std[0]
|
| 713 |
+
if not np.isfinite(detail).all():
|
| 714 |
+
raise ValueError('Native pressure contains nonfinite fields')
|
| 715 |
+
inputs = {'global_history':coarse[:9], 'regional_history':detail[:9],
|
| 716 |
+
'global_static':gs,'regional_static':rs,'detail_available':np.array([pi>=0],dtype='float32'),
|
| 717 |
+
'center':np.array(row['center'],dtype='float32'),'motion':np.array(row['motion'],dtype='float32'),
|
| 718 |
+
'issue_intensity':np.array(row['issue_intensity'],dtype='float32'),'issue_mask':np.array(row['issue_mask'],dtype='float32')}
|
| 719 |
+
targets = {'global':coarse[9:],'global_mask':np.ones_like(coarse[9:]),
|
| 720 |
+
'regional_lat_weight':np.broadcast_to(np.cos(np.deg2rad(lat))[None,None,:,None],(20,1,121,1)).astype('float32'),
|
| 721 |
+
'regional':detail[9:],'regional_mask':np.full_like(detail[9:],float(pi>=0)),
|
| 722 |
+
'center':np.array(row['target_centers'],dtype='float32'),
|
| 723 |
+
'storm_mask':np.array(row['route_mask'],dtype='float32')}
|
| 724 |
+
for k in ['pressure','pressure_mask','vmax','vmax_mask']:
|
| 725 |
+
targets[k] = np.array(row[k],dtype='float32')
|
| 726 |
+
return {k:torch.from_numpy(v.copy()) for k,v in inputs.items()},{k:torch.from_numpy(v.copy()) for k,v in targets.items()}
|
| 727 |
+
|
| 728 |
+
|
| 729 |
+
class ReuseForecaster(nn.Module):
|
| 730 |
+
def __init__(self,m,tiny=False):
|
| 731 |
+
super().__init__()
|
| 732 |
+
self.global_time = nn.Conv2d(8*9,8,1)
|
| 733 |
+
self.regional_time = nn.Conv2d(9,1,1)
|
| 734 |
+
self.global_operator = FieldOperator(8,8 if tiny else 80,False,1 if tiny else 3,tiny)
|
| 735 |
+
self.regional_operator = FieldOperator(1,8 if tiny else 88,False,0,tiny)
|
| 736 |
+
self.parent_fusion = nn.Conv2d(1+8+1,1,1)
|
| 737 |
+
self.object_score = nn.Conv2d(32+3,1,1)
|
| 738 |
+
self.intensity = nn.Sequential(nn.Linear(32+3+4,64),nn.GELU(),nn.Linear(64,2))
|
| 739 |
+
self.register_buffer('mean',torch.tensor(m['normalization']['mean'],dtype=torch.float32)[None,:,None,None])
|
| 740 |
+
self.register_buffer('std',torch.tensor(m['normalization']['std'],dtype=torch.float32)[None,:,None,None])
|
| 741 |
+
yy,xx = np.meshgrid(m['global_lat'],m['global_lon'],indexing='ij')
|
| 742 |
+
self.register_buffer('yy',torch.tensor(yy,dtype=torch.float32)[None])
|
| 743 |
+
self.register_buffer('xx',torch.tensor(xx,dtype=torch.float32)[None])
|
| 744 |
+
|
| 745 |
+
def initial(self,inputs):
|
| 746 |
+
required = {'global_history','regional_history','global_static','regional_static','detail_available','center','motion','issue_intensity','issue_mask'}
|
| 747 |
+
if set(inputs)!=required:
|
| 748 |
+
raise ValueError('Only declared causal inputs are allowed')
|
| 749 |
+
gh,rh = inputs['global_history'],inputs['regional_history']
|
| 750 |
+
rs = inputs['regional_static']
|
| 751 |
+
yy,xx = rs[:,0]*90,(rs[:,1]+1)*180
|
| 752 |
+
grid = torch.stack([(xx-100)/40-1,1-yy/30],-1)
|
| 753 |
+
return {'g':gh[:,-1],'r':rh[:,-1],'gc':self.global_time(gh.flatten(1,2)),
|
| 754 |
+
'rc':self.regional_time(rh.flatten(1,2)),'gs':inputs['global_static'],'rs':rs,
|
| 755 |
+
'grid':grid,'available':inputs['detail_available'],'center':inputs['center'],
|
| 756 |
+
'motion':inputs['motion'],'issue':inputs['issue_intensity'],'issue_mask':inputs['issue_mask']}
|
| 757 |
+
|
| 758 |
+
def step(self,s):
|
| 759 |
+
g,features = self.global_operator(s['g'],s['gc'],torch.ones_like(s['g']),s['gs'])
|
| 760 |
+
valid = (s['grid'].abs()<=1).all(-1)[:,None].to(g.dtype)
|
| 761 |
+
parent = F.grid_sample(g,s['grid'],align_corners=True,padding_mode='zeros')*valid
|
| 762 |
+
r = s['r']
|
| 763 |
+
if bool(s['available'].any()):
|
| 764 |
+
context = self.parent_fusion(torch.cat([s['rc'],parent,valid],1))
|
| 765 |
+
r,_ = self.regional_operator(r,context,torch.ones_like(r),s['rs'])
|
| 766 |
+
physical = g.float()*self.std+self.mean
|
| 767 |
+
p = physical[:,0]
|
| 768 |
+
anomaly = p-F.avg_pool2d(F.pad(p[:,None],(2,2,2,2),mode='replicate'),5,1)[:,0]
|
| 769 |
+
evidence = torch.cat([features.float(),anomaly[:,None]/10,physical[:,2:4]/30],1)
|
| 770 |
+
lat0,lon0 = s['center'][:,0,None,None],s['center'][:,1,None,None]
|
| 771 |
+
plat = lat0+s['motion'][:,1,None,None]/111.2
|
| 772 |
+
plon = lon0+s['motion'][:,0,None,None]/(111.2*lat0.deg2rad().cos().clamp_min(.2))
|
| 773 |
+
d2 = ((self.yy-plat)*111.2).square()+((self.xx-plon)*111.2*lat0.deg2rad().cos()).square()
|
| 774 |
+
score = self.object_score(evidence)[:,0]-anomaly/10-d2/(2*600**2)
|
| 775 |
+
peak = score.flatten(1).argmax(1)
|
| 776 |
+
iy,ix = torch.arange(25,device=g.device)[None,:,None],torch.arange(33,device=g.device)[None,None,:]
|
| 777 |
+
local = ((iy-(peak//33)[:,None,None]).abs()<=1)&((ix-(peak%33)[:,None,None]).abs()<=1)
|
| 778 |
+
pdf = torch.softmax(score.masked_fill(~local,-1e9).flatten(1),1).reshape_as(score)
|
| 779 |
+
center = torch.stack([(pdf*self.yy).sum((1,2)),(pdf*self.xx).sum((1,2))],1)
|
| 780 |
+
pooled = (evidence*pdf[:,None]).sum((2,3))
|
| 781 |
+
scalar = torch.cat([s['issue']/torch.tensor([100.,1000.],device=g.device),s['issue_mask']],1)
|
| 782 |
+
intensity = self.intensity(torch.cat([pooled,scalar],1))
|
| 783 |
+
grid_p = (p*pdf).sum((1,2))
|
| 784 |
+
dy,dx = s['rs'][:,0]*90,(s['rs'][:,1]+1)*180
|
| 785 |
+
sample = torch.stack([2*(center[:,1]-dx[:,0,0])/30-1,2*(dy[:,0,0]-center[:,0])/30-1],-1)[:,None,None]
|
| 786 |
+
inside = (sample.abs()<=1).all((1,2,3)) & (s['available'][:,0]>0)
|
| 787 |
+
detailed_p = F.grid_sample(r.float()*self.std[:,:1]+self.mean[:,:1],sample,align_corners=True)[:,0,0,0]
|
| 788 |
+
grid_p = torch.where(inside,detailed_p,grid_p)
|
| 789 |
+
pressure = grid_p-F.softplus(intensity[:,0])*5
|
| 790 |
+
vmax = s['issue'][:,0]*s['issue_mask'][:,0]+intensity[:,1]*10
|
| 791 |
+
meanlat = (center[:,0]+s['center'][:,0])/2
|
| 792 |
+
motion = torch.stack([(center[:,1]-s['center'][:,1])*111.2*meanlat.deg2rad().cos(),(center[:,0]-s['center'][:,0])*111.2],1)
|
| 793 |
+
return dict(s,g=g,r=r,center=center,motion=motion),{'global':g,'regional':r,'center':center,'pressure':pressure,'vmax':vmax,'grid_pressure':grid_p}
|
| 794 |
+
|
| 795 |
+
def forward(self,inputs,steps=20):
|
| 796 |
+
s = self.initial(inputs)
|
| 797 |
+
out = []
|
| 798 |
+
for _ in range(steps):
|
| 799 |
+
s,o = self.step(s); out.append(o)
|
| 800 |
+
return {k:torch.stack([o[k] for o in out],1) for k in out[0]}
|
| 801 |
+
|
| 802 |
+
|
| 803 |
+
def reuse_main():
|
| 804 |
+
global ARCHITECTURE, AtmosphereDataset
|
| 805 |
+
ARCHITECTURE,AtmosphereDataset = REUSE_ARCH,ReuseDataset
|
| 806 |
+
p = argparse.ArgumentParser(description='v1.2.64 coarse weather + native pressure training; no pilot fallback')
|
| 807 |
+
p.add_argument('batch',type=int,nargs='?',default=10)
|
| 808 |
+
p.add_argument('workers',type=int,nargs='?',default=3)
|
| 809 |
+
p.add_argument('amp',type=int,nargs='?',default=1,choices=[0,1])
|
| 810 |
+
p.add_argument('--manifest',type=Path,default=ROOT/'data/v164_reuse/manifest.json')
|
| 811 |
+
persistent = Path('/content/drive/MyDrive/typhoon_predict') if Path('/content').is_dir() else ROOT
|
| 812 |
+
p.add_argument('--output',type=Path,default=persistent/'training/v164_reuse_pressure')
|
| 813 |
+
p.add_argument('--epochs',type=int,default=12)
|
| 814 |
+
p.add_argument('--device',default='auto',choices=['auto','cpu','mps','cuda'])
|
| 815 |
+
p.add_argument('--preflight',action='store_true')
|
| 816 |
+
args = p.parse_args()
|
| 817 |
+
if args.batch<1 or args.workers<0 or args.epochs<1:
|
| 818 |
+
p.error('Invalid batch/workers/epochs')
|
| 819 |
+
m = load_reuse(args.manifest)
|
| 820 |
+
for split in ['train','validation','test']:
|
| 821 |
+
d = ReuseDataset(m,split)
|
| 822 |
+
d[0]
|
| 823 |
+
d[next(i for i,r in enumerate(d.rows) if r['patch']>=0)]
|
| 824 |
+
if args.preflight:
|
| 825 |
+
print('PASS combined non-pilot archives and native pressure data. No training started.',flush=True)
|
| 826 |
+
return
|
| 827 |
+
if Path('/content').is_dir() and not args.output.resolve().is_relative_to(persistent.resolve()):
|
| 828 |
+
raise ValueError('Colab checkpoints must persist in typhoon_predict on Drive')
|
| 829 |
+
device = device_for(args.device)
|
| 830 |
+
random.seed(164); np.random.seed(164); torch.manual_seed(164)
|
| 831 |
+
model = ReuseForecaster(m).to(device)
|
| 832 |
+
print(json.dumps({'version':VERSION,'architecture':REUSE_ARCH,'parameters':sum(x.numel() for x in model.parameters()),'device':str(device),'windows':m['counts'],'native_windows':m['patch_counts'],'output':str(args.output)}),flush=True)
|
| 833 |
+
optimizer = torch.optim.AdamW(model.parameters(),lr=1e-4,weight_decay=.01)
|
| 834 |
+
scaler = torch.amp.GradScaler('cuda',enabled=device.type=='cuda' and bool(args.amp))
|
| 835 |
+
args.output.mkdir(parents=True,exist_ok=True)
|
| 836 |
+
lock = args.output/'training.lock'
|
| 837 |
+
fd = os.open(lock,os.O_CREAT|os.O_EXCL|os.O_WRONLY)
|
| 838 |
+
os.write(fd,json.dumps({'pid':os.getpid(),'host':socket.gethostname()}).encode()); os.close(fd)
|
| 839 |
+
try:
|
| 840 |
+
train(model,optimizer,scaler,m,args,device)
|
| 841 |
+
finally:
|
| 842 |
+
lock.unlink(missing_ok=True)
|
| 843 |
+
|
| 844 |
+
if __name__ == "__main__":
|
| 845 |
+
reuse_main()
|
training/trackformer_1_2/source/data_preparation/v164_prepare_data.py
ADDED
|
@@ -0,0 +1,348 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Prepare a resumable, explicitly labelled v164 data pilot from public ERA5.
|
| 2 |
+
|
| 3 |
+
This downloads native fields. It does not upsample the old coarse archive,
|
| 4 |
+
start training, or put verification labels into model inputs.
|
| 5 |
+
"""
|
| 6 |
+
import argparse
|
| 7 |
+
import hashlib
|
| 8 |
+
import json
|
| 9 |
+
import os
|
| 10 |
+
from pathlib import Path
|
| 11 |
+
import shutil
|
| 12 |
+
import time
|
| 13 |
+
|
| 14 |
+
import numpy as np
|
| 15 |
+
import v164_global_regional_train as trainer
|
| 16 |
+
|
| 17 |
+
SOURCE = 'https://storage.googleapis.com/gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3'
|
| 18 |
+
TRACK_SHA = 'cb8d7291aa8140124ca522feb448a1d44124b018efc4ba1d54824b60a4f73451'
|
| 19 |
+
SURFACE = ['mean_sea_level_pressure', 'surface_pressure', '10m_u_component_of_wind',
|
| 20 |
+
'10m_v_component_of_wind', '2m_temperature', 'sea_surface_temperature']
|
| 21 |
+
UPPER = ['u_component_of_wind', 'v_component_of_wind', 'temperature', 'specific_humidity', 'geopotential']
|
| 22 |
+
STEP = 6 * trainer.HOUR
|
| 23 |
+
|
| 24 |
+
|
| 25 |
+
def year(ns):
|
| 26 |
+
return int(np.datetime64(int(ns), 'ns').astype('datetime64[Y]').astype(int)) + 1970
|
| 27 |
+
|
| 28 |
+
|
| 29 |
+
def build_plan(track_path, counts=(3, 1, 1), oversample=1):
|
| 30 |
+
if trainer.digest(track_path) != TRACK_SHA:
|
| 31 |
+
raise ValueError('Shared track archive checksum mismatch')
|
| 32 |
+
with np.load(track_path, allow_pickle=False) as z:
|
| 33 |
+
tracks = {k: z[k] for k in ('storm_id', 'base_time', 'base_lat', 'base_lon', 'n_leads', 'target', 'target_mask')}
|
| 34 |
+
issue_intensity = z['track'][:, -1, 4:6]*z['track_std'][4:6] + z['track_mean'][4:6]
|
| 35 |
+
sid, times = tracks['storm_id'], tracks['base_time']
|
| 36 |
+
first_year = {}
|
| 37 |
+
for storm, ns in zip(sid, times):
|
| 38 |
+
first_year[str(storm)] = min(first_year.get(str(storm), 9999), year(ns))
|
| 39 |
+
lookup = {(str(s), int(t)): i for i, (s, t) in enumerate(zip(sid, times))}
|
| 40 |
+
rows = []
|
| 41 |
+
rng = np.random.default_rng(164)
|
| 42 |
+
for split, bounds, count in zip(('train', 'validation', 'test'), ((2000, 2021), (2022, 2023), (2024, 2025)), counts):
|
| 43 |
+
candidates = {}
|
| 44 |
+
for i, (storm, ns) in enumerate(zip(sid, times)):
|
| 45 |
+
storm, ns = str(storm), int(ns)
|
| 46 |
+
lat, lon = float(tracks['base_lat'][i]), float(tracks['base_lon'][i]) % 360
|
| 47 |
+
if not (bounds[0] <= first_year[storm] <= bounds[1] and bounds[0] <= year(ns-8*STEP)
|
| 48 |
+
and year(ns+20*STEP) <= bounds[1] and ns % STEP == 0 and tracks['n_leads'][i] >= 20
|
| 49 |
+
and issue_intensity[i, 0] >= 90 and 850 < issue_intensity[i, 1] <= 960
|
| 50 |
+
and 0 < lat < 40 and 100 < lon < 180 and '1979' not in storm and storm.upper() != 'TIP'):
|
| 51 |
+
continue
|
| 52 |
+
future = [lookup.get((storm, ns + lead*STEP)) for lead in range(1, 21)]
|
| 53 |
+
if sum(j is not None for j in future) < 16:
|
| 54 |
+
continue
|
| 55 |
+
candidates.setdefault(storm, []).append((i, future))
|
| 56 |
+
storms = sorted(candidates)
|
| 57 |
+
if len(storms) < count:
|
| 58 |
+
raise ValueError(f'Not enough eligible {split} storms: {len(storms)} < {count}')
|
| 59 |
+
for storm in rng.choice(storms, size=min(len(storms), count*oversample), replace=False):
|
| 60 |
+
choices = candidates[storm]
|
| 61 |
+
i, future = min(choices, key=lambda item: float(issue_intensity[item[0], 1]))
|
| 62 |
+
p, vmax = tracks['target'][i, :, 3], tracks['target'][i, :, 2]
|
| 63 |
+
pm = tracks['target_mask'][i, :, 3] & np.isfinite(p) & (p > 800) & (p < 1100)
|
| 64 |
+
vm = tracks['target_mask'][i, :, 2] & np.isfinite(vmax) & (vmax >= 0) & (vmax < 250)
|
| 65 |
+
rows.append({'issue_ns': int(times[i]), 'split': split, 'storm_id': str(storm),
|
| 66 |
+
'verification_issue_center': [float(tracks['base_lat'][i]), float(tracks['base_lon'][i]) % 360],
|
| 67 |
+
'target_centers': [[float(tracks['base_lat'][j]), float(tracks['base_lon'][j]) % 360]
|
| 68 |
+
if j is not None else [0., 0.] for j in future],
|
| 69 |
+
'route_mask': [int(j is not None) for j in future],
|
| 70 |
+
'target_pressure_hpa': np.where(pm, p, 0).tolist(), 'pressure_mask': pm.astype(int).tolist(),
|
| 71 |
+
'target_vmax_kt': np.where(vm, vmax, 0).tolist(), 'vmax_mask': vm.astype(int).tolist()})
|
| 72 |
+
# An unconditional atmospheric sample, not a claim that no cyclone exists.
|
| 73 |
+
rows.append({'issue_ns': int(np.datetime64('2021-01-15T00', 'ns').astype('int64')), 'split': 'train',
|
| 74 |
+
'storm_id': None, 'sample_kind': 'unconditional_atmosphere'})
|
| 75 |
+
stamps = sorted({row['issue_ns'] + i*STEP for row in rows for i in range(-8, 21)})
|
| 76 |
+
indices = {t: i for i, t in enumerate(stamps)}
|
| 77 |
+
for row in rows:
|
| 78 |
+
row['issue_index'] = indices[row['issue_ns']]
|
| 79 |
+
result = {'schema': 'v164-preparation-plan-v1', 'dataset_role': 'pipeline_pilot_not_full_training',
|
| 80 |
+
'source': SOURCE, 'track_sha256': TRACK_SHA, 'rows': rows, 'time': stamps,
|
| 81 |
+
'estimated_field_bytes': len(stamps)*31*(181*360+241*321)*4,
|
| 82 |
+
'network_warning': 'ARCO upper-air chunks contain all 37 levels and the globe; transfer can greatly exceed saved size.',
|
| 83 |
+
'sampling': 'mature intense issue-time storms for association testing; not representative benchmark sampling',
|
| 84 |
+
'counts': list(counts)}
|
| 85 |
+
result['plan_sha256'] = hashlib.sha256(json.dumps(result, sort_keys=True).encode()).hexdigest()
|
| 86 |
+
return result
|
| 87 |
+
|
| 88 |
+
|
| 89 |
+
def resolve_associations(plan, source):
|
| 90 |
+
selected, rejected = [], []
|
| 91 |
+
for split, required in zip(('train', 'validation', 'test'), plan['counts']):
|
| 92 |
+
accepted = 0
|
| 93 |
+
for row in plan['rows']:
|
| 94 |
+
if row['split'] != split or not row.get('storm_id'):
|
| 95 |
+
continue
|
| 96 |
+
print('Checking issue-time analysis association:', split, row['storm_id'], flush=True)
|
| 97 |
+
pressure = source.domains(source.read('mean_sea_level_pressure', row['issue_ns'])/100)[1]
|
| 98 |
+
center = analysis_center(pressure[None], source.rl, source.ro)
|
| 99 |
+
truth = np.asarray(row['verification_issue_center'])
|
| 100 |
+
distance = float(np.linalg.norm((center-truth)*[111.2, 111.2*np.cos(np.deg2rad(truth[0]))]))
|
| 101 |
+
if distance > 400:
|
| 102 |
+
rejected.append({'storm_id': row['storm_id'], 'issue_ns': row['issue_ns'], 'distance_km': distance})
|
| 103 |
+
continue
|
| 104 |
+
selected.append(dict(row))
|
| 105 |
+
accepted += 1
|
| 106 |
+
if accepted == required:
|
| 107 |
+
break
|
| 108 |
+
if accepted < required:
|
| 109 |
+
raise ValueError(f'Insufficient unambiguous issue-time associations for {split}; no dataset was published')
|
| 110 |
+
selected.extend(dict(r) for r in plan['rows'] if not r.get('storm_id'))
|
| 111 |
+
times = sorted({r['issue_ns']+i*STEP for r in selected for i in range(-8, 21)})
|
| 112 |
+
lookup = {t: i for i, t in enumerate(times)}
|
| 113 |
+
for row in selected:
|
| 114 |
+
row['issue_index'] = lookup[row['issue_ns']]
|
| 115 |
+
result = dict(plan, rows=selected, time=times, rejected_issue_associations=rejected,
|
| 116 |
+
estimated_field_bytes=len(times)*31*(181*360+241*321)*4)
|
| 117 |
+
result.pop('plan_sha256', None)
|
| 118 |
+
result['plan_sha256'] = hashlib.sha256(json.dumps(result, sort_keys=True).encode()).hexdigest()
|
| 119 |
+
return result
|
| 120 |
+
|
| 121 |
+
|
| 122 |
+
class Era5Source:
|
| 123 |
+
def __init__(self):
|
| 124 |
+
import xarray as xr
|
| 125 |
+
self.ds = xr.open_zarr(SOURCE, consolidated=True, chunks=None)
|
| 126 |
+
missing = set(SURFACE + UPPER + ['land_sea_mask', 'geopotential_at_surface']) - set(self.ds)
|
| 127 |
+
if missing:
|
| 128 |
+
raise ValueError(f'ERA5 missing variables: {sorted(missing)}')
|
| 129 |
+
if not set(trainer.LEVELS).issubset(set(self.ds.level.values.tolist())):
|
| 130 |
+
raise ValueError('ERA5 does not supply all required pressure levels')
|
| 131 |
+
self.lat, self.lon = self.ds.latitude.values, self.ds.longitude.values
|
| 132 |
+
if not np.allclose(self.lat, np.linspace(90, -90, 721)) or not np.allclose(self.lon, np.arange(1440)*.25):
|
| 133 |
+
raise ValueError('Unexpected ERA5 native grid; no implicit regridding')
|
| 134 |
+
self.gl, self.go = self.lat[::4], self.lon[::4]
|
| 135 |
+
self.ys = np.flatnonzero((self.lat >= 0) & (self.lat <= 60))
|
| 136 |
+
self.xs = np.flatnonzero((self.lon >= 100) & (self.lon <= 180))
|
| 137 |
+
self.rl, self.ro = self.lat[self.ys], self.lon[self.xs]
|
| 138 |
+
|
| 139 |
+
def read(self, name, ns, levels=None):
|
| 140 |
+
date = np.datetime64(int(ns), 'ns')
|
| 141 |
+
stop = np.datetime64(self.ds.attrs['valid_time_stop'])
|
| 142 |
+
if date > stop or date < np.datetime64(self.ds.attrs['valid_time_start']):
|
| 143 |
+
raise ValueError(f'{date} outside final ERA5 availability; no ERA5T substitution')
|
| 144 |
+
a = self.ds[name].sel(time=date)
|
| 145 |
+
if levels is not None:
|
| 146 |
+
a = a.sel(level=list(levels)).transpose('level', 'latitude', 'longitude')
|
| 147 |
+
else:
|
| 148 |
+
a = a.transpose('latitude', 'longitude')
|
| 149 |
+
for attempt in range(3):
|
| 150 |
+
try:
|
| 151 |
+
return np.asarray(a.values, dtype='float32')
|
| 152 |
+
except (OSError, TimeoutError, ConnectionError):
|
| 153 |
+
if attempt == 2:
|
| 154 |
+
raise
|
| 155 |
+
time.sleep(2**attempt)
|
| 156 |
+
|
| 157 |
+
def domains(self, a):
|
| 158 |
+
return a[..., ::4, ::4], a[..., self.ys[:, None], self.xs[None, :]]
|
| 159 |
+
|
| 160 |
+
def frame(self, ns):
|
| 161 |
+
out = {name: np.empty((31, len(lat), len(lon)), dtype='float32')
|
| 162 |
+
for name, lat, lon in [('global', self.gl, self.go), ('regional', self.rl, self.ro)]}
|
| 163 |
+
for i, name in enumerate(SURFACE):
|
| 164 |
+
print(' reading', name, flush=True)
|
| 165 |
+
a = self.read(name, ns)
|
| 166 |
+
if i < 2:
|
| 167 |
+
a = a / 100.0
|
| 168 |
+
if i == 5:
|
| 169 |
+
a = np.where((a >= 260) & (a <= 330), a, np.nan)
|
| 170 |
+
for domain, values in zip(out, self.domains(a)):
|
| 171 |
+
out[domain][i] = values
|
| 172 |
+
for j, name in enumerate(UPPER):
|
| 173 |
+
print(' reading', name, 'levels', trainer.LEVELS, flush=True)
|
| 174 |
+
a = self.read(name, ns, trainer.LEVELS)
|
| 175 |
+
for domain, values in zip(out, self.domains(a)):
|
| 176 |
+
for level in range(5):
|
| 177 |
+
out[domain][6+level*5+j] = values[level]
|
| 178 |
+
for name, values in out.items():
|
| 179 |
+
validate_frame(values, name)
|
| 180 |
+
return out
|
| 181 |
+
|
| 182 |
+
def static(self, ns):
|
| 183 |
+
land = np.clip(self.read('land_sea_mask', ns), 0, 1)
|
| 184 |
+
altitude = self.read('geopotential_at_surface', ns)/9.80665
|
| 185 |
+
result = {}
|
| 186 |
+
for i, (name, lat, lon) in enumerate([('global', self.gl, self.go), ('regional', self.rl, self.ro)]):
|
| 187 |
+
yy, xx = np.meshgrid(lat, lon, indexing='ij')
|
| 188 |
+
result[name] = np.stack([yy/90, xx/180-1, self.domains(land)[i], self.domains(altitude)[i]/8000]).astype('float32')
|
| 189 |
+
return result
|
| 190 |
+
|
| 191 |
+
|
| 192 |
+
def validate_frame(a, name):
|
| 193 |
+
if not np.isfinite(a[[i for i in range(31) if i != 5]]).all():
|
| 194 |
+
raise ValueError(f'Nonfinite {name} atmospheric field; refusing zero-fill')
|
| 195 |
+
if not (800 < a[0].min() <= a[0].max() < 1120):
|
| 196 |
+
raise ValueError(f'{name} MSLP range/unit failure')
|
| 197 |
+
if not (150 < a[4].min() <= a[4].max() < 350):
|
| 198 |
+
raise ValueError(f'{name} temperature range/unit failure')
|
| 199 |
+
|
| 200 |
+
|
| 201 |
+
def analysis_center(field, lat, lon, previous=None):
|
| 202 |
+
pressure = field[0].copy()
|
| 203 |
+
yy, xx = np.meshgrid(lat, lon, indexing='ij')
|
| 204 |
+
allowed = (yy > 0) & (yy < 40)
|
| 205 |
+
if previous is not None:
|
| 206 |
+
distance = ((yy-previous[0])*111.2)**2 + ((xx-previous[1])*111.2*np.cos(np.deg2rad(previous[0])))**2
|
| 207 |
+
allowed &= distance < 600**2
|
| 208 |
+
index = np.argmin(np.where(allowed, pressure, np.inf))
|
| 209 |
+
return np.array([yy.flat[index], xx.flat[index]], dtype='float32')
|
| 210 |
+
|
| 211 |
+
|
| 212 |
+
def materialize(plan, root, source, max_gb=4, smoke=False):
|
| 213 |
+
root = Path(root)
|
| 214 |
+
root.mkdir(parents=True, exist_ok=True)
|
| 215 |
+
plan_path, progress_path = root/'preparation_plan.json', root/'preparation_progress.json'
|
| 216 |
+
if plan_path.exists() and json.loads(plan_path.read_text())['plan_sha256'] != plan['plan_sha256']:
|
| 217 |
+
raise ValueError('Existing preparation uses a different plan; choose a separate dataset directory')
|
| 218 |
+
needed = plan['estimated_field_bytes']
|
| 219 |
+
if needed > max_gb*1e9:
|
| 220 |
+
raise ValueError(f'Plan needs {needed/1e9:.2f} GB; exceeds explicit {max_gb} GB limit')
|
| 221 |
+
existing = sum(p.stat().st_size for p in root.glob('*.npy'))
|
| 222 |
+
if shutil.disk_usage(root).free < max(0, needed-existing) + (128 if smoke else 512)*1024**2:
|
| 223 |
+
raise OSError('Insufficient storage for data and safety margin')
|
| 224 |
+
trainer.atomic_json(plan_path, plan)
|
| 225 |
+
progress = json.loads(progress_path.read_text()) if progress_path.exists() else {'frames': {}}
|
| 226 |
+
if (root/'manifest.json').exists():
|
| 227 |
+
trainer.load_manifest(root/'manifest.json', smoke=smoke)
|
| 228 |
+
print('Prepared dataset already verified.', flush=True)
|
| 229 |
+
return
|
| 230 |
+
arrays = {}
|
| 231 |
+
for name, lat, lon in [('global', source.gl, source.go), ('regional', source.rl, source.ro)]:
|
| 232 |
+
p = root/(name+'.npy')
|
| 233 |
+
shape = (len(plan['time']), 31, len(lat), len(lon))
|
| 234 |
+
arrays[name] = np.lib.format.open_memmap(p, mode='r+' if p.exists() else 'w+', dtype='float32', shape=shape)
|
| 235 |
+
if arrays[name].shape != shape:
|
| 236 |
+
raise ValueError('Partial data shape mismatch')
|
| 237 |
+
for i, ns in enumerate(plan['time']):
|
| 238 |
+
key = str(ns)
|
| 239 |
+
def frame_digest():
|
| 240 |
+
h = hashlib.sha256()
|
| 241 |
+
for a in arrays.values():
|
| 242 |
+
h.update(np.ascontiguousarray(a[i]).tobytes())
|
| 243 |
+
return h.hexdigest()
|
| 244 |
+
if key in progress['frames'] and frame_digest() == progress['frames'][key]:
|
| 245 |
+
continue
|
| 246 |
+
print(f'ERA5 frame {i+1}/{len(plan["time"])}: {np.datetime64(ns, "ns")} (31 channels)', flush=True)
|
| 247 |
+
frame = source.frame(ns)
|
| 248 |
+
for name, a in arrays.items():
|
| 249 |
+
a[i] = frame[name]
|
| 250 |
+
a.flush()
|
| 251 |
+
progress['frames'][key] = frame_digest()
|
| 252 |
+
progress.update(completed=len(progress['frames']), total=len(plan['time']), last_completed_utc=str(np.datetime64(ns, 'ns')))
|
| 253 |
+
trainer.atomic_json(progress_path, progress)
|
| 254 |
+
np.save(root/'time.npy', np.asarray(plan['time'], dtype='int64'))
|
| 255 |
+
for name, a in source.static(plan['time'][0]).items():
|
| 256 |
+
np.save(root/(name+'_static.npy'), a)
|
| 257 |
+
sums, squares, counts = np.zeros(31), np.zeros(31), np.zeros(31)
|
| 258 |
+
for i, ns in enumerate(plan['time']):
|
| 259 |
+
if not 2000 <= year(ns) <= 2021:
|
| 260 |
+
continue
|
| 261 |
+
for a in arrays.values():
|
| 262 |
+
values = np.asarray(a[i], dtype='float64').reshape(31, -1)
|
| 263 |
+
sums += np.nansum(values, axis=1)
|
| 264 |
+
squares += np.nansum(values**2, axis=1)
|
| 265 |
+
counts += np.isfinite(values).sum(axis=1)
|
| 266 |
+
if (counts == 0).any():
|
| 267 |
+
raise ValueError('Missing training-only normalization data')
|
| 268 |
+
mean = sums/counts
|
| 269 |
+
std = np.maximum(np.sqrt(np.maximum(squares/counts-mean**2, 0)), 1e-6)
|
| 270 |
+
rows = []
|
| 271 |
+
for original in plan['rows']:
|
| 272 |
+
row = dict(original)
|
| 273 |
+
i = row['issue_index']
|
| 274 |
+
# Association is derived solely from issue/past ERA5, never best-track coordinates.
|
| 275 |
+
center = analysis_center(arrays['regional'][i], source.rl, source.ro)
|
| 276 |
+
past = analysis_center(arrays['regional'][i-1], source.rl, source.ro, center)
|
| 277 |
+
row['center'] = center.tolist()
|
| 278 |
+
row['motion_km_per_6h'] = [float((center[1]-past[1])*111.2*np.cos(np.deg2rad(center[0]))), float((center[0]-past[0])*111.2)]
|
| 279 |
+
truth = row.pop('verification_issue_center', None)
|
| 280 |
+
if truth is not None:
|
| 281 |
+
dist = np.linalg.norm((center-np.array(truth))*[111.2, 111.2*np.cos(np.deg2rad(truth[0]))])
|
| 282 |
+
row['issue_association_error_km'] = float(dist)
|
| 283 |
+
if dist > 400:
|
| 284 |
+
row['route_mask'] = row['pressure_mask'] = row['vmax_mask'] = [0]*20
|
| 285 |
+
rows.append(row)
|
| 286 |
+
(root/'samples.jsonl').write_text(''.join(json.dumps(row)+'\n' for row in rows))
|
| 287 |
+
m = {'schema': 'v164-atmosphere-v1', 'synthetic': bool(smoke), 'dataset_role': plan['dataset_role'],
|
| 288 |
+
'channels': trainer.CHANNELS, 'units': trainer.UNITS,
|
| 289 |
+
'global_lat': source.gl.tolist(), 'global_lon': source.go.tolist(),
|
| 290 |
+
'regional_lat': source.rl.tolist(), 'regional_lon': source.ro.tolist(),
|
| 291 |
+
'normalization': {'fit_split': 'train', 'end_year': 2021, 'mean': mean.tolist(), 'std': std.tolist()},
|
| 292 |
+
'future_weather_as_inputs': False, 'official_forecasts_as_inputs': False,
|
| 293 |
+
'inputs_are_analysis_valid_time_hindcasts_not_operational_release_time': True,
|
| 294 |
+
'source': SOURCE, 'track_archive_sha256': plan['track_sha256'], 'files': {}}
|
| 295 |
+
for name in ('global', 'regional', 'global_static', 'regional_static', 'time', 'samples'):
|
| 296 |
+
p = root/(name+('.jsonl' if name == 'samples' else '.npy'))
|
| 297 |
+
m['files'][name] = {'path': p.name, 'sha256': trainer.digest(p)}
|
| 298 |
+
for split in ('train', 'validation', 'test'):
|
| 299 |
+
if not any(row['split'] == split and any(row.get('route_mask', [])) for row in rows):
|
| 300 |
+
raise ValueError(f'No valid analysis-associated storm labels in {split}. Pilot is insufficient; do not train it.')
|
| 301 |
+
# Validate a temporary manifest before publishing the ready marker.
|
| 302 |
+
candidate = root/'manifest.pending.json'
|
| 303 |
+
trainer.atomic_json(candidate, m)
|
| 304 |
+
trainer.load_manifest(candidate, smoke=smoke)
|
| 305 |
+
candidate.replace(root/'manifest.json')
|
| 306 |
+
print('PILOT DATA READY. This is not the full training corpus or evidence of forecast skill.', flush=True)
|
| 307 |
+
|
| 308 |
+
|
| 309 |
+
def main():
|
| 310 |
+
p = argparse.ArgumentParser(description=__doc__)
|
| 311 |
+
p.add_argument('--track', type=Path, required=True)
|
| 312 |
+
p.add_argument('--output', type=Path, required=True)
|
| 313 |
+
p.add_argument('--prepare', action='store_true')
|
| 314 |
+
p.add_argument('--max-gb', type=float, default=4)
|
| 315 |
+
p.add_argument('--cases', type=int, nargs=3, default=[3, 1, 1], metavar=('TRAIN', 'VAL', 'TEST'))
|
| 316 |
+
args = p.parse_args()
|
| 317 |
+
if min(args.cases) < 1:
|
| 318 |
+
p.error('Each split needs at least one case')
|
| 319 |
+
plan = build_plan(args.track, args.cases, oversample=4 if args.prepare else 1)
|
| 320 |
+
if args.prepare:
|
| 321 |
+
args.output.mkdir(parents=True, exist_ok=True)
|
| 322 |
+
lock = args.output/'preparation.lock'
|
| 323 |
+
fd = os.open(lock, os.O_CREAT | os.O_EXCL | os.O_WRONLY)
|
| 324 |
+
os.close(fd)
|
| 325 |
+
try:
|
| 326 |
+
source = Era5Source()
|
| 327 |
+
previous = args.output/'preparation_plan.json'
|
| 328 |
+
if previous.exists():
|
| 329 |
+
plan = json.loads(previous.read_text())
|
| 330 |
+
if plan['track_sha256'] != TRACK_SHA or plan['counts'] != args.cases or plan['source'] != SOURCE:
|
| 331 |
+
raise ValueError('Existing preparation configuration differs; choose a new output directory')
|
| 332 |
+
saved_hash = plan['plan_sha256']
|
| 333 |
+
payload = {k: v for k, v in plan.items() if k != 'plan_sha256'}
|
| 334 |
+
if hashlib.sha256(json.dumps(payload, sort_keys=True).encode()).hexdigest() != saved_hash:
|
| 335 |
+
raise ValueError('Existing preparation plan checksum mismatch')
|
| 336 |
+
else:
|
| 337 |
+
plan = resolve_associations(plan, source)
|
| 338 |
+
print(json.dumps({k: v for k, v in plan.items() if k not in ('rows', 'time')}, indent=2), flush=True)
|
| 339 |
+
materialize(plan, args.output, source, args.max_gb)
|
| 340 |
+
finally:
|
| 341 |
+
lock.unlink(missing_ok=True)
|
| 342 |
+
else:
|
| 343 |
+
print(json.dumps({k: v for k, v in plan.items() if k not in ('rows', 'time')}, indent=2), flush=True)
|
| 344 |
+
print('Plan only. Add --prepare to download. No training is started.', flush=True)
|
| 345 |
+
|
| 346 |
+
|
| 347 |
+
if __name__ == '__main__':
|
| 348 |
+
main()
|
training/trackformer_1_2/source/data_preparation/v164_prepare_mac.py
ADDED
|
@@ -0,0 +1,36 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Mac data preparation with three concurrent source-variable reads.
|
| 2 |
+
|
| 3 |
+
Uses the same plan, fields, validation and resume records as the serial preparer.
|
| 4 |
+
This utility is not called by the Drive-only Colab notebook.
|
| 5 |
+
"""
|
| 6 |
+
from concurrent.futures import ThreadPoolExecutor
|
| 7 |
+
|
| 8 |
+
import v164_prepare_data as prep
|
| 9 |
+
|
| 10 |
+
|
| 11 |
+
class ConcurrentEra5Source(prep.Era5Source):
|
| 12 |
+
_pending = None
|
| 13 |
+
|
| 14 |
+
def read(self, name, ns, levels=None):
|
| 15 |
+
key = (name, int(ns), tuple(levels) if levels is not None else None)
|
| 16 |
+
if self._pending is not None and key in self._pending:
|
| 17 |
+
return self._pending.pop(key).result()
|
| 18 |
+
return super().read(name, ns, levels)
|
| 19 |
+
|
| 20 |
+
def frame(self, ns):
|
| 21 |
+
jobs = [(name, None) for name in prep.SURFACE]
|
| 22 |
+
jobs += [(name, tuple(prep.trainer.LEVELS)) for name in prep.UPPER]
|
| 23 |
+
with ThreadPoolExecutor(max_workers=3) as pool:
|
| 24 |
+
self._pending = {(name, int(ns), levels): pool.submit(super().read, name, ns, levels)
|
| 25 |
+
for name, levels in jobs}
|
| 26 |
+
try:
|
| 27 |
+
return super().frame(ns)
|
| 28 |
+
finally:
|
| 29 |
+
for future in self._pending.values():
|
| 30 |
+
future.cancel()
|
| 31 |
+
self._pending = None
|
| 32 |
+
|
| 33 |
+
|
| 34 |
+
if __name__ == '__main__':
|
| 35 |
+
prep.Era5Source = ConcurrentEra5Source
|
| 36 |
+
prep.main()
|
training/trackformer_1_2/source/data_preparation/v164_reuse_impl.py
ADDED
|
@@ -0,0 +1,283 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Embedded in the standalone v164 trainer: coarse archive plus native MSLP patches."""
|
| 2 |
+
REUSE_ARCH = 'wp-coarse8-native-pressure1-v164'
|
| 3 |
+
REUSE_CHANNELS = ['mslp', 'hgt500', 'uwnd850', 'vwnd850', 'uwnd500', 'vwnd500', 'uwnd200', 'vwnd200']
|
| 4 |
+
|
| 5 |
+
|
| 6 |
+
def load_reuse(path):
|
| 7 |
+
path = Path(path)
|
| 8 |
+
if not path.is_file():
|
| 9 |
+
raise FileNotFoundError('Combined dataset missing. Run the updated Drive setup cell first: '+str(path))
|
| 10 |
+
m = json.loads(path.read_text())
|
| 11 |
+
if m.get('schema') != 'v164-reuse-pressure-v1' or m.get('pilot', True):
|
| 12 |
+
raise ValueError('Combined non-pilot archive required; no pilot fallback')
|
| 13 |
+
for name, entry in m['files'].items():
|
| 14 |
+
p = path.parent/entry['path']
|
| 15 |
+
print('Checking dataset file:', name, flush=True)
|
| 16 |
+
if digest(p) != entry['sha256']:
|
| 17 |
+
raise ValueError('Dataset checksum mismatch: '+name)
|
| 18 |
+
entry['resolved'] = str(p)
|
| 19 |
+
with np.load(m['files']['basin']['resolved'], allow_pickle=False) as z:
|
| 20 |
+
if z['channels'].tolist() != REUSE_CHANNELS[1:]:
|
| 21 |
+
raise ValueError('Coarse channel order mismatch')
|
| 22 |
+
times, lat, lon = z['time'], z['lat'], z['lon']
|
| 23 |
+
q = z['q']
|
| 24 |
+
scale, offset = z['scale'], z['offset']
|
| 25 |
+
if q.shape != (len(times),7,25,33) or not np.all(np.diff(times) == 6*HOUR):
|
| 26 |
+
raise ValueError('Coarse grid/time mismatch')
|
| 27 |
+
qpath = path.parent/'basin_q_verified.npy'
|
| 28 |
+
# Refresh the compact worker mmap from the checksum-verified source every preflight.
|
| 29 |
+
temporary = qpath.with_suffix('.pending.npy')
|
| 30 |
+
np.save(temporary, q)
|
| 31 |
+
os.replace(temporary, qpath)
|
| 32 |
+
slp = np.load(m['files']['slp']['resolved'], mmap_mode='r')
|
| 33 |
+
if slp.shape != (len(times),25,33) or not np.isfinite(slp).all():
|
| 34 |
+
raise ValueError('Coarse pressure coverage invalid')
|
| 35 |
+
with np.load(m['files']['track']['resolved'], allow_pickle=False) as z:
|
| 36 |
+
track = {k:z[k] for k in ['base_time','storm_id','base_lat','base_lon','target','target_mask']}
|
| 37 |
+
issue_intensity = z['track'][:,-1,4:6]*z['track_std'][4:6]+z['track_mean'][4:6]
|
| 38 |
+
plan = json.loads(Path(m['files']['plan']['resolved']).read_text())
|
| 39 |
+
if len(plan['rows']) != 1000 or plan['shape'] != [1000,29,121,121] or plan['lead_hours'] != list(range(-48,121,6)):
|
| 40 |
+
raise ValueError('Expected the verified 1000-window native pressure package, not the six-sample pilot')
|
| 41 |
+
pressure = np.load(m['files']['pressure']['resolved'], mmap_mode='r')
|
| 42 |
+
if pressure.shape != (1000,29,121,121):
|
| 43 |
+
raise ValueError('Native pressure shape mismatch')
|
| 44 |
+
patches = {}
|
| 45 |
+
for pi,row in enumerate(plan['rows']):
|
| 46 |
+
ti = row['track_archive_row']
|
| 47 |
+
if int(track['base_time'][ti]) != row['issue_ns'] or str(track['storm_id'][ti]) != row['storm_id']:
|
| 48 |
+
raise ValueError('Native patch/track identity mismatch')
|
| 49 |
+
patches[ti] = pi
|
| 50 |
+
first_year = {}
|
| 51 |
+
years = track['base_time'].astype('datetime64[ns]').astype('datetime64[Y]').astype(int)+1970
|
| 52 |
+
for sid,y in zip(track['storm_id'],years):
|
| 53 |
+
first_year[str(sid)] = min(first_year.get(str(sid),9999),int(y))
|
| 54 |
+
lookup = {(str(s),int(t)):i for i,(s,t) in enumerate(zip(track['storm_id'],track['base_time']))}
|
| 55 |
+
rows = []
|
| 56 |
+
for i,ns in enumerate(track['base_time']):
|
| 57 |
+
sid = str(track['storm_id'][i])
|
| 58 |
+
y = first_year[sid]
|
| 59 |
+
split = 'train' if 2000 <= y <= 2021 else 'validation' if 2022 <= y <= 2023 else 'test' if 2024 <= y <= 2025 else None
|
| 60 |
+
if split is None:
|
| 61 |
+
continue
|
| 62 |
+
low,high = {'train':(2000,2021),'validation':(2022,2023),'test':(2024,2025)}[split]
|
| 63 |
+
bounds = np.array([ns-48*HOUR,ns+120*HOUR],dtype='int64').astype('datetime64[ns]').astype('datetime64[Y]').astype(int)+1970
|
| 64 |
+
ai = int(np.searchsorted(times,ns))
|
| 65 |
+
center = [float(track['base_lat'][i]),float(track['base_lon'][i])%360]
|
| 66 |
+
if bounds[0] < low or bounds[1] > high or ai < 8 or ai+20 >= len(times) or times[ai] != ns:
|
| 67 |
+
continue
|
| 68 |
+
if not np.isfinite(center).all() or not (0 < center[0] < 60 and 100 < center[1] < 180):
|
| 69 |
+
continue
|
| 70 |
+
previous = lookup.get((sid,int(ns-6*HOUR)))
|
| 71 |
+
motion = [0.,0.]
|
| 72 |
+
if previous is not None:
|
| 73 |
+
old = [float(track['base_lat'][previous]),float(track['base_lon'][previous])%360]
|
| 74 |
+
if np.isfinite(old).all():
|
| 75 |
+
motion = [(center[1]-old[1])*111.2*np.cos(np.deg2rad(center[0])),(center[0]-old[0])*111.2]
|
| 76 |
+
future = [lookup.get((sid,int(ns+j*6*HOUR))) for j in range(1,21)]
|
| 77 |
+
target_centers = [[float(track['base_lat'][j]),float(track['base_lon'][j])%360] if j is not None else [0.,0.] for j in future]
|
| 78 |
+
rm = [float(j is not None and np.isfinite(target_centers[k]).all()) for k,j in enumerate(future)]
|
| 79 |
+
target_centers = np.nan_to_num(target_centers).tolist()
|
| 80 |
+
wind,p = track['target'][i,:,2],track['target'][i,:,3]
|
| 81 |
+
wm = track['target_mask'][i,:,2].astype(bool)&np.isfinite(wind)&(wind>=0)&(wind<250)
|
| 82 |
+
pm = track['target_mask'][i,:,3].astype(bool)&np.isfinite(p)&(p>800)&(p<1100)
|
| 83 |
+
pi = patches.get(i,-1)
|
| 84 |
+
if pi >= 0 and plan['rows'][pi]['split'] != split:
|
| 85 |
+
raise ValueError('Patch split mismatch')
|
| 86 |
+
issue = issue_intensity[i]
|
| 87 |
+
iv = [bool(np.isfinite(issue[0]) and 0<=issue[0]<250),bool(np.isfinite(issue[1]) and 800<issue[1]<1100)]
|
| 88 |
+
rows.append({'split':split,'storm_id':sid,'atlas':ai,'patch':pi,'center':center,'motion':motion,
|
| 89 |
+
'issue_intensity':[float(issue[k]) if iv[k] else 0. for k in range(2)],'issue_mask':iv,
|
| 90 |
+
'target_centers':target_centers,'route_mask':rm,
|
| 91 |
+
'pressure':np.where(pm,p,0).tolist(),'pressure_mask':pm.tolist(),
|
| 92 |
+
'vmax':np.where(wm,wind,0).tolist(),'vmax_mask':wm.tolist()})
|
| 93 |
+
counts = {s:sum(r['split']==s for r in rows) for s in ['train','validation','test']}
|
| 94 |
+
patch_counts = {s:sum(r['split']==s and r['patch']>=0 for r in rows) for s in counts}
|
| 95 |
+
if counts['train'] < 10000 or patch_counts != {'train':800,'validation':100,'test':100}:
|
| 96 |
+
raise ValueError(f'Incomplete combined coverage: {counts}, native patches {patch_counts}')
|
| 97 |
+
train_times = np.flatnonzero((times.astype('datetime64[ns]').astype('datetime64[Y]').astype(int)+1970>=2000)&
|
| 98 |
+
(times.astype('datetime64[ns]').astype('datetime64[Y]').astype(int)+1970<=2021))
|
| 99 |
+
sums,squares = np.zeros(8),np.zeros(8)
|
| 100 |
+
n = 0
|
| 101 |
+
for start in range(0,len(train_times),256):
|
| 102 |
+
ii = train_times[start:start+256]
|
| 103 |
+
b = np.concatenate([np.asarray(slp[ii],dtype='float32')[:,None],q[ii].astype('float32')*scale[None,:,None,None]+offset[None,:,None,None]],1).astype('float64')
|
| 104 |
+
if not np.isfinite(b).all():
|
| 105 |
+
raise ValueError('Nonfinite coarse training fields')
|
| 106 |
+
sums += b.sum((0,2,3)); squares += (b*b).sum((0,2,3)); n += len(ii)*25*33
|
| 107 |
+
mean = sums/n
|
| 108 |
+
std = np.maximum(np.sqrt(np.maximum(squares/n-mean*mean,0)),1e-4)
|
| 109 |
+
m.update(rows=rows, plan=plan, qpath=str(qpath), scale=scale.tolist(), offset=offset.tolist(),
|
| 110 |
+
channels=REUSE_CHANNELS,units=['hPa','m','m/s','m/s','m/s','m/s','m/s','m/s'],
|
| 111 |
+
global_lat=lat.tolist(),global_lon=lon.tolist(),regional_lat=list(np.linspace(15,-15,121)),
|
| 112 |
+
regional_lon=list(np.linspace(-15,15,121)),normalization={'mean':mean.tolist(),'std':std.tolist(),'fit_split':'train','end_year':2021},
|
| 113 |
+
identity=digest(path),counts=counts,patch_counts=patch_counts)
|
| 114 |
+
print(json.dumps({'all_archive_windows':counts,'native_pressure_windows':patch_counts,'architecture':REUSE_ARCH}),flush=True)
|
| 115 |
+
return m
|
| 116 |
+
|
| 117 |
+
|
| 118 |
+
class ReuseDataset(Dataset):
|
| 119 |
+
def __init__(self,m,split,training=False):
|
| 120 |
+
self.m = m
|
| 121 |
+
self.rows = [r for r in m['rows'] if r['split']==split]
|
| 122 |
+
self.arrays = None
|
| 123 |
+
|
| 124 |
+
def __len__(self):
|
| 125 |
+
return len(self.rows)
|
| 126 |
+
|
| 127 |
+
def __getitem__(self,i):
|
| 128 |
+
m,row = self.m,self.rows[i]
|
| 129 |
+
if self.arrays is None:
|
| 130 |
+
self.arrays = {k:np.load(m['files'][k]['resolved'],mmap_mode='r') for k in ['slp','pressure']}
|
| 131 |
+
self.arrays['q'] = np.load(m['qpath'],mmap_mode='r')
|
| 132 |
+
with np.load(m['files']['geography']['resolved'],allow_pickle=False) as geo:
|
| 133 |
+
self.arrays['land'] = geo['land_fraction']
|
| 134 |
+
self.arrays['height'] = geo['elevation_m']/8000
|
| 135 |
+
ai = row['atlas']
|
| 136 |
+
ii = slice(ai-8,ai+21)
|
| 137 |
+
coarse = np.concatenate([np.array(self.arrays['slp'][ii],dtype='float32')[:,None],
|
| 138 |
+
np.array(self.arrays['q'][ii],dtype='float32')*np.array(m['scale'],dtype='float32')[None,:,None,None]+np.array(m['offset'],dtype='float32')[None,:,None,None]],1)
|
| 139 |
+
mean,std = np.array(m['normalization']['mean'],dtype='float32'),np.array(m['normalization']['std'],dtype='float32')
|
| 140 |
+
coarse = (coarse-mean[None,:,None,None])/std[None,:,None,None]
|
| 141 |
+
pi = row['patch']
|
| 142 |
+
anchor = m['plan']['rows'][pi] if pi>=0 else {'center_lat':round(row['center'][0]*4)/4,'center_lon':round(row['center'][1]*4)/4}
|
| 143 |
+
lat = anchor['center_lat']+np.linspace(15,-15,121,dtype='float32')
|
| 144 |
+
lon = anchor['center_lon']+np.linspace(-15,15,121,dtype='float32')
|
| 145 |
+
yy,xx = np.meshgrid(lat,lon,indexing='ij')
|
| 146 |
+
y = np.rint((90-yy)*4).astype(int); x = np.rint(xx*4).astype(int)%1440
|
| 147 |
+
rs = np.stack([yy/90,xx/180-1,self.arrays['land'][y,x],self.arrays['height'][y,x]]).astype('float32')
|
| 148 |
+
cy,cx = np.meshgrid(m['global_lat'],m['global_lon'],indexing='ij')
|
| 149 |
+
iy = np.rint((90-cy)*4).astype(int); ix = np.rint(cx*4).astype(int)%1440
|
| 150 |
+
gs = np.stack([cy/90,cx/180-1,self.arrays['land'][iy,ix],self.arrays['height'][iy,ix]]).astype('float32')
|
| 151 |
+
detail = np.zeros((29,1,121,121),dtype='float32')
|
| 152 |
+
if pi>=0:
|
| 153 |
+
detail[:,0] = (np.array(self.arrays['pressure'][pi],dtype='float32')-mean[0])/std[0]
|
| 154 |
+
if not np.isfinite(detail).all():
|
| 155 |
+
raise ValueError('Native pressure contains nonfinite fields')
|
| 156 |
+
inputs = {'global_history':coarse[:9], 'regional_history':detail[:9],
|
| 157 |
+
'global_static':gs,'regional_static':rs,'detail_available':np.array([pi>=0],dtype='float32'),
|
| 158 |
+
'center':np.array(row['center'],dtype='float32'),'motion':np.array(row['motion'],dtype='float32'),
|
| 159 |
+
'issue_intensity':np.array(row['issue_intensity'],dtype='float32'),'issue_mask':np.array(row['issue_mask'],dtype='float32')}
|
| 160 |
+
targets = {'global':coarse[9:],'global_mask':np.ones_like(coarse[9:]),
|
| 161 |
+
'regional_lat_weight':np.broadcast_to(np.cos(np.deg2rad(lat))[None,None,:,None],(20,1,121,1)).astype('float32'),
|
| 162 |
+
'regional':detail[9:],'regional_mask':np.full_like(detail[9:],float(pi>=0)),
|
| 163 |
+
'center':np.array(row['target_centers'],dtype='float32'),
|
| 164 |
+
'storm_mask':np.array(row['route_mask'],dtype='float32')}
|
| 165 |
+
for k in ['pressure','pressure_mask','vmax','vmax_mask']:
|
| 166 |
+
targets[k] = np.array(row[k],dtype='float32')
|
| 167 |
+
return {k:torch.from_numpy(v.copy()) for k,v in inputs.items()},{k:torch.from_numpy(v.copy()) for k,v in targets.items()}
|
| 168 |
+
|
| 169 |
+
|
| 170 |
+
class ReuseForecaster(nn.Module):
|
| 171 |
+
def __init__(self,m,tiny=False):
|
| 172 |
+
super().__init__()
|
| 173 |
+
self.global_time = nn.Conv2d(8*9,8,1)
|
| 174 |
+
self.regional_time = nn.Conv2d(9,1,1)
|
| 175 |
+
self.global_operator = FieldOperator(8,8 if tiny else 80,False,1 if tiny else 3,tiny)
|
| 176 |
+
self.regional_operator = FieldOperator(1,8 if tiny else 88,False,0,tiny)
|
| 177 |
+
self.parent_fusion = nn.Conv2d(1+8+1,1,1)
|
| 178 |
+
self.object_score = nn.Conv2d(32+3,1,1)
|
| 179 |
+
self.intensity = nn.Sequential(nn.Linear(32+3+4,64),nn.GELU(),nn.Linear(64,2))
|
| 180 |
+
self.register_buffer('mean',torch.tensor(m['normalization']['mean'],dtype=torch.float32)[None,:,None,None])
|
| 181 |
+
self.register_buffer('std',torch.tensor(m['normalization']['std'],dtype=torch.float32)[None,:,None,None])
|
| 182 |
+
yy,xx = np.meshgrid(m['global_lat'],m['global_lon'],indexing='ij')
|
| 183 |
+
self.register_buffer('yy',torch.tensor(yy,dtype=torch.float32)[None])
|
| 184 |
+
self.register_buffer('xx',torch.tensor(xx,dtype=torch.float32)[None])
|
| 185 |
+
|
| 186 |
+
def initial(self,inputs):
|
| 187 |
+
required = {'global_history','regional_history','global_static','regional_static','detail_available','center','motion','issue_intensity','issue_mask'}
|
| 188 |
+
if set(inputs)!=required:
|
| 189 |
+
raise ValueError('Only declared causal inputs are allowed')
|
| 190 |
+
gh,rh = inputs['global_history'],inputs['regional_history']
|
| 191 |
+
rs = inputs['regional_static']
|
| 192 |
+
yy,xx = rs[:,0]*90,(rs[:,1]+1)*180
|
| 193 |
+
grid = torch.stack([(xx-100)/40-1,1-yy/30],-1)
|
| 194 |
+
return {'g':gh[:,-1],'r':rh[:,-1],'gc':self.global_time(gh.flatten(1,2)),
|
| 195 |
+
'rc':self.regional_time(rh.flatten(1,2)),'gs':inputs['global_static'],'rs':rs,
|
| 196 |
+
'grid':grid,'available':inputs['detail_available'],'center':inputs['center'],
|
| 197 |
+
'motion':inputs['motion'],'issue':inputs['issue_intensity'],'issue_mask':inputs['issue_mask']}
|
| 198 |
+
|
| 199 |
+
def step(self,s):
|
| 200 |
+
g,features = self.global_operator(s['g'],s['gc'],torch.ones_like(s['g']),s['gs'])
|
| 201 |
+
valid = (s['grid'].abs()<=1).all(-1)[:,None].to(g.dtype)
|
| 202 |
+
parent = F.grid_sample(g,s['grid'],align_corners=True,padding_mode='zeros')*valid
|
| 203 |
+
r = s['r']
|
| 204 |
+
if bool(s['available'].any()):
|
| 205 |
+
context = self.parent_fusion(torch.cat([s['rc'],parent,valid],1))
|
| 206 |
+
r,_ = self.regional_operator(r,context,torch.ones_like(r),s['rs'])
|
| 207 |
+
physical = g.float()*self.std+self.mean
|
| 208 |
+
p = physical[:,0]
|
| 209 |
+
anomaly = p-F.avg_pool2d(F.pad(p[:,None],(2,2,2,2),mode='replicate'),5,1)[:,0]
|
| 210 |
+
evidence = torch.cat([features.float(),anomaly[:,None]/10,physical[:,2:4]/30],1)
|
| 211 |
+
lat0,lon0 = s['center'][:,0,None,None],s['center'][:,1,None,None]
|
| 212 |
+
plat = lat0+s['motion'][:,1,None,None]/111.2
|
| 213 |
+
plon = lon0+s['motion'][:,0,None,None]/(111.2*lat0.deg2rad().cos().clamp_min(.2))
|
| 214 |
+
d2 = ((self.yy-plat)*111.2).square()+((self.xx-plon)*111.2*lat0.deg2rad().cos()).square()
|
| 215 |
+
score = self.object_score(evidence)[:,0]-anomaly/10-d2/(2*600**2)
|
| 216 |
+
peak = score.flatten(1).argmax(1)
|
| 217 |
+
iy,ix = torch.arange(25,device=g.device)[None,:,None],torch.arange(33,device=g.device)[None,None,:]
|
| 218 |
+
local = ((iy-(peak//33)[:,None,None]).abs()<=1)&((ix-(peak%33)[:,None,None]).abs()<=1)
|
| 219 |
+
pdf = torch.softmax(score.masked_fill(~local,-1e9).flatten(1),1).reshape_as(score)
|
| 220 |
+
center = torch.stack([(pdf*self.yy).sum((1,2)),(pdf*self.xx).sum((1,2))],1)
|
| 221 |
+
pooled = (evidence*pdf[:,None]).sum((2,3))
|
| 222 |
+
scalar = torch.cat([s['issue']/torch.tensor([100.,1000.],device=g.device),s['issue_mask']],1)
|
| 223 |
+
intensity = self.intensity(torch.cat([pooled,scalar],1))
|
| 224 |
+
grid_p = (p*pdf).sum((1,2))
|
| 225 |
+
dy,dx = s['rs'][:,0]*90,(s['rs'][:,1]+1)*180
|
| 226 |
+
sample = torch.stack([2*(center[:,1]-dx[:,0,0])/30-1,2*(dy[:,0,0]-center[:,0])/30-1],-1)[:,None,None]
|
| 227 |
+
inside = (sample.abs()<=1).all((1,2,3)) & (s['available'][:,0]>0)
|
| 228 |
+
detailed_p = F.grid_sample(r.float()*self.std[:,:1]+self.mean[:,:1],sample,align_corners=True)[:,0,0,0]
|
| 229 |
+
grid_p = torch.where(inside,detailed_p,grid_p)
|
| 230 |
+
pressure = grid_p-F.softplus(intensity[:,0])*5
|
| 231 |
+
vmax = s['issue'][:,0]*s['issue_mask'][:,0]+intensity[:,1]*10
|
| 232 |
+
meanlat = (center[:,0]+s['center'][:,0])/2
|
| 233 |
+
motion = torch.stack([(center[:,1]-s['center'][:,1])*111.2*meanlat.deg2rad().cos(),(center[:,0]-s['center'][:,0])*111.2],1)
|
| 234 |
+
return dict(s,g=g,r=r,center=center,motion=motion),{'global':g,'regional':r,'center':center,'pressure':pressure,'vmax':vmax,'grid_pressure':grid_p}
|
| 235 |
+
|
| 236 |
+
def forward(self,inputs,steps=20):
|
| 237 |
+
s = self.initial(inputs)
|
| 238 |
+
out = []
|
| 239 |
+
for _ in range(steps):
|
| 240 |
+
s,o = self.step(s); out.append(o)
|
| 241 |
+
return {k:torch.stack([o[k] for o in out],1) for k in out[0]}
|
| 242 |
+
|
| 243 |
+
|
| 244 |
+
def reuse_main():
|
| 245 |
+
global ARCHITECTURE, AtmosphereDataset
|
| 246 |
+
ARCHITECTURE,AtmosphereDataset = REUSE_ARCH,ReuseDataset
|
| 247 |
+
p = argparse.ArgumentParser(description='v1.2.64 coarse weather + native pressure training; no pilot fallback')
|
| 248 |
+
p.add_argument('batch',type=int,nargs='?',default=10)
|
| 249 |
+
p.add_argument('workers',type=int,nargs='?',default=3)
|
| 250 |
+
p.add_argument('amp',type=int,nargs='?',default=1,choices=[0,1])
|
| 251 |
+
p.add_argument('--manifest',type=Path,default=ROOT/'data/v164_reuse/manifest.json')
|
| 252 |
+
persistent = Path('/content/drive/MyDrive/typhoon_predict') if Path('/content').is_dir() else ROOT
|
| 253 |
+
p.add_argument('--output',type=Path,default=persistent/'training/v164_reuse_pressure')
|
| 254 |
+
p.add_argument('--epochs',type=int,default=12)
|
| 255 |
+
p.add_argument('--device',default='auto',choices=['auto','cpu','mps','cuda'])
|
| 256 |
+
p.add_argument('--preflight',action='store_true')
|
| 257 |
+
args = p.parse_args()
|
| 258 |
+
if args.batch<1 or args.workers<0 or args.epochs<1:
|
| 259 |
+
p.error('Invalid batch/workers/epochs')
|
| 260 |
+
m = load_reuse(args.manifest)
|
| 261 |
+
for split in ['train','validation','test']:
|
| 262 |
+
d = ReuseDataset(m,split)
|
| 263 |
+
d[0]
|
| 264 |
+
d[next(i for i,r in enumerate(d.rows) if r['patch']>=0)]
|
| 265 |
+
if args.preflight:
|
| 266 |
+
print('PASS combined non-pilot archives and native pressure data. No training started.',flush=True)
|
| 267 |
+
return
|
| 268 |
+
if Path('/content').is_dir() and not args.output.resolve().is_relative_to(persistent.resolve()):
|
| 269 |
+
raise ValueError('Colab checkpoints must persist in typhoon_predict on Drive')
|
| 270 |
+
device = device_for(args.device)
|
| 271 |
+
random.seed(164); np.random.seed(164); torch.manual_seed(164)
|
| 272 |
+
model = ReuseForecaster(m).to(device)
|
| 273 |
+
print(json.dumps({'version':VERSION,'architecture':REUSE_ARCH,'parameters':sum(x.numel() for x in model.parameters()),'device':str(device),'windows':m['counts'],'native_windows':m['patch_counts'],'output':str(args.output)}),flush=True)
|
| 274 |
+
optimizer = torch.optim.AdamW(model.parameters(),lr=1e-4,weight_decay=.01)
|
| 275 |
+
scaler = torch.amp.GradScaler('cuda',enabled=device.type=='cuda' and bool(args.amp))
|
| 276 |
+
args.output.mkdir(parents=True,exist_ok=True)
|
| 277 |
+
lock = args.output/'training.lock'
|
| 278 |
+
fd = os.open(lock,os.O_CREAT|os.O_EXCL|os.O_WRONLY)
|
| 279 |
+
os.write(fd,json.dumps({'pid':os.getpid(),'host':socket.gethostname()}).encode()); os.close(fd)
|
| 280 |
+
try:
|
| 281 |
+
train(model,optimizer,scaler,m,args,device)
|
| 282 |
+
finally:
|
| 283 |
+
lock.unlink(missing_ok=True)
|
training/trackformer_1_2/source/data_preparation/v164_structure_research.md
ADDED
|
@@ -0,0 +1,142 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# TrackFormer v1.2.64: field-first atmosphere and cyclone redesign
|
| 2 |
+
|
| 3 |
+
Date: 2026-09-13. Status: investigated proposal, NOT a trained or validated v164 model.
|
| 4 |
+
No Colab interaction and no training were performed. Diagnostics ran on this Mac's MPS GPU.
|
| 5 |
+
|
| 6 |
+
## Decision
|
| 7 |
+
|
| 8 |
+
Change the state representation, atmospheric training coverage, and field/track contract before scaling the current architecture. Do not merely widen v163, remove its tracker, or add more geometric route losses.
|
| 9 |
+
|
| 10 |
+
The recommended v164 prototype is a global-to-regional, multi-variable atmospheric predictor with a storm-object decoder tied to its generated fields. It must genuinely forecast pressure, not move or paint a pressure template onto a separately predicted route. A stochastic residual extension should follow a verified deterministic baseline, not obscure an input bug.
|
| 11 |
+
|
| 12 |
+
## What the Tip investigation actually established
|
| 13 |
+
|
| 14 |
+
Case: 1979-10-12 00 UTC, nine analyses from t-48 through t0, 20 six-hour forecast steps. Checkpoint metadata says epoch 7. The model has **27,179,604 parameters**, despite the `20m` configuration/file label.
|
| 15 |
+
|
| 16 |
+
Evidence files:
|
| 17 |
+
|
| 18 |
+
- `diagnostics/v164_tip_audit/report.json`: actual baseline and two controlled inference ablations.
|
| 19 |
+
- `diagnostics/v164_tip_audit/baseline.npz`: unmodified checkpoint tensor outputs.
|
| 20 |
+
- `diagnostics/v164_tip_audit/grib_verification.json`: issue-time channel checks against original NOAA CFSR GRIB.
|
| 21 |
+
- `audit_tip_v163_for_v164.py`: repeatable local diagnostic, with labels excluded from forward inputs.
|
| 22 |
+
|
| 23 |
+
### Confirmed data/display problems
|
| 24 |
+
|
| 25 |
+
1. The demo plotted six-hour truth movements as if they were cumulative displacements. All 20 corrected coordinates now match the archived time-tagged realized route within 0.0001 degrees. The true +120 location is 21.5 N, 128.1 E, not approximately 17 N, 138 E. This was a renderer defect, not proof that the training labels were wrong: training already accumulates the step targets.
|
| 26 |
+
2. The forecast renderer used predicted latitude to convert eastward displacement, whereas model field sampling uses an issue-latitude tangent plane. The renderer now uses the model's transform. The training step-coordinate versus fixed-grid geometry remains an architectural consistency issue to remove in v164 using one tested spherical coordinate API.
|
| 27 |
+
3. The displayed v1.1 and v1.2.26 legacy routes match exactly at every lead. The v1.2.26 renderer imports a legacy v62 artifact, not a freshly identified v1.2.26 checkpoint. They are not two independently verified model replays. A provenance warning was added; this audit does not certify their independent skill or modify the global benchmark.
|
| 28 |
+
4. Tip has no SST patch in this demo, so it receives zeros. The training loader filters on `sst_ok`; it does not train this exact missing-data condition. This is a confirmed train/inference contract gap. Its causal contribution requires a matched-source SST replay, not an invented warm ocean patch.
|
| 29 |
+
|
| 30 |
+
### Checks that passed
|
| 31 |
+
|
| 32 |
+
- All forward inputs are finite. Weather ranges are plausible: pressure about 972-1031 hPa over history; height about 5116-5922 gpm; winds in m/s. There is no demonstrated Pa/hPa or height/geopotential factor error.
|
| 33 |
+
- The nine weather timestamps stop exactly at the issue time.
|
| 34 |
+
- All eight issue-time map channels were checked against the original GRIB analysis. Pressure differs by at most 0.000059 hPa after float conversion; height and wind differences are zero on the sampled native grid. GRIB time is the requested issue and lead is zero. Older snapshots were timestamp-checked, not all independently re-decoded here.
|
| 35 |
+
- The issue-pressure mask is zero. The grid-pressure fallback therefore follows the training behavior. Missing pressure is not evidence of an incorrectly supplied 870-hPa observation.
|
| 36 |
+
- ONI is intentionally masked in both training and this replay until release-dated vintages exist. An empty ONI dictionary is not an accidental missing feature for this checkpoint.
|
| 37 |
+
|
| 38 |
+
### Measured model failure
|
| 39 |
+
|
| 40 |
+
These are forecast-structure diagnostics, not future-field accuracy scores. Core minimum means the lowest grid value within 500 km of the predicted center; the ring is 500-1000 km. That region contains only about 9-13 native grid points in the core.
|
| 41 |
+
|
| 42 |
+
| Lead | Whole-grid minimum, hPa | Core minimum, hPa | Ring mean minus core minimum, hPa |
|
| 43 |
+
| --- | ---: | ---: | ---: |
|
| 44 |
+
| 0 | 980.425 | 980.425 | 24.993 |
|
| 45 |
+
| 24 | 983.234 | 983.234 | 20.906 |
|
| 46 |
+
| 48 | 987.474 | 987.474 | 16.002 |
|
| 47 |
+
| 72 | 992.260 | 994.568 | 9.411 |
|
| 48 |
+
| 120 | 999.023 | 1008.427 | 2.562 |
|
| 49 |
+
|
| 50 |
+
The generated route ends near 34.30 N, 139.66 E. The scalar central-pressure head says 992.23 hPa, but the field near that route has no comparably deep resolved low. Some difference between a subgrid scalar and coarse reanalysis is legitimate; the combination of lost core contrast, weakening gradients, and diverging route/field location is the issue.
|
| 51 |
+
|
| 52 |
+
After correcting the displayed truth and projection, descriptive Tip mean geodesic route error is **677.7 km**, with mean heading error **54.6 degrees**. The stored DeepMind route has 140.1 km mean error on those leads, but uses a different analysis pipeline and this historical case is not established as held out for its pretrained weights. Do not call that a controlled generalization benchmark.
|
| 53 |
+
|
| 54 |
+
Controlled MPS ablations (same checkpoint/inputs, local-distance approximation for this diagnostic):
|
| 55 |
+
|
| 56 |
+
| Variant | Mean route error, km | Interpretation |
|
| 57 |
+
| --- | ---: | --- |
|
| 58 |
+
| Original model | 678.1 | Baseline |
|
| 59 |
+
| Learned acceleration forced to zero | 694.9 | Does not repair the failure |
|
| 60 |
+
| Field-tracker correction disabled | 1106.4 | Tracker helps on this case |
|
| 61 |
+
|
| 62 |
+
The baseline mean tracker correction is about 21.3 km per step. These interventions do NOT isolate every possible cause, but disprove the easy claim that deleting the tracker or acceleration is the demonstrated solution. They are diagnostics, not new tuned models.
|
| 63 |
+
|
| 64 |
+
### What remains unproved
|
| 65 |
+
|
| 66 |
+
- A future, matched-source Tip atmospheric verification sequence was not evaluated here; no whole-map Tip MAE claim is made. Some real weakening is expected, so simply requiring a deep low forever would be wrong.
|
| 67 |
+
- CFSR versus the training analysis source is a distribution shift. Finite/range checks do not eliminate it; compare source-matched replays and normalized distributions on many storms.
|
| 68 |
+
- This one storm cannot prove overfitting, an optimal parameter count, or guaranteed improvements from any proposed architecture.
|
| 69 |
+
|
| 70 |
+
## Why another v163-style network is unlikely to address the main issue
|
| 71 |
+
|
| 72 |
+
The actual code uses a 25 x 33, 2.5-degree grid with eight variables: MSLP, height at 500 hPa, and winds at 850/500/200 hPa. A multi-scale encoder cannot recover missing fine-scale or thermodynamic information from that grid. The 0.5-degree source already has a relatively weak resolved low (minimum 976.435 hPa); sampling to 2.5 degrees weakens it further to 980.425 hPa. These are grid values, not observed cyclone central pressure.
|
| 73 |
+
|
| 74 |
+
It predicts bounded pointwise state increments on a fixed regional grid without a global boundary forecast. The two pressure update branches permit at most about 3.52 hPa per six-hour step at a cell (0.50 times the checkpoint pressure standard deviation). This is a potential translation/gradient limitation, not proof that saturation is the observed cause.
|
| 75 |
+
|
| 76 |
+
It lacks a prognostic humidity/temperature column and evolving ocean coupling. The current physics loss matches derivatives, vorticity, divergence and shear against targets; it is not a numerical conservation-law solver. Increasing that weight cannot supply missing thermodynamics or guarantee physical pressure evolution.
|
| 77 |
+
|
| 78 |
+
The scalar intensity head, route proposal and field tracker can disagree. Numerous route losses coexist with comparatively small field objectives. Their numerical weights alone do not establish gradient dominance: measure shared-layer gradient norms and conflict before reweighting.
|
| 79 |
+
|
| 80 |
+
Training uses storm windows that have 120-hour targets and valid SST. That removes many partial lifecycle samples and does not teach general atmospheric evolution from every available analysis time. Parameter scaling leaves these problems intact.
|
| 81 |
+
|
| 82 |
+
## Proposed architecture and data contract
|
| 83 |
+
|
| 84 |
+
### 1. Atmospheric state, not pressure-only images
|
| 85 |
+
|
| 86 |
+
Use a real global 1-degree parent state and a genuine 0.25-degree western-Pacific regional state with a surrounding halo. Start with MSLP, surface pressure, 10-m winds, 2-m temperature, SST, and u/v/T/q/geopotential at 850, 700, 500, 300 and 200 hPa. Include land fraction, orography, solar geometry, availability and age channels. Add ocean heat-content/mixed-layer information only when a causal, consistently covered archive exists.
|
| 87 |
+
|
| 88 |
+
Use physical units and one canonical regridding/normalization implementation for training, replay and benchmarking. A larger interpolated copy of the current 2.5-degree labels is NOT a high-resolution dataset. Reanalysis does not fully resolve the eye even at 0.25 degrees; retain a separately identified observation-scale intensity output.
|
| 89 |
+
|
| 90 |
+
Train from 2000 onward on all eligible atmospheric timestamps, including non-storm periods, with predeclared storm-disjoint/year splits. Older years can supplement training only when sources and masks are consistent, and Tip remains excluded. No future reanalysis, best-track values or operational forecast product may enter a rollout input. A global parent generated internally from t0 is a prediction, not a future observation.
|
| 91 |
+
|
| 92 |
+
### 2. Global-to-regional atmospheric dynamics backbone
|
| 93 |
+
|
| 94 |
+
Use a spherical multi-resolution graph or spherical neural-operator backbone for the global state, with vertically coupled tokens. It forecasts evolving ridges, troughs and steering winds, not just a global pooled context vector. A regional windowed-attention/ConvNeXt operator consumes the parent forecast, its own prior predicted fields, and lateral boundary state at each step. Decode all atmospheric variables jointly.
|
| 95 |
+
|
| 96 |
+
Start with supervised six-hour state evolution and multi-step training. A differentiable transport operator can be tested for prognostic quantities with a learned physical-process residual, but do not treat MSLP as a passive scalar or claim conservation from a simple image warp. If implementing a full dynamics solver, use an established tested solver rather than an unverified hand-written approximation.
|
| 97 |
+
|
| 98 |
+
Do not require a moving cyclone crop for the first prototype: a fixed regional field avoids using true future centers to choose inputs and avoids route-driven template placement. Predicted-center local attention is allowed, but cannot move the output field to force agreement.
|
| 99 |
+
|
| 100 |
+
### 3. One coupled storm-object and field representation
|
| 101 |
+
|
| 102 |
+
Maintain a persistent storm-identity token initialized from the observed issue-time storm. Cross-attend to generated pressure, vorticity, wind and warm-core features. Forecast a spatial association distribution over the generated field, retaining multiple candidate lows where necessary. The final route is a subgrid center associated with a generated meteorological object, not an independent trajectory plus a pasted low.
|
| 103 |
+
|
| 104 |
+
Use a motion prior only as a soft association aid. Avoid a global minimum tracker, and avoid an expectation between two distinct lows. Model cyclone existence, weakening, transition and association confidence explicitly; do not force a tropical vortex to survive after genuine dissipation.
|
| 105 |
+
|
| 106 |
+
The pressure field remains the atmospheric decoder's raw output. Grid-scale pressure, observation-scale central pressure, and maximum surface wind have different contracts. Train the subgrid intensity decoder against available observations and calibrate its relation to the resolved vortex. Do not force central pressure to equal a coarse cell or claim a head offset creates a resolved eye.
|
| 107 |
+
|
| 108 |
+
### 4. Conditional probabilistic extension, after the deterministic baseline
|
| 109 |
+
|
| 110 |
+
Add a spatially coherent stochastic residual to the atmospheric evolution, shared across variables and time, using a conditional flow or FGN-style formulation. Each ensemble member must contain its own full field, storm association, track and intensity. The purpose is to represent uncertain storm positions without averaging several locations into a diffuse low.
|
| 111 |
+
|
| 112 |
+
Keep the deterministic baseline for ablation. Score members/distributions properly; never choose the best-looking member against future truth. In the demo distinguish a fixed-seed member from the ensemble mean. A smooth ensemble-mean field is not itself proof of a bad probabilistic forecast.
|
| 113 |
+
|
| 114 |
+
### 5. Training stages and sizes
|
| 115 |
+
|
| 116 |
+
1. Pretrain atmospheric next-state prediction on non-storm and storm times, then train genuinely free-running 48/120-hour rollouts with predicted lateral boundaries.
|
| 117 |
+
2. Add the storm-object and intensity decoders and joint training. Use field accuracy/structure and storm association losses with logged gradient scales, not ten overlapping route terms by default.
|
| 118 |
+
3. Add stochastic evolution only if deterministic members lose structure under uncertainty, and evaluate calibration as well as sharpness.
|
| 119 |
+
|
| 120 |
+
Initial trainable target: **40-60M total**, for example a roughly 25-30M global core plus 15-25M regional/object components. These are architecture budgets, not measured parameter counts or confirmed T4 fits. Test the same architecture at approximately 25M, 50M and 100M under matched data/optimization before attributing gains to size. A large global dataset can require substantially more training than previous storm-only runs.
|
| 121 |
+
|
| 122 |
+
Practical alternative: freeze an existing pretrained global weather backbone and train only the regional/object refinement. This may be the fastest skillful system, but label it a hybrid and disclose its pretrained-data coverage, license and full parameter count. Do not present its accuracy as a standalone TrackFormer win over that same parent. No parent product or remote service should be silently introduced.
|
| 123 |
+
|
| 124 |
+
## Promotion gates before calling v164 better
|
| 125 |
+
|
| 126 |
+
- Shared input tests: exact timestamps, field units, channel order, coordinate round trip, longitude wrapping, missingness, normalization and source checksums. Delete future-label tensors from the forward interface and verify predictions are unchanged when verification labels are shuffled.
|
| 127 |
+
- Missing SST: either fetch a verified causal patch or use a masked encoder trained with modality dropout. No unmasked all-zero fallback on an always-present training feature.
|
| 128 |
+
- Genuine field skill: matched-grid area-weighted MSLP MAE/RMSE plus anomaly and gradient error, spectral diagnostics, cyclone-centered pressure contrast and radial profiles, low-center displacement, and co-location with vorticity/winds. Evaluate core and background separately.
|
| 129 |
+
- Route: per-step direction, turning/recurvature timing, cross-track and along-track error, and full-path mean distance. Do not substitute a single terminal-distance score or optimize only appearance.
|
| 130 |
+
- Intensity: separate observed central-pressure/wind errors from gridded MSLP error; keep wind averaging conventions explicit.
|
| 131 |
+
- Multi-storm free rollouts: tropical, recurving, landfalling and dissipating cases; report failure rate and confidence. A persistence baseline must be included.
|
| 132 |
+
- Frozen validation selection, storm-clustered bootstrap uncertainty, and a new untouched test cohort. Tip is now a development diagnostic due repeated inspection; it cannot remain an untouched test merely because it is excluded from gradient training. The repeatedly used 270 cases likewise require an independent final test.
|
| 133 |
+
- Gates must not be tuned to force the DeepMind ranking. Compare the same initialization/verification sources and disclose pretrained train/test overlap where known.
|
| 134 |
+
|
| 135 |
+
## Research basis and limits
|
| 136 |
+
|
| 137 |
+
- [GraphCast](https://arxiv.org/abs/2212.12794): useful basis for multi-scale global atmospheric forecasting rather than a storm-only latent backbone.
|
| 138 |
+
- [WeatherNext Cyclones, official explanation](https://deepmind.google/blog/weathernext-ai-model-achieves-breakthrough-in-forecasting-cyclones/): combines global atmospheric prediction with cyclone observations and coherent ensembles. Its reported Mini performance also shows that resolution alone is not a sufficient explanation for our failure. Adopting the design principles does not reproduce its skill or training corpus.
|
| 139 |
+
- [StormCast](https://arxiv.org/abs/2408.10958): supports conditioning regional probabilistic evolution on synoptic variables and evaluating realistic spectra. Its short-range kilometer-scale experiments are not proof of 120-hour western-Pacific cyclone skill.
|
| 140 |
+
- [NeuralGCM](https://www.nature.com/articles/s41586-024-07744-y): supports investigating tested hybrid dynamics and learned physics, not adding an arbitrary pressure-balance penalty and calling it physically guaranteed.
|
| 141 |
+
|
| 142 |
+
Bottom line: the defensible v164 change is stronger atmospheric evolution plus a genuinely coupled storm representation, with verified data and evaluation contracts. Tip supports investigating this direction; it does not prove that more parameters or a particular new block will win.
|
training/trackformer_1_2/source/v173/DESIGN.md
ADDED
|
@@ -0,0 +1,52 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# v1.2.73: controlled moving-core improvement
|
| 2 |
+
|
| 3 |
+
Baseline: v168 E3, SHA256
|
| 4 |
+
99ed776097e05ad6df03375465d9618741bed63a4ee57724d3ba98e295039d30.
|
| 5 |
+
The baseline is frozen as an archived reference, not frozen during v173 training.
|
| 6 |
+
|
| 7 |
+
The 21,452,595-parameter model transfers every baseline parameter. A two-scale
|
| 8 |
+
environment encoder pools physical state, recurrent memory and geographic
|
| 9 |
+
features. Two transformer blocks exchange regional context; cross attention
|
| 10 |
+
returns a correction to the environmental recurrent state. A zero-initialized
|
| 11 |
+
output makes the initial forecast equal the baseline. The moving core continues
|
| 12 |
+
to evolve pressure through relative transport and learned tendencies.
|
| 13 |
+
|
| 14 |
+
Both objectives matter: spatial pressure evolution and derived storm numbers.
|
| 15 |
+
The existing area-weighted atmospheric, outer-core pressure/gradient, track,
|
| 16 |
+
central-pressure and auxiliary wind losses remain. New masked losses supervise
|
| 17 |
+
six-hour pressure change and displacement against corresponding label changes.
|
| 18 |
+
Future fields and best tracks enter only the objective. Validation always uses
|
| 19 |
+
120 hours; training horizons are 24, 48, 72 then 120 hours. Best selection uses
|
| 20 |
+
validation, never Tip or the displayed benchmark storms.
|
| 21 |
+
|
| 22 |
+
The route is a pressure association within the moving core; central pressure is
|
| 23 |
+
sampled from that same field at that location. Maximum wind remains an auxiliary
|
| 24 |
+
scalar estimate: MSLP alone does not uniquely determine surface wind. The fixed
|
| 25 |
+
regional rendering is a resampled composite; the core is authoritative for the
|
| 26 |
+
extracted central pressure. Display +0 from the initialized core, and distinguish
|
| 27 |
+
it from raw input analysis. Do not connect raw +0 to conditioned +6.
|
| 28 |
+
|
| 29 |
+
Deliberately retained limitations: the v168 initialization includes a 100 km
|
| 30 |
+
Gaussian correction using issue-time observed pressure. It is a reconstruction,
|
| 31 |
+
not resolved truth. Fine inner geometry is underdetermined with only 800 native
|
| 32 |
+
training windows. Replacing the initializer, tracker, transport and backbone all
|
| 33 |
+
at once would destroy the controlled baseline. Learned initialization and a
|
| 34 |
+
probabilistic ensemble are later separate ablations, not implemented claims.
|
| 35 |
+
|
| 36 |
+
Only the verified full existing corpus is accepted: 13949/1041/1195 windows and
|
| 37 |
+
800/100/100 native windows. No new weather downloads. Tip is excluded, and its
|
| 38 |
+
repeatedly inspected forecasts are development diagnostics. Region is Western
|
| 39 |
+
Pacific, not a global weather model. Report native/coarse cases separately and
|
| 40 |
+
validate on untouched storms before calling this an improvement.
|
| 41 |
+
|
| 42 |
+
Mac Metal tests passed: baseline equivalence before optimization, coarse/native
|
| 43 |
+
four-step finite backward with attention gradients, field-to-pressure equality,
|
| 44 |
+
120h finite rollout, label invariance, rejected future inputs, missing intensity,
|
| 45 |
+
optimizer resume and full sampler coverage. CUDA/T4 and skill remain unverified.
|
| 46 |
+
|
| 47 |
+
Training command after the existing notebook's setup cell:
|
| 48 |
+
`python v173_pressure_train.py 10 3 1 --microbatch 4`
|
| 49 |
+
The final positional flag is accepted for compatibility; dynamics use FP32.
|
| 50 |
+
Increase physical microbatch only after measuring throughput and memory. Effective
|
| 51 |
+
batch alone does not increase GPU memory. Checkpoints and train.log persist to
|
| 52 |
+
MyDrive/typhoon_predict/training/v173; restarting resumes last.pt.
|
training/trackformer_1_2/source/v173/baseline_model.py
ADDED
|
@@ -0,0 +1,158 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Observation-initialized moving pressure state; no independent pressure head."""
|
| 2 |
+
from pathlib import Path
|
| 3 |
+
import sys
|
| 4 |
+
import numpy as np
|
| 5 |
+
import torch
|
| 6 |
+
from torch import nn
|
| 7 |
+
from torch.nn import functional as F
|
| 8 |
+
|
| 9 |
+
import v165_base as base
|
| 10 |
+
|
| 11 |
+
VERSION = '1.2.68'
|
| 12 |
+
ARCHITECTURE = 'moving-pressure-core-v168-prototype'
|
| 13 |
+
|
| 14 |
+
|
| 15 |
+
def weighted(value, mask):
|
| 16 |
+
return (value * mask).sum() / mask.expand_as(value).sum().clamp_min(1)
|
| 17 |
+
|
| 18 |
+
|
| 19 |
+
class CoreForecaster(base.TransportForecaster):
|
| 20 |
+
def __init__(self, manifest, tiny=False):
|
| 21 |
+
super().__init__(manifest, tiny)
|
| 22 |
+
del self.fine, self.fine_memory, self.wind
|
| 23 |
+
self.core_net = base.EvolutionNet(1 + 8 + 32 + 4, 3, 8 if tiny else 64, tiny)
|
| 24 |
+
self.core_memory = base.RecurrentFeatures(9)
|
| 25 |
+
self.motion_net = nn.Sequential(nn.Linear(32 + 8 + 4, 64), nn.GELU(), nn.Linear(64, 3))
|
| 26 |
+
nn.init.zeros_(self.motion_net[-1].weight)
|
| 27 |
+
nn.init.zeros_(self.motion_net[-1].bias)
|
| 28 |
+
# 20 km internal grid is a reconstruction, not new resolved observations.
|
| 29 |
+
yy, xx = torch.meshgrid(torch.linspace(640, -640, 65), torch.linspace(-640, 640, 65), indexing='ij')
|
| 30 |
+
self.register_buffer('east', xx[None])
|
| 31 |
+
self.register_buffer('north', yy[None])
|
| 32 |
+
radius = (xx.square() + yy.square()).sqrt()
|
| 33 |
+
self.register_buffer('taper', ((640-radius)/160).clamp(0, 1)[None, None])
|
| 34 |
+
|
| 35 |
+
def coordinates(self, center):
|
| 36 |
+
lat = center[:, 0, None, None] + self.north / 111.2
|
| 37 |
+
lon = center[:, 1, None, None] + self.east / (111.2 * center[:, 0, None, None].deg2rad().cos().clamp_min(.2))
|
| 38 |
+
return lat, lon
|
| 39 |
+
|
| 40 |
+
def basin_grid(self, lat, lon):
|
| 41 |
+
return torch.stack([(lon-100)/40-1, 1-lat/30], -1)
|
| 42 |
+
|
| 43 |
+
def core_grid(self, lat, lon, origin):
|
| 44 |
+
east = (lon-origin[:, 1, None, None])*111.2*origin[:, 0, None, None].deg2rad().cos().clamp_min(.2)
|
| 45 |
+
north = (lat-origin[:, 0, None, None])*111.2
|
| 46 |
+
return torch.stack([east/640, -north/640], -1)
|
| 47 |
+
|
| 48 |
+
def initial(self, inputs):
|
| 49 |
+
required = {'global_history','regional_history','global_static','regional_static','detail_available','center','motion','issue_intensity','issue_mask'}
|
| 50 |
+
if set(inputs) != required:
|
| 51 |
+
raise ValueError('Only declared past and issue-time inputs are allowed')
|
| 52 |
+
history = inputs['global_history']
|
| 53 |
+
lat, lon = self.coordinates(inputs['center'])
|
| 54 |
+
grid = self.basin_grid(lat, lon)
|
| 55 |
+
regional = inputs['regional_static']
|
| 56 |
+
rlat, rlon = regional[:, 0]*90, (regional[:, 1]+1)*180
|
| 57 |
+
native_grid = torch.stack([2*(lon-rlon[:, :1, :1])/(rlon[:, :1, -1:]-rlon[:, :1, :1])-1,
|
| 58 |
+
2*(lat-rlat[:, :1, :1])/(rlat[:, -1:, :1]-rlat[:, :1, :1])-1], -1)
|
| 59 |
+
parent = base.sample_field(history[:, -1, :1], grid)
|
| 60 |
+
available = inputs['detail_available'][:, :, None, None].bool()
|
| 61 |
+
past = []
|
| 62 |
+
for i in range(9):
|
| 63 |
+
coarse = base.sample_field(history[:, i, :1], grid)
|
| 64 |
+
native = base.sample_field(inputs['regional_history'][:, i], native_grid)
|
| 65 |
+
past.append(torch.where(available, native, coarse))
|
| 66 |
+
core = past[-1]
|
| 67 |
+
observed = (inputs['issue_intensity'][:, 1] - self.mean[0, 0, 0, 0]) / self.std[0, 0, 0, 0]
|
| 68 |
+
# Only initialization uses an explicit compact core prior. Subsequent
|
| 69 |
+
# changes are learned field tendencies, not repeated scalar insertion.
|
| 70 |
+
radius2 = self.east.square() + self.north.square()
|
| 71 |
+
prior = torch.exp(-radius2/(2*100**2))[:, None] * self.taper
|
| 72 |
+
difference = observed-core[:, 0, 32, 32]
|
| 73 |
+
core = core + difference[:, None, None, None]*prior*inputs['issue_mask'][:, 1, None, None, None]
|
| 74 |
+
return dict(g=history[:, -1], gh=self.coarse_memory.encode(history.flatten(1, 2)),
|
| 75 |
+
core=core, ch=self.core_memory.encode(torch.cat(past, 1)),
|
| 76 |
+
anomaly=core-parent, origin=inputs['center'], center=inputs['center'], motion=inputs['motion'],
|
| 77 |
+
vmax=inputs['issue_intensity'][:, 0], gs=inputs['global_static'], rs=regional,
|
| 78 |
+
issue=inputs['issue_intensity'], issue_mask=inputs['issue_mask'])
|
| 79 |
+
|
| 80 |
+
def step(self, s):
|
| 81 |
+
delta, features = self.coarse(torch.cat([s['g'], s['gh'], s['gs']], 1))
|
| 82 |
+
velocity = self.steering(s['g']) + 15*torch.tanh(delta[:, 8:])
|
| 83 |
+
departure = base.departure_grid(s['g'], velocity, self.lat, 2.5)
|
| 84 |
+
g = base.sample_field(s['g'], departure) + .2*torch.tanh(delta[:, :8])
|
| 85 |
+
gh = self.coarse_memory(base.sample_field(s['gh'], departure), features)
|
| 86 |
+
at = self.basin_grid(s['center'][:, 0, None, None], s['center'][:, 1, None, None])
|
| 87 |
+
local = base.sample_field(torch.cat([gh, g], 1), at)[:, :, 0, 0]
|
| 88 |
+
conditioning = torch.cat([s['motion']/100, s['issue_mask']], 1)
|
| 89 |
+
tendencies = self.motion_net(torch.cat([local, conditioning], 1))
|
| 90 |
+
flow = base.sample_field(velocity, at)[:, :, 0, 0]
|
| 91 |
+
motion = .7*flow*21.6 + .3*s['motion'] + 100*torch.tanh(tendencies[:, :2])
|
| 92 |
+
origin = s['center'] + torch.stack([motion[:, 1]/111.2,
|
| 93 |
+
motion[:, 0]/(111.2*s['center'][:, 0].deg2rad().cos().clamp_min(.2))], 1)
|
| 94 |
+
lat, lon = self.coordinates(origin)
|
| 95 |
+
grid = self.basin_grid(lat, lon)
|
| 96 |
+
environment = base.sample_field(g, grid)
|
| 97 |
+
static = base.sample_field(s['gs'], grid)
|
| 98 |
+
# Translate the moving frame without resampling away bulk storm motion.
|
| 99 |
+
old_lat, old_lon = self.coordinates(s['origin'])
|
| 100 |
+
old_grid = self.basin_grid(old_lat, old_lon)
|
| 101 |
+
old_environment = base.sample_field(s['g'], old_grid)
|
| 102 |
+
rd, rf = self.core_net(torch.cat([s['core'], old_environment, s['ch'], static], 1))
|
| 103 |
+
local_flow = base.sample_field(velocity, grid) + 5*torch.tanh(rd[:, 1:])
|
| 104 |
+
relative = local_flow - (motion/21.6)[:, :, None, None]
|
| 105 |
+
shift = (s['center']-s['origin'])
|
| 106 |
+
source_lat = old_lat + shift[:, 0, None, None] - relative[:, 1]*21.6/111.2
|
| 107 |
+
source_lon = old_lon + shift[:, 1, None, None] - relative[:, 0]*21.6/(111.2*old_lat.deg2rad().cos().clamp_min(.2))
|
| 108 |
+
dep = self.core_grid(source_lat, source_lon, s['origin'])
|
| 109 |
+
anomaly = base.sample_field(s['anomaly'], dep, 'zeros') + .1*rd[:, :1]
|
| 110 |
+
core = environment[:, :1] + anomaly*self.taper
|
| 111 |
+
ch = self.core_memory(base.sample_field(s['ch'], dep, 'zeros'), rf)
|
| 112 |
+
physical = core*self.std[:, :1]+self.mean[:, :1]
|
| 113 |
+
# Field-based association, constrained to this storm's local neighborhood.
|
| 114 |
+
score = -physical[:, 0]/2 - (self.east.square()+self.north.square())/(2*180**2)
|
| 115 |
+
score = score.masked_fill((self.east.square()+self.north.square()) > 300**2, -1e9)
|
| 116 |
+
pdf = score.flatten(1).softmax(1).reshape_as(score)
|
| 117 |
+
center = torch.stack([(pdf*lat).sum((1,2)), (pdf*lon).sum((1,2))], 1)
|
| 118 |
+
cg = self.core_grid(center[:, 0, None, None], center[:, 1, None, None], origin)
|
| 119 |
+
pressure = base.sample_field(physical, cg)[:, 0, 0, 0]
|
| 120 |
+
rlat, rlon = s['rs'][:, 0]*90, (s['rs'][:, 1]+1)*180
|
| 121 |
+
rg = self.basin_grid(rlat, rlon)
|
| 122 |
+
sampling = self.core_grid(rlat, rlon, origin)
|
| 123 |
+
reconstructed = base.sample_field(g[:, :1], rg) + base.sample_field(anomaly*self.taper, sampling, 'zeros')
|
| 124 |
+
# Resolved basin output remains at source resolution; reconstruction has
|
| 125 |
+
# separate provenance and is never relabeled as raw coarse analysis.
|
| 126 |
+
vmax = (s['vmax']+5*tendencies[:, 2]).clamp_min(0)
|
| 127 |
+
valid = (grid.abs()<=1).all(-1)[:, None]
|
| 128 |
+
out = dict(global_=g, regional=reconstructed, regional_valid=(rg.abs()<=1).all(-1)[:,None],
|
| 129 |
+
core=core, core_lat=lat, core_lon=lon, core_valid=valid,
|
| 130 |
+
center=center, pressure=pressure, grid_pressure=pressure, vmax=vmax,
|
| 131 |
+
track_valid=(center[:,0]>0)&(center[:,0]<60)&(center[:,1]>100)&(center[:,1]<180),
|
| 132 |
+
native_route=torch.ones_like(pressure,dtype=torch.bool),
|
| 133 |
+
pressure_tendency=pressure-base.sample_field(s['core']*self.std[:,:1]+self.mean[:,:1],
|
| 134 |
+
self.core_grid(s['center'][:,0,None,None],s['center'][:,1,None,None],s['origin']))[:,0,0,0])
|
| 135 |
+
out['global']=out.pop('global_')
|
| 136 |
+
return dict(s,g=g,gh=gh,core=core,ch=ch,anomaly=anomaly,origin=origin,center=center,motion=motion,vmax=vmax), out
|
| 137 |
+
|
| 138 |
+
|
| 139 |
+
def objective(out, target, inputs, manifest):
|
| 140 |
+
terms={}
|
| 141 |
+
area=inputs['global_static'][:,:1].mul(90).deg2rad().cos().clamp_min(0)
|
| 142 |
+
terms['environment']=weighted(F.smooth_l1_loss(out['global'],target['global'],reduction='none'),area*target['global_mask'])
|
| 143 |
+
# Reanalysis does not resolve the observed central pressure. Exclude the
|
| 144 |
+
# inner 150 km from absolute reconstruction targets instead of erasing it.
|
| 145 |
+
lat,lon=inputs['regional_static'][:,0]*90,(inputs['regional_static'][:,1]+1)*180
|
| 146 |
+
ct=target['center']
|
| 147 |
+
d2=((lat-ct[:,0,None,None])*111.2).square()+((lon-ct[:,1,None,None])*111.2*ct[:,0,None,None].deg2rad().cos()).square()
|
| 148 |
+
mask=target['regional_mask']*out['regional_valid']*(d2>=150**2)[:,None]
|
| 149 |
+
terms['regional']=weighted(F.smooth_l1_loss(out['regional'],target['regional'],reduction='none'),mask)
|
| 150 |
+
for axis in (-2,-1):
|
| 151 |
+
pred=out['regional'].diff(dim=axis);truth=target['regional'].diff(dim=axis)
|
| 152 |
+
w=mask.narrow(axis,0,pred.shape[axis])*mask.narrow(axis,1,pred.shape[axis])
|
| 153 |
+
terms['regional']=terms['regional']+.2*weighted((pred-truth).abs(),w)
|
| 154 |
+
displacement=(out['center']-ct)*torch.stack([torch.full_like(ct[:,0],111.2),111.2*ct[:,0].deg2rad().cos()],1)
|
| 155 |
+
terms['route']=weighted(F.smooth_l1_loss(displacement/100,torch.zeros_like(displacement),reduction='none'),target['storm_mask'][:,None])
|
| 156 |
+
terms['pressure']=weighted(F.smooth_l1_loss(out['pressure']/20,target['pressure']/20,reduction='none'),target['pressure_mask'])
|
| 157 |
+
terms['wind']=weighted(F.smooth_l1_loss(out['vmax']/30,target['vmax']/30,reduction='none'),target['vmax_mask'])
|
| 158 |
+
return terms['environment']+terms['regional']+.2*terms['route']+.5*terms['pressure']+.05*terms['wind'],terms
|
training/trackformer_1_2/source/v173/model.py
ADDED
|
@@ -0,0 +1,74 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""v173: moving pressure core with multiscale environmental attention.
|
| 2 |
+
|
| 3 |
+
The public track and central pressure are extracted from the evolving core.
|
| 4 |
+
Wind is an auxiliary scalar estimate, not a resolved wind field.
|
| 5 |
+
"""
|
| 6 |
+
import torch
|
| 7 |
+
from torch import nn
|
| 8 |
+
from torch.nn import functional as F
|
| 9 |
+
from baseline_model import CoreForecaster as Baseline, objective as baseline_loss, weighted, base
|
| 10 |
+
|
| 11 |
+
VERSION = '1.2.73'
|
| 12 |
+
ARCHITECTURE = 'v173-moving-core-multiscale-attention'
|
| 13 |
+
|
| 14 |
+
|
| 15 |
+
class EnvironmentalAttention(nn.Module):
|
| 16 |
+
def __init__(self):
|
| 17 |
+
super().__init__()
|
| 18 |
+
self.project = nn.Conv2d(44, 64, 1)
|
| 19 |
+
layer = nn.TransformerEncoderLayer(64, 4, 128, dropout=0., batch_first=True,
|
| 20 |
+
norm_first=True)
|
| 21 |
+
self.encoder = nn.TransformerEncoder(layer, 2, enable_nested_tensor=False)
|
| 22 |
+
self.query = nn.Linear(32, 64)
|
| 23 |
+
self.cross = nn.MultiheadAttention(64, 4, dropout=0., batch_first=True)
|
| 24 |
+
self.output = nn.Linear(64, 32)
|
| 25 |
+
nn.init.zeros_(self.output.weight)
|
| 26 |
+
nn.init.zeros_(self.output.bias)
|
| 27 |
+
|
| 28 |
+
def forward(self, state):
|
| 29 |
+
h, w = state['gh'].shape[-2:]
|
| 30 |
+
features = self.project(torch.cat([state['g'], state['gh'], state['gs']], 1))
|
| 31 |
+
tokens = torch.cat([F.avg_pool2d(features, kernel_size=size, stride=size,
|
| 32 |
+
ceil_mode=True).flatten(2)
|
| 33 |
+
for size in (3, 5)], 2).transpose(1, 2)
|
| 34 |
+
tokens = self.encoder(tokens)
|
| 35 |
+
queries = self.query(state['gh'].flatten(2).transpose(1, 2))
|
| 36 |
+
context, _ = self.cross(queries, tokens, tokens, need_weights=False)
|
| 37 |
+
return self.output(context).transpose(1, 2).reshape(-1, 32, h, w)
|
| 38 |
+
|
| 39 |
+
|
| 40 |
+
class CoreForecaster(Baseline):
|
| 41 |
+
def __init__(self, manifest, tiny=False):
|
| 42 |
+
super().__init__(manifest, tiny)
|
| 43 |
+
# Meshgrid views overlap on CPU; checkpoint loading requires owned storage.
|
| 44 |
+
self.east = self.east.clone()
|
| 45 |
+
self.north = self.north.clone()
|
| 46 |
+
self.context = EnvironmentalAttention()
|
| 47 |
+
|
| 48 |
+
def step(self, state):
|
| 49 |
+
conditioned = dict(state, gh=state['gh'] + self.context(state))
|
| 50 |
+
next_state, out = super().step(conditioned)
|
| 51 |
+
out['previous_center'] = state['center']
|
| 52 |
+
out['previous_pressure'] = out['pressure'] - out['pressure_tendency']
|
| 53 |
+
return next_state, out
|
| 54 |
+
|
| 55 |
+
|
| 56 |
+
def objective(out, target, inputs, manifest):
|
| 57 |
+
loss, terms = baseline_loss(out, target, inputs, manifest)
|
| 58 |
+
if 'previous_pressure' in target:
|
| 59 |
+
mask = target['pressure_mask'] * target['previous_pressure_mask']
|
| 60 |
+
true_delta = target['pressure'] - target['previous_pressure']
|
| 61 |
+
terms['pressure_change'] = weighted(F.smooth_l1_loss(
|
| 62 |
+
out['pressure_tendency']/10, true_delta/10, reduction='none'), mask)
|
| 63 |
+
loss = loss + .25 * terms['pressure_change']
|
| 64 |
+
if 'previous_center' in target:
|
| 65 |
+
center = target['center']
|
| 66 |
+
scale = torch.stack([torch.full_like(center[:, 0], 111.2),
|
| 67 |
+
111.2*center[:, 0].deg2rad().cos()], 1)
|
| 68 |
+
predicted = (out['center']-out['previous_center'])*scale
|
| 69 |
+
truth = (center-target['previous_center'])*scale
|
| 70 |
+
mask = target['storm_mask']*target['previous_storm_mask']
|
| 71 |
+
terms['motion_change'] = weighted(F.smooth_l1_loss(
|
| 72 |
+
predicted/100, truth/100, reduction='none'), mask[:, None])
|
| 73 |
+
loss = loss + .1 * terms['motion_change']
|
| 74 |
+
return loss, terms
|
training/trackformer_1_2/source/v173/test_model.py
ADDED
|
@@ -0,0 +1,90 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Real-data Metal preflight, pressure coupling, gradients and resume."""
|
| 2 |
+
from pathlib import Path
|
| 3 |
+
import tempfile
|
| 4 |
+
import torch
|
| 5 |
+
from model import CoreForecaster, objective, base
|
| 6 |
+
from v173_pressure_train import save, restore, warm_start, batches
|
| 7 |
+
from baseline_model import CoreForecaster as Baseline
|
| 8 |
+
|
| 9 |
+
|
| 10 |
+
def main():
|
| 11 |
+
root=Path(__file__).resolve().parents[1]
|
| 12 |
+
torch.set_num_threads(4);torch.manual_seed(168)
|
| 13 |
+
base.EvolutionNet.block=lambda self,module,x:module(x)
|
| 14 |
+
m=base.load_reuse(root/'data/v164_reuse/manifest.json')
|
| 15 |
+
assert m['counts']==dict(train=13949,validation=1041,test=1195)
|
| 16 |
+
assert not any('1979275N06159'==r['storm_id'] for r in m['rows'])
|
| 17 |
+
model=CoreForecaster(m).to('mps');warm_start(model,root/'v168/runs/mac/epoch_003.pt',m)
|
| 18 |
+
reference=Baseline(m).to('mps')
|
| 19 |
+
reference.load_state_dict(torch.load(root/'v168/runs/mac/epoch_003.pt',map_location='cpu',weights_only=False)['model'])
|
| 20 |
+
reference.eval();model.eval()
|
| 21 |
+
optimizer=torch.optim.AdamW(model.parameters(),lr=3e-5,weight_decay=.01)
|
| 22 |
+
data=base.ReuseDataset(m,'train')
|
| 23 |
+
choices=[next(i for i,r in enumerate(data.rows) if r['patch']<0 and r['issue_mask'][1]),
|
| 24 |
+
next(i for i,r in enumerate(data.rows) if r['patch']>=0 and r['issue_mask'][1])]
|
| 25 |
+
for index in choices:
|
| 26 |
+
x,y=data[index];x={k:v[None].to('mps') for k,v in x.items()}
|
| 27 |
+
if index==choices[0]:
|
| 28 |
+
with torch.no_grad():
|
| 29 |
+
expected=reference(x,1);actual=model(x,1)
|
| 30 |
+
for key in expected: assert torch.allclose(actual[key],expected[key],atol=1e-5,rtol=1e-5), key
|
| 31 |
+
print('PASS zero-correction warm start equals v168 E3 field and route',flush=True)
|
| 32 |
+
model.train()
|
| 33 |
+
optimizer.zero_grad(set_to_none=True);state=model.initial(x)
|
| 34 |
+
initial_p=state['core'][:,0,32,32]*model.std[0,0,0,0]+model.mean[0,0,0,0]
|
| 35 |
+
assert torch.allclose(initial_p,x['issue_intensity'][:,1],atol=1e-3)
|
| 36 |
+
total=0
|
| 37 |
+
for lead in range(4):
|
| 38 |
+
state,out=model.step(state)
|
| 39 |
+
target={k:v[lead:lead+1].to('mps') for k,v in y.items()}
|
| 40 |
+
target['previous_pressure']=y['pressure'][lead-1:lead].to('mps') if lead else x['issue_intensity'][:,1]
|
| 41 |
+
target['previous_pressure_mask']=y['pressure_mask'][lead-1:lead].to('mps') if lead else x['issue_mask'][:,1]
|
| 42 |
+
target['previous_center']=y['center'][lead-1:lead].to('mps') if lead else x['center']
|
| 43 |
+
target['previous_storm_mask']=y['storm_mask'][lead-1:lead].to('mps') if lead else torch.ones_like(target['storm_mask'])
|
| 44 |
+
loss,_=objective(out,target,x,m);total=total+loss/4
|
| 45 |
+
assert all(bool(torch.isfinite(v).all()) for v in out.values())
|
| 46 |
+
grid=model.core_grid(out['center'][:,0,None,None],out['center'][:,1,None,None],state['origin'])
|
| 47 |
+
field_pressure=base.sample_field(out['core']*model.std[:,:1]+model.mean[:,:1],grid)[:,0,0,0]
|
| 48 |
+
assert torch.equal(field_pressure,out['pressure'])
|
| 49 |
+
total.backward()
|
| 50 |
+
assert all(bool(torch.isfinite(p.grad).all()) for p in model.parameters() if p.grad is not None)
|
| 51 |
+
for layer in [model.core_net.out,model.coarse.out,model.motion_net[-1]]:
|
| 52 |
+
assert layer.weight.grad is not None and layer.weight.grad.abs().sum()>0
|
| 53 |
+
assert model.context.output.weight.grad.abs().sum()>0
|
| 54 |
+
torch.nn.utils.clip_grad_norm_(model.parameters(),1.);optimizer.step()
|
| 55 |
+
print('PASS real 4-step backward, anchored pressure and field-derived output; native=',bool(x['detail_available'].item()),flush=True)
|
| 56 |
+
model.eval()
|
| 57 |
+
with torch.no_grad():
|
| 58 |
+
out=model(x,20)
|
| 59 |
+
assert all(bool(torch.isfinite(v).all()) for v in out.values())
|
| 60 |
+
first=model(x,1)
|
| 61 |
+
altered={k:v.clone() for k,v in target.items()};altered['center']+=10
|
| 62 |
+
objective({k:v[:,0] for k,v in first.items()},altered,x,m)
|
| 63 |
+
again=model(x,1)
|
| 64 |
+
assert all(torch.equal(first[k],again[k]) for k in first)
|
| 65 |
+
try:model.initial(dict(x,future_pressure=target['regional']))
|
| 66 |
+
except ValueError:pass
|
| 67 |
+
else:raise AssertionError('Future input accepted')
|
| 68 |
+
missing={k:v.clone() for k,v in x.items()};missing['issue_mask'].zero_();missing['issue_intensity'].zero_()
|
| 69 |
+
absent=model(missing,1)
|
| 70 |
+
assert all(bool(torch.isfinite(v).all()) for v in absent.values())
|
| 71 |
+
print('PASS 120h rollout, label invariance, missing intensity and future-input rejection',flush=True)
|
| 72 |
+
with tempfile.TemporaryDirectory(dir=root/'v173') as tmp:
|
| 73 |
+
path=Path(tmp)/'resume.pt';save(path,model,optimizer,m,0,[],100.,dict(window=1,totals=[[1.,1],[1.,1]],batch=10))
|
| 74 |
+
def update():
|
| 75 |
+
model.train();optimizer.zero_grad(set_to_none=True)
|
| 76 |
+
o=model(x,1);loss,_=objective({k:v[:,0] for k,v in o.items()},target,x,m)
|
| 77 |
+
loss.backward();torch.nn.utils.clip_grad_norm_(model.parameters(),1.);optimizer.step()
|
| 78 |
+
return model.core_net.out.weight.detach().clone()
|
| 79 |
+
expected=update();payload=restore(path,model,optimizer,m);actual=update()
|
| 80 |
+
assert torch.allclose(expected,actual,atol=1e-6,rtol=1e-5)
|
| 81 |
+
assert payload['progress']['window']==1
|
| 82 |
+
indices=base.training_indices(data.rows,1)
|
| 83 |
+
grouped=list(batches(indices,data.rows,10,2))
|
| 84 |
+
assert sorted(i for group,_ in grouped for part in group for i in part)==sorted(indices)
|
| 85 |
+
assert set(indices)==set(range(len(data)))
|
| 86 |
+
print('PASS checkpoint optimizer resume and full dataset coverage',flush=True)
|
| 87 |
+
print('parameters=',sum(p.numel() for p in model.parameters()),flush=True)
|
| 88 |
+
|
| 89 |
+
|
| 90 |
+
if __name__=='__main__':main()
|
training/trackformer_1_2/source/v173/v165_base.py
ADDED
|
@@ -0,0 +1,657 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#!/usr/bin/env python3
|
| 2 |
+
"""Standalone v1.2.65. Existing full archives; no web downloads or pilot fallback."""
|
| 3 |
+
from __future__ import annotations
|
| 4 |
+
|
| 5 |
+
import argparse
|
| 6 |
+
import hashlib
|
| 7 |
+
import json
|
| 8 |
+
import math
|
| 9 |
+
import os
|
| 10 |
+
from pathlib import Path
|
| 11 |
+
import random
|
| 12 |
+
import socket
|
| 13 |
+
import time
|
| 14 |
+
from contextlib import nullcontext
|
| 15 |
+
|
| 16 |
+
import numpy as np
|
| 17 |
+
import torch
|
| 18 |
+
from torch import nn
|
| 19 |
+
from torch.nn import functional as F
|
| 20 |
+
from torch.utils.checkpoint import checkpoint
|
| 21 |
+
from torch.utils.data import Dataset, DataLoader
|
| 22 |
+
|
| 23 |
+
ROOT = Path(__file__).resolve().parent
|
| 24 |
+
HOUR = 3_600_000_000_000
|
| 25 |
+
|
| 26 |
+
def digest(path):
|
| 27 |
+
h = hashlib.sha256()
|
| 28 |
+
with Path(path).open("rb") as f:
|
| 29 |
+
for block in iter(lambda: f.read(8 * 1024**2), b""):
|
| 30 |
+
h.update(block)
|
| 31 |
+
return h.hexdigest()
|
| 32 |
+
|
| 33 |
+
|
| 34 |
+
def atomic_json(path, data):
|
| 35 |
+
path = Path(path)
|
| 36 |
+
tmp = path.with_suffix(path.suffix + ".tmp")
|
| 37 |
+
tmp.write_text(json.dumps(data, indent=2, allow_nan=False) + "\n")
|
| 38 |
+
tmp.replace(path)
|
| 39 |
+
|
| 40 |
+
|
| 41 |
+
def atomic_checkpoint(path, data):
|
| 42 |
+
path = Path(path)
|
| 43 |
+
tmp = path.with_suffix(".tmp")
|
| 44 |
+
with tmp.open("wb") as f:
|
| 45 |
+
torch.save(data, f)
|
| 46 |
+
f.flush()
|
| 47 |
+
os.fsync(f.fileno())
|
| 48 |
+
tmp.replace(path)
|
| 49 |
+
|
| 50 |
+
|
| 51 |
+
def device_for(name):
|
| 52 |
+
if name == "auto":
|
| 53 |
+
name = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
|
| 54 |
+
if name == "cuda" and not torch.cuda.is_available():
|
| 55 |
+
raise RuntimeError("CUDA unavailable in this Python. Install a CUDA PyTorch wheel before training.")
|
| 56 |
+
if name == "mps" and not torch.backends.mps.is_available():
|
| 57 |
+
raise RuntimeError("MPS requested but unavailable")
|
| 58 |
+
return torch.device(name)
|
| 59 |
+
|
| 60 |
+
|
| 61 |
+
"""Embedded in the standalone v164 trainer: coarse archive plus native MSLP patches."""
|
| 62 |
+
REUSE_ARCH = 'wp-coarse8-native-pressure1-v164'
|
| 63 |
+
REUSE_CHANNELS = ['mslp', 'hgt500', 'uwnd850', 'vwnd850', 'uwnd500', 'vwnd500', 'uwnd200', 'vwnd200']
|
| 64 |
+
|
| 65 |
+
|
| 66 |
+
def load_reuse(path):
|
| 67 |
+
path = Path(path)
|
| 68 |
+
if not path.is_file():
|
| 69 |
+
raise FileNotFoundError('Combined dataset missing. Run the updated Drive setup cell first: '+str(path))
|
| 70 |
+
m = json.loads(path.read_text())
|
| 71 |
+
if m.get('schema') != 'v164-reuse-pressure-v1' or m.get('pilot', True):
|
| 72 |
+
raise ValueError('Combined non-pilot archive required; no pilot fallback')
|
| 73 |
+
for name, entry in m['files'].items():
|
| 74 |
+
p = path.parent/entry['path']
|
| 75 |
+
print('Checking dataset file:', name, flush=True)
|
| 76 |
+
if digest(p) != entry['sha256']:
|
| 77 |
+
raise ValueError('Dataset checksum mismatch: '+name)
|
| 78 |
+
entry['resolved'] = str(p)
|
| 79 |
+
with np.load(m['files']['basin']['resolved'], allow_pickle=False) as z:
|
| 80 |
+
if z['channels'].tolist() != REUSE_CHANNELS[1:]:
|
| 81 |
+
raise ValueError('Coarse channel order mismatch')
|
| 82 |
+
times, lat, lon = z['time'], z['lat'], z['lon']
|
| 83 |
+
q = z['q']
|
| 84 |
+
scale, offset = z['scale'], z['offset']
|
| 85 |
+
if q.shape != (len(times),7,25,33) or not np.all(np.diff(times) == 6*HOUR):
|
| 86 |
+
raise ValueError('Coarse grid/time mismatch')
|
| 87 |
+
qpath = path.parent/'basin_q_verified.npy'
|
| 88 |
+
# Refresh the compact worker mmap from the checksum-verified source every preflight.
|
| 89 |
+
temporary = qpath.with_suffix('.pending.npy')
|
| 90 |
+
np.save(temporary, q)
|
| 91 |
+
os.replace(temporary, qpath)
|
| 92 |
+
slp = np.load(m['files']['slp']['resolved'], mmap_mode='r')
|
| 93 |
+
if slp.shape != (len(times),25,33) or not np.isfinite(slp).all():
|
| 94 |
+
raise ValueError('Coarse pressure coverage invalid')
|
| 95 |
+
with np.load(m['files']['track']['resolved'], allow_pickle=False) as z:
|
| 96 |
+
track = {k:z[k] for k in ['base_time','storm_id','base_lat','base_lon','target','target_mask']}
|
| 97 |
+
issue_intensity = z['track'][:,-1,4:6]*z['track_std'][4:6]+z['track_mean'][4:6]
|
| 98 |
+
plan = json.loads(Path(m['files']['plan']['resolved']).read_text())
|
| 99 |
+
if len(plan['rows']) != 1000 or plan['shape'] != [1000,29,121,121] or plan['lead_hours'] != list(range(-48,121,6)):
|
| 100 |
+
raise ValueError('Expected the verified 1000-window native pressure package, not the six-sample pilot')
|
| 101 |
+
pressure = np.load(m['files']['pressure']['resolved'], mmap_mode='r')
|
| 102 |
+
if pressure.shape != (1000,29,121,121):
|
| 103 |
+
raise ValueError('Native pressure shape mismatch')
|
| 104 |
+
patches = {}
|
| 105 |
+
for pi,row in enumerate(plan['rows']):
|
| 106 |
+
ti = row['track_archive_row']
|
| 107 |
+
if int(track['base_time'][ti]) != row['issue_ns'] or str(track['storm_id'][ti]) != row['storm_id']:
|
| 108 |
+
raise ValueError('Native patch/track identity mismatch')
|
| 109 |
+
patches[ti] = pi
|
| 110 |
+
first_year = {}
|
| 111 |
+
years = track['base_time'].astype('datetime64[ns]').astype('datetime64[Y]').astype(int)+1970
|
| 112 |
+
for sid,y in zip(track['storm_id'],years):
|
| 113 |
+
first_year[str(sid)] = min(first_year.get(str(sid),9999),int(y))
|
| 114 |
+
lookup = {(str(s),int(t)):i for i,(s,t) in enumerate(zip(track['storm_id'],track['base_time']))}
|
| 115 |
+
rows = []
|
| 116 |
+
for i,ns in enumerate(track['base_time']):
|
| 117 |
+
sid = str(track['storm_id'][i])
|
| 118 |
+
y = first_year[sid]
|
| 119 |
+
split = 'train' if 2000 <= y <= 2021 else 'validation' if 2022 <= y <= 2023 else 'test' if 2024 <= y <= 2025 else None
|
| 120 |
+
if split is None:
|
| 121 |
+
continue
|
| 122 |
+
low,high = {'train':(2000,2021),'validation':(2022,2023),'test':(2024,2025)}[split]
|
| 123 |
+
bounds = np.array([ns-48*HOUR,ns+120*HOUR],dtype='int64').astype('datetime64[ns]').astype('datetime64[Y]').astype(int)+1970
|
| 124 |
+
ai = int(np.searchsorted(times,ns))
|
| 125 |
+
center = [float(track['base_lat'][i]),float(track['base_lon'][i])%360]
|
| 126 |
+
if bounds[0] < low or bounds[1] > high or ai < 8 or ai+20 >= len(times) or times[ai] != ns:
|
| 127 |
+
continue
|
| 128 |
+
if not np.isfinite(center).all() or not (0 < center[0] < 60 and 100 < center[1] < 180):
|
| 129 |
+
continue
|
| 130 |
+
previous = lookup.get((sid,int(ns-6*HOUR)))
|
| 131 |
+
motion = [0.,0.]
|
| 132 |
+
if previous is not None:
|
| 133 |
+
old = [float(track['base_lat'][previous]),float(track['base_lon'][previous])%360]
|
| 134 |
+
if np.isfinite(old).all():
|
| 135 |
+
motion = [(center[1]-old[1])*111.2*np.cos(np.deg2rad(center[0])),(center[0]-old[0])*111.2]
|
| 136 |
+
future = [lookup.get((sid,int(ns+j*6*HOUR))) for j in range(1,21)]
|
| 137 |
+
target_centers = [[float(track['base_lat'][j]),float(track['base_lon'][j])%360] if j is not None else [0.,0.] for j in future]
|
| 138 |
+
rm = [float(j is not None and np.isfinite(target_centers[k]).all()) for k,j in enumerate(future)]
|
| 139 |
+
target_centers = np.nan_to_num(target_centers).tolist()
|
| 140 |
+
wind,p = track['target'][i,:,2],track['target'][i,:,3]
|
| 141 |
+
wm = track['target_mask'][i,:,2].astype(bool)&np.isfinite(wind)&(wind>=0)&(wind<250)
|
| 142 |
+
pm = track['target_mask'][i,:,3].astype(bool)&np.isfinite(p)&(p>800)&(p<1100)
|
| 143 |
+
pi = patches.get(i,-1)
|
| 144 |
+
if pi >= 0 and plan['rows'][pi]['split'] != split:
|
| 145 |
+
raise ValueError('Patch split mismatch')
|
| 146 |
+
issue = issue_intensity[i]
|
| 147 |
+
iv = [bool(np.isfinite(issue[0]) and 0<=issue[0]<250),bool(np.isfinite(issue[1]) and 800<issue[1]<1100)]
|
| 148 |
+
rows.append({'split':split,'storm_id':sid,'atlas':ai,'patch':pi,'center':center,'motion':motion,
|
| 149 |
+
'issue_intensity':[float(issue[k]) if iv[k] else 0. for k in range(2)],'issue_mask':iv,
|
| 150 |
+
'target_centers':target_centers,'route_mask':rm,
|
| 151 |
+
'pressure':np.where(pm,p,0).tolist(),'pressure_mask':pm.tolist(),
|
| 152 |
+
'vmax':np.where(wm,wind,0).tolist(),'vmax_mask':wm.tolist()})
|
| 153 |
+
counts = {s:sum(r['split']==s for r in rows) for s in ['train','validation','test']}
|
| 154 |
+
patch_counts = {s:sum(r['split']==s and r['patch']>=0 for r in rows) for s in counts}
|
| 155 |
+
if counts['train'] < 10000 or patch_counts != {'train':800,'validation':100,'test':100}:
|
| 156 |
+
raise ValueError(f'Incomplete combined coverage: {counts}, native patches {patch_counts}')
|
| 157 |
+
train_times = np.flatnonzero((times.astype('datetime64[ns]').astype('datetime64[Y]').astype(int)+1970>=2000)&
|
| 158 |
+
(times.astype('datetime64[ns]').astype('datetime64[Y]').astype(int)+1970<=2021))
|
| 159 |
+
sums,squares = np.zeros(8),np.zeros(8)
|
| 160 |
+
n = 0
|
| 161 |
+
for start in range(0,len(train_times),256):
|
| 162 |
+
ii = train_times[start:start+256]
|
| 163 |
+
b = np.concatenate([np.asarray(slp[ii],dtype='float32')[:,None],q[ii].astype('float32')*scale[None,:,None,None]+offset[None,:,None,None]],1).astype('float64')
|
| 164 |
+
if not np.isfinite(b).all():
|
| 165 |
+
raise ValueError('Nonfinite coarse training fields')
|
| 166 |
+
sums += b.sum((0,2,3)); squares += (b*b).sum((0,2,3)); n += len(ii)*25*33
|
| 167 |
+
mean = sums/n
|
| 168 |
+
std = np.maximum(np.sqrt(np.maximum(squares/n-mean*mean,0)),1e-4)
|
| 169 |
+
m.update(rows=rows, plan=plan, qpath=str(qpath), scale=scale.tolist(), offset=offset.tolist(),
|
| 170 |
+
channels=REUSE_CHANNELS,units=['hPa','m','m/s','m/s','m/s','m/s','m/s','m/s'],
|
| 171 |
+
global_lat=lat.tolist(),global_lon=lon.tolist(),regional_lat=list(np.linspace(15,-15,121)),
|
| 172 |
+
regional_lon=list(np.linspace(-15,15,121)),normalization={'mean':mean.tolist(),'std':std.tolist(),'fit_split':'train','end_year':2021},
|
| 173 |
+
identity=digest(path),counts=counts,patch_counts=patch_counts)
|
| 174 |
+
print(json.dumps({'all_archive_windows':counts,'native_pressure_windows':patch_counts,'architecture':REUSE_ARCH}),flush=True)
|
| 175 |
+
return m
|
| 176 |
+
|
| 177 |
+
|
| 178 |
+
class ReuseDataset(Dataset):
|
| 179 |
+
def __init__(self,m,split,training=False):
|
| 180 |
+
self.m = m
|
| 181 |
+
self.rows = [r for r in m['rows'] if r['split']==split]
|
| 182 |
+
self.arrays = None
|
| 183 |
+
|
| 184 |
+
def __len__(self):
|
| 185 |
+
return len(self.rows)
|
| 186 |
+
|
| 187 |
+
def __getitem__(self,i):
|
| 188 |
+
m,row = self.m,self.rows[i]
|
| 189 |
+
if self.arrays is None:
|
| 190 |
+
self.arrays = {k:np.load(m['files'][k]['resolved'],mmap_mode='r') for k in ['slp','pressure']}
|
| 191 |
+
self.arrays['q'] = np.load(m['qpath'],mmap_mode='r')
|
| 192 |
+
with np.load(m['files']['geography']['resolved'],allow_pickle=False) as geo:
|
| 193 |
+
self.arrays['land'] = geo['land_fraction']
|
| 194 |
+
self.arrays['height'] = geo['elevation_m']/8000
|
| 195 |
+
ai = row['atlas']
|
| 196 |
+
ii = slice(ai-8,ai+21)
|
| 197 |
+
coarse = np.concatenate([np.array(self.arrays['slp'][ii],dtype='float32')[:,None],
|
| 198 |
+
np.array(self.arrays['q'][ii],dtype='float32')*np.array(m['scale'],dtype='float32')[None,:,None,None]+np.array(m['offset'],dtype='float32')[None,:,None,None]],1)
|
| 199 |
+
mean,std = np.array(m['normalization']['mean'],dtype='float32'),np.array(m['normalization']['std'],dtype='float32')
|
| 200 |
+
coarse = (coarse-mean[None,:,None,None])/std[None,:,None,None]
|
| 201 |
+
pi = row['patch']
|
| 202 |
+
anchor = m['plan']['rows'][pi] if pi>=0 else {'center_lat':round(row['center'][0]*4)/4,'center_lon':round(row['center'][1]*4)/4}
|
| 203 |
+
lat = anchor['center_lat']+np.linspace(15,-15,121,dtype='float32')
|
| 204 |
+
lon = anchor['center_lon']+np.linspace(-15,15,121,dtype='float32')
|
| 205 |
+
yy,xx = np.meshgrid(lat,lon,indexing='ij')
|
| 206 |
+
y = np.rint((90-yy)*4).astype(int); x = np.rint(xx*4).astype(int)%1440
|
| 207 |
+
rs = np.stack([yy/90,xx/180-1,self.arrays['land'][y,x],self.arrays['height'][y,x]]).astype('float32')
|
| 208 |
+
cy,cx = np.meshgrid(m['global_lat'],m['global_lon'],indexing='ij')
|
| 209 |
+
iy = np.rint((90-cy)*4).astype(int); ix = np.rint(cx*4).astype(int)%1440
|
| 210 |
+
gs = np.stack([cy/90,cx/180-1,self.arrays['land'][iy,ix],self.arrays['height'][iy,ix]]).astype('float32')
|
| 211 |
+
detail = np.zeros((29,1,121,121),dtype='float32')
|
| 212 |
+
if pi>=0:
|
| 213 |
+
detail[:,0] = (np.array(self.arrays['pressure'][pi],dtype='float32')-mean[0])/std[0]
|
| 214 |
+
if not np.isfinite(detail).all():
|
| 215 |
+
raise ValueError('Native pressure contains nonfinite fields')
|
| 216 |
+
inputs = {'global_history':coarse[:9], 'regional_history':detail[:9],
|
| 217 |
+
'global_static':gs,'regional_static':rs,'detail_available':np.array([pi>=0],dtype='float32'),
|
| 218 |
+
'center':np.array(row['center'],dtype='float32'),'motion':np.array(row['motion'],dtype='float32'),
|
| 219 |
+
'issue_intensity':np.array(row['issue_intensity'],dtype='float32'),'issue_mask':np.array(row['issue_mask'],dtype='float32')}
|
| 220 |
+
targets = {'global':coarse[9:],'global_mask':np.ones_like(coarse[9:]),
|
| 221 |
+
'regional_lat_weight':np.broadcast_to(np.cos(np.deg2rad(lat))[None,None,:,None],(20,1,121,1)).astype('float32'),
|
| 222 |
+
'regional':detail[9:],'regional_mask':np.full_like(detail[9:],float(pi>=0)),
|
| 223 |
+
'center':np.array(row['target_centers'],dtype='float32'),
|
| 224 |
+
'storm_mask':np.array(row['route_mask'],dtype='float32')}
|
| 225 |
+
for k in ['pressure','pressure_mask','vmax','vmax_mask']:
|
| 226 |
+
targets[k] = np.array(row[k],dtype='float32')
|
| 227 |
+
return {k:torch.from_numpy(v.copy()) for k,v in inputs.items()},{k:torch.from_numpy(v.copy()) for k,v in targets.items()}
|
| 228 |
+
|
| 229 |
+
|
| 230 |
+
"""v165 implementation, embedded with verified archive readers by the builder."""
|
| 231 |
+
VERSION = '1.2.65'
|
| 232 |
+
ARCHITECTURE = REUSE_ARCH = 'wp-causal-transport-evolution-field-tracker-v165'
|
| 233 |
+
TRAIN_HORIZONS = [4, 8, 12, 20]
|
| 234 |
+
|
| 235 |
+
|
| 236 |
+
def sample_field(field, grid, padding='border'):
|
| 237 |
+
# Geometry and interpolation stay FP32, including under CUDA autocast.
|
| 238 |
+
with torch.autocast(device_type=field.device.type, enabled=False):
|
| 239 |
+
return F.grid_sample(field.float(), grid.float(), align_corners=True,
|
| 240 |
+
mode='bilinear', padding_mode=padding)
|
| 241 |
+
|
| 242 |
+
|
| 243 |
+
def identity_grid(field):
|
| 244 |
+
h, w = field.shape[-2:]
|
| 245 |
+
yy, xx = torch.meshgrid(torch.linspace(-1, 1, h, device=field.device),
|
| 246 |
+
torch.linspace(-1, 1, w, device=field.device), indexing='ij')
|
| 247 |
+
return torch.stack([xx, yy], -1)[None].expand(field.shape[0], -1, -1, -1)
|
| 248 |
+
|
| 249 |
+
|
| 250 |
+
def departure_grid(field, velocity, latitude, spacing):
|
| 251 |
+
# East/north m/s -> backward departure coordinates on a north-to-south grid.
|
| 252 |
+
h, w = field.shape[-2:]
|
| 253 |
+
u, v = velocity.float().unbind(1)
|
| 254 |
+
dx = u * 21.6 / (111.2 * latitude.float().deg2rad().cos().clamp_min(.2) * spacing)
|
| 255 |
+
dy = -v * 21.6 / (111.2 * spacing)
|
| 256 |
+
return identity_grid(field) - torch.stack([2*dx/(w-1), 2*dy/(h-1)], -1)
|
| 257 |
+
|
| 258 |
+
|
| 259 |
+
class TransportBlock(nn.Module):
|
| 260 |
+
def __init__(self, width):
|
| 261 |
+
super().__init__()
|
| 262 |
+
self.net = nn.Sequential(nn.Conv2d(width, width, 5, padding=2, groups=width),
|
| 263 |
+
nn.GroupNorm(8, width), nn.Conv2d(width, width*4, 1),
|
| 264 |
+
nn.GELU(), nn.Conv2d(width*4, width, 1))
|
| 265 |
+
self.scale = nn.Parameter(torch.full((1, width, 1, 1), .01))
|
| 266 |
+
|
| 267 |
+
def forward(self, x):
|
| 268 |
+
return x + self.scale * self.net(x)
|
| 269 |
+
|
| 270 |
+
|
| 271 |
+
class EvolutionNet(nn.Module):
|
| 272 |
+
def __init__(self, channels, outputs, base=72, tiny=False):
|
| 273 |
+
super().__init__()
|
| 274 |
+
widths = [base, base*2, base*4, base*6]
|
| 275 |
+
depths = [1]*4 if tiny else [2, 2, 4, 4]
|
| 276 |
+
self.stem = nn.Conv2d(channels, base, 3, padding=1)
|
| 277 |
+
self.enc = nn.ModuleList(nn.Sequential(*[TransportBlock(w) for _ in range(d)])
|
| 278 |
+
for w, d in zip(widths, depths))
|
| 279 |
+
self.down = nn.ModuleList(nn.Conv2d(a, b, 3, stride=2, padding=1)
|
| 280 |
+
for a, b in zip(widths, widths[1:]))
|
| 281 |
+
self.up = nn.ModuleList(nn.Conv2d(widths[i+1], widths[i], 1) for i in (2, 1, 0))
|
| 282 |
+
self.dec = nn.ModuleList(TransportBlock(widths[i]) for i in (2, 1, 0))
|
| 283 |
+
self.out = nn.Conv2d(base, outputs, 1)
|
| 284 |
+
self.features = nn.Conv2d(base, 32, 1)
|
| 285 |
+
nn.init.zeros_(self.out.weight)
|
| 286 |
+
nn.init.zeros_(self.out.bias)
|
| 287 |
+
|
| 288 |
+
def block(self, module, x):
|
| 289 |
+
return checkpoint(module, x, use_reentrant=False) if self.training and torch.is_grad_enabled() else module(x)
|
| 290 |
+
|
| 291 |
+
def forward(self, x):
|
| 292 |
+
x = self.stem(x)
|
| 293 |
+
skips = []
|
| 294 |
+
for i, layer in enumerate(self.enc):
|
| 295 |
+
x = self.block(layer, x)
|
| 296 |
+
skips.append(x)
|
| 297 |
+
if i < 3:
|
| 298 |
+
x = self.down[i](x)
|
| 299 |
+
for up, layer, skip in zip(self.up, self.dec, reversed(skips[:-1])):
|
| 300 |
+
x = F.interpolate(up(x), size=skip.shape[-2:], mode='bilinear', align_corners=False)+skip
|
| 301 |
+
x = self.block(layer, x)
|
| 302 |
+
return self.out(x).float(), self.features(x).float()
|
| 303 |
+
|
| 304 |
+
|
| 305 |
+
class RecurrentFeatures(nn.Module):
|
| 306 |
+
def __init__(self, history_channels):
|
| 307 |
+
super().__init__()
|
| 308 |
+
self.encode = nn.Sequential(nn.Conv2d(history_channels, 32, 3, padding=1),
|
| 309 |
+
nn.GELU(), TransportBlock(32))
|
| 310 |
+
self.gates = nn.Conv2d(64, 64, 3, padding=1)
|
| 311 |
+
self.candidate = nn.Conv2d(64, 32, 3, padding=1)
|
| 312 |
+
|
| 313 |
+
def forward(self, old, features):
|
| 314 |
+
reset, update = torch.sigmoid(self.gates(torch.cat([old, features], 1))).chunk(2, 1)
|
| 315 |
+
proposal = torch.tanh(self.candidate(torch.cat([old*reset, features], 1)))
|
| 316 |
+
return ((1-update)*old + update*proposal).float()
|
| 317 |
+
|
| 318 |
+
|
| 319 |
+
def track_field(pressure, lat, lon, prior, domain=None):
|
| 320 |
+
"""Associate a local pressure minimum; no independent learned route head."""
|
| 321 |
+
lat0, lon0 = prior[:, 0, None, None], prior[:, 1, None, None]
|
| 322 |
+
distance2 = ((lat-lat0)*111.2).square()+((lon-lon0)*111.2*lat0.deg2rad().cos()).square()
|
| 323 |
+
candidate = distance2 <= 650**2
|
| 324 |
+
if domain is not None:
|
| 325 |
+
candidate = candidate & domain
|
| 326 |
+
score = -pressure/3 - distance2/(2*650**2)
|
| 327 |
+
peak = score.masked_fill(~candidate, -1e9).flatten(1).argmax(1)
|
| 328 |
+
h, w = pressure.shape[-2:]
|
| 329 |
+
y, x = torch.arange(h, device=pressure.device)[None, :, None], torch.arange(w, device=pressure.device)[None, None, :]
|
| 330 |
+
local = ((y-(peak//w)[:, None, None]).abs() <= 1) & ((x-(peak%w)[:, None, None]).abs() <= 1) & candidate
|
| 331 |
+
pdf = torch.softmax(score.masked_fill(~local, -1e9).flatten(1), 1).reshape_as(score)
|
| 332 |
+
center = torch.stack([(pdf*lat).sum((1, 2)), (pdf*lon).sum((1, 2))], 1)
|
| 333 |
+
grid = torch.stack([2*(center[:, 1]-lon[:, 0, 0])/(lon[:, 0, -1]-lon[:, 0, 0])-1,
|
| 334 |
+
2*(center[:, 0]-lat[:, 0, 0])/(lat[:, -1, 0]-lat[:, 0, 0])-1], -1)[:, None, None]
|
| 335 |
+
central = sample_field(pressure[:, None], grid)[:, 0, 0, 0]
|
| 336 |
+
annulus = (distance2 >= 300**2) & (distance2 <= 650**2)
|
| 337 |
+
background = (pressure*annulus).sum((1, 2))/annulus.sum((1, 2)).clamp_min(1)
|
| 338 |
+
interior = (peak//w > 0) & (peak//w < h-1) & (peak%w > 0) & (peak%w < w-1)
|
| 339 |
+
valid = candidate.flatten(1).any(1) & interior & (background-central > 1.)
|
| 340 |
+
return center, central, valid, background-central
|
| 341 |
+
|
| 342 |
+
|
| 343 |
+
class TransportForecaster(nn.Module):
|
| 344 |
+
def __init__(self, manifest, tiny=False):
|
| 345 |
+
super().__init__()
|
| 346 |
+
base = 8 if tiny else 72
|
| 347 |
+
self.coarse_memory, self.fine_memory = RecurrentFeatures(72), RecurrentFeatures(9)
|
| 348 |
+
self.coarse = EvolutionNet(8+32+4, 8+2, base, tiny)
|
| 349 |
+
self.fine = EvolutionNet(1+8+32+4+1, 1+2, base, tiny)
|
| 350 |
+
self.steering_logits = nn.Parameter(torch.tensor([.4, .4, .2]).log())
|
| 351 |
+
# Wind intensity is explicitly a scalar estimate, not a resolved 10-m wind map.
|
| 352 |
+
self.wind = nn.Sequential(nn.Linear(32+2, 64), nn.GELU(), nn.Linear(64, 1))
|
| 353 |
+
self.register_buffer('mean', torch.tensor(manifest['normalization']['mean'], dtype=torch.float32)[None, :, None, None])
|
| 354 |
+
self.register_buffer('std', torch.tensor(manifest['normalization']['std'], dtype=torch.float32)[None, :, None, None])
|
| 355 |
+
yy, xx = np.meshgrid(manifest['global_lat'], manifest['global_lon'], indexing='ij')
|
| 356 |
+
self.register_buffer('lat', torch.tensor(yy, dtype=torch.float32)[None])
|
| 357 |
+
self.register_buffer('lon', torch.tensor(xx, dtype=torch.float32)[None])
|
| 358 |
+
|
| 359 |
+
def initial(self, inputs):
|
| 360 |
+
required = {'global_history', 'regional_history', 'global_static', 'regional_static',
|
| 361 |
+
'detail_available', 'center', 'motion', 'issue_intensity', 'issue_mask'}
|
| 362 |
+
if set(inputs) != required:
|
| 363 |
+
raise ValueError('Only declared issue-time and historical inputs are allowed')
|
| 364 |
+
gh, rh, rs = inputs['global_history'], inputs['regional_history'], inputs['regional_static']
|
| 365 |
+
grid = torch.stack([((rs[:, 1]+1)*180-100)/40-1, 1-rs[:, 0]*3], -1)
|
| 366 |
+
parent = sample_field(gh[:, -1, :1], grid)
|
| 367 |
+
available = inputs['detail_available'][:, :, None, None]
|
| 368 |
+
r = torch.where(available.bool(), rh[:, -1], parent)
|
| 369 |
+
return dict(g=gh[:, -1], r=r, gh=self.coarse_memory.encode(gh.flatten(1, 2)).float(),
|
| 370 |
+
rh=self.fine_memory.encode(rh.flatten(1, 2)).float(), gs=inputs['global_static'], rs=rs,
|
| 371 |
+
grid=grid, available=available, anchor=(r-parent)*available,
|
| 372 |
+
coordinates=identity_grid(r).permute(0, 3, 1, 2), source=torch.zeros_like(r),
|
| 373 |
+
center=inputs['center'], motion=inputs['motion'], issue=inputs['issue_intensity'],
|
| 374 |
+
issue_mask=inputs['issue_mask'])
|
| 375 |
+
|
| 376 |
+
def steering(self, g):
|
| 377 |
+
physical = g.float()*self.std+self.mean
|
| 378 |
+
winds = physical[:, 2:].reshape(g.shape[0], 3, 2, *g.shape[-2:])
|
| 379 |
+
return (winds*torch.softmax(self.steering_logits, 0)[None, :, None, None, None]).sum(1)
|
| 380 |
+
|
| 381 |
+
def step(self, s):
|
| 382 |
+
delta, features = self.coarse(torch.cat([s['g'], s['gh'], s['gs']], 1))
|
| 383 |
+
velocity = self.steering(s['g']) + 15*torch.tanh(delta[:, 8:])
|
| 384 |
+
departure = departure_grid(s['g'], velocity, self.lat, 2.5)
|
| 385 |
+
g = sample_field(s['g'], departure) + .2*torch.tanh(delta[:, :8])
|
| 386 |
+
gh = self.coarse_memory(sample_field(s['gh'], departure), features)
|
| 387 |
+
parent = sample_field(g, s['grid'])
|
| 388 |
+
fine_velocity = sample_field(velocity, s['grid'])
|
| 389 |
+
r, rh, coordinates, source = s['r'], s['rh'], s['coordinates'], s['source']
|
| 390 |
+
parent_valid = (s['grid'].abs() <= 1).all(-1)[:, None].float()
|
| 391 |
+
if bool(s['available'].any()):
|
| 392 |
+
rd, rf = self.fine(torch.cat([r, parent, rh, s['rs'], parent_valid], 1))
|
| 393 |
+
fine_velocity = fine_velocity + 10*torch.tanh(rd[:, 1:])
|
| 394 |
+
dep = departure_grid(r, fine_velocity, s['rs'][:, 0]*90, .25)
|
| 395 |
+
# Compose material coordinates and sample the ORIGINAL anomaly once per lead.
|
| 396 |
+
# This avoids repeatedly interpolating away its core. Learned source terms
|
| 397 |
+
# can strengthen, weaken or deform it; survival is not hard-coded.
|
| 398 |
+
displacement = coordinates-identity_grid(r).permute(0, 3, 1, 2)
|
| 399 |
+
coordinates = dep.permute(0, 3, 1, 2)+sample_field(displacement, dep)
|
| 400 |
+
source = sample_field(source, dep, 'zeros') + .2*torch.tanh(rd[:, :1])
|
| 401 |
+
anomaly = sample_field(s['anchor'], coordinates.permute(0, 2, 3, 1), 'zeros')+source
|
| 402 |
+
proposed = parent[:, :1]+anomaly*parent_valid
|
| 403 |
+
r = torch.where(s['available'].bool(), proposed, parent[:, :1])
|
| 404 |
+
rh = self.fine_memory(sample_field(rh, dep), rf)
|
| 405 |
+
else:
|
| 406 |
+
r = parent[:, :1]
|
| 407 |
+
# Association prior uses only predicted circulation, never a future center.
|
| 408 |
+
at = torch.stack([(s['center'][:, 1]-100)/40-1, 1-s['center'][:, 0]/30], -1)[:, None, None]
|
| 409 |
+
steering = sample_field(velocity, at)[:, :, 0, 0]
|
| 410 |
+
displacement_km = .7*steering*21.6 + .3*s['motion']
|
| 411 |
+
prior = s['center']+torch.stack([displacement_km[:, 1]/111.2,
|
| 412 |
+
displacement_km[:, 0]/(111.2*s['center'][:, 0].deg2rad().cos().clamp_min(.2))], 1)
|
| 413 |
+
gp, rp = (g*self.std+self.mean)[:, 0], (r*self.std[:, :1]+self.mean[:, :1])[:, 0]
|
| 414 |
+
coarse_center, coarse_p, coarse_ok, coarse_depth = track_field(gp, self.lat, self.lon, prior)
|
| 415 |
+
fine_lat, fine_lon = s['rs'][:, 0]*90, (s['rs'][:, 1]+1)*180
|
| 416 |
+
fine_center, fine_p, fine_ok, fine_depth = track_field(rp, fine_lat, fine_lon, prior, parent_valid[:, 0].bool())
|
| 417 |
+
inside = (prior[:, 0] < fine_lat[:, 0, 0]) & (prior[:, 0] > fine_lat[:, -1, 0]) & \
|
| 418 |
+
(prior[:, 1] > fine_lon[:, 0, 0]) & (prior[:, 1] < fine_lon[:, 0, -1])
|
| 419 |
+
inside = inside & (prior[:, 0] >= 0) & (prior[:, 0] <= 60) & (prior[:, 1] >= 100) & (prior[:, 1] <= 180)
|
| 420 |
+
use_fine = s['available'][:, 0, 0, 0].bool() & inside
|
| 421 |
+
center = torch.where(use_fine[:, None], fine_center, coarse_center)
|
| 422 |
+
pressure = torch.where(use_fine, fine_p, coarse_p)
|
| 423 |
+
ok = torch.where(use_fine, fine_ok, coarse_ok)
|
| 424 |
+
wind_features = torch.cat([gh.mean((2, 3)), s['issue'][:, :1]/100, s['issue_mask'][:, :1]], 1)
|
| 425 |
+
vmax = F.softplus(self.wind(wind_features)[:, 0])*30
|
| 426 |
+
motion = torch.stack([(center[:, 1]-s['center'][:, 1])*111.2*((center[:, 0]+s['center'][:, 0])/2).deg2rad().cos(),
|
| 427 |
+
(center[:, 0]-s['center'][:, 0])*111.2], 1)
|
| 428 |
+
out = dict(global_=g, regional=r, center=center, pressure=pressure, grid_pressure=pressure,
|
| 429 |
+
vmax=vmax, track_valid=ok, native_route=use_fine, regional_valid=parent_valid.bool(),
|
| 430 |
+
pressure_depth=torch.where(use_fine, fine_depth, coarse_depth),
|
| 431 |
+
global_velocity=velocity, regional_velocity=fine_velocity)
|
| 432 |
+
out['global'] = out.pop('global_')
|
| 433 |
+
return dict(s, g=g, r=r, gh=gh, rh=rh, coordinates=coordinates, source=source,
|
| 434 |
+
center=center, motion=motion), out
|
| 435 |
+
|
| 436 |
+
def forward(self, inputs, steps=20):
|
| 437 |
+
state, outputs = self.initial(inputs), []
|
| 438 |
+
for _ in range(steps):
|
| 439 |
+
state, output = self.step(state)
|
| 440 |
+
outputs.append(output)
|
| 441 |
+
return {k: torch.stack([o[k] for o in outputs], 1) for k in outputs[0]}
|
| 442 |
+
|
| 443 |
+
|
| 444 |
+
def weighted_mean(value, weight):
|
| 445 |
+
return (value*weight).sum()/weight.expand_as(value).sum().clamp_min(1)
|
| 446 |
+
|
| 447 |
+
|
| 448 |
+
def objective165(output, target, inputs, manifest):
|
| 449 |
+
terms = {}
|
| 450 |
+
for key in ('global', 'regional'):
|
| 451 |
+
pred, truth, mask = output[key].float(), target[key].float(), target[key+'_mask'].float()
|
| 452 |
+
if key == 'regional':
|
| 453 |
+
mask = mask*output['regional_valid']
|
| 454 |
+
latitude = inputs['global_static' if key == 'global' else 'regional_static'][:, :1]*90
|
| 455 |
+
area = latitude.deg2rad().cos().clamp_min(0)*mask
|
| 456 |
+
terms[key] = weighted_mean(F.smooth_l1_loss(pred, truth, reduction='none'), area)
|
| 457 |
+
for axis in (-2, -1):
|
| 458 |
+
a, b = pred.diff(dim=axis), truth.diff(dim=axis)
|
| 459 |
+
w = area.narrow(axis, 0, a.shape[axis])
|
| 460 |
+
terms[key] = terms[key]+.2*weighted_mean((a-b).abs(), w)
|
| 461 |
+
# Labels localize supervision only. They are never passed to initial()/step().
|
| 462 |
+
lat, lon = inputs['regional_static'][:, 0]*90, (inputs['regional_static'][:, 1]+1)*180
|
| 463 |
+
ct = target['center']
|
| 464 |
+
dist2 = ((lat-ct[:, 0, None, None])*111.2).square()+ \
|
| 465 |
+
((lon-ct[:, 1, None, None])*111.2*ct[:, 0, None, None].deg2rad().cos()).square()
|
| 466 |
+
region = (dist2 < 600**2)[:, None]*target['regional_mask']*output['regional_valid']*target['storm_mask'][:, None, None, None]
|
| 467 |
+
p, truth = output['regional'].float(), target['regional'].float()
|
| 468 |
+
terms['storm_field'] = weighted_mean((p-truth).abs(), region)
|
| 469 |
+
# Compare resolved native pressure minima against the SAME resolved target field,
|
| 470 |
+
# not subgrid best-track pressure (e.g. Tip's 875 hPa is absent in this archive).
|
| 471 |
+
local = region.flatten(1).bool()
|
| 472 |
+
pmin = p.flatten(1).masked_fill(~local, 1e4).min(1).values
|
| 473 |
+
tmin = truth.flatten(1).masked_fill(~local, 1e4).min(1).values
|
| 474 |
+
terms['resolved_core'] = weighted_mean((pmin-tmin).abs(), local.any(1).float())
|
| 475 |
+
# Radial-band pressure contrasts retain size/asymmetry information without
|
| 476 |
+
# imposing an idealized vortex or forcing a storm to survive.
|
| 477 |
+
terms['radial_structure'] = 0
|
| 478 |
+
for low, high in ((0, 150), (150, 300), (300, 600)):
|
| 479 |
+
band = ((dist2 >= low**2) & (dist2 < high**2))[:, None]*region
|
| 480 |
+
den = band.sum((1, 2, 3)).clamp_min(1)
|
| 481 |
+
contrast = ((p-truth)*band).sum((1, 2, 3))/den
|
| 482 |
+
terms['radial_structure'] = terms['radial_structure']+weighted_mean(contrast.abs(), (den > 1).float())/3
|
| 483 |
+
displacement = (output['center']-ct)*torch.stack([torch.ones_like(ct[:, 0])*111.2,
|
| 484 |
+
111.2*ct[:, 0].deg2rad().cos()], 1)
|
| 485 |
+
terms['route'] = weighted_mean(F.smooth_l1_loss(displacement/100, torch.zeros_like(displacement), reduction='none'),
|
| 486 |
+
target['storm_mask'][:, None])
|
| 487 |
+
terms['wind'] = weighted_mean(F.smooth_l1_loss(output['vmax']/30, target['vmax']/30, reduction='none'), target['vmax_mask'])
|
| 488 |
+
terms['flow_smooth'] = sum(v.diff(dim=a).abs().mean()/30 for v in
|
| 489 |
+
[output['global_velocity'], output['regional_velocity']] for a in [-2, -1])
|
| 490 |
+
loss = terms['global']+terms['regional']+.5*terms['storm_field']+.2*terms['resolved_core']+ \
|
| 491 |
+
.1*terms['radial_structure']+.1*terms['route']+.05*terms['wind']+.002*terms['flow_smooth']
|
| 492 |
+
return loss, terms
|
| 493 |
+
|
| 494 |
+
|
| 495 |
+
def training_indices(rows, epoch):
|
| 496 |
+
"""Every coarse window once, plus storm-balanced native replays (~25% native)."""
|
| 497 |
+
rng = np.random.default_rng(165000+epoch)
|
| 498 |
+
indices = list(range(len(rows)))
|
| 499 |
+
groups = {}
|
| 500 |
+
for i, row in enumerate(rows):
|
| 501 |
+
if row['patch'] >= 0:
|
| 502 |
+
groups.setdefault(row['storm_id'], []).append(i)
|
| 503 |
+
native = sum(len(v) for v in groups.values())
|
| 504 |
+
extra = max(0, math.ceil((.25*len(rows)-native)/.75))
|
| 505 |
+
keys = sorted(groups)
|
| 506 |
+
for j in range(extra):
|
| 507 |
+
# Round-robin storms, random window within each storm.
|
| 508 |
+
indices.append(int(rng.choice(groups[keys[j % len(keys)]])))
|
| 509 |
+
rng.shuffle(indices)
|
| 510 |
+
return indices
|
| 511 |
+
|
| 512 |
+
|
| 513 |
+
def save_training_state(path, model, optimizer, scaler, manifest, epoch, history, best):
|
| 514 |
+
atomic_checkpoint(path, dict(version=VERSION, architecture=ARCHITECTURE, dataset_sha256=manifest['identity'],
|
| 515 |
+
data_contract={k:manifest[k] for k in ('channels', 'units', 'global_lat', 'global_lon', 'regional_lat', 'regional_lon', 'normalization')},
|
| 516 |
+
epoch=epoch, model=model.state_dict(), optimizer=optimizer.state_dict(), scaler=scaler.state_dict(),
|
| 517 |
+
history=history, best_value=best, python_rng=random.getstate(), numpy_rng=np.random.get_state(),
|
| 518 |
+
torch_rng=torch.get_rng_state(), cuda_rng=torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None))
|
| 519 |
+
|
| 520 |
+
|
| 521 |
+
def restore_training_state(path, model, optimizer, scaler, manifest):
|
| 522 |
+
saved = torch.load(path, map_location='cpu', weights_only=False)
|
| 523 |
+
if saved['architecture'] != ARCHITECTURE or saved['dataset_sha256'] != manifest['identity']:
|
| 524 |
+
raise ValueError('Resume requires the identical v165 architecture and dataset; v164 is not compatible')
|
| 525 |
+
model.load_state_dict(saved['model'])
|
| 526 |
+
optimizer.load_state_dict(saved['optimizer'])
|
| 527 |
+
scaler.load_state_dict(saved['scaler'])
|
| 528 |
+
random.setstate(saved['python_rng']); np.random.set_state(saved['numpy_rng']); torch.set_rng_state(saved['torch_rng'])
|
| 529 |
+
if torch.cuda.is_available() and saved.get('cuda_rng') is not None:
|
| 530 |
+
torch.cuda.set_rng_state_all(saved['cuda_rng'])
|
| 531 |
+
return saved['epoch']+1, saved['history'], saved['best_value']
|
| 532 |
+
|
| 533 |
+
|
| 534 |
+
def train165(model, m, args, device):
|
| 535 |
+
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=.01)
|
| 536 |
+
scaler = torch.amp.GradScaler('cuda', enabled=device.type == 'cuda' and bool(args.amp))
|
| 537 |
+
last, best_path = args.output/'last.pt', args.output/'best.pt'
|
| 538 |
+
start, history, best = 1, [], float('inf')
|
| 539 |
+
if last.is_file():
|
| 540 |
+
start, history, best = restore_training_state(last, model, optimizer, scaler, m)
|
| 541 |
+
print('Resuming completed epoch', start-1, flush=True)
|
| 542 |
+
atomic_json(args.output/'manifest.json', dict(version=VERSION, architecture=ARCHITECTURE,
|
| 543 |
+
parameters=sum(p.numel() for p in model.parameters()), dataset_sha256=m['identity'],
|
| 544 |
+
script_sha256=digest(__file__), samples=m['counts'], native_samples=m['patch_counts'],
|
| 545 |
+
future_weather_as_inputs=False, official_forecasts_as_inputs=False, input_hours=list(range(-48, 1, 6)),
|
| 546 |
+
forecast_hours=list(range(6, 121, 6)), pressure='resolved field at the reported field-derived center',
|
| 547 |
+
native_sampling='all coarse windows plus storm-balanced native replays to 25 percent',
|
| 548 |
+
validation='fixed 120h joint loss, balanced native/coarse strata; no test-set selection',
|
| 549 |
+
physical_microbatch=1, effective_batch=args.batch, tbptt_steps=4,
|
| 550 |
+
limitations=['regional WP, not global', 'native MSLP only on 1000 windows',
|
| 551 |
+
'no fine upper-air variables', 'not an exact conservation solver', 'Tip is a development case']))
|
| 552 |
+
for epoch in range(start, args.epochs+1):
|
| 553 |
+
horizon = TRAIN_HORIZONS[min(epoch-1, len(TRAIN_HORIZONS)-1)]
|
| 554 |
+
metrics = {}
|
| 555 |
+
for split in ('train', 'validation'):
|
| 556 |
+
training = split == 'train'
|
| 557 |
+
model.train(training)
|
| 558 |
+
data = ReuseDataset(m, split)
|
| 559 |
+
sampler = training_indices(data.rows, epoch) if training else None
|
| 560 |
+
loader = DataLoader(data, batch_size=1, sampler=sampler, num_workers=min(args.workers, 2),
|
| 561 |
+
**({'prefetch_factor':1} if args.workers else {}))
|
| 562 |
+
groups, failures = {False:[], True:[]}, []
|
| 563 |
+
optimizer.zero_grad(set_to_none=True)
|
| 564 |
+
began = time.time()
|
| 565 |
+
for index, (inputs, targets) in enumerate(loader, 1):
|
| 566 |
+
inputs = {k:v.to(device) for k, v in inputs.items()}
|
| 567 |
+
steps = horizon if training else 20
|
| 568 |
+
total, block_loss = 0., 0
|
| 569 |
+
with torch.set_grad_enabled(training):
|
| 570 |
+
state = model.initial(inputs)
|
| 571 |
+
for step in range(steps):
|
| 572 |
+
with torch.autocast('cuda', dtype=torch.float16) if scaler.is_enabled() else nullcontext():
|
| 573 |
+
state, output = model.step(state)
|
| 574 |
+
frame = {k:v[:, step].to(device) for k, v in targets.items()}
|
| 575 |
+
loss, terms = objective165(output, frame, inputs, m)
|
| 576 |
+
if not torch.isfinite(loss):
|
| 577 |
+
raise RuntimeError(f'Nonfinite loss epoch={epoch} sample={index} step={step}; no checkpoint advanced')
|
| 578 |
+
total += float(loss.detach())/steps
|
| 579 |
+
failures.append(float((~output['track_valid']).float().mean()))
|
| 580 |
+
if training:
|
| 581 |
+
block_loss = block_loss+loss/(steps*args.batch)
|
| 582 |
+
if (step+1) % 4 == 0 or step+1 == steps:
|
| 583 |
+
scaler.scale(block_loss).backward()
|
| 584 |
+
state = {k:v.detach() for k, v in state.items()}
|
| 585 |
+
block_loss = 0
|
| 586 |
+
if training and (index % args.batch == 0 or index == len(loader)):
|
| 587 |
+
scaler.unscale_(optimizer)
|
| 588 |
+
remainder = index % args.batch
|
| 589 |
+
if index == len(loader) and remainder:
|
| 590 |
+
for parameter in model.parameters():
|
| 591 |
+
if parameter.grad is not None:
|
| 592 |
+
parameter.grad.mul_(args.batch/remainder)
|
| 593 |
+
norm = nn.utils.clip_grad_norm_(model.parameters(), 1.)
|
| 594 |
+
if not torch.isfinite(norm):
|
| 595 |
+
raise RuntimeError('Nonfinite gradient; checkpoint unchanged. Retry with AMP=0.')
|
| 596 |
+
scaler.step(optimizer); scaler.update(); optimizer.zero_grad(set_to_none=True)
|
| 597 |
+
groups[bool(inputs['detail_available'][0, 0])].append(total)
|
| 598 |
+
if index == 1 or index % 10 == 0:
|
| 599 |
+
print(f'{split} epoch={epoch}/{args.epochs} sample={index}/{len(loader)} horizon={steps} loss={total:.5f} elapsed_s={time.time()-began:.0f}', flush=True)
|
| 600 |
+
# Keep rare native examples visible in selection, rather than 5.7% dilution.
|
| 601 |
+
metrics[split+'_native'] = float(np.mean(groups[True]))
|
| 602 |
+
metrics[split+'_coarse'] = float(np.mean(groups[False]))
|
| 603 |
+
metrics[split] = (metrics[split+'_native']+metrics[split+'_coarse'])/2
|
| 604 |
+
metrics[split+'_association_failure_fraction'] = float(np.mean(failures))
|
| 605 |
+
promote = metrics['validation'] < best
|
| 606 |
+
if promote:
|
| 607 |
+
best = metrics['validation']
|
| 608 |
+
history.append(dict(epoch=epoch, horizon=horizon, **metrics))
|
| 609 |
+
save_training_state(last, model, optimizer, scaler, m, epoch, history, best)
|
| 610 |
+
save_training_state(args.output/f'epoch_{epoch:03d}.pt', model, optimizer, scaler, m, epoch, history, best)
|
| 611 |
+
if promote:
|
| 612 |
+
save_training_state(best_path, model, optimizer, scaler, m, epoch, history, best)
|
| 613 |
+
atomic_json(args.output/'history.json', history)
|
| 614 |
+
print(f'Saved epoch {epoch} directly to {args.output}; best={promote}', flush=True)
|
| 615 |
+
print('Training complete. Test set untouched; best.pt selected by full-120h validation.', flush=True)
|
| 616 |
+
|
| 617 |
+
|
| 618 |
+
def main165():
|
| 619 |
+
parser = argparse.ArgumentParser(description='v1.2.65 causal pressure transport and evolution; Drive-only data')
|
| 620 |
+
parser.add_argument('batch', type=int, nargs='?', default=10)
|
| 621 |
+
parser.add_argument('workers', type=int, nargs='?', default=3)
|
| 622 |
+
parser.add_argument('amp', type=int, nargs='?', default=1, choices=[0, 1])
|
| 623 |
+
parser.add_argument('--manifest', type=Path, default=ROOT/'data/v164_reuse/manifest.json')
|
| 624 |
+
persistent = Path('/content/drive/MyDrive/typhoon_predict') if Path('/content').is_dir() else ROOT
|
| 625 |
+
parser.add_argument('--output', type=Path, default=persistent/'training/v165_transport_pressure')
|
| 626 |
+
parser.add_argument('--epochs', type=int, default=12)
|
| 627 |
+
parser.add_argument('--device', choices=['auto', 'cpu', 'mps', 'cuda'], default='auto')
|
| 628 |
+
parser.add_argument('--preflight', action='store_true')
|
| 629 |
+
args = parser.parse_args()
|
| 630 |
+
if args.batch < 1 or args.workers < 0 or args.epochs < 1:
|
| 631 |
+
parser.error('Invalid batch, workers or epochs')
|
| 632 |
+
if Path('/content').is_dir() and (not (persistent.parent).is_dir() or not args.output.resolve().is_relative_to(persistent.resolve())):
|
| 633 |
+
raise ValueError('Mount Google Drive first; Colab checkpoints must persist under MyDrive/typhoon_predict')
|
| 634 |
+
manifest = load_reuse(args.manifest)
|
| 635 |
+
for split in ('train', 'validation', 'test'):
|
| 636 |
+
data = ReuseDataset(manifest, split)
|
| 637 |
+
data[0]; data[next(i for i, r in enumerate(data.rows) if r['patch'] >= 0)]
|
| 638 |
+
if args.preflight:
|
| 639 |
+
print('PASS full existing archive checksums, splits and native samples. No training started.', flush=True)
|
| 640 |
+
return
|
| 641 |
+
device = device_for(args.device)
|
| 642 |
+
random.seed(165); np.random.seed(165); torch.manual_seed(165)
|
| 643 |
+
model = TransportForecaster(manifest).to(device)
|
| 644 |
+
print(json.dumps(dict(version=VERSION, architecture=ARCHITECTURE, parameters=sum(p.numel() for p in model.parameters()),
|
| 645 |
+
device=str(device), samples=manifest['counts'], native_samples=manifest['patch_counts'], output=str(args.output))), flush=True)
|
| 646 |
+
args.output.mkdir(parents=True, exist_ok=True)
|
| 647 |
+
lock = args.output/'training.lock'
|
| 648 |
+
fd = os.open(lock, os.O_CREAT | os.O_EXCL | os.O_WRONLY)
|
| 649 |
+
os.write(fd, json.dumps(dict(pid=os.getpid(), host=socket.gethostname())).encode()); os.close(fd)
|
| 650 |
+
try:
|
| 651 |
+
train165(model, manifest, args, device)
|
| 652 |
+
finally:
|
| 653 |
+
lock.unlink(missing_ok=True)
|
| 654 |
+
|
| 655 |
+
|
| 656 |
+
if __name__ == '__main__':
|
| 657 |
+
main165()
|
training/trackformer_1_2/source/v173/v173_pressure_train.py
ADDED
|
@@ -0,0 +1,176 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Full-dataset, resumable Metal training for the v168 moving-core prototype."""
|
| 2 |
+
import argparse
|
| 3 |
+
import fcntl
|
| 4 |
+
import importlib.util
|
| 5 |
+
import json
|
| 6 |
+
import os
|
| 7 |
+
from pathlib import Path
|
| 8 |
+
import random
|
| 9 |
+
import signal
|
| 10 |
+
import socket
|
| 11 |
+
import time
|
| 12 |
+
import sys
|
| 13 |
+
import numpy as np
|
| 14 |
+
import torch
|
| 15 |
+
from torch.utils.data import DataLoader
|
| 16 |
+
from model import CoreForecaster, objective, base, VERSION, ARCHITECTURE
|
| 17 |
+
|
| 18 |
+
ROOT=Path(__file__).resolve().parents[1]
|
| 19 |
+
HORIZONS=[4,8,12,20]
|
| 20 |
+
|
| 21 |
+
|
| 22 |
+
class Tee:
|
| 23 |
+
def __init__(self, terminal, log): self.terminal,self.log=terminal,log
|
| 24 |
+
def write(self, value):
|
| 25 |
+
self.terminal.write(value);self.log.write(value);self.log.flush()
|
| 26 |
+
def flush(self): self.terminal.flush();self.log.flush()
|
| 27 |
+
|
| 28 |
+
|
| 29 |
+
def batches(indices, rows, effective, physical):
|
| 30 |
+
for start in range(0,len(indices),effective):
|
| 31 |
+
window=indices[start:start+effective];parts=[]
|
| 32 |
+
for native in (True,False):
|
| 33 |
+
group=[i for i in window if (rows[i]['patch']>=0)==native]
|
| 34 |
+
parts.extend(group[j:j+physical] for j in range(0,len(group),physical))
|
| 35 |
+
yield parts,len(window)
|
| 36 |
+
|
| 37 |
+
|
| 38 |
+
def save(path,model,optimizer,m,epoch,history,best,progress=None):
|
| 39 |
+
base.atomic_checkpoint(path,dict(version=VERSION,architecture=ARCHITECTURE,
|
| 40 |
+
dataset_sha256=m['identity'],model=model.state_dict(),optimizer=optimizer.state_dict(),
|
| 41 |
+
data_contract={k:m[k] for k in ('channels','units','global_lat','global_lon','regional_lat','regional_lon','normalization')},
|
| 42 |
+
epoch=epoch,history=history,best=best,progress=progress,
|
| 43 |
+
python_rng=random.getstate(),numpy_rng=np.random.get_state(),torch_rng=torch.get_rng_state()))
|
| 44 |
+
|
| 45 |
+
|
| 46 |
+
def restore(path,model,optimizer,m):
|
| 47 |
+
p=torch.load(path,map_location='cpu',weights_only=False)
|
| 48 |
+
if p['architecture']!=ARCHITECTURE or p['dataset_sha256']!=m['identity']:
|
| 49 |
+
raise ValueError('Architecture/dataset mismatch')
|
| 50 |
+
model.load_state_dict(p['model'],strict=True);optimizer.load_state_dict(p['optimizer'])
|
| 51 |
+
random.setstate(p['python_rng']);np.random.set_state(p['numpy_rng']);torch.set_rng_state(p['torch_rng'])
|
| 52 |
+
return p
|
| 53 |
+
|
| 54 |
+
|
| 55 |
+
def warm_start(model,path,m):
|
| 56 |
+
if base.digest(path)!='99ed776097e05ad6df03375465d9618741bed63a4ee57724d3ba98e295039d30':
|
| 57 |
+
raise ValueError('Expected the verified v168 epoch 3 baseline checkpoint')
|
| 58 |
+
p=torch.load(path,map_location='cpu',weights_only=False)
|
| 59 |
+
if p['version']!='1.2.68' or p['epoch']!=3 or p['dataset_sha256']!=m['identity']:
|
| 60 |
+
raise ValueError('Warm-start identity mismatch')
|
| 61 |
+
weights=p['model']
|
| 62 |
+
missing,extra=model.load_state_dict(weights,strict=False)
|
| 63 |
+
if extra or any(not k.startswith('context.') for k in missing):
|
| 64 |
+
raise ValueError('Incomplete environmental transfer')
|
| 65 |
+
|
| 66 |
+
|
| 67 |
+
def main():
|
| 68 |
+
p=argparse.ArgumentParser()
|
| 69 |
+
for name,default in [('batch',10),('workers',2),('amp',0)]:p.add_argument(name,type=int,nargs='?',default=default)
|
| 70 |
+
p.add_argument('--microbatch',type=int,default=2)
|
| 71 |
+
p.add_argument('--epochs',type=int,default=12)
|
| 72 |
+
p.add_argument('--device',default='cuda' if torch.cuda.is_available() else 'mps',choices=['mps','cuda','cpu'])
|
| 73 |
+
here=Path(__file__).resolve().parent
|
| 74 |
+
colab=Path('/content/drive/MyDrive/typhoon_predict')
|
| 75 |
+
p.add_argument('--manifest',type=Path,default=(here if colab.exists() else ROOT)/'data/v164_reuse/manifest.json')
|
| 76 |
+
p.add_argument('--output',type=Path,default=colab/'training/v173' if colab.exists() else here/'runs/mac')
|
| 77 |
+
p.add_argument('--warm-start',type=Path,default=here/'baseline_epoch_003.pt' if colab.exists() else ROOT/'v168/runs/mac/epoch_003.pt')
|
| 78 |
+
p.add_argument('--preflight',action='store_true')
|
| 79 |
+
a=p.parse_args()
|
| 80 |
+
if min(a.batch,a.microbatch,a.epochs)<1 or a.workers<0 or a.amp not in (0,1):p.error('Positive batches/epochs required; amp must be 0 or 1')
|
| 81 |
+
if a.amp: print('AMP request recorded; pressure dynamics use FP32 for numerical stability.',flush=True)
|
| 82 |
+
if a.preflight:
|
| 83 |
+
m=base.load_reuse(a.manifest)
|
| 84 |
+
assert m['counts']==dict(train=13949,validation=1041,test=1195)
|
| 85 |
+
assert m['patch_counts']==dict(train=800,validation=100,test=100)
|
| 86 |
+
model=CoreForecaster(m);warm_start(model,a.warm_start,m)
|
| 87 |
+
print('Verified full dataset and v168 E3 warm start; parameters=',sum(p.numel() for p in model.parameters()))
|
| 88 |
+
return
|
| 89 |
+
torch.set_num_threads(4);random.seed(168);np.random.seed(168);torch.manual_seed(168)
|
| 90 |
+
base.EvolutionNet.block=lambda self,module,x:module(x)
|
| 91 |
+
device=base.device_for(a.device);a.output.mkdir(parents=True,exist_ok=True)
|
| 92 |
+
handle=(a.output/'training.lock').open('a+')
|
| 93 |
+
try:fcntl.flock(handle,fcntl.LOCK_EX|fcntl.LOCK_NB)
|
| 94 |
+
except BlockingIOError:raise RuntimeError('Training already active for this output')
|
| 95 |
+
handle.seek(0);handle.truncate();json.dump(dict(pid=os.getpid(),host=socket.gethostname()),handle);handle.flush()
|
| 96 |
+
log=(a.output/'train.log').open('a',buffering=1)
|
| 97 |
+
sys.stdout=Tee(sys.stdout,log);sys.stderr=Tee(sys.stderr,log)
|
| 98 |
+
stop=[False]
|
| 99 |
+
for sig in (signal.SIGTERM,signal.SIGINT):signal.signal(sig,lambda *_:stop.__setitem__(0,True))
|
| 100 |
+
m=base.load_reuse(a.manifest)
|
| 101 |
+
assert m['counts']==dict(train=13949,validation=1041,test=1195)
|
| 102 |
+
assert m['patch_counts']==dict(train=800,validation=100,test=100)
|
| 103 |
+
model=CoreForecaster(m).to(device)
|
| 104 |
+
optimizer=torch.optim.AdamW(model.parameters(),lr=3e-5,weight_decay=.01)
|
| 105 |
+
last=a.output/'last.pt';history=[];best=float('inf');start=1;progress=None
|
| 106 |
+
if last.exists():
|
| 107 |
+
saved=restore(last,model,optimizer,m)
|
| 108 |
+
history,best,start,progress=saved['history'],saved['best'],saved['epoch']+1,saved['progress']
|
| 109 |
+
if progress and progress['batch']!=a.batch:raise ValueError('Mid-epoch resume needs identical effective batch')
|
| 110 |
+
else:warm_start(model,a.warm_start,m)
|
| 111 |
+
manifest=dict(version=VERSION,architecture=ARCHITECTURE,parameters=sum(x.numel() for x in model.parameters()),
|
| 112 |
+
samples=m['counts'],native_samples=m['patch_counts'],dataset_sha256=m['identity'],device=str(device),
|
| 113 |
+
effective_batch=a.batch,microbatch=a.microbatch,precision='float32',future_weather_as_inputs=False,
|
| 114 |
+
future_best_track_as_inputs=False,official_forecasts_as_inputs=False,
|
| 115 |
+
pressure_map='moving learned core plus resolved environment; observation-conditioned initialization, not raw analysis',
|
| 116 |
+
internal_core_spacing_km=20,initial_core_prior_sigma_km=100,inner_absolute_reanalysis_loss_exclusion_km=150,
|
| 117 |
+
core_structure_status='prototype; unresolved inner structure is not spatially validated',
|
| 118 |
+
pressure_change_supervision=True,motion_change_supervision=True,
|
| 119 |
+
environmental_attention='two spatial pooling scales, two transformer blocks, cross attention; zero initialized correction',
|
| 120 |
+
route='pressure-field association; no independent displayed route or pressure head',
|
| 121 |
+
warm_start=str(a.warm_start),warm_start_sha256=base.digest(a.warm_start),
|
| 122 |
+
model_sha256=base.digest(Path(__file__).with_name('model.py')),trainer_sha256=base.digest(__file__))
|
| 123 |
+
base.atomic_json(a.output/'manifest.json',manifest)
|
| 124 |
+
print(json.dumps(manifest),flush=True)
|
| 125 |
+
for epoch in range(start,a.epochs+1):
|
| 126 |
+
metrics={}
|
| 127 |
+
for split in ('train','validation'):
|
| 128 |
+
training=split=='train';model.train(training);data=base.ReuseDataset(m,split)
|
| 129 |
+
indices=base.training_indices(data.rows,epoch) if training else list(range(len(data)))
|
| 130 |
+
windows=list(batches(indices,data.rows,a.batch,a.microbatch))
|
| 131 |
+
offset=progress['window'] if training and progress else 0
|
| 132 |
+
totals=progress['totals'] if training and progress else [[0.,0],[0.,0]]
|
| 133 |
+
parts=[part for group,_ in windows[offset:] for part in group]
|
| 134 |
+
loader=iter(DataLoader(data,batch_sampler=parts,num_workers=a.workers,**({'prefetch_factor':1} if a.workers else {})))
|
| 135 |
+
began=time.time();steps=HORIZONS[min(epoch-1,3)] if training else 20
|
| 136 |
+
for wi,(group,count) in enumerate(windows[offset:],offset):
|
| 137 |
+
optimizer.zero_grad(set_to_none=True)
|
| 138 |
+
for _ in group:
|
| 139 |
+
x,y=next(loader);x={k:v.to(device) for k,v in x.items()};size=x['center'].shape[0]
|
| 140 |
+
with torch.set_grad_enabled(training):
|
| 141 |
+
state=model.initial(x);block=0;total=0.
|
| 142 |
+
for lead in range(steps):
|
| 143 |
+
state,out=model.step(state);target={k:v[:,lead].to(device) for k,v in y.items()}
|
| 144 |
+
target['previous_pressure']=y['pressure'][:,lead-1].to(device) if lead else x['issue_intensity'][:,1]
|
| 145 |
+
target['previous_pressure_mask']=y['pressure_mask'][:,lead-1].to(device) if lead else x['issue_mask'][:,1]
|
| 146 |
+
target['previous_center']=y['center'][:,lead-1].to(device) if lead else x['center']
|
| 147 |
+
target['previous_storm_mask']=y['storm_mask'][:,lead-1].to(device) if lead else torch.ones_like(target['storm_mask'])
|
| 148 |
+
loss,terms=objective(out,target,x,m)
|
| 149 |
+
if not bool(torch.isfinite(loss)):raise RuntimeError('Nonfinite loss; previous checkpoint retained')
|
| 150 |
+
total+=float(loss.detach())/steps
|
| 151 |
+
if training:
|
| 152 |
+
block=block+loss*size/(count*steps)
|
| 153 |
+
if (lead+1)%4==0 or lead+1==steps:
|
| 154 |
+
block.backward();block=0;state={k:v.detach() for k,v in state.items()}
|
| 155 |
+
native=int(bool(x['detail_available'][0,0]));totals[native][0]+=total*size;totals[native][1]+=size
|
| 156 |
+
if training:
|
| 157 |
+
norm=torch.nn.utils.clip_grad_norm_(model.parameters(),1.)
|
| 158 |
+
if not bool(torch.isfinite(norm)):raise RuntimeError('Nonfinite gradient; previous checkpoint retained')
|
| 159 |
+
optimizer.step()
|
| 160 |
+
if (wi+1)%50==0 or wi+1==len(windows) or stop[0]:
|
| 161 |
+
save(last,model,optimizer,m,epoch-1,history,best,dict(window=wi+1,totals=totals,batch=a.batch))
|
| 162 |
+
if wi%5==0 or wi+1==len(windows):
|
| 163 |
+
print(f'{split} epoch={epoch}/{a.epochs} window={wi+1}/{len(windows)} horizon={steps} loss={total:.5f} pressure={float(terms["pressure"].detach()):.5f} route={float(terms["route"].detach()):.5f} elapsed_s={time.time()-began:.0f}',flush=True)
|
| 164 |
+
if stop[0]:print('Stopped at optimizer boundary; resume from last.pt',flush=True);return
|
| 165 |
+
metrics[split]=sum(v/n for v,n in totals if n)/sum(n>0 for _,n in totals);progress=None
|
| 166 |
+
promote=metrics['validation']<best
|
| 167 |
+
if promote:best=metrics['validation']
|
| 168 |
+
history.append(dict(epoch=epoch,horizon=HORIZONS[min(epoch-1,3)],**metrics))
|
| 169 |
+
for path in [a.output/f'epoch_{epoch:03d}.pt',last]+([a.output/'best.pt'] if promote else []):
|
| 170 |
+
save(path,model,optimizer,m,epoch,history,best)
|
| 171 |
+
base.atomic_json(a.output/'history.json',history)
|
| 172 |
+
print(f'Saved completed epoch {epoch}; best={promote}',flush=True)
|
| 173 |
+
print('Training complete; best selected on fixed 120h validation.',flush=True)
|
| 174 |
+
|
| 175 |
+
|
| 176 |
+
if __name__=='__main__':main()
|
training/trackformer_1_2/train.py
ADDED
|
@@ -0,0 +1,33 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Portable launcher for the unchanged archived v173 trainer; no auto-training."""
|
| 2 |
+
import argparse
|
| 3 |
+
from pathlib import Path
|
| 4 |
+
import runpy
|
| 5 |
+
import sys
|
| 6 |
+
|
| 7 |
+
|
| 8 |
+
def main():
|
| 9 |
+
here = Path(__file__).resolve().parent
|
| 10 |
+
p = argparse.ArgumentParser(description=__doc__)
|
| 11 |
+
p.add_argument('--data-root', type=Path, default=here / 'data')
|
| 12 |
+
p.add_argument('--warm-start', type=Path, default=here / 'checkpoints/v168_epoch_003.pt')
|
| 13 |
+
p.add_argument('--output', type=Path, required=True)
|
| 14 |
+
p.add_argument('--device', choices=['cpu', 'mps', 'cuda'], required=True)
|
| 15 |
+
p.add_argument('--epochs', type=int, default=4)
|
| 16 |
+
p.add_argument('--batch', type=int, default=10)
|
| 17 |
+
p.add_argument('--microbatch', type=int, default=4)
|
| 18 |
+
p.add_argument('--workers', type=int, default=3)
|
| 19 |
+
p.add_argument('--preflight', action='store_true')
|
| 20 |
+
a = p.parse_args()
|
| 21 |
+
source = here / 'source/v173'
|
| 22 |
+
sys.path.insert(0, str(source))
|
| 23 |
+
sys.argv = [str(source / 'v173_pressure_train.py'), str(a.batch), str(a.workers), '0',
|
| 24 |
+
'--microbatch', str(a.microbatch), '--epochs', str(a.epochs),
|
| 25 |
+
'--device', a.device, '--manifest', str(a.data_root.resolve() / 'manifest.json'),
|
| 26 |
+
'--warm-start', str(a.warm_start.resolve()), '--output', str(a.output.resolve())]
|
| 27 |
+
if a.preflight:
|
| 28 |
+
sys.argv.append('--preflight')
|
| 29 |
+
runpy.run_path(str(source / 'v173_pressure_train.py'), run_name='__main__')
|
| 30 |
+
|
| 31 |
+
|
| 32 |
+
if __name__ == '__main__':
|
| 33 |
+
main()
|