zhangrenchao commited on
Commit
2f3c9e4
·
verified ·
1 Parent(s): 8737219

Upload folder using huggingface_hub

Browse files
.gitattributes CHANGED
@@ -1,35 +1,3 @@
1
- *.7z filter=lfs diff=lfs merge=lfs -text
2
- *.arrow filter=lfs diff=lfs merge=lfs -text
3
- *.bin filter=lfs diff=lfs merge=lfs -text
4
- *.bz2 filter=lfs diff=lfs merge=lfs -text
5
- *.ckpt filter=lfs diff=lfs merge=lfs -text
6
- *.ftz filter=lfs diff=lfs merge=lfs -text
7
- *.gz filter=lfs diff=lfs merge=lfs -text
8
- *.h5 filter=lfs diff=lfs merge=lfs -text
9
- *.joblib filter=lfs diff=lfs merge=lfs -text
10
- *.lfs.* filter=lfs diff=lfs merge=lfs -text
11
- *.mlmodel filter=lfs diff=lfs merge=lfs -text
12
- *.model filter=lfs diff=lfs merge=lfs -text
13
- *.msgpack filter=lfs diff=lfs merge=lfs -text
14
  *.npy filter=lfs diff=lfs merge=lfs -text
15
  *.npz filter=lfs diff=lfs merge=lfs -text
16
- *.onnx filter=lfs diff=lfs merge=lfs -text
17
- *.ot filter=lfs diff=lfs merge=lfs -text
18
- *.parquet filter=lfs diff=lfs merge=lfs -text
19
- *.pb filter=lfs diff=lfs merge=lfs -text
20
- *.pickle filter=lfs diff=lfs merge=lfs -text
21
- *.pkl filter=lfs diff=lfs merge=lfs -text
22
- *.pt filter=lfs diff=lfs merge=lfs -text
23
- *.pth filter=lfs diff=lfs merge=lfs -text
24
- *.rar filter=lfs diff=lfs merge=lfs -text
25
- *.safetensors filter=lfs diff=lfs merge=lfs -text
26
- saved_model/**/* filter=lfs diff=lfs merge=lfs -text
27
- *.tar.* filter=lfs diff=lfs merge=lfs -text
28
- *.tar filter=lfs diff=lfs merge=lfs -text
29
- *.tflite filter=lfs diff=lfs merge=lfs -text
30
- *.tgz filter=lfs diff=lfs merge=lfs -text
31
- *.wasm filter=lfs diff=lfs merge=lfs -text
32
- *.xz filter=lfs diff=lfs merge=lfs -text
33
- *.zip filter=lfs diff=lfs merge=lfs -text
34
- *.zst filter=lfs diff=lfs merge=lfs -text
35
- *tfevents* filter=lfs diff=lfs merge=lfs -text
 
1
+ *.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