OneScience commited on
Commit
5b87d5f
·
verified ·
1 Parent(s): a3f2220

Upload folder using huggingface_hub

Browse files
.gitattributes CHANGED
@@ -1,35 +1,50 @@
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
  *.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
+ *.bin.* filter=lfs diff=lfs merge=lfs -text
5
  *.bz2 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
  *.model filter=lfs diff=lfs merge=lfs -text
12
  *.msgpack filter=lfs diff=lfs merge=lfs -text
 
 
13
  *.onnx filter=lfs diff=lfs merge=lfs -text
14
  *.ot filter=lfs diff=lfs merge=lfs -text
15
  *.parquet filter=lfs diff=lfs merge=lfs -text
16
  *.pb filter=lfs diff=lfs merge=lfs -text
17
+
 
 
18
  *.pth filter=lfs diff=lfs merge=lfs -text
19
  *.rar filter=lfs diff=lfs merge=lfs -text
 
20
  saved_model/**/* filter=lfs diff=lfs merge=lfs -text
21
  *.tar.* filter=lfs diff=lfs merge=lfs -text
 
22
  *.tflite filter=lfs diff=lfs merge=lfs -text
23
  *.tgz filter=lfs diff=lfs merge=lfs -text
 
24
  *.xz filter=lfs diff=lfs merge=lfs -text
25
  *.zip filter=lfs diff=lfs merge=lfs -text
