OneScience commited on
Commit
1f85b04
·
verified ·
1 Parent(s): b490d16

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,133 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ frameworks:
3
+ - pytorch
4
+ language:
5
+ - en
6
+ license: apache-2.0
7
+ tags:
8
+ - OneScience
9
+ - SA-PINNs
10
+ - self-adaptive-physics-informed-neural-networks
11
+ - soft-attention
12
+ - partial-differential-equations
13
+ - Laplace
14
+ - Helmholtz
15
+ - Burgers
16
+ tasks:
17
+ - pde-solving
18
+ ---
19
+ <p align="center">
20
+ <strong>
21
+ <span style="font-size: 30px;">SA-PINNs</span>
22
+ </strong>
23
+ </p>
24
+
25
+ # Model Overview
26
+
27
+ SA-PINNs (Self-Adaptive Physics-Informed Neural Networks) assign learnable positive attention weights to observation points, boundary points, and PDE collocation points, enabling the model to focus automatically on regions that are difficult to fit during training.
28
+
29
+ This model package provides examples for the one-dimensional Laplace equation, the two-dimensional Helmholtz equation, and the one-dimensional time-dependent Burgers equation.
30
+
31
+ Paper: Self-Adaptive Physics-Informed Neural Networks using a Soft Attention Mechanism
32
+ https://arxiv.org/abs/2009.04544
33
+
34
+ # Model Description
35
+
36
+ SA-PINNs use min–max saddle-point optimization: network parameters reduce the weighted physics loss through gradient descent, while attention parameters increase the weights assigned to difficult locations through gradient ascent. The attention weights are constructed from normalized `exp(alpha)`, ensuring that they remain positive with a mean close to 1. The Adam stage updates both the network and attention parameters, whereas the L-BFGS stage freezes the attention parameters and refines the network.
37
+
38
+ # Use Cases
39
+
40
+ | Use Case | Description |
41
+ | :---: | :--- |
42
+ | One-dimensional Laplace equation | Examine how attention focuses on collocation points with high residuals |
43
+ | Two-dimensional Helmholtz equation | Solve a two-dimensional analytical solution with high-frequency spatial variation |
44
+ | Burgers equation | Train jointly on initial conditions, boundary conditions, and nonlinear PDE residuals |
45
+ | Adaptive-weight research | Compare a standard PINN with pointwise soft-attention weighting |
46
+
47
+ # Usage
48
+
49
+ ## 1. OneCode
50
+
51
+ Use the online OneCode environment for an intelligent, one-click AI for Science (AI4S) programming experience:
52
+
53
+ [Launch OneCode for one-click AI4S programming](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)
54
+
55
+ ## 2. Manual Setup
56
+
57
+ **Hardware Requirements**
58
+
59
+ - A GPU or DCU is recommended for training.
60
+ - A CPU can be used for small-scale pipeline validation.
61
+ - DCU users must install DTK and a PyTorch environment compatible with the target cluster.
62
+
63
+ ### Download the Model Package
64
+
65
+ ```bash
66
+ modelscope download --model OneScience/SA-PINNs --local_dir ./SA-PINNs
67
+ cd SA-PINNs
68
+ ```
69
+
70
+ ### Set Up the Runtime Environment
71
+
72
+ **DCU Environment**
73
+
74
+ ```bash
75
+ # Activate DTK and Conda first
76
+ conda create -n onescience311 python=3.11 -y
77
+ conda activate onescience311
78
+ pip install onescience[cfd-dcu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
79
+ ```
80
+
81
+ **GPU Environment**
82
+
83
+ ```bash
84
+ # Activate Conda first
85
+ conda create -n onescience311 python=3.11 -y libstdcxx-ng=12 libgcc-ng=12 gcc_linux-64=12 gxx_linux-64=12
86
+ conda activate onescience311
87
+ pip install onescience[cfd-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
88
+ ```
89
+
90
+ ### Training Data
91
+
92
+ All three examples generate observation points, boundary points, and PDE collocation points from analytical equations and random sampling, with no dependency on external data files. The Burgers example does not include a reference solution and therefore outputs only the predicted field and attention distribution.
93
+
94
+ ### Training
95
+
96
+ Select an example with `--case`:
97
+
98
+ ```bash
99
+ python scripts/train.py --case laplace
100
+ python scripts/train.py --case helmholtz
101
+ python scripts/train.py --case burgers
102
+ ```
103
+
104
+ Checkpoints and training histories are saved to `weight/` and `result/` by default. Command-line arguments such as `--epochs`, `--lbfgs-iters`, `--n-pde`, and `--device` can override values in `conf/config.yaml`.
105
+
106
+ ### Model Weights
107
+
108
+ This repository provides weights trained for all three examples in the `weight/` directory.
109
+
110
+
111
+ ### Inference, Evaluation, and Visualization
112
+
113
+ After training the selected example, run:
114
+
115
+ ```bash
116
+ python scripts/inference.py --case laplace
117
+ python scripts/inference.py --case helmholtz
118
+ python scripts/inference.py --case burgers
119
+ ```
120
+
121
+ Predictions, error plots, and PDE attention distributions are saved to the `result/` directory. The Laplace and Helmholtz examples report relative L2 errors. Default model and training parameters can be modified in `conf/config.yaml`.
122
+
123
+ # Official OneScience Resources
124
+
125
+ | Platform | OneScience 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
+ # Citations and License
131
+
132
+ - McClenny, L. and Braga-Neto, U. Self-Adaptive Physics-Informed Neural Networks using a Soft Attention Mechanism. arXiv:2009.04544, 2020.
133
+ - This model package is released under the Apache-2.0 license and retains attribution to the original paper.
conf/config.yaml ADDED
@@ -0,0 +1,103 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ root:
2
+ common:
3
+ device: "auto"
4
+ dtype: "float64"
5
+ seed: 42
6
+ weight_dir: "weight"
7
+ result_dir: "result"
8
+
9
+ cases:
10
+ laplace:
11
+ model:
12
+ input_dim: 1
13
+ n_layers: 4
14
+ n_neurons: 50
15
+ activation: "tanh"
16
+ attention:
17
+ enabled: true
18
+ lr: 0.01
19
+ training:
20
+ epochs: 10000
21
+ lr: 0.001
22
+ lbfgs_iters: 2000
23
+ log_interval: 1000
24
+ loss:
25
+ data: 10.0
26
+ pde: 1.0
27
+ boundary: 10.0
28
+ data:
29
+ domain: [0.0, 1.0]
30
+ n_sol: 200
31
+ n_pde: 1000
32
+ n_bnd: 2
33
+ noise_std: 0.0
34
+ test_res: 200
35
+ output:
36
+ checkpoint_name: "sapinn_laplace.pt"
37
+ predictions_name: "sapinn_laplace.npz"
38
+ figure_name: "sapinn_laplace.png"
39
+ history_name: "sapinn_laplace_history.npz"
40
+
41
+ helmholtz:
42
+ model:
43
+ input_dim: 2
44
+ n_layers: 4
45
+ n_neurons: 50
46
+ activation: "tanh"
47
+ attention:
48
+ enabled: true
49
+ lr: 0.001
50
+ training:
51
+ epochs: 10000
52
+ lr: 0.001
53
+ lbfgs_iters: 5000
54
+ log_interval: 1000
55
+ loss:
56
+ data: 0.0
57
+ pde: 1.0
58
+ boundary: 10.0
59
+ data:
60
+ domain_x: [-1.0, 1.0]
61
+ domain_y: [-1.0, 1.0]
62
+ n_sol: 200
63
+ n_pde: 10000
64
+ n_bnd: 400
65
+ test_res: 100
66
+ wave_number: 1.0
67
+ output:
68
+ checkpoint_name: "sapinn_helmholtz.pt"
69
+ predictions_name: "sapinn_helmholtz.npz"
70
+ figure_name: "sapinn_helmholtz.png"
71
+ history_name: "sapinn_helmholtz_history.npz"
72
+
73
+ burgers:
74
+ model:
75
+ input_dim: 2
76
+ n_layers: 8
77
+ n_neurons: 20
78
+ activation: "tanh"
79
+ attention:
80
+ enabled: true
81
+ lr: 0.005
82
+ training:
83
+ epochs: 100
84
+ lr: 0.005
85
+ lbfgs_iters: 5000
86
+ log_interval: 10
87
+ loss:
88
+ data: 10.0
89
+ pde: 1.0
90
+ boundary: 10.0
91
+ data:
92
+ domain_x: [-1.0, 1.0]
93
+ domain_t: [0.0, 1.0]
94
+ n_sol: 100
95
+ n_pde: 10000
96
+ n_bnd: 50
97
+ test_res: 100
98
+ viscosity: 0.003183098861837907
99
+ output:
100
+ checkpoint_name: "sapinn_burgers.pt"
101
+ predictions_name: "sapinn_burgers.npz"
102
+ figure_name: "sapinn_burgers.png"
103
+ history_name: "sapinn_burgers_history.npz"
configuration.json ADDED
@@ -0,0 +1 @@
 
 
1
+ {"framework":"Pytorch","task":"other"}
data/.gitkeep ADDED
File without changes
model/__pycache__/sa_pinn.cpython-311.pyc ADDED
Binary file (16.7 kB). View file
 
