Upload folder using huggingface_hub
Browse files- .gitattributes +1 -33
- README.md +134 -0
- conf/config.yaml +60 -0
- config.json +55 -0
- configuration.json +12 -0
- model/pinn-tc.py +92 -0
- scripts/fake_data.py +82 -0
- scripts/inference.py +77 -0
- scripts/result.py +108 -0
- scripts/train.py +137 -0
- weight/.gitkeep +0 -0
.gitattributes
CHANGED
|
@@ -1,35 +1,3 @@
|
|
| 1 |
-
*.
|
| 2 |
-
*.arrow filter=lfs diff=lfs merge=lfs -text
|
| 3 |
-
*.bin filter=lfs diff=lfs merge=lfs -text
|
| 4 |
-
*.bz2 filter=lfs diff=lfs merge=lfs -text
|
| 5 |
-
*.ckpt filter=lfs diff=lfs merge=lfs -text
|
| 6 |
-
*.ftz filter=lfs diff=lfs merge=lfs -text
|
| 7 |
-
*.gz filter=lfs diff=lfs merge=lfs -text
|
| 8 |
-
*.h5 filter=lfs diff=lfs merge=lfs -text
|
| 9 |
-
*.joblib filter=lfs diff=lfs merge=lfs -text
|
| 10 |
-
*.lfs.* filter=lfs diff=lfs merge=lfs -text
|
| 11 |
-
*.mlmodel filter=lfs diff=lfs merge=lfs -text
|
| 12 |
-
*.model filter=lfs diff=lfs merge=lfs -text
|
| 13 |
-
*.msgpack filter=lfs diff=lfs merge=lfs -text
|
| 14 |
*.npy filter=lfs diff=lfs merge=lfs -text
|
| 15 |
*.npz filter=lfs diff=lfs merge=lfs -text
|
| 16 |
-
*.onnx filter=lfs diff=lfs merge=lfs -text
|
| 17 |
-
*.ot filter=lfs diff=lfs merge=lfs -text
|
| 18 |
-
*.parquet filter=lfs diff=lfs merge=lfs -text
|
| 19 |
-
*.pb filter=lfs diff=lfs merge=lfs -text
|
| 20 |
-
*.pickle filter=lfs diff=lfs merge=lfs -text
|
| 21 |
-
*.pkl filter=lfs diff=lfs merge=lfs -text
|
| 22 |
-
*.pt filter=lfs diff=lfs merge=lfs -text
|
| 23 |
-
*.pth filter=lfs diff=lfs merge=lfs -text
|
| 24 |
-
*.rar filter=lfs diff=lfs merge=lfs -text
|
| 25 |
-
*.safetensors filter=lfs diff=lfs merge=lfs -text
|
| 26 |
-
saved_model/**/* filter=lfs diff=lfs merge=lfs -text
|
| 27 |
-
*.tar.* filter=lfs diff=lfs merge=lfs -text
|
| 28 |
-
*.tar filter=lfs diff=lfs merge=lfs -text
|
| 29 |
-
*.tflite filter=lfs diff=lfs merge=lfs -text
|
| 30 |
-
*.tgz filter=lfs diff=lfs merge=lfs -text
|
| 31 |
-
*.wasm filter=lfs diff=lfs merge=lfs -text
|
| 32 |
-
*.xz filter=lfs diff=lfs merge=lfs -text
|
| 33 |
-
*.zip filter=lfs diff=lfs merge=lfs -text
|
| 34 |
-
*.zst filter=lfs diff=lfs merge=lfs -text
|
| 35 |
-
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
|
|
|
| 1 |
+
*.pt filter=lfs diff=lfs merge=lfs -text
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 2 |
*.npy filter=lfs diff=lfs merge=lfs -text
|
| 3 |
*.npz filter=lfs diff=lfs merge=lfs -text
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
README.md
ADDED
|
@@ -0,0 +1,134 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
---
|
| 2 |
+
license: cc-by-4.0
|
| 3 |
+
language:
|
| 4 |
+
- en
|
| 5 |
+
tags:
|
| 6 |
+
- OneScience
|
| 7 |
+
- Earth Science
|
| 8 |
+
- Physics-Informed Neural Network
|
| 9 |
+
frameworks: PyTorch
|
| 10 |
+
---
|
| 11 |
+
|
| 12 |
+
<p align="center">
|
| 13 |
+
<strong><span style="font-size: 30px;">PINN-TC</span></strong>
|
| 14 |
+
</p>
|
| 15 |
+
|
| 16 |
+
# Model Introduction
|
| 17 |
+
|
| 18 |
+
PINN-TC reconstructs three-dimensional tropical-cyclone wind and pressure-related fields from sparse flight, dropsonde, or simulated observations, using physical equations to constrain unobserved regions for vortex initialization and data-assimilation research.
|
| 19 |
+
|
| 20 |
+
Paper: Realistic tropical cyclone wind and pressure fields can be reconstructed from sparse data using deep learning
|
| 21 |
+
https://doi.org/10.1038/s43247-023-01144-2
|
| 22 |
+
|
| 23 |
+
# Model Description
|
| 24 |
+
|
| 25 |
+
The method was proposed by research teams at Princeton University, the California Institute of Technology, Stanford University, and NOAA's Geophysical Fluid Dynamics Laboratory. The paper trains and validates the method with Hurricane Ida T-SHiELD simulations and NOAA/AOML flight-level, dropsonde, and Tail Doppler Radar observations. It is intended for sparse tropical-cyclone wind and geopotential reconstruction, physics-consistent constraints, and vortex-initialization research.
|
| 26 |
+
|
| 27 |
+
# Use Cases
|
| 28 |
+
|
| 29 |
+
| Use Case | Description |
|
| 30 |
+
| :---: | :--- |
|
| 31 |
+
| Tropical-cyclone reconstruction | Reconstruct wind, geopotential, and pressure-vertical-velocity fields from sparse coordinate observations. |
|
| 32 |
+
| Physics-informed constraints | Constrain unobserved regions with horizontal momentum and pressure-continuity equations. |
|
| 33 |
+
| ModelScope/OneCode execution | Validate sparse sampling, training, full-grid inference, scientific metrics, and visualization in ModelScope or OneCode environments. |
|
| 34 |
+
| Multi-GPU training | Launch distributed data-parallel training through `torchrun`. |
|
| 35 |
+
|
| 36 |
+
# Usage Instructions
|
| 37 |
+
|
| 38 |
+
## 1.OneCode
|
| 39 |
+
|
| 40 |
+
Experience intelligent, one-click AI4S programming through the OneCode online environment:
|
| 41 |
+
|
| 42 |
+
[Try intelligent, one-click AI4S programming](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)
|
| 43 |
+
|
| 44 |
+
## 2. Download and Installation
|
| 45 |
+
|
| 46 |
+
```bash
|
| 47 |
+
hf download OneScience-Group/PINN-TC --local-dir ./PINN-TC
|
| 48 |
+
cd PINN-TC
|
| 49 |
+
```
|
| 50 |
+
|
| 51 |
+
### Environment Dependencies
|
| 52 |
+
|
| 53 |
+
**Hardware Requirements**
|
| 54 |
+
|
| 55 |
+
- A GPU or DCU is recommended.
|
| 56 |
+
- A CPU can be used to validate the workflow with the default small-sample configuration.
|
| 57 |
+
- DCU users must install DTK in advance. DTK 25.04.2 or later, or the OneScience-recommended version matching the current cluster, is recommended.
|
| 58 |
+
|
| 59 |
+
**DCU Environment**
|
| 60 |
+
|
| 61 |
+
```bash
|
| 62 |
+
# Activate DTK and Conda first
|
| 63 |
+
conda create -n onescience311 python=3.11 -y
|
| 64 |
+
conda activate onescience311
|
| 65 |
+
pip install onescience[earth-dcu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
|
| 66 |
+
pip install numpy pyyaml matplotlib
|
| 67 |
+
```
|
| 68 |
+
|
| 69 |
+
**GPU Environment**
|
| 70 |
+
|
| 71 |
+
```bash
|
| 72 |
+
# Activate Conda first
|
| 73 |
+
conda create -n onescience311 python=3.11 -y libstdcxx-ng=12 libgcc-ng=12 gcc_linux-64=12 gxx_linux-64=12
|
| 74 |
+
conda activate onescience311
|
| 75 |
+
pip install onescience[earth-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
|
| 76 |
+
pip install numpy pyyaml matplotlib
|
| 77 |
+
```
|
| 78 |
+
|
| 79 |
+
### Training Data
|
| 80 |
+
|
| 81 |
+
Training data consist of tropical-cyclone simulation fields and sparse observations containing storm-centered coordinates, horizontal winds, and geopotential. Model inputs are `[y,x,t,p]` coordinates, and targets are winds and geopotential at the same locations, while vertical velocity is constrained by the physical equations. The data cover a `±400 km` horizontal domain, `150-900 hPa`, and observation times at `-3/0/+3 h`. This repository uses a small number of synthetic samples to validate training, inference, and evaluation; they do not represent the official T-SHiELD or NOAA distribution, training scale, or formal paper performance.
|
| 82 |
+
|
| 83 |
+
```bash
|
| 84 |
+
python scripts/fake_data.py --force
|
| 85 |
+
```
|
| 86 |
+
|
| 87 |
+
### Training
|
| 88 |
+
|
| 89 |
+
For single-device training, use:
|
| 90 |
+
|
| 91 |
+
```bash
|
| 92 |
+
python scripts/train.py
|
| 93 |
+
```
|
| 94 |
+
|
| 95 |
+
For multi-GPU training, use:
|
| 96 |
+
|
| 97 |
+
```bash
|
| 98 |
+
torchrun --nproc_per_node=8 --nnodes=1 --rdzv_id=1000 --rdzv_backend=c10d --max_restarts=0 --master_addr="localhost" --master_port=29500 scripts/train.py
|
| 99 |
+
```
|
| 100 |
+
|
| 101 |
+
Training jointly optimizes observational data loss and PDE residuals from the beta-plane horizontal momentum equations and the pressure-coordinate continuity equation. The checkpoint at `result/checkpoints/pinn-tc.pt` contains model state, model configuration, data format version, Adam/L-BFGS states, training history, and completed-step state.
|
| 102 |
+
|
| 103 |
+
### Trained Weights
|
| 104 |
+
|
| 105 |
+
This repository does not bundle official paper weights, and no confirmed official checkpoint is available in the public resources. Running the training script produces a local engineering checkpoint, but no compatibility with external weights or equivalence to paper training results is claimed.
|
| 106 |
+
|
| 107 |
+
### Inference
|
| 108 |
+
|
| 109 |
+
```bash
|
| 110 |
+
python scripts/inference.py
|
| 111 |
+
```
|
| 112 |
+
|
| 113 |
+
Inference loads the trained checkpoint and generates complete three-dimensional wind, geopotential, and pressure-vertical-velocity fields from sparse tropical-cyclone observation coordinates and physical constraints. Complete numerical results are saved to `result/output/predictions.npz`.
|
| 114 |
+
|
| 115 |
+
### Evaluation and Visualization
|
| 116 |
+
|
| 117 |
+
```bash
|
| 118 |
+
python scripts/result.py
|
| 119 |
+
```
|
| 120 |
+
|
| 121 |
+
Evaluation computes wind-speed RMSE, Pearson correlation, PDE residuals, and radial RMSE as aggregates over the three observation times; it saves neither per-class metrics nor separate per-time metrics. It also generates a four-panel comparison of sparse observations, target wind speed, predicted wind speed, and error. Synthetic-data results validate the engineering workflow only and do not represent formal paper performance.
|
| 122 |
+
|
| 123 |
+
# Official OneScience Information
|
| 124 |
+
|
| 125 |
+
| Platform | OneScience Main Repository | Skills Repository |
|
| 126 |
+
| --- | --- | --- |
|
| 127 |
+
| Gitee | https://gitee.com/onescience-ai/onescience | https://gitee.com/onescience-ai/oneskills |
|
| 128 |
+
| GitHub | https://github.com/onescience-ai/OneScience | https://github.com/onescience-ai/oneskills |
|
| 129 |
+
|
| 130 |
+
# Citation and License
|
| 131 |
+
|
| 132 |
+
This repository is an independent engineering reproduction of the public PINN-TC specifications.
|
| 133 |
+
|
| 134 |
+
Use of this repository's code, official model weights, and data remains subject to the licenses and terms of their respective projects.
|
conf/config.yaml
ADDED
|
@@ -0,0 +1,60 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
seed: 42
|
| 2 |
+
domain:
|
| 3 |
+
horizontal_extent_m: 400000
|
| 4 |
+
time_extent_s: 21600
|
| 5 |
+
pressure_min_pa: 15000
|
| 6 |
+
pressure_max_pa: 90000
|
| 7 |
+
data:
|
| 8 |
+
root: data
|
| 9 |
+
format_version: pinn_tc_v1
|
| 10 |
+
input_order: [y, x, t, p]
|
| 11 |
+
supervised_order: [u, v, geopotential]
|
| 12 |
+
output_order: [u, v, geopotential, omega]
|
| 13 |
+
horizontal_extent_m: 400000
|
| 14 |
+
time_extent_s: 21600
|
| 15 |
+
pressure_min_pa: 15000
|
| 16 |
+
pressure_max_pa: 90000
|
| 17 |
+
grid_points: 251
|
| 18 |
+
pressure_levels: 44
|
| 19 |
+
time_steps: 3
|
| 20 |
+
observation_times_s: [-10800, 0, 10800]
|
| 21 |
+
observation_mode: Switch
|
| 22 |
+
observation_line_points: 17
|
| 23 |
+
observation_pressure_stride: 11
|
| 24 |
+
boundary_points_per_edge: 5
|
| 25 |
+
collocation_points: 512
|
| 26 |
+
model:
|
| 27 |
+
width: 16
|
| 28 |
+
depth: 2
|
| 29 |
+
physics:
|
| 30 |
+
f0: 0.00005
|
| 31 |
+
beta: 0.00000000002
|
| 32 |
+
train:
|
| 33 |
+
optimizer_sequence: [Adam, L-BFGS]
|
| 34 |
+
adam_steps: 2
|
| 35 |
+
adam_lr: 0.001
|
| 36 |
+
lbfgs_steps: 1
|
| 37 |
+
lbfgs_lr: 0.5
|
| 38 |
+
lbfgs_history_size: 10
|
| 39 |
+
observation_batch: 128
|
| 40 |
+
collocation_batch: 128
|
| 41 |
+
wind_speed_weight: 1.0
|
| 42 |
+
gamma: 0.1
|
| 43 |
+
momentum_scale: 0.01
|
| 44 |
+
continuity_scale: 0.0001
|
| 45 |
+
runtime:
|
| 46 |
+
device: auto
|
| 47 |
+
ddp: torchrun-compatible
|
| 48 |
+
inference_chunk_size: 131072
|
| 49 |
+
evaluation:
|
| 50 |
+
pde_points: 256
|
| 51 |
+
paths:
|
| 52 |
+
checkpoint: result/checkpoints/pinn-tc.pt
|
| 53 |
+
training_metrics: result/training/metrics.json
|
| 54 |
+
prediction: result/output/predictions.npz
|
| 55 |
+
evaluation_dir: result/evaluation
|
| 56 |
+
paper_model:
|
| 57 |
+
2D: {hidden_layers: 8, width: 100}
|
| 58 |
+
3D: {hidden_layers: 4, width: 100}
|
| 59 |
+
real: {hidden_layers: 4, width: 50}
|
| 60 |
+
optimizer_sequence: [Adam, L-BFGS]
|
config.json
ADDED
|
@@ -0,0 +1,55 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"model_name": "PINN-TC",
|
| 3 |
+
"model_type": "pinn_tc",
|
| 4 |
+
"architectures": [
|
| 5 |
+
"PINNTC"
|
| 6 |
+
],
|
| 7 |
+
"framework": "PyTorch",
|
| 8 |
+
"domain": "tropical cyclone dynamics",
|
| 9 |
+
"task": "continuous-domain tropical-cyclone field reconstruction",
|
| 10 |
+
"implementation": {
|
| 11 |
+
"entry_point": "model/pinn-tc.py",
|
| 12 |
+
"scope": "Core-method and full-data-dimension reduced-scale end-to-end engineering reproduction",
|
| 13 |
+
"train_script": "scripts/train.py",
|
| 14 |
+
"inference_script": "scripts/inference.py",
|
| 15 |
+
"evaluation_script": "scripts/result.py",
|
| 16 |
+
"synthetic_data_script": "scripts/fake_data.py"
|
| 17 |
+
},
|
| 18 |
+
"architecture": {
|
| 19 |
+
"inputs": ["y", "x", "t", "p"],
|
| 20 |
+
"outputs": ["u", "v", "geopotential", "omega"],
|
| 21 |
+
"supervised_outputs": ["u", "v", "geopotential"],
|
| 22 |
+
"pde_constrained_output": "omega",
|
| 23 |
+
"activation": "tanh",
|
| 24 |
+
"engineering_model": {"width": 16, "depth": 2},
|
| 25 |
+
"paper_models": {
|
| 26 |
+
"2D": {"hidden_layers": 8, "width": 100},
|
| 27 |
+
"3D": {"hidden_layers": 4, "width": 100},
|
| 28 |
+
"real": {"hidden_layers": 4, "width": 50}
|
| 29 |
+
},
|
| 30 |
+
"physics_residuals": ["zonal_momentum", "meridional_momentum", "pressure_continuity"]
|
| 31 |
+
},
|
| 32 |
+
"data": {
|
| 33 |
+
"dataset": "T-SHiELD Hurricane Ida case and NOAA/AOML observations",
|
| 34 |
+
"protocol": "Sparse storm-centered coordinate observations with collocation points over the full continuous domain",
|
| 35 |
+
"format_version": "pinn_tc_v1",
|
| 36 |
+
"input_shape": ["N", 4],
|
| 37 |
+
"output_shape": ["N", 4],
|
| 38 |
+
"dense_output_shape": [251, 251, 44, 3, 4],
|
| 39 |
+
"dense_output_layout": "YXPTV",
|
| 40 |
+
"horizontal_extent_m": [-400000, 400000],
|
| 41 |
+
"horizontal_resolution_km": 3.2,
|
| 42 |
+
"pressure_levels": 44,
|
| 43 |
+
"pressure_range_pa": [15000, 90000],
|
| 44 |
+
"observation_times_s": [-10800, 0, 10800],
|
| 45 |
+
"units": {"u_v": "m s-1", "geopotential": "m2 s-2", "omega": "Pa s-1"}
|
| 46 |
+
},
|
| 47 |
+
"configuration_sources": [
|
| 48 |
+
"conf/config.yaml",
|
| 49 |
+
"model/pinn-tc.py",
|
| 50 |
+
"scripts/fake_data.py",
|
| 51 |
+
"scripts/train.py",
|
| 52 |
+
"scripts/inference.py",
|
| 53 |
+
"scripts/result.py"
|
| 54 |
+
]
|
| 55 |
+
}
|
configuration.json
ADDED
|
@@ -0,0 +1,12 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"framework": "PyTorch",
|
| 3 |
+
"task": "continuous-domain tropical-cyclone field reconstruction",
|
| 4 |
+
"model": "PINN-TC",
|
| 5 |
+
"input_format": "coordinates [N,4] ordered as [y,x,t,p]",
|
| 6 |
+
"protocol": "Sparse-observation and PDE-constrained reconstruction on a 251x251x44x3 evaluation grid",
|
| 7 |
+
"default_config": "conf/config.yaml",
|
| 8 |
+
"training": "scripts/train.py",
|
| 9 |
+
"inference": "scripts/inference.py",
|
| 10 |
+
"evaluation": "scripts/result.py",
|
| 11 |
+
"visualization": "scripts/result.py"
|
| 12 |
+
}
|
model/pinn-tc.py
ADDED
|
@@ -0,0 +1,92 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Continuous-coordinate physics-informed neural network for a tropical cyclone."""
|
| 2 |
+
|
| 3 |
+
from __future__ import annotations
|
| 4 |
+
|
| 5 |
+
import math
|
| 6 |
+
|
| 7 |
+
import numpy as np
|
| 8 |
+
import torch
|
| 9 |
+
from torch import nn
|
| 10 |
+
|
| 11 |
+
|
| 12 |
+
INPUT_ORDER = ("y", "x", "t", "p")
|
| 13 |
+
OUTPUT_ORDER = ("u", "v", "geopotential", "omega")
|
| 14 |
+
|
| 15 |
+
|
| 16 |
+
class PINNTC(nn.Module):
|
| 17 |
+
"""Tanh MLP mapping physical [y, x, t, p] coordinates to four fields."""
|
| 18 |
+
|
| 19 |
+
def __init__(self, width: int = 16, depth: int = 2):
|
| 20 |
+
super().__init__()
|
| 21 |
+
if width < 1 or depth < 1:
|
| 22 |
+
raise ValueError("width and depth must be positive")
|
| 23 |
+
layers: list[nn.Module] = []
|
| 24 |
+
dimensions = [4] + [width] * depth + [4]
|
| 25 |
+
for index, (source, target) in enumerate(zip(dimensions[:-1], dimensions[1:])):
|
| 26 |
+
layer = nn.Linear(source, target)
|
| 27 |
+
nn.init.xavier_normal_(layer.weight, gain=nn.init.calculate_gain("tanh") if index < depth else 1.0)
|
| 28 |
+
nn.init.zeros_(layer.bias)
|
| 29 |
+
layers.append(layer)
|
| 30 |
+
if index < depth:
|
| 31 |
+
layers.append(nn.Tanh())
|
| 32 |
+
self.network = nn.Sequential(*layers)
|
| 33 |
+
self.register_buffer("coordinate_scale", torch.tensor([400_000.0, 400_000.0, 21_600.0, 75_000.0]))
|
| 34 |
+
self.register_buffer("coordinate_center", torch.tensor([0.0, 0.0, 0.0, 52_500.0]))
|
| 35 |
+
self.register_buffer("output_scale", torch.tensor([45.0, 45.0, 6_000.0, 0.5]))
|
| 36 |
+
self.register_buffer("output_center", torch.tensor([0.0, 0.0, 50_000.0, 0.0]))
|
| 37 |
+
|
| 38 |
+
def forward(self, coordinates: torch.Tensor) -> torch.Tensor:
|
| 39 |
+
if coordinates.ndim != 2 or coordinates.shape[1] != 4:
|
| 40 |
+
raise ValueError("coordinates must have shape [N,4] in [y,x,t,p] order")
|
| 41 |
+
normalized = (coordinates - self.coordinate_center) / self.coordinate_scale
|
| 42 |
+
return self.network(normalized) * self.output_scale + self.output_center
|
| 43 |
+
|
| 44 |
+
|
| 45 |
+
def pde_residuals(
|
| 46 |
+
model: nn.Module,
|
| 47 |
+
coordinates: torch.Tensor,
|
| 48 |
+
f0: float = 5.0e-5,
|
| 49 |
+
beta: float = 2.0e-11,
|
| 50 |
+
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
| 51 |
+
"""Return beta-plane pressure-coordinate momentum and continuity residuals."""
|
| 52 |
+
if not coordinates.requires_grad:
|
| 53 |
+
coordinates = coordinates.detach().requires_grad_(True)
|
| 54 |
+
fields = model(coordinates)
|
| 55 |
+
u, v, phi, omega = (fields[:, index] for index in range(4))
|
| 56 |
+
|
| 57 |
+
def gradient(value: torch.Tensor) -> torch.Tensor:
|
| 58 |
+
return torch.autograd.grad(value, coordinates, torch.ones_like(value), create_graph=True, retain_graph=True)[0]
|
| 59 |
+
|
| 60 |
+
grad_u, grad_v, grad_phi, grad_omega = map(gradient, (u, v, phi, omega))
|
| 61 |
+
y = coordinates[:, 0]
|
| 62 |
+
coriolis = f0 + beta * y
|
| 63 |
+
# Coordinate columns are [y, x, t, p]. Geopotential has units m2 s-2.
|
| 64 |
+
ru = grad_u[:, 2] + u * grad_u[:, 1] + v * grad_u[:, 0] + omega * grad_u[:, 3] - coriolis * v + grad_phi[:, 1]
|
| 65 |
+
rv = grad_v[:, 2] + u * grad_v[:, 1] + v * grad_v[:, 0] + omega * grad_v[:, 3] + coriolis * u + grad_phi[:, 0]
|
| 66 |
+
rc = grad_u[:, 1] + grad_v[:, 0] + grad_omega[:, 3]
|
| 67 |
+
return ru, rv, rc
|
| 68 |
+
|
| 69 |
+
|
| 70 |
+
def analytic_vortex_numpy(coordinates: np.ndarray) -> np.ndarray:
|
| 71 |
+
"""Structured, noise-free translating and pressure-decaying analytic cyclone."""
|
| 72 |
+
coordinates = np.asarray(coordinates, dtype=np.float64)
|
| 73 |
+
if coordinates.ndim != 2 or coordinates.shape[1] != 4:
|
| 74 |
+
raise ValueError("coordinates must have shape [N,4]")
|
| 75 |
+
y, x, time, pressure = coordinates.T
|
| 76 |
+
center_x = 45_000.0 * np.sin(time / 21_600.0 * np.pi / 2.0)
|
| 77 |
+
center_y = 30_000.0 * np.cos(time / 21_600.0 * np.pi / 2.0)
|
| 78 |
+
dx, dy = x - center_x, y - center_y
|
| 79 |
+
radius = np.sqrt(dx * dx + dy * dy + 4.0e6)
|
| 80 |
+
radial_scale = 95_000.0
|
| 81 |
+
vertical = np.exp(-((pressure - 55_000.0) / 27_000.0) ** 2)
|
| 82 |
+
tangential = 42.0 * (radius / radial_scale) * np.exp(1.0 - radius / radial_scale) * vertical
|
| 83 |
+
radial = -3.0 * np.exp(-radius / 180_000.0) * vertical
|
| 84 |
+
u = -tangential * dy / radius + radial * dx / radius + 2.5
|
| 85 |
+
v = tangential * dx / radius + radial * dy / radius + 1.0
|
| 86 |
+
geopotential = 50_000.0 - 5_500.0 * np.exp(-(radius / 130_000.0) ** 2) * vertical
|
| 87 |
+
omega = -0.32 * np.exp(-(radius / 110_000.0) ** 2) * np.sin(np.pi * (pressure - 15_000.0) / 75_000.0)
|
| 88 |
+
return np.stack((u, v, geopotential, omega), axis=1).astype(np.float32)
|
| 89 |
+
|
| 90 |
+
|
| 91 |
+
def parameter_count(model: nn.Module) -> int:
|
| 92 |
+
return sum(parameter.numel() for parameter in model.parameters())
|
scripts/fake_data.py
ADDED
|
@@ -0,0 +1,82 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Generate sparse structured observations and collocation coordinates lazily."""
|
| 2 |
+
|
| 3 |
+
import argparse
|
| 4 |
+
import importlib.util
|
| 5 |
+
from pathlib import Path
|
| 6 |
+
|
| 7 |
+
import numpy as np
|
| 8 |
+
import yaml
|
| 9 |
+
|
| 10 |
+
|
| 11 |
+
ROOT = Path(__file__).resolve().parents[1]
|
| 12 |
+
|
| 13 |
+
|
| 14 |
+
def load_model_module():
|
| 15 |
+
spec = importlib.util.spec_from_file_location("pinn_tc_model", ROOT / "model/pinn-tc.py")
|
| 16 |
+
module = importlib.util.module_from_spec(spec)
|
| 17 |
+
spec.loader.exec_module(module)
|
| 18 |
+
return module
|
| 19 |
+
|
| 20 |
+
|
| 21 |
+
def structured_observations(config, mode):
|
| 22 |
+
data = config["data"]
|
| 23 |
+
extent = float(data["horizontal_extent_m"])
|
| 24 |
+
levels = np.linspace(float(data["pressure_min_pa"]), float(data["pressure_max_pa"]), int(data["pressure_levels"]))
|
| 25 |
+
times = np.asarray(data["observation_times_s"], dtype=np.float64)
|
| 26 |
+
if times.shape != (int(data["time_steps"]),) or not np.array_equal(times, [-10800.0, 0.0, 10800.0]):
|
| 27 |
+
raise ValueError("paper observation times must be [-3h, 0h, +3h]")
|
| 28 |
+
line = np.linspace(-extent, extent, int(data["observation_line_points"]))
|
| 29 |
+
coordinates = []
|
| 30 |
+
for time_index, time in enumerate(times):
|
| 31 |
+
active_mode = ("Cross" if time_index % 2 == 0 else "Plus") if mode == "Switch" else mode
|
| 32 |
+
for pressure in levels[:: int(data["observation_pressure_stride"])]:
|
| 33 |
+
if active_mode == "Plus":
|
| 34 |
+
coordinates.extend((y, x, time, pressure) for y, x in zip(np.zeros_like(line), line))
|
| 35 |
+
coordinates.extend((y, x, time, pressure) for y, x in zip(line, np.zeros_like(line)))
|
| 36 |
+
elif active_mode == "Cross":
|
| 37 |
+
coordinates.extend((y, x, time, pressure) for y, x in zip(line, line))
|
| 38 |
+
coordinates.extend((y, x, time, pressure) for y, x in zip(line, -line))
|
| 39 |
+
else:
|
| 40 |
+
raise ValueError("observation mode must be Cross, Plus, or Switch")
|
| 41 |
+
edge = np.linspace(-extent, extent, int(data["boundary_points_per_edge"]))
|
| 42 |
+
coordinates.extend((-extent, value, time, pressure) for value in edge)
|
| 43 |
+
coordinates.extend((extent, value, time, pressure) for value in edge)
|
| 44 |
+
coordinates.extend((value, -extent, time, pressure) for value in edge[1:-1])
|
| 45 |
+
coordinates.extend((value, extent, time, pressure) for value in edge[1:-1])
|
| 46 |
+
return np.unique(np.asarray(coordinates, dtype=np.float32), axis=0)
|
| 47 |
+
|
| 48 |
+
|
| 49 |
+
def main():
|
| 50 |
+
parser = argparse.ArgumentParser()
|
| 51 |
+
parser.add_argument("--mode", choices=("Cross", "Plus", "Switch"))
|
| 52 |
+
parser.add_argument("--force", action="store_true")
|
| 53 |
+
args = parser.parse_args()
|
| 54 |
+
config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
|
| 55 |
+
mode = args.mode or config["data"]["observation_mode"]
|
| 56 |
+
output = ROOT / config["data"]["root"]
|
| 57 |
+
output.mkdir(parents=True, exist_ok=True)
|
| 58 |
+
path = output / "training_points.npz"
|
| 59 |
+
if path.exists() and not args.force:
|
| 60 |
+
print(f"exists={path.relative_to(ROOT)} (use --force to regenerate)")
|
| 61 |
+
return
|
| 62 |
+
module = load_model_module()
|
| 63 |
+
observations = structured_observations(config, mode)
|
| 64 |
+
targets = module.analytic_vortex_numpy(observations)[:, :3]
|
| 65 |
+
rng = np.random.default_rng(int(config["seed"]))
|
| 66 |
+
count = int(config["data"]["collocation_points"])
|
| 67 |
+
collocation = np.column_stack((
|
| 68 |
+
rng.uniform(-config["data"]["horizontal_extent_m"], config["data"]["horizontal_extent_m"], (count, 2)),
|
| 69 |
+
rng.uniform(-config["data"]["time_extent_s"], config["data"]["time_extent_s"], count),
|
| 70 |
+
rng.uniform(config["data"]["pressure_min_pa"], config["data"]["pressure_max_pa"], count),
|
| 71 |
+
)).astype(np.float32)
|
| 72 |
+
# The first two random columns already follow the required [y, x] order.
|
| 73 |
+
np.savez(path, observation_coordinates=observations, observation_targets=targets,
|
| 74 |
+
collocation_coordinates=collocation, input_order=np.asarray(module.INPUT_ORDER),
|
| 75 |
+
supervised_order=np.asarray(module.OUTPUT_ORDER[:3]), observation_mode=np.asarray(mode),
|
| 76 |
+
dense_shape=np.asarray([data_dim := int(config["data"]["grid_points"]), data_dim,
|
| 77 |
+
int(config["data"]["pressure_levels"]), int(config["data"]["time_steps"])]))
|
| 78 |
+
print(f"generated={path.relative_to(ROOT)} observations={len(observations)} collocation={len(collocation)} mode={mode}")
|
| 79 |
+
|
| 80 |
+
|
| 81 |
+
if __name__ == "__main__":
|
| 82 |
+
main()
|
scripts/inference.py
ADDED
|
@@ -0,0 +1,77 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Run chunked inference on the full 251 x 251 x 44 x 3 paper-resolution domain."""
|
| 2 |
+
|
| 3 |
+
import importlib.util
|
| 4 |
+
from pathlib import Path
|
| 5 |
+
|
| 6 |
+
import numpy as np
|
| 7 |
+
import torch
|
| 8 |
+
import yaml
|
| 9 |
+
|
| 10 |
+
|
| 11 |
+
ROOT = Path(__file__).resolve().parents[1]
|
| 12 |
+
|
| 13 |
+
|
| 14 |
+
def load_model_module():
|
| 15 |
+
spec = importlib.util.spec_from_file_location("pinn_tc_model", ROOT / "model/pinn-tc.py")
|
| 16 |
+
module = importlib.util.module_from_spec(spec)
|
| 17 |
+
spec.loader.exec_module(module)
|
| 18 |
+
return module
|
| 19 |
+
|
| 20 |
+
|
| 21 |
+
def coordinates_for_indices(indices, y_axis, x_axis, times, pressures):
|
| 22 |
+
x_count, p_count, t_count = len(x_axis), len(pressures), len(times)
|
| 23 |
+
y_index = indices // (x_count * p_count * t_count)
|
| 24 |
+
remainder = indices % (x_count * p_count * t_count)
|
| 25 |
+
x_index = remainder // (p_count * t_count)
|
| 26 |
+
remainder %= p_count * t_count
|
| 27 |
+
p_index = remainder // t_count
|
| 28 |
+
t_index = remainder % t_count
|
| 29 |
+
return np.column_stack((y_axis[y_index], x_axis[x_index], times[t_index], pressures[p_index])).astype(np.float32)
|
| 30 |
+
|
| 31 |
+
|
| 32 |
+
def main():
|
| 33 |
+
config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
|
| 34 |
+
module = load_model_module()
|
| 35 |
+
requested = config["runtime"]["device"]
|
| 36 |
+
device = torch.device("cuda" if requested == "auto" and torch.cuda.is_available() else ("cpu" if requested == "auto" else requested))
|
| 37 |
+
checkpoint = torch.load(ROOT / config["paths"]["checkpoint"], map_location=device, weights_only=False)
|
| 38 |
+
if checkpoint.get("format_version") != config["data"]["format_version"]:
|
| 39 |
+
raise ValueError("checkpoint format_version mismatch")
|
| 40 |
+
if checkpoint["input_order"] != list(module.INPUT_ORDER) or checkpoint["output_order"] != list(module.OUTPUT_ORDER):
|
| 41 |
+
raise ValueError("checkpoint coordinate or variable protocol mismatch")
|
| 42 |
+
model = module.PINNTC(**checkpoint["model_config"]).to(device)
|
| 43 |
+
model.load_state_dict(checkpoint["model_state_dict"]); model.eval()
|
| 44 |
+
data = config["data"]
|
| 45 |
+
y_axis = np.linspace(-data["horizontal_extent_m"], data["horizontal_extent_m"], data["grid_points"], dtype=np.float32)
|
| 46 |
+
x_axis = y_axis.copy()
|
| 47 |
+
times = np.asarray(data["observation_times_s"], dtype=np.float32)
|
| 48 |
+
pressures = np.linspace(data["pressure_min_pa"], data["pressure_max_pa"], data["pressure_levels"], dtype=np.float32)
|
| 49 |
+
shape = (len(y_axis), len(x_axis), len(pressures), len(times), 4)
|
| 50 |
+
expected = (251, 251, 44, 3, 4)
|
| 51 |
+
if shape != expected:
|
| 52 |
+
raise ValueError(f"dense grid must remain {expected}, got {shape}")
|
| 53 |
+
output = ROOT / config["paths"]["prediction"]
|
| 54 |
+
output.parent.mkdir(parents=True, exist_ok=True)
|
| 55 |
+
temporary = output.with_suffix(".work.npy")
|
| 56 |
+
prediction = np.lib.format.open_memmap(temporary, mode="w+", dtype=np.float32, shape=shape)
|
| 57 |
+
flat = prediction.reshape(-1, 4)
|
| 58 |
+
chunk_size = int(config["runtime"]["inference_chunk_size"])
|
| 59 |
+
with torch.inference_mode():
|
| 60 |
+
for start in range(0, len(flat), chunk_size):
|
| 61 |
+
stop = min(start + chunk_size, len(flat))
|
| 62 |
+
coordinates = coordinates_for_indices(np.arange(start, stop, dtype=np.int64), y_axis, x_axis, times, pressures)
|
| 63 |
+
values = model(torch.from_numpy(coordinates).to(device)).cpu().numpy().astype(np.float32)
|
| 64 |
+
if not np.isfinite(values).all():
|
| 65 |
+
raise FloatingPointError(f"non-finite prediction in rows {start}:{stop}")
|
| 66 |
+
flat[start:stop] = values
|
| 67 |
+
prediction.flush()
|
| 68 |
+
np.savez(output, predictions=prediction, y=y_axis, x=x_axis, time=times, pressure=pressures,
|
| 69 |
+
input_order=np.asarray(module.INPUT_ORDER), output_order=np.asarray(module.OUTPUT_ORDER),
|
| 70 |
+
format_version=np.asarray(config["data"]["format_version"]), layout=np.asarray("YXPTV"))
|
| 71 |
+
del prediction
|
| 72 |
+
temporary.unlink()
|
| 73 |
+
print(f"predictions={output.relative_to(ROOT)} shape={shape} bytes={output.stat().st_size}")
|
| 74 |
+
|
| 75 |
+
|
| 76 |
+
if __name__ == "__main__":
|
| 77 |
+
main()
|
scripts/result.py
ADDED
|
@@ -0,0 +1,108 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Evaluate dense PINN-TC predictions and create sparse/target/prediction/error panels."""
|
| 2 |
+
|
| 3 |
+
import importlib.util
|
| 4 |
+
import json
|
| 5 |
+
from pathlib import Path
|
| 6 |
+
|
| 7 |
+
import matplotlib
|
| 8 |
+
matplotlib.use("Agg")
|
| 9 |
+
import matplotlib.pyplot as plt
|
| 10 |
+
import numpy as np
|
| 11 |
+
import torch
|
| 12 |
+
import yaml
|
| 13 |
+
|
| 14 |
+
|
| 15 |
+
ROOT = Path(__file__).resolve().parents[1]
|
| 16 |
+
|
| 17 |
+
|
| 18 |
+
def load_model_module():
|
| 19 |
+
spec = importlib.util.spec_from_file_location("pinn_tc_model", ROOT / "model/pinn-tc.py")
|
| 20 |
+
module = importlib.util.module_from_spec(spec)
|
| 21 |
+
spec.loader.exec_module(module)
|
| 22 |
+
return module
|
| 23 |
+
|
| 24 |
+
|
| 25 |
+
def main():
|
| 26 |
+
config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
|
| 27 |
+
module = load_model_module()
|
| 28 |
+
archive = np.load(ROOT / config["paths"]["prediction"])
|
| 29 |
+
if str(archive["format_version"]) != config["data"]["format_version"]:
|
| 30 |
+
raise ValueError("prediction format_version mismatch")
|
| 31 |
+
prediction = archive["predictions"]
|
| 32 |
+
expected = (251, 251, 44, 3, 4)
|
| 33 |
+
if prediction.shape != expected or prediction.dtype != np.float32:
|
| 34 |
+
raise ValueError(f"prediction must be float32 with shape {expected}")
|
| 35 |
+
y_axis, x_axis, pressures, times = archive["y"], archive["x"], archive["pressure"], archive["time"]
|
| 36 |
+
sum_speed_error = sum_pred = sum_true = sum_pred2 = sum_true2 = sum_cross = 0.0
|
| 37 |
+
count = 0
|
| 38 |
+
radial_edges = np.linspace(0.0, np.sqrt(2.0) * config["data"]["horizontal_extent_m"], 13)
|
| 39 |
+
radial_sse = np.zeros(len(radial_edges) - 1); radial_count = np.zeros(len(radial_edges) - 1, dtype=np.int64)
|
| 40 |
+
yy, xx = np.meshgrid(y_axis, x_axis, indexing="ij")
|
| 41 |
+
radius = np.sqrt(xx * xx + yy * yy)
|
| 42 |
+
for pressure_index, pressure in enumerate(pressures):
|
| 43 |
+
for time_index, time in enumerate(times):
|
| 44 |
+
coordinates = np.column_stack((yy.ravel(), xx.ravel(), np.full(yy.size, time), np.full(yy.size, pressure))).astype(np.float32)
|
| 45 |
+
target = module.analytic_vortex_numpy(coordinates).reshape(251, 251, 4)
|
| 46 |
+
predicted_speed = np.linalg.norm(prediction[:, :, pressure_index, time_index, :2].astype(np.float64), axis=-1)
|
| 47 |
+
true_speed = np.linalg.norm(target[:, :, :2].astype(np.float64), axis=-1)
|
| 48 |
+
difference = predicted_speed - true_speed
|
| 49 |
+
sum_speed_error += np.square(difference).sum(); count += difference.size
|
| 50 |
+
sum_pred += predicted_speed.sum(); sum_true += true_speed.sum()
|
| 51 |
+
sum_pred2 += np.square(predicted_speed).sum(); sum_true2 += np.square(true_speed).sum()
|
| 52 |
+
sum_cross += (predicted_speed * true_speed).sum()
|
| 53 |
+
bins = np.clip(np.digitize(radius.ravel(), radial_edges) - 1, 0, len(radial_sse) - 1)
|
| 54 |
+
radial_sse += np.bincount(bins, weights=np.square(difference).ravel(), minlength=len(radial_sse))
|
| 55 |
+
radial_count += np.bincount(bins, minlength=len(radial_count))
|
| 56 |
+
covariance = sum_cross - sum_pred * sum_true / count
|
| 57 |
+
variance_pred = sum_pred2 - sum_pred * sum_pred / count
|
| 58 |
+
variance_true = sum_true2 - sum_true * sum_true / count
|
| 59 |
+
pearson = covariance / np.sqrt(max(variance_pred * variance_true, 1e-30))
|
| 60 |
+
|
| 61 |
+
checkpoint = torch.load(ROOT / config["paths"]["checkpoint"], map_location="cpu", weights_only=False)
|
| 62 |
+
if checkpoint.get("format_version") != config["data"]["format_version"]:
|
| 63 |
+
raise ValueError("checkpoint format_version mismatch")
|
| 64 |
+
model = module.PINNTC(**checkpoint["model_config"]); model.load_state_dict(checkpoint["model"]); model.eval()
|
| 65 |
+
rng = np.random.default_rng(int(config["seed"]) + 99)
|
| 66 |
+
residual_count = int(config["evaluation"]["pde_points"])
|
| 67 |
+
residual_coordinates = np.column_stack((rng.uniform(-400_000, 400_000, (residual_count, 2)),
|
| 68 |
+
rng.uniform(-21_600, 21_600, residual_count),
|
| 69 |
+
rng.uniform(15_000, 90_000, residual_count))).astype(np.float32)
|
| 70 |
+
tensor = torch.from_numpy(residual_coordinates).requires_grad_(True)
|
| 71 |
+
ru, rv, rc = module.pde_residuals(model, tensor, **config["physics"])
|
| 72 |
+
report = {"dense_shape": list(expected), "dense_points": int(np.prod(expected[:-1])),
|
| 73 |
+
"wind_speed_rmse_m_s": float(np.sqrt(sum_speed_error / count)), "wind_speed_pearson": float(pearson),
|
| 74 |
+
"pde_residual_rms": {"ru_m_s2": float(torch.sqrt(torch.mean(ru.square())).detach()),
|
| 75 |
+
"rv_m_s2": float(torch.sqrt(torch.mean(rv.square())).detach()),
|
| 76 |
+
"rc_s-1": float(torch.sqrt(torch.mean(rc.square())).detach())},
|
| 77 |
+
"radial_rmse": [{"radius_mid_km": float((left + right) / 2000.0), "rmse_m_s": float(np.sqrt(error / max(samples, 1))), "count": int(samples)}
|
| 78 |
+
for left, right, error, samples in zip(radial_edges[:-1], radial_edges[1:], radial_sse, radial_count)],
|
| 79 |
+
"scope": "Structured analytic engineering validation, not a paper-performance claim."}
|
| 80 |
+
values = [report["wind_speed_rmse_m_s"], report["wind_speed_pearson"], *report["pde_residual_rms"].values()]
|
| 81 |
+
if not np.isfinite(values).all():
|
| 82 |
+
raise FloatingPointError("evaluation produced a non-finite metric")
|
| 83 |
+
output = ROOT / config["paths"]["evaluation_dir"]
|
| 84 |
+
output.mkdir(parents=True, exist_ok=True)
|
| 85 |
+
(output / "metrics.json").write_text(json.dumps(report, indent=2) + "\n")
|
| 86 |
+
|
| 87 |
+
pressure_index, time_index = int(np.argmin(abs(pressures - 50_000))), int(np.argmin(abs(times)))
|
| 88 |
+
coordinates = np.column_stack((yy.ravel(), xx.ravel(), np.full(yy.size, times[time_index]),
|
| 89 |
+
np.full(yy.size, pressures[pressure_index]))).astype(np.float32)
|
| 90 |
+
target_speed = np.linalg.norm(module.analytic_vortex_numpy(coordinates)[:, :2], axis=1).reshape(251, 251)
|
| 91 |
+
predicted_speed = np.linalg.norm(prediction[:, :, pressure_index, time_index, :2], axis=-1)
|
| 92 |
+
sparse = np.load(ROOT / config["data"]["root"] / "training_points.npz")["observation_coordinates"]
|
| 93 |
+
selected = sparse[(abs(sparse[:, 2] - times[time_index]) < 1) & (abs(sparse[:, 3] - pressures[pressure_index]) < 1000)]
|
| 94 |
+
figure, axes_plot = plt.subplots(1, 4, figsize=(16, 4), constrained_layout=True)
|
| 95 |
+
axes_plot[0].scatter(selected[:, 1] / 1000, selected[:, 0] / 1000, s=5, c="black"); axes_plot[0].set_title("Sparse observations")
|
| 96 |
+
panels = ((target_speed, "Target wind speed", "viridis"), (predicted_speed, "Predicted wind speed", "viridis"),
|
| 97 |
+
(predicted_speed - target_speed, "Prediction error", "coolwarm"))
|
| 98 |
+
for axis, (field, title, color_map) in zip(axes_plot[1:], panels):
|
| 99 |
+
image = axis.imshow(field, origin="lower", extent=(-400, 400, -400, 400), cmap=color_map)
|
| 100 |
+
axis.set_title(title); figure.colorbar(image, ax=axis, shrink=0.75)
|
| 101 |
+
for axis in axes_plot:
|
| 102 |
+
axis.set_xlabel("x (km)"); axis.set_ylabel("y (km)")
|
| 103 |
+
figure.savefig(output / "comparison.png", dpi=140); plt.close(figure)
|
| 104 |
+
print(f"evaluation={output.relative_to(ROOT)} wind_rmse={report['wind_speed_rmse_m_s']:.6g} pearson={pearson:.6g}")
|
| 105 |
+
|
| 106 |
+
|
| 107 |
+
if __name__ == "__main__":
|
| 108 |
+
main()
|
scripts/train.py
ADDED
|
@@ -0,0 +1,137 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Train PINN-TC with Adam and an optional L-BFGS refinement under torchrun DDP."""
|
| 2 |
+
|
| 3 |
+
import argparse
|
| 4 |
+
import importlib.util
|
| 5 |
+
import json
|
| 6 |
+
import os
|
| 7 |
+
import random
|
| 8 |
+
from pathlib import Path
|
| 9 |
+
|
| 10 |
+
import numpy as np
|
| 11 |
+
import torch
|
| 12 |
+
import yaml
|
| 13 |
+
from torch.nn.parallel import DistributedDataParallel
|
| 14 |
+
|
| 15 |
+
|
| 16 |
+
ROOT = Path(__file__).resolve().parents[1]
|
| 17 |
+
|
| 18 |
+
|
| 19 |
+
def load_model_module():
|
| 20 |
+
spec = importlib.util.spec_from_file_location("pinn_tc_model", ROOT / "model/pinn-tc.py")
|
| 21 |
+
module = importlib.util.module_from_spec(spec)
|
| 22 |
+
spec.loader.exec_module(module)
|
| 23 |
+
return module
|
| 24 |
+
|
| 25 |
+
|
| 26 |
+
def choose_device(config, local_rank):
|
| 27 |
+
requested = config["runtime"]["device"]
|
| 28 |
+
if requested == "auto":
|
| 29 |
+
return torch.device("cuda", local_rank) if torch.cuda.is_available() else torch.device("cpu")
|
| 30 |
+
return torch.device(requested)
|
| 31 |
+
|
| 32 |
+
|
| 33 |
+
def main():
|
| 34 |
+
parser = argparse.ArgumentParser()
|
| 35 |
+
parser.add_argument("--resume", action="store_true", help="resume optimizer and model state from the configured checkpoint")
|
| 36 |
+
args = parser.parse_args()
|
| 37 |
+
config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
|
| 38 |
+
module = load_model_module()
|
| 39 |
+
seed = int(config["seed"])
|
| 40 |
+
random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
|
| 41 |
+
distributed = int(os.environ.get("WORLD_SIZE", "1")) > 1
|
| 42 |
+
local_rank = int(os.environ.get("LOCAL_RANK", "0"))
|
| 43 |
+
if distributed:
|
| 44 |
+
torch.distributed.init_process_group("nccl" if torch.cuda.is_available() else "gloo")
|
| 45 |
+
rank = torch.distributed.get_rank() if distributed else 0
|
| 46 |
+
device = choose_device(config, local_rank)
|
| 47 |
+
if device.type == "cuda":
|
| 48 |
+
torch.cuda.set_device(device); torch.cuda.manual_seed_all(seed)
|
| 49 |
+
|
| 50 |
+
source = np.load(ROOT / config["data"]["root"] / "training_points.npz")
|
| 51 |
+
observations = torch.from_numpy(source["observation_coordinates"]).to(device)
|
| 52 |
+
targets = torch.from_numpy(source["observation_targets"]).to(device)
|
| 53 |
+
collocation = torch.from_numpy(source["collocation_coordinates"]).to(device)
|
| 54 |
+
model = module.PINNTC(**config["model"]).to(device)
|
| 55 |
+
wrapped = DistributedDataParallel(model, device_ids=[local_rank] if device.type == "cuda" else None) if distributed else model
|
| 56 |
+
bare = wrapped.module if distributed else wrapped
|
| 57 |
+
gamma = float(config["train"]["gamma"])
|
| 58 |
+
batch_observations = int(config["train"]["observation_batch"])
|
| 59 |
+
batch_collocation = int(config["train"]["collocation_batch"])
|
| 60 |
+
|
| 61 |
+
def loss_for(obs_index, col_index):
|
| 62 |
+
prediction = wrapped(observations[obs_index])[:, :3]
|
| 63 |
+
speed = torch.sqrt(targets[obs_index, 0] ** 2 + targets[obs_index, 1] ** 2)
|
| 64 |
+
wind_weight = 1.0 + float(config["train"]["wind_speed_weight"]) * speed / 50.0
|
| 65 |
+
normalized_error = (prediction - targets[obs_index]) / torch.tensor([45.0, 45.0, 6000.0], device=device)
|
| 66 |
+
data_loss = (normalized_error.square() * wind_weight[:, None]).mean()
|
| 67 |
+
points = collocation[col_index].detach().requires_grad_(True)
|
| 68 |
+
ru, rv, rc = module.pde_residuals(wrapped, points, **config["physics"])
|
| 69 |
+
pde_loss = (ru.square().mean() + rv.square().mean()) / float(config["train"]["momentum_scale"]) ** 2
|
| 70 |
+
pde_loss = pde_loss + rc.square().mean() / float(config["train"]["continuity_scale"]) ** 2
|
| 71 |
+
return data_loss + gamma * pde_loss, data_loss, pde_loss
|
| 72 |
+
|
| 73 |
+
optimizer = torch.optim.Adam(wrapped.parameters(), lr=float(config["train"]["adam_lr"]))
|
| 74 |
+
history = []
|
| 75 |
+
checkpoint_path = ROOT / config["paths"]["checkpoint"]
|
| 76 |
+
completed_adam_steps = 0
|
| 77 |
+
checkpoint = None
|
| 78 |
+
if args.resume:
|
| 79 |
+
checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
|
| 80 |
+
if checkpoint["model_config"] != config["model"]:
|
| 81 |
+
raise ValueError("checkpoint model configuration mismatch")
|
| 82 |
+
bare.load_state_dict(checkpoint["model_state_dict"])
|
| 83 |
+
optimizer.load_state_dict(checkpoint["adam_optimizer_state_dict"])
|
| 84 |
+
history = list(checkpoint.get("history", []))
|
| 85 |
+
completed_adam_steps = int(checkpoint.get("completed_adam_steps", 0))
|
| 86 |
+
generator = torch.Generator(device="cpu").manual_seed(seed + rank)
|
| 87 |
+
for step in range(completed_adam_steps, completed_adam_steps + int(config["train"]["adam_steps"])):
|
| 88 |
+
obs_index = torch.randperm(len(observations), generator=generator)[:batch_observations].to(device)
|
| 89 |
+
col_index = torch.randperm(len(collocation), generator=generator)[:batch_collocation].to(device)
|
| 90 |
+
optimizer.zero_grad(set_to_none=True)
|
| 91 |
+
loss, data_loss, pde_loss = loss_for(obs_index, col_index)
|
| 92 |
+
loss.backward(); optimizer.step()
|
| 93 |
+
history.append({"stage": "Adam", "step": step + 1, "loss": float(loss.detach()),
|
| 94 |
+
"data_loss": float(data_loss.detach()), "pde_loss": float(pde_loss.detach())})
|
| 95 |
+
|
| 96 |
+
lbfgs_steps = int(config["train"]["lbfgs_steps"])
|
| 97 |
+
optimizer_lbfgs = None
|
| 98 |
+
if lbfgs_steps > 0:
|
| 99 |
+
# Every DDP rank runs the same closure count and deterministic subset so gradient collectives remain aligned.
|
| 100 |
+
obs_index = torch.arange(min(batch_observations, len(observations)), device=device)
|
| 101 |
+
col_index = torch.arange(min(batch_collocation, len(collocation)), device=device)
|
| 102 |
+
optimizer_lbfgs = torch.optim.LBFGS(wrapped.parameters(), lr=float(config["train"]["lbfgs_lr"]),
|
| 103 |
+
max_iter=lbfgs_steps, history_size=int(config["train"]["lbfgs_history_size"]),
|
| 104 |
+
line_search_fn=None)
|
| 105 |
+
if checkpoint is not None and checkpoint.get("lbfgs_optimizer_state_dict") is not None:
|
| 106 |
+
optimizer_lbfgs.load_state_dict(checkpoint["lbfgs_optimizer_state_dict"])
|
| 107 |
+
latest = {}
|
| 108 |
+
|
| 109 |
+
def closure():
|
| 110 |
+
optimizer_lbfgs.zero_grad(set_to_none=True)
|
| 111 |
+
loss, data_loss, pde_loss = loss_for(obs_index, col_index)
|
| 112 |
+
loss.backward()
|
| 113 |
+
latest.update(loss=float(loss.detach()), data_loss=float(data_loss.detach()), pde_loss=float(pde_loss.detach()))
|
| 114 |
+
return loss
|
| 115 |
+
|
| 116 |
+
optimizer_lbfgs.step(closure)
|
| 117 |
+
history.append({"stage": "L-BFGS", "step": 1, **latest})
|
| 118 |
+
|
| 119 |
+
if rank == 0:
|
| 120 |
+
metrics_path = ROOT / config["paths"]["training_metrics"]
|
| 121 |
+
checkpoint_path.parent.mkdir(parents=True, exist_ok=True); metrics_path.parent.mkdir(parents=True, exist_ok=True)
|
| 122 |
+
torch.save({"model": bare.state_dict(), "model_state_dict": bare.state_dict(),
|
| 123 |
+
"model_config": config["model"], "format_version": config["data"]["format_version"],
|
| 124 |
+
"adam_optimizer_state_dict": optimizer.state_dict(),
|
| 125 |
+
"lbfgs_optimizer_state_dict": optimizer_lbfgs.state_dict() if optimizer_lbfgs is not None else None,
|
| 126 |
+
"completed_adam_steps": completed_adam_steps + int(config["train"]["adam_steps"]),
|
| 127 |
+
"history": history,
|
| 128 |
+
"input_order": list(module.INPUT_ORDER), "output_order": list(module.OUTPUT_ORDER),
|
| 129 |
+
"domain": config["domain"], "physics": config["physics"], "seed": seed}, checkpoint_path)
|
| 130 |
+
metrics_path.write_text(json.dumps({"history": history, "distributed_world_size": int(os.environ.get("WORLD_SIZE", "1"))}, indent=2) + "\n")
|
| 131 |
+
print(f"checkpoint={checkpoint_path.relative_to(ROOT)} parameters={module.parameter_count(bare)} final_loss={history[-1]['loss']:.6g}")
|
| 132 |
+
if distributed:
|
| 133 |
+
torch.distributed.destroy_process_group()
|
| 134 |
+
|
| 135 |
+
|
| 136 |
+
if __name__ == "__main__":
|
| 137 |
+
main()
|
weight/.gitkeep
ADDED
|
File without changes
|