26
+ *.zstandard filter=lfs diff=lfs merge=lfs -text
27
+ *.tfevents* filter=lfs diff=lfs merge=lfs -text
28
+ *.db* filter=lfs diff=lfs merge=lfs -text
29
+ *.ark* filter=lfs diff=lfs merge=lfs -text
30
+ **/*ckpt*data* filter=lfs diff=lfs merge=lfs -text
31
+ **/*ckpt*.meta filter=lfs diff=lfs merge=lfs -text
32
+ **/*ckpt*.index filter=lfs diff=lfs merge=lfs -text
33
+ *.safetensors filter=lfs diff=lfs merge=lfs -text
34
+ *.ckpt filter=lfs diff=lfs merge=lfs -text
35
+ *.gguf* filter=lfs diff=lfs merge=lfs -text
36
+ *.ggml filter=lfs diff=lfs merge=lfs -text
37
+ *.llamafile* filter=lfs diff=lfs merge=lfs -text
38
+ *.pt2 filter=lfs diff=lfs merge=lfs -text
39
+ *.mlmodel filter=lfs diff=lfs merge=lfs -text
40
+ *.npy filter=lfs diff=lfs merge=lfs -text
41
+
42
+ *.pickle filter=lfs diff=lfs merge=lfs -text
43
+ *.pkl filter=lfs diff=lfs merge=lfs -text
44
+ *.tar filter=lfs diff=lfs merge=lfs -text
45
+ *.wasm filter=lfs diff=lfs merge=lfs -text
46
  *.zst filter=lfs diff=lfs merge=lfs -text
47
  *tfevents* filter=lfs diff=lfs merge=lfs -text
48
+
49
+ *.npz filter=lfs diff=lfs merge=lfs -text
50
+ *.pt filter=lfs diff=lfs merge=lfs -text
README.md ADDED
@@ -0,0 +1,126 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ frameworks:
3
+ - pytorch
4
+ language:
5
+ - en
6
+ license: apache-2.0
7
+ tags:
8
+ - OneScience
9
+ - gPINNs
10
+ - physics-informed-neural-networks
11
+ - partial-differential-equations
12
+ - Poisson
13
+ - Burgers
14
+ tasks:
15
+ - pde-solving
16
+ ---
17
+ <p align="center">
18
+ <strong>
19
+ <span style="font-size: 30px;">gPINNs</span>
20
+ </strong>
21
+ </p>
22
+
23
+ # Model Overview
24
+
25
+ gPINNs (Gradient-enhanced Physics-Informed Neural Networks) augment the PDE residual constraints of conventional PINNs with constraints on residual gradients with respect to the input coordinates. This additional higher-order derivative information can improve sample efficiency and solution accuracy for selected forward and inverse problems.
26
+
27
+ This model package provides one-dimensional and two-dimensional Poisson examples, together with a Burgers equation example using residual-based adaptive refinement (RAR).
28
+
29
+ Paper: Gradient-enhanced physics-informed neural networks for forward and inverse PDE problems
30
+ https://doi.org/10.1016/j.cma.2022.114823
31
+
32
+ # Model Description
33
+
34
+ gPINNs use automatic differentiation to compute both PDE residuals and their coordinate gradients, incorporating both as training constraints. The two-dimensional Poisson example applies a hard-boundary output transformation, while the Burgers example uses RAR to add collocation points progressively in high-residual regions.
35
+
36
+ # Use Cases
37
+
38
+ | Use Case | Description |
39
+ | :---: | :--- |
40
+ | One-dimensional Poisson equation | Validate joint training with residual-gradient constraints and boundary loss |
41
+ | Two-dimensional Poisson equation | Solve a homogeneous Dirichlet boundary-value problem using a hard-boundary output transformation |
42
+ | Burgers equation | Combine higher-order automatic differentiation with RAR in high-residual regions |
43
+ | Pipeline validation | Validate training, inference, and plotting with the bundled data and pretrained weights |
44
+
45
+ # Usage
46
+
47
+ ## 1. OneCode
48
+
49
+ Use the online OneCode environment for an intelligent, one-click AI for Science (AI4S) programming experience:
50
+
51
+ [Launch OneCode for one-click AI4S programming](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)
52
+
53
+ ## 2. Manual Setup
54
+
55
+ **Hardware Requirements**
56
+
57
+ - A GPU or DCU is recommended for training.
58
+ - A CPU can be used for inference and small-scale pipeline validation, but full training will be slow.
59
+ - DCU users must install DTK and a PyTorch environment compatible with the target cluster.
60
+
61
+ ### Download the Model Package
62
+
63
+ ```bash
64
+ modelscope download --model OneScience/gPINNs --local_dir ./gPINNs
65
+ cd gPINNs
66
+ ```
67
+
68
+ ### Set Up the Runtime Environment
69
+
70
+ **DCU Environment**
71
+
72
+ ```bash
73
+ # Activate DTK and Conda first
74
+ conda create -n onescience311 python=3.11 -y
75
+ conda activate onescience311
76
+ pip install onescience[cfd-dcu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
77
+ ```
78
+
79
+ **GPU Environment**
80
+
81
+ ```bash
82
+ # Activate Conda first
83
+ conda create -n onescience311 python=3.11 -y libstdcxx-ng=12 libgcc-ng=12 gcc_linux-64=12 gxx_linux-64=12
84
+ conda activate onescience311
85
+ pip install onescience[cfd-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
86
+ ```
87
+
88
+ ### Training Data
89
+
90
+ The one- and two-dimensional Poisson examples construct training and evaluation data from analytical solutions; the Burgers example uses the bundled reference solution `data/Burgers.npz`. Data and training parameters can be adjusted in `conf/config.yaml`.
91
+
92
+ ### Training
93
+
94
+ To train all examples:
95
+
96
+ ```bash
97
+ python scripts/train.py --case all
98
+ ```
99
+
100
+ Alternatively, train an individual example with `--case 1d`, `--case 2d`, or `--case burgers`. The `--epochs` and `--nf` options override the defaults for the selected example; `--lbfgs-iters` applies only to the one-dimensional Poisson case, `--rar-rounds` controls the number of RAR rounds for the Burgers case, and `--quick` skips RAR for rapid pipeline validation.
101
+
102
+ ### Model Weights
103
+
104
+ This model package provides pretrained weights for the one-dimensional Poisson, two-dimensional Poisson, and Burgers equation examples in the `weight/` directory.
105
+
106
+ ### Inference, Evaluation, and Visualization
107
+
108
+ The model package provides weights for validating all three examples. Run:
109
+
110
+ ```bash
111
+ python scripts/inference.py --case all
112
+ ```
113
+
114
+ The script reports the relative L2 error for each example and saves prediction visualizations under `result/`. Use `--case` to select an individual example. Model, data, and training parameters can all be modified in `conf/config.yaml`.
115
+
116
+ # Official OneScience Resources
117
+
118
+ | Platform | OneScience Repository | Skills Repository |
119
+ | --- | --- | --- |
120
+ | Gitee | https://gitee.com/onescience-ai/onescience | https://gitee.com/onescience-ai/oneskills |
121
+ | GitHub | https://github.com/onescience-ai/OneScience | https://github.com/onescience-ai/oneskills |
122
+
123
+ # Citations and License
124
+
125
+ - Yu, J., Lu, L., Meng, X., and Karniadakis, G. E. Gradient-enhanced physics-informed neural networks for forward and inverse PDE problems. Computer Methods in Applied Mechanics and Engineering, 393, 114823, 2022.
126
+ - This model package is released under the Apache-2.0 license and retains attribution to the original paper and data sources.
conf/config.yaml ADDED
@@ -0,0 +1,39 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ root:
2
+ common:
3
+ device: "auto"
4
+ dtype: "float64"
5
+ seed: 1234
6
+ weight_dir: "weight"
7
+ result_dir: "result"
8
+
9
+ poisson1d:
10
+ nf: 30
11
+ w_g: 0.1
12
+ epochs: 30000
13
+ lbfgs_iters: 2000
14
+ lr: 0.001
15
+ log_every: 1000
16
+ layers: [1, 30, 30, 30, 30, 1]
17
+
18
+ poisson2d:
19
+ a: 2.0
20
+ nf: 400
21
+ w_g: 0.001
22
+ epochs: 15000
23
+ lr: 0.001
24
+ log_every: 2000
25
+ layers: [2, 30, 30, 30, 1]
26
+
27
+ burgers:
28
+ data: "data/Burgers.npz"
29
+ nf: 1500
30
+ w_g: 0.0001
31
+ epochs: 20000
32
+ lr: 0.001
33
+ log_every: 2000
34
+ layers: [2, 32, 32, 32, 1]
35
+ rar_rounds: 5
36
+ rar_epochs: 5000
37
+ rar_candidate_points: 100000
38
+ rar_candidate_batch_size: 10000
39
+ rar_add_points: 10
configuration.json ADDED
@@ -0,0 +1 @@
 
 
1
+ {"framework":"Pytorch","task":"other"}
data/Burgers.npz ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:e9d57a5824cd54b8345ac0c7390dad93ccb1bebec85d4f9d71f5407c99874315
3
+ size 208318
model/gpinn.py ADDED
@@ -0,0 +1,220 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from __future__ import annotations
2
+
3
+ from collections.abc import Sequence
4
+
5
+ import numpy as np
6
+ import torch
7
+ import torch.autograd as autograd
8
+ import torch.nn as nn
9
+
10
+
11
+ class SubNet(nn.Module):
12
+ """Fully connected network with trainable adaptive activations."""
13
+
14
+ def __init__(self, layers: Sequence[int], activation: str = "tanh") -> None:
15
+ super().__init__()
16
+ if len(layers) < 2:
17
+ raise ValueError("layers must contain at least an input and an output size")
18
+ if activation not in {"tanh", "sin", "cos"}:
19
+ raise ValueError(f"unsupported activation: {activation}")
20
+
21
+ self.activation = activation
22
+ self.linears = nn.ModuleList(
23
+ nn.Linear(layers[index], layers[index + 1], dtype=torch.float64)
24
+ for index in range(len(layers) - 1)
25
+ )
26
+ # Keep one value per layer for compatibility with the published checkpoints.
27
+ self.a = nn.ParameterList(
28
+ nn.Parameter(torch.tensor(0.05, dtype=torch.float64))
29
+ for _ in range(len(layers) - 1)
30
+ )
31
+
32
+ def forward(self, inputs: torch.Tensor) -> torch.Tensor:
33
+ hidden = inputs
34
+ for index, linear in enumerate(self.linears[:-1]):
35
+ hidden = 20.0 * self.a[index] * linear(hidden)
36
+ hidden = getattr(torch, self.activation)(hidden)
37
+ return self.linears[-1](hidden)
38
+
39
+
40
+ class gPINN(nn.Module):
41
+ """Gradient-enhanced physics-informed neural network backbone."""
42
+
43
+ def __init__(self, layers: Sequence[int], activation: str = "tanh") -> None:
44
+ super().__init__()
45
+ self.net = SubNet(layers, activation)
46
+
47
+ def forward(self, inputs: torch.Tensor) -> torch.Tensor:
48
+ return self.net(inputs)
49
+
50
+
51
+ class gPINNPoisson2D(nn.Module):
52
+ """2D Poisson model whose output transform enforces zero boundaries."""
53
+
54
+ def __init__(self, layers: Sequence[int]) -> None:
55
+ super().__init__()
56
+ self.net = gPINN(layers)
57
+
58
+ def forward(self, coordinates: torch.Tensor) -> torch.Tensor:
59
+ x = coordinates[:, 0:1]
60
+ y = coordinates[:, 1:2]
61
+ return x * y * (1.0 - x) * (1.0 - y) * self.net(coordinates)
62
+
63
+
64
+ def exact_poisson1d(x: np.ndarray) -> np.ndarray:
65
+ solution = x + np.sin(8.0 * x) / 8.0
66
+ for frequency in range(1, 5):
67
+ solution += np.sin(frequency * x) / frequency
68
+ return solution
69
+
70
+
71
+ def gpinn_loss_poisson1d(
72
+ model: nn.Module,
73
+ interior: torch.Tensor,
74
+ boundary: torch.Tensor,
75
+ boundary_values: torch.Tensor,
76
+ gradient_weight: float,
77
+ ) -> tuple[torch.Tensor, dict[str, float]]:
78
+ x = interior.detach().requires_grad_(True)
79
+ prediction = model(x)
80
+ first = autograd.grad(prediction, x, torch.ones_like(prediction), create_graph=True)[0]
81
+ second = autograd.grad(first, x, torch.ones_like(first), create_graph=True)[0]
82
+ third = autograd.grad(second, x, torch.ones_like(second), create_graph=True)[0]
83
+
84
+ source = 8.0 * torch.sin(8.0 * x) + sum(
85
+ frequency * torch.sin(frequency * x) for frequency in range(1, 5)
86
+ )
87
+ source_gradient = (
88
+ torch.cos(x)
89
+ + 4.0 * torch.cos(2.0 * x)
90
+ + 9.0 * torch.cos(3.0 * x)
91
+ + 16.0 * torch.cos(4.0 * x)
92
+ + 64.0 * torch.cos(8.0 * x)
93
+ )
94
+
95
+ residual = -second - source
96
+ residual_gradient = -third - source_gradient
97
+ residual_loss = torch.mean(residual.square())
98
+ boundary_loss = torch.mean((model(boundary) - boundary_values).square())
99
+ gradient_loss = torch.mean(residual_gradient.square())
100
+ loss = residual_loss + boundary_loss + gradient_weight * gradient_loss
101
+ return loss, {
102
+ "residual": residual_loss.item(),
103
+ "boundary": boundary_loss.item(),
104
+ "gradient": gradient_loss.item(),
105
+ }
106
+
107
+
108
+ def exact_poisson2d(x: np.ndarray, y: np.ndarray, exponent: float) -> np.ndarray:
109
+ return (16.0 * x * y * (1.0 - x) * (1.0 - y)) ** exponent
110
+
111
+
112
+ def _poisson2d_source(coordinates: torch.Tensor, exponent: float) -> torch.Tensor:
113
+ x = coordinates[:, 0:1]
114
+ y = coordinates[:, 1:2]
115
+ exact = (16.0 * x * y * (1.0 - x) * (1.0 - y)) ** exponent
116
+ exact_gradient = autograd.grad(
117
+ exact, coordinates, torch.ones_like(exact), create_graph=True
118
+ )[0]
119
+ exact_xx = autograd.grad(
120
+ exact_gradient[:, 0:1],
121
+ coordinates,
122
+ torch.ones_like(exact_gradient[:, 0:1]),
123
+ create_graph=True,
124
+ )[0][:, 0:1]
125
+ exact_yy = autograd.grad(
126
+ exact_gradient[:, 1:2],
127
+ coordinates,
128
+ torch.ones_like(exact_gradient[:, 1:2]),
129
+ create_graph=True,
130
+ )[0][:, 1:2]
131
+ return -(exact_xx + exact_yy)
132
+
133
+
134
+ def gpinn_loss_poisson2d(
135
+ model: nn.Module,
136
+ interior: torch.Tensor,
137
+ gradient_weight: float,
138
+ exponent: float,
139
+ ) -> tuple[torch.Tensor, dict[str, float]]:
140
+ coordinates = interior.detach().requires_grad_(True)
141
+ prediction = model(coordinates)
142
+ prediction_gradient = autograd.grad(
143
+ prediction, coordinates, torch.ones_like(prediction), create_graph=True
144
+ )[0]
145
+ prediction_xx = autograd.grad(
146
+ prediction_gradient[:, 0:1],
147
+ coordinates,
148
+ torch.ones_like(prediction_gradient[:, 0:1]),
149
+ create_graph=True,
150
+ )[0][:, 0:1]
151
+ prediction_yy = autograd.grad(
152
+ prediction_gradient[:, 1:2],
153
+ coordinates,
154
+ torch.ones_like(prediction_gradient[:, 1:2]),
155
+ create_graph=True,
156
+ )[0][:, 1:2]
157
+
158
+ residual = prediction_xx + prediction_yy + _poisson2d_source(coordinates, exponent)
159
+ residual_gradient = autograd.grad(
160
+ residual, coordinates, torch.ones_like(residual), create_graph=True
161
+ )[0]
162
+ residual_loss = torch.mean(residual.square())
163
+ gradient_loss = torch.mean(residual_gradient[:, 0:1].square()) + torch.mean(
164
+ residual_gradient[:, 1:2].square()
165
+ )
166
+ loss = residual_loss + gradient_weight * gradient_loss
167
+ return loss, {"residual": residual_loss.item(), "gradient": gradient_loss.item()}
168
+
169
+
170
+ def output_transform_burgers(
171
+ coordinates: torch.Tensor, raw_prediction: torch.Tensor
172
+ ) -> torch.Tensor:
173
+ x = coordinates[:, 0:1]
174
+ t = coordinates[:, 1:2]
175
+ return (
176
+ (1.0 - x) * (1.0 + x) * (1.0 - torch.exp(-t)) * raw_prediction
177
+ - torch.sin(torch.pi * x)
178
+ )
179
+
180
+
181
+ class gPINNBurgers(nn.Module):
182
+ """Burgers model with hard initial and boundary constraints."""
183
+
184
+ def __init__(self, layers: Sequence[int]) -> None:
185
+ super().__init__()
186
+ self.net = gPINN(layers)
187
+
188
+ def forward(self, coordinates: torch.Tensor) -> torch.Tensor:
189
+ return output_transform_burgers(coordinates, self.net(coordinates))
190
+
191
+
192
+ def burgers_residual(
193
+ model: nn.Module, coordinates: torch.Tensor, create_graph: bool = True
194
+ ) -> tuple[torch.Tensor, torch.Tensor]:
195
+ inputs = coordinates.detach().requires_grad_(True)
196
+ prediction = model(inputs)
197
+ prediction_gradient = autograd.grad(
198
+ prediction, inputs, torch.ones_like(prediction), create_graph=True
199
+ )[0]
200
+ prediction_x = prediction_gradient[:, 0:1]
201
+ prediction_t = prediction_gradient[:, 1:2]
202
+ prediction_xx = autograd.grad(
203
+ prediction_x,
204
+ inputs,
205
+ torch.ones_like(prediction_x),
206
+ create_graph=create_graph,
207
+ )[0][:, 0:1]
208
+ viscosity = 0.01 / torch.pi
209
+ residual = prediction_t + prediction * prediction_x - viscosity * prediction_xx
210
+ return residual, inputs
211
+
212
+
213
+ def burgers_gpinn_terms(
214
+ model: nn.Module, coordinates: torch.Tensor
215
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
216
+ residual, inputs = burgers_residual(model, coordinates, create_graph=True)
217
+ residual_gradient = autograd.grad(
218
+ residual, inputs, torch.ones_like(residual), create_graph=True
219
+ )[0]
220
+ return residual, residual_gradient[:, 0:1], residual_gradient[:, 1:2]
scripts/inference.py ADDED
@@ -0,0 +1,270 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from __future__ import annotations
2
+
3
+ import argparse
4
+ import os
5
+ import sys
6
+ from pathlib import Path
7
+
8
+ import matplotlib
9
+
10
+ matplotlib.use("Agg")
11
+ import matplotlib.pyplot as plt
12
+ import numpy as np
13
+ import torch
14
+ import yaml
15
+
16
+
17
+ PROJECT_ROOT = Path(__file__).resolve().parents[1]
18
+ sys.path.insert(0, str(PROJECT_ROOT))
19
+
20
+ from model.gpinn import ( # noqa: E402
21
+ exact_poisson1d,
22
+ exact_poisson2d,
23
+ gPINN,
24
+ gPINNBurgers,
25
+ gPINNPoisson2D,
26
+ )
27
+
28
+
29
+ DEFAULT_CONFIG = PROJECT_ROOT / "conf" / "config.yaml"
30
+ CASES = ("1d", "2d", "burgers")
31
+
32
+
33
+ def load_config(path: Path) -> dict:
34
+ with path.open("r", encoding="utf-8") as stream:
35
+ config = yaml.safe_load(stream)
36
+ if not isinstance(config, dict) or "root" not in config:
37
+ raise ValueError(f"config must contain a 'root' mapping: {path}")
38
+ return config["root"]
39
+
40
+
41
+ def project_path(value: str | Path) -> Path:
42
+ path = Path(value).expanduser()
43
+ return path if path.is_absolute() else PROJECT_ROOT / path
44
+
45
+
46
+ def resolve_device(requested: str) -> torch.device:
47
+ if requested == "auto":
48
+ return torch.device("cuda" if torch.cuda.is_available() else "cpu")
49
+ device = torch.device(requested)
50
+ if device.type == "cuda" and not torch.cuda.is_available():
51
+ raise RuntimeError("CUDA/DCU was requested but torch.cuda.is_available() is false")
52
+ return device
53
+
54
+
55
+ def resolve_dtype(name: str) -> torch.dtype:
56
+ try:
57
+ return {"float32": torch.float32, "float64": torch.float64}[name]
58
+ except KeyError as error:
59
+ raise ValueError(f"unsupported dtype: {name}") from error
60
+
61
+
62
+ def load_checkpoint(weight_dir: Path, filename: str) -> dict:
63
+ checkpoint_path = weight_dir / filename
64
+ if not checkpoint_path.is_file():
65
+ raise FileNotFoundError(
66
+ f"checkpoint not found: {checkpoint_path}. Run scripts/train.py first."
67
+ )
68
+ checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=True)
69
+ if not isinstance(checkpoint, dict) or "model_state" not in checkpoint:
70
+ raise ValueError(f"invalid checkpoint format: {checkpoint_path}")
71
+ return checkpoint
72
+
73
+
74
+ def infer_poisson1d(
75
+ weight_dir: Path,
76
+ result_dir: Path,
77
+ device: torch.device,
78
+ dtype: torch.dtype,
79
+ ) -> float:
80
+ checkpoint = load_checkpoint(weight_dir, "gpinn_poisson1d.pt")
81
+ layers = checkpoint.get("layers", [1, 30, 30, 30, 30, 1])
82
+ model = gPINN(layers).to(device=device, dtype=dtype)
83
+ model.load_state_dict(checkpoint["model_state"], strict=True)
84
+ model.eval()
85
+
86
+ x = torch.linspace(0.0, np.pi, 1000, dtype=dtype, device=device).unsqueeze(-1)
87
+ exact = exact_poisson1d(x.cpu().numpy().reshape(-1))
88
+ with torch.no_grad():
89
+ prediction = model(x).cpu().numpy().reshape(-1)
90
+ relative_l2 = float(np.linalg.norm(prediction - exact) / np.linalg.norm(exact))
91
+
92
+ figure, axis = plt.subplots(figsize=(8, 4))
93
+ axis.plot(x.cpu().numpy(), exact, "k-", label="Exact")
94
+ axis.plot(x.cpu().numpy(), prediction, "r--", label="gPINN")
95
+ axis.set_xlabel("x")
96
+ axis.set_ylabel("u")
97
+ axis.legend()
98
+ figure.tight_layout()
99
+ output_path = result_dir / "gpinn_poisson1d.png"
100
+ figure.savefig(output_path, dpi=150)
101
+ plt.close(figure)
102
+ print(f"[1D] relative L2={relative_l2:.6e}, plot={output_path}")
103
+ return relative_l2
104
+
105
+
106
+ def infer_poisson2d(
107
+ weight_dir: Path,
108
+ result_dir: Path,
109
+ device: torch.device,
110
+ dtype: torch.dtype,
111
+ default_exponent: float,
112
+ ) -> float:
113
+ checkpoint = load_checkpoint(weight_dir, "gpinn_poisson2d.pt")
114
+ layers = checkpoint.get("layers", [2, 30, 30, 30, 1])
115
+ architecture = checkpoint.get("architecture", "gpinn")
116
+ if architecture == "poisson2d_hard_bc":
117
+ model = gPINNPoisson2D(layers)
118
+ elif architecture == "gpinn":
119
+ # The bundled legacy checkpoint predates the hard boundary transform.
120
+ model = gPINN(layers)
121
+ else:
122
+ raise ValueError(f"unsupported 2D Poisson architecture: {architecture}")
123
+ model = model.to(device=device, dtype=dtype)
124
+ model.load_state_dict(checkpoint["model_state"], strict=True)
125
+ model.eval()
126
+
127
+ exponent = float(checkpoint.get("a", default_exponent))
128
+ axis = np.linspace(0.0, 1.0, 100)
129
+ grid_x, grid_y = np.meshgrid(axis, axis)
130
+ coordinates = torch.as_tensor(
131
+ np.column_stack([grid_x.ravel(), grid_y.ravel()]), dtype=dtype, device=device
132
+ )
133
+ with torch.no_grad():
134
+ prediction = model(coordinates).cpu().numpy().reshape(grid_x.shape)
135
+ exact = exact_poisson2d(grid_x, grid_y, exponent)
136
+ absolute_error = np.abs(prediction - exact)
137
+ relative_l2 = float(
138
+ np.linalg.norm(prediction.ravel() - exact.ravel()) / np.linalg.norm(exact.ravel())
139
+ )
140
+
141
+ figure, axes = plt.subplots(1, 3, figsize=(15, 4))
142
+ for plot_axis, title, field in zip(
143
+ axes,
144
+ ("Exact", "gPINN", "Absolute error"),
145
+ (exact, prediction, absolute_error),
146
+ strict=True,
147
+ ):
148
+ image = plot_axis.imshow(field, extent=(0, 1, 0, 1), origin="lower", cmap="jet")
149
+ plot_axis.set_title(title)
150
+ plot_axis.set_xlabel("x")
151
+ plot_axis.set_ylabel("y")
152
+ figure.colorbar(image, ax=plot_axis)
153
+ figure.tight_layout()
154
+ output_path = result_dir / "gpinn_poisson2d.png"
155
+ figure.savefig(output_path, dpi=150)
156
+ plt.close(figure)
157
+ print(f"[2D] relative L2={relative_l2:.6e}, plot={output_path}")
158
+ return relative_l2
159
+
160
+
161
+ def load_burgers_grid(data_path: Path) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
162
+ if not data_path.is_file():
163
+ raise FileNotFoundError(f"Burgers data not found: {data_path}")
164
+ with np.load(data_path) as data:
165
+ missing = {"t", "x", "usol"}.difference(data.files)
166
+ if missing:
167
+ raise ValueError(f"Burgers data is missing arrays: {sorted(missing)}")
168
+ time_axis = np.asarray(data["t"]).reshape(-1)
169
+ space_axis = np.asarray(data["x"]).reshape(-1)
170
+ exact = np.asarray(data["usol"])
171
+ expected_shape = (space_axis.size, time_axis.size)
172
+ if exact.shape != expected_shape:
173
+ raise ValueError(f"usol shape must be {expected_shape}, got {exact.shape}")
174
+ return time_axis, space_axis, exact
175
+
176
+
177
+ def infer_burgers(
178
+ weight_dir: Path,
179
+ result_dir: Path,
180
+ data_path: Path,
181
+ device: torch.device,
182
+ dtype: torch.dtype,
183
+ ) -> float:
184
+ checkpoint = load_checkpoint(weight_dir, "gpinn_burgers.pt")
185
+ layers = checkpoint.get("layers", [2, 32, 32, 32, 1])
186
+ model = gPINNBurgers(layers).to(device=device, dtype=dtype)
187
+ model.load_state_dict(checkpoint["model_state"], strict=True)
188
+ model.eval()
189
+
190
+ time_axis, space_axis, exact = load_burgers_grid(data_path)
191
+ time_grid, space_grid = np.meshgrid(time_axis, space_axis)
192
+ coordinates = torch.as_tensor(
193
+ np.column_stack([space_grid.ravel(), time_grid.ravel()]),
194
+ dtype=dtype,
195
+ device=device,
196
+ )
197
+ with torch.no_grad():
198
+ prediction = model(coordinates).cpu().numpy().reshape(exact.shape)
199
+ absolute_error = np.abs(prediction - exact)
200
+ relative_l2 = float(
201
+ np.linalg.norm(prediction.ravel() - exact.ravel()) / np.linalg.norm(exact.ravel())
202
+ )
203
+
204
+ figure, axes = plt.subplots(1, 3, figsize=(15, 4))
205
+ for plot_axis, title, field in zip(
206
+ axes,
207
+ ("Exact", "gPINN", "Absolute error"),
208
+ (exact, prediction, absolute_error),
209
+ strict=True,
210
+ ):
211
+ image = plot_axis.contourf(time_axis, space_axis, field, 100, cmap="jet")
212
+ plot_axis.set_title(title)
213
+ plot_axis.set_xlabel("t")
214
+ plot_axis.set_ylabel("x")
215
+ figure.colorbar(image, ax=plot_axis)
216
+ figure.tight_layout()
217
+ output_path = result_dir / "gpinn_burgers.png"
218
+ figure.savefig(output_path, dpi=150)
219
+ plt.close(figure)
220
+ print(f"[Burgers] relative L2={relative_l2:.6e}, plot={output_path}")
221
+ return relative_l2
222
+
223
+
224
+ def parse_args() -> argparse.Namespace:
225
+ parser = argparse.ArgumentParser(description="Run gPINN inference and visualization")
226
+ parser.add_argument(
227
+ "--config",
228
+ type=Path,
229
+ default=Path(os.environ.get("GPINN_CONFIG", DEFAULT_CONFIG)),
230
+ help="YAML configuration file",
231
+ )
232
+ parser.add_argument("--case", choices=(*CASES, "all"), default="all")
233
+ parser.add_argument("--device", help="Override common.device, for example cpu or cuda:0")
234
+ parser.add_argument("--data", type=Path, help="Override Burgers.npz path")
235
+ parser.add_argument("--weight-dir", type=Path, help="Override checkpoint input directory")
236
+ parser.add_argument("--result-dir", type=Path, help="Override plot output directory")
237
+ return parser.parse_args()
238
+
239
+
240
+ def main() -> None:
241
+ args = parse_args()
242
+ config_path = args.config.expanduser().resolve()
243
+ config = load_config(config_path)
244
+ common = config["common"]
245
+ device = resolve_device(args.device or str(common["device"]))
246
+ dtype = resolve_dtype(str(common["dtype"]))
247
+ weight_dir = project_path(args.weight_dir or common["weight_dir"])
248
+ result_dir = project_path(args.result_dir or common["result_dir"])
249
+ data_path = project_path(args.data or config["burgers"]["data"])
250
+ result_dir.mkdir(parents=True, exist_ok=True)
251
+
252
+ print(f"Config: {config_path}")
253
+ print(f"Device: {device}")
254
+ selected_cases = CASES if args.case == "all" else (args.case,)
255
+ if "1d" in selected_cases:
256
+ infer_poisson1d(weight_dir, result_dir, device, dtype)
257
+ if "2d" in selected_cases:
258
+ infer_poisson2d(
259
+ weight_dir,
260
+ result_dir,
261
+ device,
262
+ dtype,
263
+ float(config["poisson2d"]["a"]),
264
+ )
265
+ if "burgers" in selected_cases:
266
+ infer_burgers(weight_dir, result_dir, data_path, device, dtype)
267
+
268
+
269
+ if __name__ == "__main__":
270
+ main()
scripts/train.py ADDED
@@ -0,0 +1,495 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from __future__ import annotations
2
+
3
+ import argparse
4
+ import os
5
+ import sys
6
+ import time
7
+ from pathlib import Path
8
+
9
+ import numpy as np
10
+ import torch
11
+ import yaml
12
+
13
+
14
+ PROJECT_ROOT = Path(__file__).resolve().parents[1]
15
+ sys.path.insert(0, str(PROJECT_ROOT))
16
+
17
+ from model.gpinn import ( # noqa: E402
18
+ burgers_gpinn_terms,
19
+ burgers_residual,
20
+ exact_poisson1d,
21
+ exact_poisson2d,
22
+ gPINN,
23
+ gPINNBurgers,
24
+ gPINNPoisson2D,
25
+ gpinn_loss_poisson1d,
26
+ gpinn_loss_poisson2d,
27
+ )
28
+
29
+
30
+ DEFAULT_CONFIG = PROJECT_ROOT / "conf" / "config.yaml"
31
+ CASES = ("1d", "2d", "burgers")
32
+
33
+
34
+ def load_config(path: Path) -> dict:
35
+ with path.open("r", encoding="utf-8") as stream:
36
+ config = yaml.safe_load(stream)
37
+ if not isinstance(config, dict) or "root" not in config:
38
+ raise ValueError(f"config must contain a 'root' mapping: {path}")
39
+ return config["root"]
40
+
41
+
42
+ def project_path(value: str | Path) -> Path:
43
+ path = Path(value).expanduser()
44
+ return path if path.is_absolute() else PROJECT_ROOT / path
45
+
46
+
47
+ def resolve_device(requested: str) -> torch.device:
48
+ if requested == "auto":
49
+ return torch.device("cuda" if torch.cuda.is_available() else "cpu")
50
+ device = torch.device(requested)
51
+ if device.type == "cuda" and not torch.cuda.is_available():
52
+ raise RuntimeError("CUDA/DCU was requested but torch.cuda.is_available() is false")
53
+ return device
54
+
55
+
56
+ def resolve_dtype(name: str) -> torch.dtype:
57
+ try:
58
+ return {"float32": torch.float32, "float64": torch.float64}[name]
59
+ except KeyError as error:
60
+ raise ValueError(f"unsupported dtype: {name}") from error
61
+
62
+
63
+ def seed_everything(seed: int) -> None:
64
+ np.random.seed(seed)
65
+ torch.manual_seed(seed)
66
+ if torch.cuda.is_available():
67
+ torch.cuda.manual_seed_all(seed)
68
+
69
+
70
+ def ensure_finite(loss: torch.Tensor, case: str, step: int) -> None:
71
+ if not torch.isfinite(loss):
72
+ raise FloatingPointError(f"{case} loss became non-finite at step {step}")
73
+
74
+
75
+ def save_checkpoint(weight_dir: Path, filename: str, payload: dict) -> Path:
76
+ weight_dir.mkdir(parents=True, exist_ok=True)
77
+ checkpoint_path = weight_dir / filename
78
+ torch.save(payload, checkpoint_path)
79
+ print(f"Saved checkpoint: {checkpoint_path}")
80
+ return checkpoint_path
81
+
82
+
83
+ def train_poisson1d(
84
+ config: dict,
85
+ device: torch.device,
86
+ dtype: torch.dtype,
87
+ weight_dir: Path,
88
+ epochs_override: int | None,
89
+ nf_override: int | None,
90
+ lbfgs_override: int | None,
91
+ ) -> None:
92
+ layers = [int(value) for value in config["layers"]]
93
+ nf = nf_override if nf_override is not None else int(config["nf"])
94
+ epochs = epochs_override if epochs_override is not None else int(config["epochs"])
95
+ lbfgs_iters = (
96
+ lbfgs_override if lbfgs_override is not None else int(config["lbfgs_iters"])
97
+ )
98
+ log_every = int(config["log_every"])
99
+
100
+ interior = torch.linspace(0.0, np.pi, nf, dtype=dtype, device=device).unsqueeze(-1)
101
+ boundary = torch.tensor([[0.0], [np.pi]], dtype=dtype, device=device)
102
+ boundary_values = boundary.clone()
103
+ test_x = torch.linspace(0.0, np.pi, 1000, dtype=dtype, device=device).unsqueeze(-1)
104
+ exact = torch.as_tensor(
105
+ exact_poisson1d(test_x.detach().cpu().numpy().reshape(-1)),
106
+ dtype=dtype,
107
+ device=device,
108
+ ).unsqueeze(-1)
109
+
110
+ model = gPINN(layers).to(device=device, dtype=dtype)
111
+ optimizer = torch.optim.Adam(model.parameters(), lr=float(config["lr"]))
112
+ scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=5000, gamma=0.5)
113
+ started = time.time()
114
+ print(
115
+ f"[1D] device={device} nf={nf} epochs={epochs} "
116
+ f"lbfgs_iters={lbfgs_iters}"
117
+ )
118
+
119
+ for step in range(1, epochs + 1):
120
+ loss, parts = gpinn_loss_poisson1d(
121
+ model, interior, boundary, boundary_values, float(config["w_g"])
122
+ )
123
+ ensure_finite(loss, "1D Poisson", step)
124
+ optimizer.zero_grad(set_to_none=True)
125
+ loss.backward()
126
+ optimizer.step()
127
+ scheduler.step()
128
+ if step == 1 or step % log_every == 0 or step == epochs:
129
+ with torch.no_grad():
130
+ relative_l2 = torch.linalg.vector_norm(model(test_x) - exact) / torch.linalg.vector_norm(
131
+ exact
132
+ )
133
+ print(
134
+ f"[1D] step={step:5d} loss={loss.item():.3e} "
135
+ f"residual={parts['residual']:.3e} boundary={parts['boundary']:.3e} "
136
+ f"gradient={parts['gradient']:.3e} l2={relative_l2.item():.3e}"
137
+ )
138
+
139
+ if lbfgs_iters > 0:
140
+ lbfgs = torch.optim.LBFGS(
141
+ model.parameters(),
142
+ lr=1.0,
143
+ max_iter=lbfgs_iters,
144
+ line_search_fn="strong_wolfe",
145
+ )
146
+
147
+ def closure() -> torch.Tensor:
148
+ lbfgs.zero_grad(set_to_none=True)
149
+ closure_loss, _ = gpinn_loss_poisson1d(
150
+ model, interior, boundary, boundary_values, float(config["w_g"])
151
+ )
152
+ ensure_finite(closure_loss, "1D Poisson L-BFGS", 0)
153
+ closure_loss.backward()
154
+ return closure_loss
155
+
156
+ lbfgs.step(closure)
157
+
158
+ with torch.no_grad():
159
+ relative_l2 = torch.linalg.vector_norm(model(test_x) - exact) / torch.linalg.vector_norm(exact)
160
+ print(f"[1D] finished in {time.time() - started:.1f}s, relative L2={relative_l2.item():.6e}")
161
+ save_checkpoint(
162
+ weight_dir,
163
+ "gpinn_poisson1d.pt",
164
+ {
165
+ "case": "1d",
166
+ "architecture": "gpinn",
167
+ "model_state": model.state_dict(),
168
+ "layers": layers,
169
+ "nf": nf,
170
+ "w_g": float(config["w_g"]),
171
+ },
172
+ )
173
+
174
+
175
+ def train_poisson2d(
176
+ config: dict,
177
+ device: torch.device,
178
+ dtype: torch.dtype,
179
+ weight_dir: Path,
180
+ epochs_override: int | None,
181
+ nf_override: int | None,
182
+ ) -> None:
183
+ layers = [int(value) for value in config["layers"]]
184
+ nf = nf_override if nf_override is not None else int(config["nf"])
185
+ epochs = epochs_override if epochs_override is not None else int(config["epochs"])
186
+ exponent = float(config["a"])
187
+ log_every = int(config["log_every"])
188
+
189
+ interior = torch.rand(nf, 2, dtype=dtype, device=device)
190
+ model = gPINNPoisson2D(layers).to(device=device, dtype=dtype)
191
+ optimizer = torch.optim.Adam(model.parameters(), lr=float(config["lr"]))
192
+ scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=5000, gamma=0.5)
193
+ started = time.time()
194
+ print(f"[2D] device={device} nf={nf} epochs={epochs} a={exponent:g}")
195
+
196
+ for step in range(1, epochs + 1):
197
+ loss, parts = gpinn_loss_poisson2d(
198
+ model, interior, float(config["w_g"]), exponent
199
+ )
200
+ ensure_finite(loss, "2D Poisson", step)
201
+ optimizer.zero_grad(set_to_none=True)
202
+ loss.backward()
203
+ optimizer.step()
204
+ scheduler.step()
205
+ if step == 1 or step % log_every == 0 or step == epochs:
206
+ print(
207
+ f"[2D] step={step:5d} loss={loss.item():.3e} "
208
+ f"residual={parts['residual']:.3e} gradient={parts['gradient']:.3e}"
209
+ )
210
+
211
+ axis = np.linspace(0.0, 1.0, 100)
212
+ grid_x, grid_y = np.meshgrid(axis, axis)
213
+ test_coordinates = torch.as_tensor(
214
+ np.column_stack([grid_x.ravel(), grid_y.ravel()]), dtype=dtype, device=device
215
+ )
216
+ with torch.no_grad():
217
+ prediction = model(test_coordinates).cpu().numpy().reshape(grid_x.shape)
218
+ exact = exact_poisson2d(grid_x, grid_y, exponent)
219
+ relative_l2 = np.linalg.norm(prediction.ravel() - exact.ravel()) / np.linalg.norm(
220
+ exact.ravel()
221
+ )
222
+ print(f"[2D] finished in {time.time() - started:.1f}s, relative L2={relative_l2:.6e}")
223
+ save_checkpoint(
224
+ weight_dir,
225
+ "gpinn_poisson2d.pt",
226
+ {
227
+ "case": "2d",
228
+ "architecture": "poisson2d_hard_bc",
229
+ "model_state": model.state_dict(),
230
+ "layers": layers,
231
+ "a": exponent,
232
+ "nf": nf,
233
+ "w_g": float(config["w_g"]),
234
+ },
235
+ )
236
+
237
+
238
+ def load_burgers_reference(
239
+ data_path: Path, device: torch.device, dtype: torch.dtype
240
+ ) -> tuple[torch.Tensor, torch.Tensor]:
241
+ if not data_path.is_file():
242
+ raise FileNotFoundError(f"Burgers data not found: {data_path}")
243
+ with np.load(data_path) as data:
244
+ missing = {"t", "x", "usol"}.difference(data.files)
245
+ if missing:
246
+ raise ValueError(f"Burgers data is missing arrays: {sorted(missing)}")
247
+ time_axis = np.asarray(data["t"]).reshape(-1)
248
+ space_axis = np.asarray(data["x"]).reshape(-1)
249
+ exact = np.asarray(data["usol"])
250
+ expected_shape = (space_axis.size, time_axis.size)
251
+ if exact.shape != expected_shape:
252
+ raise ValueError(f"usol shape must be {expected_shape}, got {exact.shape}")
253
+ time_grid, space_grid = np.meshgrid(time_axis, space_axis)
254
+ coordinates = torch.as_tensor(
255
+ np.column_stack([space_grid.ravel(), time_grid.ravel()]),
256
+ dtype=dtype,
257
+ device=device,
258
+ )
259
+ values = torch.as_tensor(exact.reshape(-1, 1), dtype=dtype, device=device)
260
+ return coordinates, values
261
+
262
+
263
+ def burgers_boundary_points(
264
+ device: torch.device, dtype: torch.dtype
265
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
266
+ boundary = torch.cat(
267
+ [
268
+ torch.cat(
269
+ [
270
+ torch.full((100, 1), -1.0, dtype=dtype, device=device),
271
+ torch.rand(100, 1, dtype=dtype, device=device),
272
+ ],
273
+ dim=1,
274
+ ),
275
+ torch.cat(
276
+ [
277
+ torch.full((100, 1), 1.0, dtype=dtype, device=device),
278
+ torch.rand(100, 1, dtype=dtype, device=device),
279
+ ],
280
+ dim=1,
281
+ ),
282
+ ]
283
+ )
284
+ initial = torch.cat(
285
+ [
286
+ torch.linspace(-1.0, 1.0, 200, dtype=dtype, device=device).unsqueeze(-1),
287
+ torch.zeros(200, 1, dtype=dtype, device=device),
288
+ ],
289
+ dim=1,
290
+ )
291
+ initial_values = -torch.sin(torch.pi * initial[:, 0:1])
292
+ return boundary, initial, initial_values
293
+
294
+
295
+ def add_rar_points(
296
+ model: torch.nn.Module,
297
+ config: dict,
298
+ device: torch.device,
299
+ dtype: torch.dtype,
300
+ ) -> torch.Tensor:
301
+ candidate_count = int(config["rar_candidate_points"])
302
+ batch_size = int(config["rar_candidate_batch_size"])
303
+ add_count = int(config["rar_add_points"])
304
+ if min(candidate_count, batch_size, add_count) <= 0:
305
+ raise ValueError("RAR point counts must be positive")
306
+
307
+ shortlisted_points = []
308
+ shortlisted_residuals = []
309
+ was_training = model.training
310
+ model.eval()
311
+ for start in range(0, candidate_count, batch_size):
312
+ current_size = min(batch_size, candidate_count - start)
313
+ candidates = torch.rand(current_size, 2, dtype=dtype, device=device)
314
+ candidates[:, 0] = 2.0 * candidates[:, 0] - 1.0
315
+ residual, inputs = burgers_residual(model, candidates, create_graph=False)
316
+ count = min(add_count, current_size)
317
+ values, indices = torch.topk(residual.detach().abs().reshape(-1), count)
318
+ shortlisted_points.append(inputs.detach()[indices])
319
+ shortlisted_residuals.append(values)
320
+ if was_training:
321
+ model.train()
322
+
323
+ residuals = torch.cat(shortlisted_residuals)
324
+ points = torch.cat(shortlisted_points)
325
+ final_count = min(add_count, residuals.numel())
326
+ indices = torch.topk(residuals, final_count).indices
327
+ return points[indices].detach()
328
+
329
+
330
+ def train_burgers(
331
+ config: dict,
332
+ device: torch.device,
333
+ dtype: torch.dtype,
334
+ weight_dir: Path,
335
+ data_path: Path,
336
+ epochs_override: int | None,
337
+ nf_override: int | None,
338
+ rar_override: int | None,
339
+ quick: bool,
340
+ ) -> None:
341
+ layers = [int(value) for value in config["layers"]]
342
+ nf = nf_override if nf_override is not None else int(config["nf"])
343
+ epochs = epochs_override if epochs_override is not None else int(config["epochs"])
344
+ rar_rounds = rar_override if rar_override is not None else int(config["rar_rounds"])
345
+ if quick:
346
+ rar_rounds = 0
347
+ log_every = int(config["log_every"])
348
+
349
+ collocation = torch.rand(nf, 2, dtype=dtype, device=device)
350
+ collocation[:, 0] = 2.0 * collocation[:, 0] - 1.0
351
+ reference_coordinates, reference_values = load_burgers_reference(data_path, device, dtype)
352
+ boundary, initial, initial_values = burgers_boundary_points(device, dtype)
353
+ model = gPINNBurgers(layers).to(device=device, dtype=dtype)
354
+ optimizer = torch.optim.Adam(model.parameters(), lr=float(config["lr"]))
355
+ scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=5000, gamma=0.5)
356
+ anchors: list[torch.Tensor] = []
357
+
358
+ def evaluate() -> float:
359
+ with torch.no_grad():
360
+ return (
361
+ torch.linalg.vector_norm(model(reference_coordinates) - reference_values)
362
+ / torch.linalg.vector_norm(reference_values)
363
+ ).item()
364
+
365
+ def optimization_step(points: torch.Tensor, step: int) -> tuple[torch.Tensor, dict[str, float]]:
366
+ residual, residual_x, residual_t = burgers_gpinn_terms(model, points)
367
+ residual_loss = torch.mean(residual.square())
368
+ boundary_loss = torch.mean(model(boundary).square()) + torch.mean(
369
+ (model(initial) - initial_values).square()
370
+ )
371
+ gradient_loss = torch.mean(residual_x.square()) + torch.mean(residual_t.square())
372
+ loss = residual_loss + boundary_loss + float(config["w_g"]) * gradient_loss
373
+ ensure_finite(loss, "Burgers", step)
374
+ optimizer.zero_grad(set_to_none=True)
375
+ loss.backward()
376
+ optimizer.step()
377
+ scheduler.step()
378
+ return loss, {
379
+ "residual": residual_loss.item(),
380
+ "boundary": boundary_loss.item(),
381
+ "gradient": gradient_loss.item(),
382
+ }
383
+
384
+ started = time.time()
385
+ print(f"[Burgers] device={device} nf={nf} epochs={epochs} rar_rounds={rar_rounds}")
386
+ for step in range(1, epochs + 1):
387
+ points = torch.cat([collocation, *anchors], dim=0)
388
+ loss, parts = optimization_step(points, step)
389
+ if step == 1 or step % log_every == 0 or step == epochs:
390
+ print(
391
+ f"[Burgers] step={step:5d} loss={loss.item():.3e} "
392
+ f"residual={parts['residual']:.3e} boundary={parts['boundary']:.3e} "
393
+ f"gradient={parts['gradient']:.3e} l2={evaluate():.3e}"
394
+ )
395
+
396
+ rar_epochs = int(config["rar_epochs"])
397
+ for round_index in range(1, rar_rounds + 1):
398
+ new_points = add_rar_points(model, config, device, dtype)
399
+ anchors.append(new_points)
400
+ points = torch.cat([collocation, *anchors], dim=0)
401
+ for rar_step in range(1, rar_epochs + 1):
402
+ optimization_step(points, epochs + (round_index - 1) * rar_epochs + rar_step)
403
+ print(
404
+ f"[Burgers] RAR round={round_index}/{rar_rounds} "
405
+ f"anchors={sum(item.shape[0] for item in anchors)} l2={evaluate():.3e}"
406
+ )
407
+
408
+ relative_l2 = evaluate()
409
+ print(
410
+ f"[Burgers] finished in {time.time() - started:.1f}s, "
411
+ f"relative L2={relative_l2:.6e}"
412
+ )
413
+ save_checkpoint(
414
+ weight_dir,
415
+ "gpinn_burgers.pt",
416
+ {
417
+ "case": "burgers",
418
+ "architecture": "burgers_hard_bc",
419
+ "model_state": model.state_dict(),
420
+ "layers": layers,
421
+ "nf": nf,
422
+ "w_g": float(config["w_g"]),
423
+ "rar_rounds": rar_rounds,
424
+ },
425
+ )
426
+
427
+
428
+ def parse_args() -> argparse.Namespace:
429
+ parser = argparse.ArgumentParser(description="Train gPINN PDE examples")
430
+ parser.add_argument(
431
+ "--config",
432
+ type=Path,
433
+ default=Path(os.environ.get("GPINN_CONFIG", DEFAULT_CONFIG)),
434
+ help="YAML configuration file",
435
+ )
436
+ parser.add_argument("--case", choices=(*CASES, "all"), default="all")
437
+ parser.add_argument("--device", help="Override common.device, for example cpu or cuda:0")
438
+ parser.add_argument("--epochs", type=int, help="Override Adam iterations for selected cases")
439
+ parser.add_argument("--nf", type=int, help="Override collocation point count")
440
+ parser.add_argument("--lbfgs-iters", type=int, help="Override 1D Poisson L-BFGS iterations")
441
+ parser.add_argument("--rar-rounds", type=int, help="Override Burgers RAR rounds")
442
+ parser.add_argument("--quick", action="store_true", help="Skip Burgers RAR rounds")
443
+ parser.add_argument("--data", type=Path, help="Override Burgers.npz path")
444
+ parser.add_argument("--weight-dir", type=Path, help="Override checkpoint output directory")
445
+ return parser.parse_args()
446
+
447
+
448
+ def main() -> None:
449
+ args = parse_args()
450
+ config_path = args.config.expanduser().resolve()
451
+ config = load_config(config_path)
452
+ common = config["common"]
453
+ device = resolve_device(args.device or str(common["device"]))
454
+ dtype = resolve_dtype(str(common["dtype"]))
455
+ weight_dir = project_path(args.weight_dir or common["weight_dir"])
456
+ data_path = project_path(args.data or config["burgers"]["data"])
457
+ seed_everything(int(common["seed"]))
458
+
459
+ print(f"Config: {config_path}")
460
+ selected_cases = CASES if args.case == "all" else (args.case,)
461
+ if "1d" in selected_cases:
462
+ train_poisson1d(
463
+ config["poisson1d"],
464
+ device,
465
+ dtype,
466
+ weight_dir,
467
+ args.epochs,
468
+ args.nf,
469
+ args.lbfgs_iters,
470
+ )
471
+ if "2d" in selected_cases:
472
+ train_poisson2d(
473
+ config["poisson2d"],
474
+ device,
475
+ dtype,
476
+ weight_dir,
477
+ args.epochs,
478
+ args.nf,
479
+ )
480
+ if "burgers" in selected_cases:
481
+ train_burgers(
482
+ config["burgers"],
483
+ device,
484
+ dtype,
485
+ weight_dir,
486
+ data_path,
487
+ args.epochs,
488
+ args.nf,
489
+ args.rar_rounds,
490
+ args.quick,
491
+ )
492
+
493
+
494
+ if __name__ == "__main__":
495
+ main()
weight/gpinn_burgers.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:aeb01c61bba1d304e8d986605b18b9a99b113dc6cdbfe8abfa1d1d060af73b55
3
+ size 22674
weight/gpinn_poisson1d.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:e7d934849cd64c65ed87016db2175aef801e59113563ae21aa3e33b95636e403
3
+ size 28161
weight/gpinn_poisson2d.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:6aeb343877399cfa2f48ee4b48965d880009311bb4188471eb739a2bb2fbe8d9
3
+ size 20210