model/sa_pinn.py ADDED
@@ -0,0 +1,267 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from __future__ import annotations
2
+
3
+ import math
4
+ from collections.abc import Mapping, Sequence
5
+
6
+ import torch
7
+ import torch.nn as nn
8
+
9
+
10
+ class FCN(nn.Module):
11
+ """Fully connected network used by each SA-PINN case."""
12
+
13
+ def __init__(
14
+ self,
15
+ layer_sizes: Sequence[int],
16
+ activation: str = "tanh",
17
+ dtype: torch.dtype = torch.float64,
18
+ ) -> None:
19
+ super().__init__()
20
+ if len(layer_sizes) < 2:
21
+ raise ValueError("layer_sizes must include input and output widths")
22
+ if activation not in {"tanh", "sin", "sine", "relu", "gelu"}:
23
+ raise ValueError(f"unsupported activation: {activation}")
24
+ self.activation = "sin" if activation == "sine" else activation
25
+ stddev = math.sqrt(50.0 / layer_sizes[1])
26
+ self.linears = nn.ModuleList(
27
+ nn.Linear(layer_sizes[index], layer_sizes[index + 1], dtype=dtype)
28
+ for index in range(len(layer_sizes) - 1)
29
+ )
30
+ for linear in self.linears:
31
+ nn.init.trunc_normal_(
32
+ linear.weight,
33
+ mean=0.0,
34
+ std=stddev,
35
+ a=-2.0 * stddev,
36
+ b=2.0 * stddev,
37
+ )
38
+ nn.init.zeros_(linear.bias)
39
+
40
+ def _activate(self, values: torch.Tensor) -> torch.Tensor:
41
+ if self.activation == "gelu":
42
+ return torch.nn.functional.gelu(values)
43
+ return getattr(torch, self.activation)(values)
44
+
45
+ def forward(self, inputs: torch.Tensor) -> torch.Tensor:
46
+ hidden = inputs
47
+ for linear in self.linears[:-1]:
48
+ hidden = self._activate(linear(hidden))
49
+ return self.linears[-1](hidden)
50
+
51
+
52
+ class AttentionWeights(nn.Module):
53
+ """Positive, normalized, per-point self-adaptive loss weights."""
54
+
55
+ def __init__(
56
+ self,
57
+ n_points: int,
58
+ dtype: torch.dtype = torch.float64,
59
+ initial_value: float = 0.0,
60
+ ) -> None:
61
+ super().__init__()
62
+ if n_points <= 0:
63
+ raise ValueError("attention requires a positive point count")
64
+ self.alpha = nn.Parameter(
65
+ torch.full((n_points, 1), initial_value, dtype=dtype)
66
+ )
67
+
68
+ def forward(self) -> torch.Tensor:
69
+ shifted = self.alpha - self.alpha.detach().max()
70
+ weights = torch.exp(shifted)
71
+ weights = weights / (weights.mean() + torch.finfo(weights.dtype).eps)
72
+ return torch.clamp(weights, min=1.0e-2, max=100.0)
73
+
74
+
75
+ class SAPINN(nn.Module):
76
+ """Self-adaptive PINN with separate attention weights for each loss group."""
77
+
78
+ def __init__(
79
+ self,
80
+ input_dim: int,
81
+ n_layers: int,
82
+ n_neurons: int,
83
+ activation: str,
84
+ n_pde: int,
85
+ n_data: int,
86
+ n_boundary: int,
87
+ attention_enabled: bool = True,
88
+ dtype: torch.dtype = torch.float64,
89
+ ) -> None:
90
+ super().__init__()
91
+ if min(input_dim, n_layers, n_neurons) <= 0:
92
+ raise ValueError("model dimensions must be positive")
93
+ self.net = FCN(
94
+ [input_dim, *([n_neurons] * n_layers), 1],
95
+ activation=activation,
96
+ dtype=dtype,
97
+ )
98
+ self.att_pde = (
99
+ AttentionWeights(n_pde, dtype=dtype) if attention_enabled and n_pde else None
100
+ )
101
+ self.att_data = (
102
+ AttentionWeights(n_data, dtype=dtype) if attention_enabled and n_data else None
103
+ )
104
+ self.att_boundary = (
105
+ AttentionWeights(n_boundary, dtype=dtype)
106
+ if attention_enabled and n_boundary
107
+ else None
108
+ )
109
+
110
+ def forward(self, inputs: torch.Tensor) -> torch.Tensor:
111
+ return self.net(inputs)
112
+
113
+ def network_parameters(self):
114
+ return self.net.parameters()
115
+
116
+ def attention_parameters(self):
117
+ for attention in (self.att_pde, self.att_data, self.att_boundary):
118
+ if attention is not None:
119
+ yield from attention.parameters()
120
+
121
+ def set_attention_trainable(self, trainable: bool) -> None:
122
+ for parameter in self.attention_parameters():
123
+ parameter.requires_grad_(trainable)
124
+
125
+
126
+ def build_model(
127
+ model_config: Mapping,
128
+ point_counts: Mapping[str, int],
129
+ attention_enabled: bool,
130
+ dtype: torch.dtype,
131
+ ) -> SAPINN:
132
+ return SAPINN(
133
+ input_dim=int(model_config["input_dim"]),
134
+ n_layers=int(model_config["n_layers"]),
135
+ n_neurons=int(model_config["n_neurons"]),
136
+ activation=str(model_config["activation"]),
137
+ n_pde=int(point_counts["pde"]),
138
+ n_data=int(point_counts["data"]),
139
+ n_boundary=int(point_counts["boundary"]),
140
+ attention_enabled=attention_enabled,
141
+ dtype=dtype,
142
+ )
143
+
144
+
145
+ class Equation:
146
+ def residual(
147
+ self, model: SAPINN, coordinates: torch.Tensor
148
+ ) -> torch.Tensor:
149
+ raise NotImplementedError
150
+
151
+
152
+ class Laplace1D(Equation):
153
+ def __init__(self, diffusion: float = 1.0) -> None:
154
+ self.diffusion = diffusion
155
+
156
+ def residual(self, model: SAPINN, coordinates: torch.Tensor) -> torch.Tensor:
157
+ inputs = coordinates.detach().requires_grad_(True)
158
+ prediction = model(inputs)
159
+ first = torch.autograd.grad(
160
+ prediction, inputs, torch.ones_like(prediction), create_graph=True
161
+ )[0]
162
+ second = torch.autograd.grad(
163
+ first, inputs, torch.ones_like(first), create_graph=True
164
+ )[0]
165
+ source = torch.pi**2 * torch.sin(torch.pi * inputs)
166
+ return self.diffusion * second + source
167
+
168
+
169
+ class Helmholtz2D(Equation):
170
+ def __init__(self, wave_number: float = 1.0) -> None:
171
+ self.wave_number = wave_number
172
+
173
+ def residual(self, model: SAPINN, coordinates: torch.Tensor) -> torch.Tensor:
174
+ inputs = coordinates.detach().requires_grad_(True)
175
+ prediction = model(inputs)
176
+ gradient = torch.autograd.grad(
177
+ prediction, inputs, torch.ones_like(prediction), create_graph=True
178
+ )[0]
179
+ prediction_xx = torch.autograd.grad(
180
+ gradient[:, 0:1],
181
+ inputs,
182
+ torch.ones_like(gradient[:, 0:1]),
183
+ create_graph=True,
184
+ )[0][:, 0:1]
185
+ prediction_yy = torch.autograd.grad(
186
+ gradient[:, 1:2],
187
+ inputs,
188
+ torch.ones_like(gradient[:, 1:2]),
189
+ create_graph=True,
190
+ )[0][:, 1:2]
191
+ exact = torch.sin(torch.pi * inputs[:, 0:1]) * torch.sin(
192
+ 4.0 * torch.pi * inputs[:, 1:2]
193
+ )
194
+ forcing = (
195
+ -torch.pi**2 - (4.0 * torch.pi) ** 2 + self.wave_number**2
196
+ ) * exact
197
+ return (
198
+ prediction_xx
199
+ + prediction_yy
200
+ + self.wave_number**2 * prediction
201
+ - forcing
202
+ )
203
+
204
+
205
+ class Burgers2D(Equation):
206
+ def __init__(self, viscosity: float = 0.01 / math.pi) -> None:
207
+ self.viscosity = viscosity
208
+
209
+ def residual(self, model: SAPINN, coordinates: torch.Tensor) -> torch.Tensor:
210
+ inputs = coordinates.detach().requires_grad_(True)
211
+ prediction = model(inputs)
212
+ gradient = torch.autograd.grad(
213
+ prediction, inputs, torch.ones_like(prediction), create_graph=True
214
+ )[0]
215
+ prediction_x = gradient[:, 0:1]
216
+ prediction_t = gradient[:, 1:2]
217
+ prediction_xx = torch.autograd.grad(
218
+ prediction_x,
219
+ inputs,
220
+ torch.ones_like(prediction_x),
221
+ create_graph=True,
222
+ )[0][:, 0:1]
223
+ return (
224
+ prediction_t
225
+ + prediction * prediction_x
226
+ - self.viscosity * prediction_xx
227
+ )
228
+
229
+
230
+ def weighted_mean_square(
231
+ residual: torch.Tensor, attention: AttentionWeights | None
232
+ ) -> torch.Tensor:
233
+ if attention is None:
234
+ return torch.mean(residual.square())
235
+ weights = attention()
236
+ if weights.shape != residual.shape:
237
+ raise ValueError(
238
+ f"attention shape {tuple(weights.shape)} does not match residual {tuple(residual.shape)}"
239
+ )
240
+ return torch.mean(weights * residual.square())
241
+
242
+
243
+ def loss_components(
244
+ model: SAPINN,
245
+ equation: Equation,
246
+ tensors: Mapping[str, torch.Tensor],
247
+ ) -> dict[str, torch.Tensor]:
248
+ pde_residual = equation.residual(model, tensors["x_pde"])
249
+ boundary_residual = model(tensors["x_boundary"]) - tensors["u_boundary"]
250
+ if tensors.get("x_data") is not None:
251
+ data_residual = model(tensors["x_data"]) - tensors["u_data"]
252
+ data_loss = weighted_mean_square(data_residual, model.att_data)
253
+ else:
254
+ data_loss = torch.zeros((), dtype=pde_residual.dtype, device=pde_residual.device)
255
+ return {
256
+ "data": data_loss,
257
+ "boundary": weighted_mean_square(boundary_residual, model.att_boundary),
258
+ "pde": weighted_mean_square(pde_residual, model.att_pde),
259
+ }
260
+
261
+
262
+ def weighted_loss(components: Mapping[str, torch.Tensor], weights: Mapping) -> torch.Tensor:
263
+ return (
264
+ float(weights["data"]) * components["data"]
265
+ + float(weights["boundary"]) * components["boundary"]
266
+ + float(weights["pde"]) * components["pde"]
267
+ )
result/sapinn_burgers_history.npz ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:02eaf5639d4c2439486ac0bbc22c216551a9c607d07dae3eed225b8484174299
3
+ size 1796
result/sapinn_helmholtz_history.npz ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:f6d505c41e5372afe8296d3ba3fde39c8a668cfb23a860bb51dfa2f73fa185eb
3
+ size 76130
result/sapinn_laplace_history.npz ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:ec6972364605b30147a566d6348f3fe6dd20d46c12a6749c79558978939ba38c
3
+ size 76445
scripts/__pycache__/common.cpython-311.pyc ADDED
Binary file (4.22 kB). View file
 
scripts/__pycache__/problems.cpython-311.pyc ADDED
Binary file (10.6 kB). View file
 
scripts/common.py ADDED
@@ -0,0 +1,54 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from __future__ import annotations
2
+
3
+ import random
4
+ from collections.abc import Mapping
5
+ from pathlib import Path
6
+
7
+ import numpy as np
8
+ import torch
9
+ import yaml
10
+
11
+
12
+ def load_config(path: Path) -> dict:
13
+ with path.open("r", encoding="utf-8") as stream:
14
+ config = yaml.safe_load(stream)
15
+ if not isinstance(config, dict) or "root" not in config:
16
+ raise ValueError(f"config must contain a 'root' mapping: {path}")
17
+ return config["root"]
18
+
19
+
20
+ def project_path(value: str | Path, project_root: Path) -> Path:
21
+ path = Path(value).expanduser()
22
+ return path if path.is_absolute() else project_root / path
23
+
24
+
25
+ def resolve_device(requested: str) -> torch.device:
26
+ if requested == "auto":
27
+ return torch.device("cuda" if torch.cuda.is_available() else "cpu")
28
+ device = torch.device(requested)
29
+ if device.type == "cuda" and not torch.cuda.is_available():
30
+ raise RuntimeError("CUDA/DCU was requested but torch.cuda.is_available() is false")
31
+ return device
32
+
33
+
34
+ def resolve_dtype(name: str) -> torch.dtype:
35
+ try:
36
+ return {"float32": torch.float32, "float64": torch.float64}[name]
37
+ except KeyError as error:
38
+ raise ValueError(f"unsupported dtype: {name}") from error
39
+
40
+
41
+ def seed_everything(seed: int) -> None:
42
+ random.seed(seed)
43
+ np.random.seed(seed)
44
+ torch.manual_seed(seed)
45
+ if torch.cuda.is_available():
46
+ torch.cuda.manual_seed_all(seed)
47
+
48
+
49
+ def checkpoint_state(checkpoint: Mapping) -> tuple[Mapping[str, torch.Tensor], dict]:
50
+ if "model_state" in checkpoint:
51
+ return checkpoint["model_state"], dict(checkpoint)
52
+ if checkpoint and all(torch.is_tensor(value) for value in checkpoint.values()):
53
+ return checkpoint, {}
54
+ raise ValueError("checkpoint contains no valid SA-PINN model state")
scripts/inference.py ADDED
@@ -0,0 +1,241 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from __future__ import annotations
2
+
3
+ import argparse
4
+ import os
5
+ import sys
6
+ from collections.abc import Mapping
7
+ from pathlib import Path
8
+
9
+ import matplotlib
10
+
11
+ matplotlib.use("Agg")
12
+ import matplotlib.pyplot as plt
13
+ import numpy as np
14
+ import torch
15
+
16
+
17
+ PROJECT_ROOT = Path(__file__).resolve().parents[1]
18
+ sys.path.insert(0, str(PROJECT_ROOT))
19
+ sys.path.insert(0, str(Path(__file__).resolve().parent))
20
+
21
+ from common import ( # noqa: E402
22
+ checkpoint_state,
23
+ load_config,
24
+ project_path,
25
+ resolve_device,
26
+ resolve_dtype,
27
+ )
28
+ from model.sa_pinn import build_model # noqa: E402
29
+ from problems import CASES, generate_data, point_counts, relative_l2, to_tensors # noqa: E402
30
+
31
+
32
+ DEFAULT_CONFIG = PROJECT_ROOT / "conf" / "config.yaml"
33
+
34
+
35
+ def parse_args() -> argparse.Namespace:
36
+ parser = argparse.ArgumentParser(description="Run SA-PINNs inference and visualization")
37
+ parser.add_argument(
38
+ "--config",
39
+ type=Path,
40
+ default=Path(os.environ.get("SAPINN_CONFIG", DEFAULT_CONFIG)),
41
+ help="YAML configuration file",
42
+ )
43
+ parser.add_argument("--case", choices=CASES, default="laplace")
44
+ parser.add_argument("--device", help="Override common.device")
45
+ parser.add_argument("--checkpoint", type=Path, help="Checkpoint path")
46
+ parser.add_argument("--weight-dir", type=Path, help="Override weight directory")
47
+ parser.add_argument("--result-dir", type=Path, help="Override result directory")
48
+ parser.add_argument("--test-res", type=int, help="Override test resolution per axis")
49
+ return parser.parse_args()
50
+
51
+
52
+ def save_figure(
53
+ case: str,
54
+ data: Mapping,
55
+ prediction: np.ndarray,
56
+ attention: np.ndarray | None,
57
+ output_path: Path,
58
+ error: float | None,
59
+ ) -> None:
60
+ if case == "laplace":
61
+ exact = data["u_exact"].reshape(-1)
62
+ predicted = prediction.reshape(-1)
63
+ x_axis = data["x_test"].reshape(-1)
64
+ figure, axes = plt.subplots(1, 3, figsize=(15, 4))
65
+ axes[0].plot(x_axis, exact, "k-", label="Exact")
66
+ axes[0].plot(x_axis, predicted, "r--", label="SA-PINN")
67
+ axes[0].set_xlabel("x")
68
+ axes[0].set_ylabel("u")
69
+ axes[0].legend()
70
+ axes[1].semilogy(
71
+ x_axis, np.maximum(np.abs(predicted - exact), 1.0e-16), "r-"
72
+ )
73
+ axes[1].set_xlabel("x")
74
+ axes[1].set_ylabel("Absolute error")
75
+ if attention is not None:
76
+ order = np.argsort(data["x_pde"].reshape(-1))
77
+ axes[2].plot(
78
+ data["x_pde"].reshape(-1)[order],
79
+ attention.reshape(-1)[order],
80
+ "b-",
81
+ )
82
+ axes[2].set_xlabel("x")
83
+ axes[2].set_ylabel("Attention weight")
84
+ else:
85
+ axes[2].text(0.5, 0.5, "Attention disabled", ha="center", va="center")
86
+ axes[2].set_axis_off()
87
+ elif case == "helmholtz":
88
+ shape = data["test_shape"]
89
+ exact = data["u_exact"].reshape(shape)
90
+ predicted = prediction.reshape(shape)
91
+ figure, axes = plt.subplots(1, 3, figsize=(15, 4))
92
+ for axis, title, field in zip(
93
+ axes,
94
+ ("Exact", "SA-PINN", "Absolute error"),
95
+ (exact, predicted, np.abs(predicted - exact)),
96
+ strict=True,
97
+ ):
98
+ image = axis.imshow(field, extent=(-1, 1, -1, 1), origin="lower", cmap="jet")
99
+ axis.set_title(title)
100
+ axis.set_xlabel("x")
101
+ axis.set_ylabel("y")
102
+ figure.colorbar(image, ax=axis)
103
+ else:
104
+ shape = data["test_shape"]
105
+ predicted = prediction.reshape(shape)
106
+ figure, axes = plt.subplots(1, 2, figsize=(12, 4))
107
+ image = axes[0].imshow(
108
+ predicted,
109
+ extent=(-1, 1, 0, 1),
110
+ origin="lower",
111
+ aspect="auto",
112
+ cmap="jet",
113
+ )
114
+ axes[0].set_title("SA-PINN Burgers prediction")
115
+ axes[0].set_xlabel("x")
116
+ axes[0].set_ylabel("t")
117
+ figure.colorbar(image, ax=axes[0])
118
+ if attention is not None:
119
+ scatter = axes[1].scatter(
120
+ data["x_pde"][:, 0],
121
+ data["x_pde"][:, 1],
122
+ c=attention.reshape(-1),
123
+ s=8,
124
+ cmap="viridis",
125
+ )
126
+ axes[1].set_title("PDE attention weights")
127
+ axes[1].set_xlabel("x")
128
+ axes[1].set_ylabel("t")
129
+ figure.colorbar(scatter, ax=axes[1])
130
+ else:
131
+ axes[1].text(0.5, 0.5, "Attention disabled", ha="center", va="center")
132
+ axes[1].set_axis_off()
133
+ figure.suptitle(f"{case} | L2={'N/A' if error is None else f'{error:.3e}'}")
134
+ figure.tight_layout()
135
+ figure.savefig(output_path, dpi=150)
136
+ plt.close(figure)
137
+
138
+
139
+ def main() -> None:
140
+ args = parse_args()
141
+ config_path = args.config.expanduser().resolve()
142
+ config = load_config(config_path)
143
+ common = config["common"]
144
+ case_config = config["cases"][args.case]
145
+ device = resolve_device(args.device or str(common["device"]))
146
+ dtype = resolve_dtype(str(common["dtype"]))
147
+ weight_dir = project_path(args.weight_dir or common["weight_dir"], PROJECT_ROOT)
148
+ result_dir = project_path(args.result_dir or common["result_dir"], PROJECT_ROOT)
149
+ checkpoint_path = (
150
+ project_path(args.checkpoint, PROJECT_ROOT)
151
+ if args.checkpoint
152
+ else weight_dir / case_config["output"]["checkpoint_name"]
153
+ )
154
+ if not checkpoint_path.is_file():
155
+ raise FileNotFoundError(
156
+ f"checkpoint not found: {checkpoint_path}. Run scripts/train.py first."
157
+ )
158
+ checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=True)
159
+ if not isinstance(checkpoint, Mapping):
160
+ raise ValueError(f"invalid SA-PINN checkpoint: {checkpoint_path}")
161
+ state, metadata = checkpoint_state(checkpoint)
162
+ checkpoint_case = metadata.get("case", args.case)
163
+ if checkpoint_case != args.case:
164
+ raise ValueError(
165
+ f"checkpoint case is {checkpoint_case}, but --case requested {args.case}"
166
+ )
167
+ model_config = metadata.get("model_config", case_config["model"])
168
+ data_config = dict(metadata.get("data_config", case_config["data"]))
169
+ if args.test_res is not None:
170
+ data_config["test_res"] = args.test_res
171
+ seed = int(metadata.get("seed", common["seed"]))
172
+ stored_counts = metadata.get("point_counts")
173
+ if stored_counts is None:
174
+ stored_counts = {
175
+ "pde": int(state["att_pde.alpha"].shape[0])
176
+ if "att_pde.alpha" in state
177
+ else int(data_config["n_pde"]),
178
+ "data": int(state["att_data.alpha"].shape[0])
179
+ if "att_data.alpha" in state
180
+ else int(data_config["n_sol"]),
181
+ "boundary": int(state["att_boundary.alpha"].shape[0])
182
+ if "att_boundary.alpha" in state
183
+ else int(data_config["n_bnd"]),
184
+ }
185
+ data_config.update(
186
+ {
187
+ "n_pde": stored_counts["pde"],
188
+ "n_sol": stored_counts["data"],
189
+ "n_bnd": stored_counts["boundary"],
190
+ }
191
+ )
192
+ data = generate_data(args.case, data_config, seed)
193
+ counts = stored_counts or point_counts(data)
194
+ attention_enabled = bool(
195
+ metadata.get(
196
+ "attention_enabled", any(key.startswith("att_") for key in state)
197
+ )
198
+ )
199
+ model = build_model(model_config, counts, attention_enabled, dtype).to(
200
+ device=device, dtype=dtype
201
+ )
202
+ model.load_state_dict(state, strict=True)
203
+ model.eval()
204
+ tensors = to_tensors(data, device, dtype)
205
+ with torch.no_grad():
206
+ prediction = model(tensors["x_test"]).cpu().numpy()
207
+ error = relative_l2(prediction, data["u_exact"])
208
+ attention = (
209
+ model.att_pde().detach().cpu().numpy() if model.att_pde is not None else None
210
+ )
211
+
212
+ result_dir.mkdir(parents=True, exist_ok=True)
213
+ predictions_path = result_dir / case_config["output"]["predictions_name"]
214
+ figure_path = result_dir / case_config["output"]["figure_name"]
215
+ np.savez_compressed(
216
+ predictions_path,
217
+ coordinates=data["x_test"],
218
+ prediction=prediction,
219
+ exact=np.array([]) if data["u_exact"] is None else data["u_exact"],
220
+ relative_l2=np.nan if error is None else error,
221
+ pde_attention=np.array([]) if attention is None else attention,
222
+ )
223
+ save_figure(args.case, data, prediction, attention, figure_path, error)
224
+
225
+ print(f"Config: {config_path}")
226
+ print(f"Case: {args.case}")
227
+ print(f"Checkpoint: {checkpoint_path}")
228
+ print(f"Device: {device}")
229
+ print(f"Relative L2={'N/A' if error is None else f'{error:.6e}'}")
230
+ if attention is not None:
231
+ print(
232
+ f"PDE attention: mean={attention.mean():.6f}, "
233
+ f"std={attention.std():.6f}, min={attention.min():.6f}, "
234
+ f"max={attention.max():.6f}"
235
+ )
236
+ print(f"Predictions: {predictions_path}")
237
+ print(f"Plot: {figure_path}")
238
+
239
+
240
+ if __name__ == "__main__":
241
+ main()
scripts/problems.py ADDED
@@ -0,0 +1,167 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from __future__ import annotations
2
+
3
+ from collections.abc import Mapping
4
+
5
+ import numpy as np
6
+ import torch
7
+
8
+ from model.sa_pinn import Burgers2D, Equation, Helmholtz2D, Laplace1D
9
+
10
+
11
+ CASES = ("laplace", "helmholtz", "burgers")
12
+
13
+
14
+ def exact_solution(case: str, coordinates: np.ndarray) -> np.ndarray | None:
15
+ if case == "laplace":
16
+ return np.sin(np.pi * coordinates[:, 0:1])
17
+ if case == "helmholtz":
18
+ return np.sin(np.pi * coordinates[:, 0:1]) * np.sin(
19
+ 4.0 * np.pi * coordinates[:, 1:2]
20
+ )
21
+ if case == "burgers":
22
+ return None
23
+ raise ValueError(f"unsupported case: {case}")
24
+
25
+
26
+ def build_equation(case: str, data_config: Mapping) -> Equation:
27
+ if case == "laplace":
28
+ return Laplace1D()
29
+ if case == "helmholtz":
30
+ return Helmholtz2D(float(data_config["wave_number"]))
31
+ if case == "burgers":
32
+ return Burgers2D(float(data_config["viscosity"]))
33
+ raise ValueError(f"unsupported case: {case}")
34
+
35
+
36
+ def generate_data(case: str, config: Mapping, seed: int) -> dict[str, np.ndarray | None]:
37
+ generator = np.random.default_rng(seed)
38
+ n_pde = int(config["n_pde"])
39
+ n_solution = int(config["n_sol"])
40
+ n_boundary = int(config["n_bnd"])
41
+ test_resolution = int(config["test_res"])
42
+ if min(n_pde, n_solution, n_boundary, test_resolution) <= 0:
43
+ raise ValueError("all data counts must be positive")
44
+
45
+ if case == "laplace":
46
+ lower, upper = map(float, config["domain"])
47
+ x_pde = generator.uniform(lower, upper, (n_pde, 1))
48
+ x_boundary = np.array([[lower], [upper]], dtype=np.float64)
49
+ u_boundary = exact_solution(case, x_boundary)
50
+ x_data = np.linspace(lower, upper, n_solution)[:, None]
51
+ u_data = exact_solution(case, x_data)
52
+ noise_std = float(config.get("noise_std", 0.0))
53
+ if noise_std:
54
+ u_data = u_data + noise_std * generator.standard_normal(u_data.shape)
55
+ x_test = np.linspace(lower, upper, test_resolution)[:, None]
56
+ test_shape = (test_resolution,)
57
+ elif case == "helmholtz":
58
+ x_lower, x_upper = map(float, config["domain_x"])
59
+ y_lower, y_upper = map(float, config["domain_y"])
60
+ x_pde = np.column_stack(
61
+ (
62
+ generator.uniform(x_lower, x_upper, n_pde),
63
+ generator.uniform(y_lower, y_upper, n_pde),
64
+ )
65
+ )
66
+ if n_boundary < 4 or n_boundary % 4:
67
+ raise ValueError("Helmholtz n_bnd must be divisible by 4 and at least 4")
68
+ per_edge = n_boundary // 4
69
+ x_axis = np.linspace(x_lower, x_upper, per_edge)
70
+ y_axis = np.linspace(y_lower, y_upper, per_edge)
71
+ x_boundary = np.vstack(
72
+ (
73
+ np.column_stack((x_axis, np.full(per_edge, y_lower))),
74
+ np.column_stack((x_axis, np.full(per_edge, y_upper))),
75
+ np.column_stack((np.full(per_edge, x_lower), y_axis)),
76
+ np.column_stack((np.full(per_edge, x_upper), y_axis)),
77
+ )
78
+ )
79
+ u_boundary = exact_solution(case, x_boundary)
80
+ x_data = np.column_stack(
81
+ (
82
+ generator.uniform(x_lower, x_upper, n_solution),
83
+ generator.uniform(y_lower, y_upper, n_solution),
84
+ )
85
+ )
86
+ u_data = exact_solution(case, x_data)
87
+ grid_x = np.linspace(x_lower, x_upper, test_resolution)
88
+ grid_y = np.linspace(y_lower, y_upper, test_resolution)
89
+ mesh_x, mesh_y = np.meshgrid(grid_x, grid_y, indexing="xy")
90
+ x_test = np.column_stack((mesh_x.ravel(), mesh_y.ravel()))
91
+ test_shape = mesh_x.shape
92
+ elif case == "burgers":
93
+ x_lower, x_upper = map(float, config["domain_x"])
94
+ t_lower, t_upper = map(float, config["domain_t"])
95
+ x_pde = np.column_stack(
96
+ (
97
+ generator.uniform(x_lower, x_upper, n_pde),
98
+ generator.uniform(t_lower, t_upper, n_pde),
99
+ )
100
+ )
101
+ if n_boundary < 2:
102
+ raise ValueError("Burgers n_bnd must be at least 2")
103
+ left_count = n_boundary // 2
104
+ right_count = n_boundary - left_count
105
+ left_time = generator.uniform(t_lower, t_upper, left_count)
106
+ right_time = generator.uniform(t_lower, t_upper, right_count)
107
+ x_boundary = np.vstack(
108
+ (
109
+ np.column_stack((np.full(left_count, x_lower), left_time)),
110
+ np.column_stack((np.full(right_count, x_upper), right_time)),
111
+ )
112
+ )
113
+ u_boundary = np.zeros((n_boundary, 1))
114
+ initial_x = generator.uniform(x_lower, x_upper, n_solution)
115
+ x_data = np.column_stack((initial_x, np.full(n_solution, t_lower)))
116
+ u_data = -np.sin(np.pi * initial_x)[:, None]
117
+ grid_x = np.linspace(x_lower, x_upper, test_resolution)
118
+ grid_t = np.linspace(t_lower, t_upper, test_resolution)
119
+ mesh_x, mesh_t = np.meshgrid(grid_x, grid_t, indexing="xy")
120
+ x_test = np.column_stack((mesh_x.ravel(), mesh_t.ravel()))
121
+ test_shape = mesh_x.shape
122
+ else:
123
+ raise ValueError(f"unsupported case: {case}")
124
+
125
+ return {
126
+ "x_pde": np.asarray(x_pde, dtype=np.float64),
127
+ "x_boundary": np.asarray(x_boundary, dtype=np.float64),
128
+ "u_boundary": np.asarray(u_boundary, dtype=np.float64),
129
+ "x_data": np.asarray(x_data, dtype=np.float64),
130
+ "u_data": np.asarray(u_data, dtype=np.float64),
131
+ "x_test": np.asarray(x_test, dtype=np.float64),
132
+ "u_exact": exact_solution(case, np.asarray(x_test, dtype=np.float64)),
133
+ "test_shape": test_shape,
134
+ }
135
+
136
+
137
+ def to_tensors(
138
+ data: Mapping[str, np.ndarray | None],
139
+ device: torch.device,
140
+ dtype: torch.dtype,
141
+ ) -> dict[str, torch.Tensor | None]:
142
+ return {
143
+ key: (
144
+ torch.as_tensor(value, dtype=dtype, device=device)
145
+ if isinstance(value, np.ndarray)
146
+ else None
147
+ )
148
+ for key, value in data.items()
149
+ if key != "test_shape"
150
+ }
151
+
152
+
153
+ def point_counts(data: Mapping[str, np.ndarray | None]) -> dict[str, int]:
154
+ return {
155
+ "pde": int(data["x_pde"].shape[0]),
156
+ "boundary": int(data["x_boundary"].shape[0]),
157
+ "data": int(data["x_data"].shape[0]) if data.get("x_data") is not None else 0,
158
+ }
159
+
160
+
161
+ def relative_l2(prediction: np.ndarray, reference: np.ndarray | None) -> float | None:
162
+ if reference is None:
163
+ return None
164
+ return float(
165
+ np.linalg.norm(prediction.reshape(-1) - reference.reshape(-1))
166
+ / (np.linalg.norm(reference.reshape(-1)) + 1.0e-12)
167
+ )
scripts/train.py ADDED
@@ -0,0 +1,234 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
+
12
+
13
+ PROJECT_ROOT = Path(__file__).resolve().parents[1]
14
+ sys.path.insert(0, str(PROJECT_ROOT))
15
+ sys.path.insert(0, str(Path(__file__).resolve().parent))
16
+
17
+ from common import ( # noqa: E402
18
+ load_config,
19
+ project_path,
20
+ resolve_device,
21
+ resolve_dtype,
22
+ seed_everything,
23
+ )
24
+ from model.sa_pinn import build_model, loss_components, weighted_loss # noqa: E402
25
+ from problems import ( # noqa: E402
26
+ CASES,
27
+ build_equation,
28
+ generate_data,
29
+ point_counts,
30
+ relative_l2,
31
+ to_tensors,
32
+ )
33
+
34
+
35
+ DEFAULT_CONFIG = PROJECT_ROOT / "conf" / "config.yaml"
36
+
37
+
38
+ def parse_args() -> argparse.Namespace:
39
+ parser = argparse.ArgumentParser(description="Train a self-adaptive PINN")
40
+ parser.add_argument(
41
+ "--config",
42
+ type=Path,
43
+ default=Path(os.environ.get("SAPINN_CONFIG", DEFAULT_CONFIG)),
44
+ help="YAML configuration file",
45
+ )
46
+ parser.add_argument("--case", choices=CASES, default="laplace")
47
+ parser.add_argument("--device", help="Override common.device, for example cpu or cuda:0")
48
+ parser.add_argument("--seed", type=int, help="Override common.seed")
49
+ parser.add_argument("--epochs", type=int, help="Override Adam iterations")
50
+ parser.add_argument("--lr", type=float, help="Override network learning rate")
51
+ parser.add_argument("--attention-lr", type=float, help="Override attention learning rate")
52
+ parser.add_argument("--lbfgs-iters", type=int, help="Override L-BFGS iterations")
53
+ parser.add_argument("--n-sol", type=int, help="Override data or initial-condition points")
54
+ parser.add_argument("--n-pde", type=int, help="Override PDE collocation points")
55
+ parser.add_argument("--n-bnd", type=int, help="Override boundary points")
56
+ parser.add_argument("--test-res", type=int, help="Override test resolution per axis")
57
+ parser.add_argument("--weight-dir", type=Path, help="Override checkpoint output directory")
58
+ parser.add_argument("--result-dir", type=Path, help="Override history output directory")
59
+ return parser.parse_args()
60
+
61
+
62
+ def main() -> None:
63
+ args = parse_args()
64
+ config_path = args.config.expanduser().resolve()
65
+ config = load_config(config_path)
66
+ common = config["common"]
67
+ case_config = config["cases"][args.case]
68
+ device = resolve_device(args.device or str(common["device"]))
69
+ dtype = resolve_dtype(str(common["dtype"]))
70
+ seed = args.seed if args.seed is not None else int(common["seed"])
71
+ epochs = (
72
+ args.epochs if args.epochs is not None else int(case_config["training"]["epochs"])
73
+ )
74
+ lbfgs_iters = (
75
+ args.lbfgs_iters
76
+ if args.lbfgs_iters is not None
77
+ else int(case_config["training"]["lbfgs_iters"])
78
+ )
79
+ network_lr = (
80
+ args.lr if args.lr is not None else float(case_config["training"]["lr"])
81
+ )
82
+ attention_lr = (
83
+ args.attention_lr
84
+ if args.attention_lr is not None
85
+ else float(case_config["attention"]["lr"])
86
+ )
87
+ if min(epochs, lbfgs_iters) < 0 or epochs + lbfgs_iters == 0:
88
+ raise ValueError("at least one optimizer iteration count must be positive")
89
+ if min(network_lr, attention_lr) <= 0:
90
+ raise ValueError("network and attention learning rates must be positive")
91
+
92
+ data_config = dict(case_config["data"])
93
+ for key, value in (
94
+ ("n_sol", args.n_sol),
95
+ ("n_pde", args.n_pde),
96
+ ("n_bnd", args.n_bnd),
97
+ ("test_res", args.test_res),
98
+ ):
99
+ if value is not None:
100
+ data_config[key] = value
101
+ weight_dir = project_path(args.weight_dir or common["weight_dir"], PROJECT_ROOT)
102
+ result_dir = project_path(args.result_dir or common["result_dir"], PROJECT_ROOT)
103
+ checkpoint_path = weight_dir / case_config["output"]["checkpoint_name"]
104
+ history_path = result_dir / case_config["output"]["history_name"]
105
+ seed_everything(seed)
106
+
107
+ data = generate_data(args.case, data_config, seed)
108
+ counts = point_counts(data)
109
+ tensors = to_tensors(data, device, dtype)
110
+ attention_enabled = bool(case_config["attention"]["enabled"])
111
+ model = build_model(
112
+ case_config["model"], counts, attention_enabled, dtype
113
+ ).to(device=device, dtype=dtype)
114
+ equation = build_equation(args.case, data_config)
115
+ parameter_count = sum(parameter.numel() for parameter in model.parameters())
116
+
117
+ print(f"Config: {config_path}")
118
+ print(f"Case: {args.case}")
119
+ print(f"Device: {device}")
120
+ print(f"Points: {counts}")
121
+ print(f"Parameters: {parameter_count:,}")
122
+
123
+ def loss_value() -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
124
+ components = loss_components(model, equation, tensors)
125
+ return weighted_loss(components, case_config["loss"]), components
126
+
127
+ def evaluate() -> float | None:
128
+ with torch.no_grad():
129
+ prediction = model(tensors["x_test"]).cpu().numpy()
130
+ return relative_l2(prediction, data["u_exact"])
131
+
132
+ network_optimizer = torch.optim.Adam(model.network_parameters(), lr=network_lr)
133
+ attention_parameters = list(model.attention_parameters())
134
+ attention_optimizer = (
135
+ torch.optim.Adam(attention_parameters, lr=attention_lr)
136
+ if attention_parameters
137
+ else None
138
+ )
139
+ history = []
140
+ l2_history = []
141
+ log_interval = int(case_config["training"]["log_interval"])
142
+ started = time.time()
143
+ for epoch in range(1, epochs + 1):
144
+ network_optimizer.zero_grad(set_to_none=True)
145
+ if attention_optimizer is not None:
146
+ attention_optimizer.zero_grad(set_to_none=True)
147
+ loss, components = loss_value()
148
+ if not torch.isfinite(loss):
149
+ raise FloatingPointError(f"SA-PINN loss became non-finite at epoch {epoch}")
150
+ loss.backward()
151
+ network_optimizer.step()
152
+ if attention_optimizer is not None:
153
+ for parameter in attention_parameters:
154
+ if parameter.grad is not None:
155
+ parameter.grad.neg_()
156
+ attention_optimizer.step()
157
+ history.append(loss.item())
158
+
159
+ if epoch == 1 or epoch % log_interval == 0 or epoch == epochs:
160
+ error = evaluate()
161
+ if error is not None:
162
+ l2_history.append((epoch, error))
163
+ attention = model.att_pde()
164
+ error_text = "N/A" if error is None else f"{error:.3e}"
165
+ print(
166
+ f"epoch={epoch:6d} loss={loss.item():.3e} "
167
+ f"data={components['data'].item():.3e} "
168
+ f"boundary={components['boundary'].item():.3e} "
169
+ f"pde={components['pde'].item():.3e} l2={error_text} "
170
+ f"attention_std={attention.std(unbiased=False).item():.3e}"
171
+ )
172
+
173
+ if lbfgs_iters > 0:
174
+ model.set_attention_trainable(False)
175
+ lbfgs = torch.optim.LBFGS(
176
+ model.network_parameters(),
177
+ lr=0.8,
178
+ max_iter=lbfgs_iters,
179
+ max_eval=max(1, 2 * lbfgs_iters),
180
+ tolerance_grad=1.0e-7,
181
+ tolerance_change=1.0e-9,
182
+ history_size=50,
183
+ line_search_fn="strong_wolfe",
184
+ )
185
+
186
+ def closure() -> torch.Tensor:
187
+ lbfgs.zero_grad(set_to_none=True)
188
+ closure_loss, _ = loss_value()
189
+ if not torch.isfinite(closure_loss):
190
+ raise FloatingPointError("SA-PINN L-BFGS loss became non-finite")
191
+ closure_loss.backward()
192
+ return closure_loss
193
+
194
+ lbfgs.step(closure)
195
+ model.set_attention_trainable(True)
196
+ error = evaluate()
197
+ if error is not None:
198
+ l2_history.append((epochs + lbfgs_iters, error))
199
+ print(f"L-BFGS L2={'N/A' if error is None else f'{error:.6e}'}")
200
+
201
+ final_l2 = evaluate()
202
+ elapsed = time.time() - started
203
+ weight_dir.mkdir(parents=True, exist_ok=True)
204
+ result_dir.mkdir(parents=True, exist_ok=True)
205
+ checkpoint = {
206
+ "case": args.case,
207
+ "architecture": "sa_pinn",
208
+ "model_state": model.state_dict(),
209
+ "model_config": case_config["model"],
210
+ "attention_enabled": attention_enabled,
211
+ "point_counts": counts,
212
+ "data_config": data_config,
213
+ "seed": seed,
214
+ "epochs": epochs,
215
+ "lbfgs_iters": lbfgs_iters,
216
+ "final_l2": final_l2,
217
+ }
218
+ torch.save(checkpoint, checkpoint_path)
219
+ l2_array = np.asarray(l2_history, dtype=np.float64).reshape(-1, 2)
220
+ np.savez_compressed(
221
+ history_path,
222
+ loss=np.asarray(history, dtype=np.float64),
223
+ l2_steps=l2_array[:, 0] if l2_array.size else np.array([]),
224
+ l2_values=l2_array[:, 1] if l2_array.size else np.array([]),
225
+ elapsed_seconds=elapsed,
226
+ final_l2=np.nan if final_l2 is None else final_l2,
227
+ )
228
+ print(f"Final L2={'N/A' if final_l2 is None else f'{final_l2:.6e}'}")
229
+ print(f"Saved checkpoint: {checkpoint_path}")
230
+ print(f"Saved history: {history_path}")
231
+
232
+
233
+ if __name__ == "__main__":
234
+ main()
weight/.gitkeep ADDED
File without changes
weight/sapinn_burgers.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:18c0b61281076df8690fc39d21badfff9f73fca810288b23a43ccb0a081d2cfc
3
+ size 112902
weight/sapinn_helmholtz.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:678f2468f00521fe1a470119ce19df6d93c03c44fdc75aff6f7dfce6c554b3fe
3
+ size 152968
weight/sapinn_laplace.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:bf7233b874e92551f4d0fd001ce559fe78bff97e9b39c8d9760c2dcb44bad330
3
+ size 77286