yushuang88 commited on
Commit
6eed390
·
verified ·
1 Parent(s): d5c3cf2

Upload folder using huggingface_hub

Browse files
.gitattributes CHANGED
@@ -1,35 +1,49 @@
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
+ *.npz filter=lfs diff=lfs merge=lfs -text
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
+ *.pt filter=lfs diff=lfs merge=lfs -text
README.md ADDED
@@ -0,0 +1,125 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ frameworks:
3
+ - pytorch
4
+ language:
5
+ - en
6
+ license: apache-2.0
7
+ tags:
8
+ - OneScience
9
+ - hp-VPINN
10
+ - variational-physics-informed-neural-networks
11
+ - partial-differential-equations
12
+ - Poisson
13
+ tasks:
14
+ - pde-solving
15
+ ---
16
+ <p align="center">
17
+ <strong>
18
+ <span style="font-size: 30px;">hp-VPINN</span>
19
+ </strong>
20
+ </p>
21
+
22
+ # Model Overview
23
+
24
+ hp-VPINN (Variational Physics-Informed Neural Network with h- and p-Refinement) trains a neural network using the variational weak form of a partial differential equation. Test functions impose integral constraints on the equation residual, while domain decomposition (h-refinement) and higher-order polynomial test functions (p-refinement) improve the representation and solution of complex fields.
25
+
26
+ This model package provides one- and two-dimensional Poisson equation examples.
27
+
28
+ Paper: hp-VPINNs: Variational physics-informed neural networks with domain decomposition
29
+ https://doi.org/10.1016/j.cma.2020.113547
30
+
31
+ # Model Description
32
+
33
+ hp-VPINN constructs its variational loss with Gauss–Lobatto–Jacobi quadrature and represents the equation solution with a fully connected network. The one-dimensional example supports piecewise integration domains, while the two-dimensional example supports non-square subdomain grids. The number of subdomains and the order of the test functions can be adjusted independently to study the effects of h- and p-refinement on solution quality.
34
+
35
+ # Use Cases
36
+
37
+ | Use Case | Description |
38
+ | :---: | :--- |
39
+ | One-dimensional Poisson equation | Constrain the weak-form residual using piecewise integration domains and higher-order test functions |
40
+ | Two-dimensional Poisson equation | Perform tensor-product quadrature on a two-dimensional domain-decomposition grid |
41
+ | h-refinement research | Vary the number of subdomains to compare spatial decomposition strategies |
42
+ | p-refinement research | Vary the test-function order to compare weak-form constraint capacity |
43
+
44
+ # Usage
45
+
46
+ ## 1. OneCode
47
+
48
+ Use the online OneCode environment for an intelligent, one-click AI for Science (AI4S) programming experience:
49
+
50
+ [Launch OneCode for one-click AI4S programming](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)
51
+
52
+ ## 2. Manual Setup
53
+
54
+ **Hardware Requirements**
55
+
56
+ - A GPU or DCU is recommended for training.
57
+ - A CPU can be used for inference and small-scale pipeline validation, but full training will be slow.
58
+ - DCU users must install DTK and a PyTorch environment compatible with the target cluster.
59
+
60
+ ### Download the Model Package
61
+
62
+ ```bash
63
+ hf download --model OneScience-Sugon/VPINNs --local-dir ./VPINNs
64
+ cd VPINNs
65
+ ```
66
+
67
+ ### Set Up the Runtime Environment
68
+
69
+ **DCU Environment**
70
+
71
+ ```bash
72
+ # Activate DTK and Conda first
73
+ conda create -n onescience311 python=3.11 -y
74
+ conda activate onescience311
75
+ pip install onescience[cfd-dcu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
76
+ ```
77
+
78
+ **GPU Environment**
79
+
80
+ ```bash
81
+ # Activate Conda first
82
+ conda create -n onescience311 python=3.11 -y libstdcxx-ng=12 libgcc-ng=12 gcc_linux-64=12 gxx_linux-64=12
83
+ conda activate onescience311
84
+ pip install onescience[cfd-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
85
+ ```
86
+
87
+ ### Training Data
88
+
89
+ Both the one- and two-dimensional Poisson examples construct source terms, boundary values, and evaluation data from analytical solutions, with no dependency on external data files. The numbers of subdomains, test functions, and quadrature points can be configured in `conf/config.yaml`.
90
+
91
+ ### Training
92
+
93
+ The training example is controlled by `common.case` in `conf/config.yaml`, which supports `1d`, `2d`, and `all`:
94
+
95
+ ```bash
96
+ python scripts/train.py
97
+ ```
98
+
99
+ Training weights are saved to the `weight/` directory by default.
100
+
101
+ ### Model Weights
102
+
103
+ This repository provides weights trained on the one- and two-dimensional Poisson datasets in the `weight/` directory.
104
+
105
+ ### Inference, Evaluation, and Visualization
106
+
107
+ The model package provides one-dimensional Poisson weights for validation. Run:
108
+
109
+ ```bash
110
+ python scripts/inference.py
111
+ ```
112
+
113
+ For the two-dimensional example, first train the corresponding weights, then set `common.case` to `2d` or `all` before running inference. The script reports the relative L2 error and saves result plots to `result/`. Model, training, and inference parameters can all be modified in `conf/config.yaml`.
114
+
115
+ # Official OneScience Resources
116
+
117
+ | Platform | OneScience Repository | Skills Repository |
118
+ | --- | --- | --- |
119
+ | Gitee | https://gitee.com/onescience-ai/onescience | https://gitee.com/onescience-ai/oneskills |
120
+ | GitHub | https://github.com/onescience-ai/OneScience | https://github.com/onescience-ai/oneskills |
121
+
122
+ # Citations and License
123
+
124
+ - Kharazmi, E., Zhang, Z., and Karniadakis, G. E. hp-VPINNs: Variational physics-informed neural networks with domain decomposition. Computer Methods in Applied Mechanics and Engineering, 374, 113547, 2021.
125
+ - 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,31 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ root:
2
+ common:
3
+ case: "1d"
4
+ device: "auto"
5
+ dtype: "float64"
6
+ seed: 1234
7
+ weight_dir: "weight"
8
+ result_dir: "result"
9
+
10
+ poisson1d:
11
+ n_element: 8
12
+ n_test: 40
13
+ n_quad: 60
14
+ epochs: 10000
15
+ lbfgs_iters: 1000
16
+ lr: 0.001
17
+ log_every: 2000
18
+ layers: [1, 20, 20, 20, 1]
19
+
20
+ poisson2d:
21
+ n_el_x: 4
22
+ n_el_y: 4
23
+ n_test: 8
24
+ n_quad: 12
25
+ n_bound: 80
26
+ boundary_weight: 50.0
27
+ epochs: 10000
28
+ lbfgs_iters: 1000
29
+ lr: 0.001
30
+ log_every: 2000
31
+ layers: [2, 30, 30, 30, 30, 1]
configuration.json ADDED
@@ -0,0 +1 @@
 
 
1
+ {"framework":"Pytorch","task":"other"}
data/.gitkeep ADDED
File without changes
model/__pycache__/vpinn.cpython-311.pyc ADDED
Binary file (13.4 kB). View file
 
model/vpinn.py ADDED
@@ -0,0 +1,203 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from __future__ import annotations
2
+
3
+ from collections.abc import Mapping, Sequence
4
+ from pathlib import Path
5
+
6
+ import numpy as np
7
+ import torch
8
+ import torch.nn as nn
9
+ from scipy.special import gamma, jacobi, roots_jacobi
10
+
11
+
12
+ def jacobi_poly(
13
+ degree: int, alpha: float, beta: float, points: np.ndarray | float
14
+ ) -> np.ndarray:
15
+ return jacobi(degree, alpha, beta)(np.asarray(points, dtype=float))
16
+
17
+
18
+ def gauss_lobatto_jacobi_weights(
19
+ order: int, alpha: float = 0.0, beta: float = 0.0
20
+ ) -> tuple[np.ndarray, np.ndarray]:
21
+ """Return Gauss-Lobatto-Jacobi nodes and weights, including both endpoints."""
22
+ if order < 3:
23
+ raise ValueError("Gauss-Lobatto-Jacobi quadrature order must be at least 3")
24
+
25
+ interior = roots_jacobi(order - 2, alpha + 1, beta + 1)[0]
26
+ polynomial = jacobi_poly(order - 1, alpha, beta, interior)
27
+ if alpha == 0 and beta == 0:
28
+ weights = 2.0 / ((order - 1) * order * np.square(polynomial))
29
+ left_weight = 2.0 / (
30
+ (order - 1)
31
+ * order
32
+ * np.square(jacobi_poly(order - 1, 0, 0, -1.0))
33
+ )
34
+ right_weight = 2.0 / (
35
+ (order - 1)
36
+ * order
37
+ * np.square(jacobi_poly(order - 1, 0, 0, 1.0))
38
+ )
39
+ else:
40
+ denominator = (
41
+ (order - 1)
42
+ * gamma(order)
43
+ * gamma(alpha + beta + order + 1)
44
+ * np.square(polynomial)
45
+ )
46
+ numerator = (
47
+ 2.0 ** (alpha + beta + 1)
48
+ * gamma(alpha + order)
49
+ * gamma(beta + order)
50
+ )
51
+ weights = numerator / denominator
52
+ left_weight = (
53
+ (beta + 1)
54
+ * numerator
55
+ / (
56
+ (order - 1)
57
+ * gamma(order)
58
+ * gamma(alpha + beta + order + 1)
59
+ * np.square(jacobi_poly(order - 1, alpha, beta, -1.0))
60
+ )
61
+ )
62
+ right_weight = (
63
+ (alpha + 1)
64
+ * numerator
65
+ / (
66
+ (order - 1)
67
+ * gamma(order)
68
+ * gamma(alpha + beta + order + 1)
69
+ * np.square(jacobi_poly(order - 1, alpha, beta, 1.0))
70
+ )
71
+ )
72
+
73
+ nodes = np.concatenate(([-1.0], interior, [1.0]))
74
+ all_weights = np.concatenate(([left_weight], weights, [right_weight]))
75
+ return nodes, all_weights
76
+
77
+
78
+ def test_function_jacobi(
79
+ degree: int, points: np.ndarray, alpha: float = 0.0, beta: float = 0.0
80
+ ) -> np.ndarray:
81
+ """Evaluate phi_n = P_(n+1) - P_(n-1), which vanishes at both endpoints."""
82
+ if degree < 1:
83
+ raise ValueError("test function degree must be at least 1")
84
+ return jacobi_poly(degree + 1, alpha, beta, points) - jacobi_poly(
85
+ degree - 1, alpha, beta, points
86
+ )
87
+
88
+
89
+ def d_test_function_jacobi(
90
+ degree: int, points: np.ndarray
91
+ ) -> tuple[np.ndarray, np.ndarray]:
92
+ """Evaluate the first and second derivatives of a Jacobi test function."""
93
+ if degree < 1:
94
+ raise ValueError("test function degree must be at least 1")
95
+ points = np.asarray(points, dtype=float)
96
+ if degree == 1:
97
+ first = ((degree + 2) / 2) * jacobi_poly(degree, 1, 1, points)
98
+ second = ((degree + 2) * (degree + 3) / 4) * jacobi_poly(
99
+ degree - 1, 2, 2, points
100
+ )
101
+ elif degree == 2:
102
+ first = ((degree + 2) / 2) * jacobi_poly(
103
+ degree, 1, 1, points
104
+ ) - (degree / 2) * jacobi_poly(degree - 2, 1, 1, points)
105
+ second = ((degree + 2) * (degree + 3) / 4) * jacobi_poly(
106
+ degree - 1, 2, 2, points
107
+ )
108
+ else:
109
+ first = ((degree + 2) / 2) * jacobi_poly(
110
+ degree, 1, 1, points
111
+ ) - (degree / 2) * jacobi_poly(degree - 2, 1, 1, points)
112
+ second = ((degree + 2) * (degree + 3) / 4) * jacobi_poly(
113
+ degree - 1, 2, 2, points
114
+ ) - (degree * (degree + 1) / 4) * jacobi_poly(
115
+ degree - 3, 2, 2, points
116
+ )
117
+ return np.asarray(first), np.asarray(second)
118
+
119
+
120
+ class VPINN(nn.Module):
121
+ """Fully connected VPINN backbone with trainable adaptive activations."""
122
+
123
+ def __init__(
124
+ self,
125
+ layers: Sequence[int],
126
+ activation: str = "tanh",
127
+ dtype: torch.dtype = torch.float64,
128
+ ) -> None:
129
+ super().__init__()
130
+ if len(layers) < 2:
131
+ raise ValueError("layers must contain at least an input and an output size")
132
+ if activation not in {"tanh", "sin", "cos"}:
133
+ raise ValueError(f"unsupported activation: {activation}")
134
+
135
+ self.activation = activation
136
+ self.linears = nn.ModuleList(
137
+ nn.Linear(layers[index], layers[index + 1], dtype=dtype)
138
+ for index in range(len(layers) - 1)
139
+ )
140
+ # Keep one value per layer for compatibility with the bundled checkpoint.
141
+ self.a = nn.ParameterList(
142
+ nn.Parameter(torch.tensor(0.05, dtype=dtype))
143
+ for _ in range(len(layers) - 1)
144
+ )
145
+ for linear in self.linears:
146
+ nn.init.xavier_normal_(linear.weight)
147
+ nn.init.zeros_(linear.bias)
148
+
149
+ def forward(self, inputs: torch.Tensor) -> torch.Tensor:
150
+ hidden = inputs
151
+ activation = getattr(torch, self.activation)
152
+ for index, linear in enumerate(self.linears[:-1]):
153
+ hidden = activation(20.0 * self.a[index] * linear(hidden))
154
+ return self.linears[-1](hidden)
155
+
156
+
157
+ def unpack_checkpoint(checkpoint: Mapping) -> tuple[Mapping[str, torch.Tensor], dict]:
158
+ """Support both legacy raw state dictionaries and metadata checkpoints."""
159
+ if "model_state" in checkpoint:
160
+ return checkpoint["model_state"], dict(checkpoint)
161
+ if checkpoint and all(torch.is_tensor(value) for value in checkpoint.values()):
162
+ return checkpoint, {}
163
+ raise ValueError("checkpoint does not contain a valid VPINN state dictionary")
164
+
165
+
166
+ def infer_layers(state: Mapping[str, torch.Tensor]) -> list[int]:
167
+ weight_keys = sorted(
168
+ (key for key in state if key.startswith("linears.") and key.endswith(".weight")),
169
+ key=lambda key: int(key.split(".")[1]),
170
+ )
171
+ if not weight_keys:
172
+ raise ValueError("checkpoint contains no VPINN linear-layer weights")
173
+ layers = [int(state[weight_keys[0]].shape[1])]
174
+ layers.extend(int(state[key].shape[0]) for key in weight_keys)
175
+ return layers
176
+
177
+
178
+ class VPINNWrapper(nn.Module):
179
+ """Load a VPINN checkpoint and expose it as an inference module."""
180
+
181
+ def __init__(
182
+ self,
183
+ weight_path: str | Path,
184
+ device: str | torch.device | None = None,
185
+ dtype: torch.dtype = torch.float64,
186
+ ) -> None:
187
+ super().__init__()
188
+ resolved_device = torch.device(
189
+ device or ("cuda" if torch.cuda.is_available() else "cpu")
190
+ )
191
+ checkpoint = torch.load(weight_path, map_location="cpu", weights_only=True)
192
+ if not isinstance(checkpoint, Mapping):
193
+ raise ValueError(f"invalid VPINN checkpoint: {weight_path}")
194
+ state, metadata = unpack_checkpoint(checkpoint)
195
+ layers = metadata.get("layers") or infer_layers(state)
196
+ self.device = resolved_device
197
+ self.dtype = dtype
198
+ self.net = VPINN(layers, dtype=dtype).to(device=resolved_device, dtype=dtype)
199
+ self.net.load_state_dict(state, strict=True)
200
+ self.net.eval()
201
+
202
+ def forward(self, inputs: torch.Tensor) -> torch.Tensor:
203
+ return self.net(inputs.to(device=self.device, dtype=self.dtype))
scripts/inference.py ADDED
@@ -0,0 +1,186 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from __future__ import annotations
2
+
3
+ import sys
4
+ from collections.abc import Mapping
5
+ from pathlib import Path
6
+
7
+ import matplotlib
8
+
9
+ matplotlib.use("Agg")
10
+ import matplotlib.pyplot as plt
11
+ import numpy as np
12
+ import torch
13
+ import yaml
14
+
15
+
16
+ PROJECT_ROOT = Path(__file__).resolve().parents[1]
17
+ sys.path.insert(0, str(PROJECT_ROOT))
18
+
19
+ from model.vpinn import VPINN, infer_layers, unpack_checkpoint # noqa: E402
20
+
21
+
22
+ DEFAULT_CONFIG = PROJECT_ROOT / "conf" / "config.yaml"
23
+ CASES = ("1d", "2d")
24
+
25
+
26
+ def load_config(path: Path) -> dict:
27
+ with path.open("r", encoding="utf-8") as stream:
28
+ config = yaml.safe_load(stream)
29
+ if not isinstance(config, dict) or "root" not in config:
30
+ raise ValueError(f"config must contain a 'root' mapping: {path}")
31
+ return config["root"]
32
+
33
+
34
+ def project_path(value: str | Path) -> Path:
35
+ path = Path(value).expanduser()
36
+ return path if path.is_absolute() else PROJECT_ROOT / path
37
+
38
+
39
+ def resolve_device(requested: str) -> torch.device:
40
+ if requested == "auto":
41
+ return torch.device("cuda" if torch.cuda.is_available() else "cpu")
42
+ device = torch.device(requested)
43
+ if device.type == "cuda" and not torch.cuda.is_available():
44
+ raise RuntimeError("CUDA/DCU was requested but torch.cuda.is_available() is false")
45
+ return device
46
+
47
+
48
+ def resolve_dtype(name: str) -> torch.dtype:
49
+ try:
50
+ return {"float32": torch.float32, "float64": torch.float64}[name]
51
+ except KeyError as error:
52
+ raise ValueError(f"unsupported dtype: {name}") from error
53
+
54
+
55
+ def load_model(
56
+ checkpoint_path: Path, device: torch.device, dtype: torch.dtype
57
+ ) -> tuple[VPINN, dict]:
58
+ if not checkpoint_path.is_file():
59
+ raise FileNotFoundError(
60
+ f"checkpoint not found: {checkpoint_path}. Run scripts/train.py first."
61
+ )
62
+ checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=True)
63
+ if not isinstance(checkpoint, Mapping):
64
+ raise ValueError(f"invalid VPINN checkpoint: {checkpoint_path}")
65
+ state, metadata = unpack_checkpoint(checkpoint)
66
+ layers = metadata.get("layers") or infer_layers(state)
67
+ model = VPINN(layers, dtype=dtype).to(device=device, dtype=dtype)
68
+ model.load_state_dict(state, strict=True)
69
+ model.eval()
70
+ return model, metadata
71
+
72
+
73
+ def exact_poisson1d(x: np.ndarray) -> np.ndarray:
74
+ return 0.1 * np.sin(8.0 * np.pi * x) + np.tanh(80.0 * x)
75
+
76
+
77
+ def infer_poisson1d(
78
+ weight_dir: Path,
79
+ result_dir: Path,
80
+ device: torch.device,
81
+ dtype: torch.dtype,
82
+ ) -> float:
83
+ checkpoint_path = weight_dir / "hpvpinn_poisson1d.pt"
84
+ model, _ = load_model(checkpoint_path, device, dtype)
85
+ points = torch.linspace(-1.0, 1.0, 2001, dtype=dtype, device=device).unsqueeze(-1)
86
+ exact = exact_poisson1d(points.cpu().numpy().reshape(-1))
87
+ with torch.no_grad():
88
+ prediction = model(points).cpu().numpy().reshape(-1)
89
+ absolute_error = np.abs(prediction - exact)
90
+ relative_l2 = float(np.linalg.norm(prediction - exact) / np.linalg.norm(exact))
91
+
92
+ figure, axes = plt.subplots(1, 2, figsize=(12, 4))
93
+ coordinates = points.cpu().numpy().reshape(-1)
94
+ axes[0].plot(coordinates, exact, "k-", linewidth=1.5, label="Exact")
95
+ axes[0].plot(coordinates, prediction, "r--", linewidth=1.0, label="hp-VPINN")
96
+ axes[0].set_xlabel("x")
97
+ axes[0].set_ylabel("u")
98
+ axes[0].legend()
99
+ axes[1].semilogy(coordinates, np.maximum(absolute_error, 1.0e-16), "r-")
100
+ axes[1].set_xlabel("x")
101
+ axes[1].set_ylabel("Absolute error")
102
+ figure.tight_layout()
103
+ output_path = result_dir / "hpvpinn_poisson1d.png"
104
+ figure.savefig(output_path, dpi=150)
105
+ plt.close(figure)
106
+ print(f"[1D] relative L2={relative_l2:.6e}, plot={output_path}")
107
+ return relative_l2
108
+
109
+
110
+ def exact_poisson2d(x: np.ndarray, y: np.ndarray) -> np.ndarray:
111
+ return (0.1 * np.sin(2.0 * np.pi * x) + np.tanh(10.0 * x)) * np.sin(
112
+ 2.0 * np.pi * y
113
+ )
114
+
115
+
116
+ def infer_poisson2d(
117
+ weight_dir: Path,
118
+ result_dir: Path,
119
+ device: torch.device,
120
+ dtype: torch.dtype,
121
+ ) -> float:
122
+ checkpoint_path = weight_dir / "hpvpinn_poisson2d.pt"
123
+ model, _ = load_model(checkpoint_path, device, dtype)
124
+ axis = np.linspace(-1.0, 1.0, 100)
125
+ mesh_x, mesh_y = np.meshgrid(axis, axis, indexing="ij")
126
+ points = torch.as_tensor(
127
+ np.column_stack((mesh_x.ravel(), mesh_y.ravel())), dtype=dtype, device=device
128
+ )
129
+ with torch.no_grad():
130
+ prediction = model(points).cpu().numpy().reshape(mesh_x.shape)
131
+ exact = exact_poisson2d(mesh_x, mesh_y)
132
+ absolute_error = np.abs(prediction - exact)
133
+ relative_l2 = float(
134
+ np.linalg.norm(prediction.ravel() - exact.ravel()) / np.linalg.norm(exact.ravel())
135
+ )
136
+
137
+ figure, axes = plt.subplots(1, 3, figsize=(15, 4))
138
+ for plot_axis, title, field in zip(
139
+ axes,
140
+ ("Exact", "hp-VPINN", "Absolute error"),
141
+ (exact, prediction, absolute_error),
142
+ strict=True,
143
+ ):
144
+ image = plot_axis.imshow(
145
+ field.T,
146
+ extent=(-1, 1, -1, 1),
147
+ origin="lower",
148
+ aspect="auto",
149
+ cmap="jet",
150
+ )
151
+ plot_axis.set_title(title)
152
+ plot_axis.set_xlabel("x")
153
+ plot_axis.set_ylabel("y")
154
+ figure.colorbar(image, ax=plot_axis)
155
+ figure.tight_layout()
156
+ output_path = result_dir / "hpvpinn_poisson2d.png"
157
+ figure.savefig(output_path, dpi=150)
158
+ plt.close(figure)
159
+ print(f"[2D] relative L2={relative_l2:.6e}, plot={output_path}")
160
+ return relative_l2
161
+
162
+
163
+ def main() -> None:
164
+ config_path = DEFAULT_CONFIG.resolve()
165
+ config = load_config(config_path)
166
+ common = config["common"]
167
+ device = resolve_device(str(common["device"]))
168
+ dtype = resolve_dtype(str(common["dtype"]))
169
+ weight_dir = project_path(common["weight_dir"])
170
+ result_dir = project_path(common["result_dir"])
171
+ result_dir.mkdir(parents=True, exist_ok=True)
172
+
173
+ print(f"Config: {config_path}")
174
+ print(f"Device: {device}")
175
+ selected_case = str(common["case"]).lower()
176
+ if selected_case not in (*CASES, "all"):
177
+ raise ValueError("common.case must be one of: 1d, 2d, all")
178
+ selected_cases = CASES if selected_case == "all" else (selected_case,)
179
+ if "1d" in selected_cases:
180
+ infer_poisson1d(weight_dir, result_dir, device, dtype)
181
+ if "2d" in selected_cases:
182
+ infer_poisson2d(weight_dir, result_dir, device, dtype)
183
+
184
+
185
+ if __name__ == "__main__":
186
+ main()
scripts/train.py ADDED
@@ -0,0 +1,509 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from __future__ import annotations
2
+
3
+ import sys
4
+ import time
5
+ from pathlib import Path
6
+
7
+ import numpy as np
8
+ import torch
9
+ import yaml
10
+
11
+
12
+ PROJECT_ROOT = Path(__file__).resolve().parents[1]
13
+ sys.path.insert(0, str(PROJECT_ROOT))
14
+
15
+ from model.vpinn import ( # noqa: E402
16
+ VPINN,
17
+ d_test_function_jacobi,
18
+ gauss_lobatto_jacobi_weights,
19
+ test_function_jacobi,
20
+ )
21
+
22
+
23
+ DEFAULT_CONFIG = PROJECT_ROOT / "conf" / "config.yaml"
24
+ CASES = ("1d", "2d")
25
+ OMEGA_1D = 8.0 * np.pi
26
+ STEEPNESS_1D = 80.0
27
+
28
+
29
+ def load_config(path: Path) -> dict:
30
+ with path.open("r", encoding="utf-8") as stream:
31
+ config = yaml.safe_load(stream)
32
+ if not isinstance(config, dict) or "root" not in config:
33
+ raise ValueError(f"config must contain a 'root' mapping: {path}")
34
+ return config["root"]
35
+
36
+
37
+ def project_path(value: str | Path) -> Path:
38
+ path = Path(value).expanduser()
39
+ return path if path.is_absolute() else PROJECT_ROOT / path
40
+
41
+
42
+ def resolve_device(requested: str) -> torch.device:
43
+ if requested == "auto":
44
+ return torch.device("cuda" if torch.cuda.is_available() else "cpu")
45
+ device = torch.device(requested)
46
+ if device.type == "cuda" and not torch.cuda.is_available():
47
+ raise RuntimeError("CUDA/DCU was requested but torch.cuda.is_available() is false")
48
+ return device
49
+
50
+
51
+ def resolve_dtype(name: str) -> torch.dtype:
52
+ try:
53
+ return {"float32": torch.float32, "float64": torch.float64}[name]
54
+ except KeyError as error:
55
+ raise ValueError(f"unsupported dtype: {name}") from error
56
+
57
+
58
+ def seed_everything(seed: int) -> None:
59
+ np.random.seed(seed)
60
+ torch.manual_seed(seed)
61
+ if torch.cuda.is_available():
62
+ torch.cuda.manual_seed_all(seed)
63
+
64
+
65
+ def ensure_finite(loss: torch.Tensor, case: str, step: int) -> None:
66
+ if not torch.isfinite(loss):
67
+ raise FloatingPointError(f"{case} loss became non-finite at step {step}")
68
+
69
+
70
+ def positive(value: int, name: str) -> int:
71
+ if value <= 0:
72
+ raise ValueError(f"{name} must be positive")
73
+ return value
74
+
75
+
76
+ def save_checkpoint(weight_dir: Path, filename: str, payload: dict) -> Path:
77
+ weight_dir.mkdir(parents=True, exist_ok=True)
78
+ output_path = weight_dir / filename
79
+ torch.save(payload, output_path)
80
+ print(f"Saved checkpoint: {output_path}")
81
+ return output_path
82
+
83
+
84
+ def exact_poisson1d(x: np.ndarray) -> np.ndarray:
85
+ return 0.1 * np.sin(OMEGA_1D * x) + np.tanh(STEEPNESS_1D * x)
86
+
87
+
88
+ def source_poisson1d(x: np.ndarray) -> np.ndarray:
89
+ oscillatory = -0.1 * OMEGA_1D**2 * np.sin(OMEGA_1D * x)
90
+ transition = (
91
+ -2.0
92
+ * STEEPNESS_1D**2
93
+ * np.tanh(STEEPNESS_1D * x)
94
+ / np.cosh(STEEPNESS_1D * x) ** 2
95
+ )
96
+ return -(oscillatory + transition)
97
+
98
+
99
+ def build_rhs_1d(
100
+ element_grid: np.ndarray,
101
+ quadrature_points: np.ndarray,
102
+ quadrature_weights: np.ndarray,
103
+ n_test: int,
104
+ ) -> np.ndarray:
105
+ test_values = np.stack(
106
+ [
107
+ test_function_jacobi(index + 1, quadrature_points)
108
+ for index in range(n_test)
109
+ ]
110
+ )
111
+ rhs = np.empty((element_grid.size - 1, n_test), dtype=float)
112
+ for element in range(element_grid.size - 1):
113
+ left, right = element_grid[element : element + 2]
114
+ jacobian = (right - left) / 2.0
115
+ physical_points = left + jacobian * (quadrature_points + 1.0)
116
+ weighted_source = quadrature_weights * source_poisson1d(physical_points)
117
+ rhs[element] = jacobian * (test_values @ weighted_source)
118
+ return rhs
119
+
120
+
121
+ def train_poisson1d(
122
+ config: dict,
123
+ device: torch.device,
124
+ dtype: torch.dtype,
125
+ weight_dir: Path,
126
+ ) -> None:
127
+ n_element = positive(int(config["n_element"]), "n_element")
128
+ n_test = positive(int(config["n_test"]), "n_test")
129
+ n_quad = positive(int(config["n_quad"]), "n_quad")
130
+ epochs = int(config["epochs"])
131
+ lbfgs_iters = int(config["lbfgs_iters"])
132
+ if epochs < 0 or lbfgs_iters < 0:
133
+ raise ValueError("epochs and lbfgs_iters must be non-negative")
134
+ layers = [int(value) for value in config["layers"]]
135
+ if layers[0] != 1 or layers[-1] != 1:
136
+ raise ValueError("1D Poisson layers must start and end with width 1")
137
+
138
+ element_grid = np.linspace(-1.0, 1.0, n_element + 1)
139
+ quadrature_points, quadrature_weights = gauss_lobatto_jacobi_weights(n_quad)
140
+ rhs = build_rhs_1d(element_grid, quadrature_points, quadrature_weights, n_test)
141
+ test_derivatives = np.stack(
142
+ [
143
+ d_test_function_jacobi(index + 1, quadrature_points)[0]
144
+ for index in range(n_test)
145
+ ]
146
+ )
147
+
148
+ model = VPINN(layers, dtype=dtype).to(device=device, dtype=dtype)
149
+ quadrature_tensor = torch.as_tensor(quadrature_points, dtype=dtype, device=device)
150
+ weight_tensor = torch.as_tensor(quadrature_weights, dtype=dtype, device=device)
151
+ derivative_tensor = torch.as_tensor(test_derivatives, dtype=dtype, device=device)
152
+ rhs_tensor = torch.as_tensor(rhs, dtype=dtype, device=device)
153
+ boundary = torch.tensor([[-1.0], [1.0]], dtype=dtype, device=device)
154
+ boundary_values = torch.as_tensor(
155
+ exact_poisson1d(np.array([-1.0, 1.0]))[:, None], dtype=dtype, device=device
156
+ )
157
+ test_points = torch.linspace(-1.0, 1.0, 2001, dtype=dtype, device=device).unsqueeze(-1)
158
+ exact_values = torch.as_tensor(
159
+ exact_poisson1d(test_points.cpu().numpy().reshape(-1))[:, None],
160
+ dtype=dtype,
161
+ device=device,
162
+ )
163
+
164
+ def variational_loss() -> torch.Tensor:
165
+ loss = torch.zeros((), dtype=dtype, device=device)
166
+ for element in range(n_element):
167
+ left, right = element_grid[element : element + 2]
168
+ jacobian = (right - left) / 2.0
169
+ physical_points = (
170
+ left + jacobian * (quadrature_tensor + 1.0)
171
+ ).unsqueeze(-1)
172
+ physical_points.requires_grad_(True)
173
+ prediction = model(physical_points)
174
+ prediction_x = torch.autograd.grad(
175
+ prediction,
176
+ physical_points,
177
+ torch.ones_like(prediction),
178
+ create_graph=True,
179
+ )[0]
180
+ weak_prediction = derivative_tensor @ (
181
+ weight_tensor * prediction_x.reshape(-1)
182
+ )
183
+ loss = loss + torch.mean((weak_prediction - rhs_tensor[element]).square())
184
+ boundary_loss = torch.mean((model(boundary) - boundary_values).square())
185
+ return loss + boundary_loss
186
+
187
+ def relative_l2() -> float:
188
+ with torch.no_grad():
189
+ return (
190
+ torch.linalg.vector_norm(model(test_points) - exact_values)
191
+ / torch.linalg.vector_norm(exact_values)
192
+ ).item()
193
+
194
+ optimizer = torch.optim.Adam(model.parameters(), lr=float(config["lr"]))
195
+ log_every = int(config["log_every"])
196
+ started = time.time()
197
+ print(
198
+ f"[1D] device={device} elements={n_element} n_test={n_test} "
199
+ f"n_quad={n_quad} epochs={epochs} lbfgs_iters={lbfgs_iters}"
200
+ )
201
+ for step in range(1, epochs + 1):
202
+ loss = variational_loss()
203
+ ensure_finite(loss, "1D Poisson", step)
204
+ optimizer.zero_grad(set_to_none=True)
205
+ loss.backward()
206
+ optimizer.step()
207
+ if step == 1 or step % log_every == 0 or step == epochs:
208
+ print(
209
+ f"[1D] step={step:5d} loss={loss.item():.3e} "
210
+ f"relative_l2={relative_l2():.3e}"
211
+ )
212
+
213
+ if lbfgs_iters > 0:
214
+ lbfgs = torch.optim.LBFGS(
215
+ model.parameters(),
216
+ lr=1.0,
217
+ max_iter=lbfgs_iters,
218
+ max_eval=max(1, 2 * lbfgs_iters),
219
+ history_size=50,
220
+ line_search_fn="strong_wolfe",
221
+ )
222
+
223
+ def closure() -> torch.Tensor:
224
+ lbfgs.zero_grad(set_to_none=True)
225
+ closure_loss = variational_loss()
226
+ ensure_finite(closure_loss, "1D Poisson L-BFGS", 0)
227
+ closure_loss.backward()
228
+ return closure_loss
229
+
230
+ lbfgs.step(closure)
231
+
232
+ error = relative_l2()
233
+ print(f"[1D] finished in {time.time() - started:.1f}s, relative L2={error:.6e}")
234
+ save_checkpoint(
235
+ weight_dir,
236
+ "hpvpinn_poisson1d.pt",
237
+ {
238
+ "case": "1d",
239
+ "architecture": "vpinn",
240
+ "model_state": model.state_dict(),
241
+ "layers": layers,
242
+ "n_element": n_element,
243
+ "n_test": n_test,
244
+ "n_quad": n_quad,
245
+ },
246
+ )
247
+
248
+
249
+ def exact_poisson2d(x: np.ndarray, y: np.ndarray) -> np.ndarray:
250
+ return (0.1 * np.sin(2.0 * np.pi * x) + np.tanh(10.0 * x)) * np.sin(
251
+ 2.0 * np.pi * y
252
+ )
253
+
254
+
255
+ def source_poisson2d(x: np.ndarray, y: np.ndarray) -> np.ndarray:
256
+ x_component = (
257
+ -0.4 * np.pi**2 * np.sin(2.0 * np.pi * x)
258
+ - 200.0 * np.tanh(10.0 * x) / np.cosh(10.0 * x) ** 2
259
+ ) * np.sin(2.0 * np.pi * y)
260
+ y_component = (
261
+ 0.1 * np.sin(2.0 * np.pi * x) + np.tanh(10.0 * x)
262
+ ) * (-4.0 * np.pi**2 * np.sin(2.0 * np.pi * y))
263
+ return x_component + y_component
264
+
265
+
266
+ def build_rhs_2d(
267
+ grid_x: np.ndarray,
268
+ grid_y: np.ndarray,
269
+ quadrature_points: np.ndarray,
270
+ quadrature_weights: np.ndarray,
271
+ n_test: int,
272
+ ) -> np.ndarray:
273
+ test_values = np.stack(
274
+ [
275
+ test_function_jacobi(index + 1, quadrature_points)
276
+ for index in range(n_test)
277
+ ]
278
+ )
279
+ n_elements_x = grid_x.size - 1
280
+ n_elements_y = grid_y.size - 1
281
+ rhs = np.empty((n_elements_x, n_elements_y, n_test, n_test), dtype=float)
282
+ for element_x in range(n_elements_x):
283
+ left_x, right_x = grid_x[element_x : element_x + 2]
284
+ jacobian_x = (right_x - left_x) / 2.0
285
+ physical_x = left_x + jacobian_x * (quadrature_points + 1.0)
286
+ for element_y in range(n_elements_y):
287
+ left_y, right_y = grid_y[element_y : element_y + 2]
288
+ jacobian_y = (right_y - left_y) / 2.0
289
+ physical_y = left_y + jacobian_y * (quadrature_points + 1.0)
290
+ mesh_x, mesh_y = np.meshgrid(physical_x, physical_y, indexing="ij")
291
+ weighted_source = (
292
+ quadrature_weights[:, None]
293
+ * source_poisson2d(mesh_x, mesh_y)
294
+ * quadrature_weights[None, :]
295
+ )
296
+ rhs[element_x, element_y] = (
297
+ jacobian_x
298
+ * jacobian_y
299
+ * (test_values @ weighted_source @ test_values.T)
300
+ )
301
+ return rhs
302
+
303
+
304
+ def sample_boundary_2d(n_bound: int) -> tuple[np.ndarray, np.ndarray]:
305
+ points = []
306
+ values = []
307
+ random_x = 2.0 * np.random.rand(n_bound, 1) - 1.0
308
+ for boundary_y in (1.0, -1.0):
309
+ y = np.full_like(random_x, boundary_y)
310
+ points.append(np.hstack((random_x, y)))
311
+ values.append(exact_poisson2d(random_x, y))
312
+ random_y = 2.0 * np.random.rand(n_bound, 1) - 1.0
313
+ for boundary_x in (1.0, -1.0):
314
+ x = np.full_like(random_y, boundary_x)
315
+ points.append(np.hstack((x, random_y)))
316
+ values.append(exact_poisson2d(x, random_y))
317
+ return np.vstack(points), np.vstack(values)
318
+
319
+
320
+ def train_poisson2d(
321
+ config: dict,
322
+ device: torch.device,
323
+ dtype: torch.dtype,
324
+ weight_dir: Path,
325
+ ) -> None:
326
+ n_el_x = positive(int(config["n_el_x"]), "n_el_x")
327
+ n_el_y = positive(int(config["n_el_y"]), "n_el_y")
328
+ n_test = positive(int(config["n_test"]), "n_test")
329
+ n_quad = positive(int(config["n_quad"]), "n_quad")
330
+ n_bound = positive(int(config["n_bound"]), "n_bound")
331
+ epochs = int(config["epochs"])
332
+ lbfgs_iters = int(config["lbfgs_iters"])
333
+ if epochs < 0 or lbfgs_iters < 0:
334
+ raise ValueError("epochs and lbfgs_iters must be non-negative")
335
+ layers = [int(value) for value in config["layers"]]
336
+ if layers[0] != 2 or layers[-1] != 1:
337
+ raise ValueError("2D Poisson layers must start with width 2 and end with width 1")
338
+
339
+ grid_x = np.linspace(-1.0, 1.0, n_el_x + 1)
340
+ grid_y = np.linspace(-1.0, 1.0, n_el_y + 1)
341
+ quadrature_points, quadrature_weights = gauss_lobatto_jacobi_weights(n_quad)
342
+ rhs = build_rhs_2d(
343
+ grid_x, grid_y, quadrature_points, quadrature_weights, n_test
344
+ )
345
+ test_values = np.stack(
346
+ [
347
+ test_function_jacobi(index + 1, quadrature_points)
348
+ for index in range(n_test)
349
+ ]
350
+ )
351
+ test_derivatives = np.stack(
352
+ [
353
+ d_test_function_jacobi(index + 1, quadrature_points)[0]
354
+ for index in range(n_test)
355
+ ]
356
+ )
357
+ boundary_points, boundary_values = sample_boundary_2d(n_bound)
358
+
359
+ model = VPINN(layers, dtype=dtype).to(device=device, dtype=dtype)
360
+ quadrature_tensor = torch.as_tensor(quadrature_points, dtype=dtype, device=device)
361
+ weight_tensor = torch.as_tensor(quadrature_weights, dtype=dtype, device=device)
362
+ test_tensor = torch.as_tensor(test_values, dtype=dtype, device=device)
363
+ derivative_tensor = torch.as_tensor(test_derivatives, dtype=dtype, device=device)
364
+ rhs_tensor = torch.as_tensor(rhs, dtype=dtype, device=device)
365
+ boundary_tensor = torch.as_tensor(boundary_points, dtype=dtype, device=device)
366
+ boundary_value_tensor = torch.as_tensor(boundary_values, dtype=dtype, device=device)
367
+
368
+ evaluation_axis = torch.linspace(-1.0, 1.0, 100, dtype=dtype, device=device)
369
+ evaluation_x, evaluation_y = torch.meshgrid(
370
+ evaluation_axis, evaluation_axis, indexing="ij"
371
+ )
372
+ evaluation_points = torch.stack(
373
+ (evaluation_x.reshape(-1), evaluation_y.reshape(-1)), dim=-1
374
+ )
375
+ exact_values = torch.as_tensor(
376
+ exact_poisson2d(
377
+ evaluation_points[:, 0].cpu().numpy(),
378
+ evaluation_points[:, 1].cpu().numpy(),
379
+ )[:, None],
380
+ dtype=dtype,
381
+ device=device,
382
+ )
383
+
384
+ def variational_loss() -> torch.Tensor:
385
+ loss = torch.zeros((), dtype=dtype, device=device)
386
+ for element_x in range(n_el_x):
387
+ left_x, right_x = grid_x[element_x : element_x + 2]
388
+ jacobian_x = (right_x - left_x) / 2.0
389
+ physical_x = left_x + jacobian_x * (quadrature_tensor + 1.0)
390
+ for element_y in range(n_el_y):
391
+ left_y, right_y = grid_y[element_y : element_y + 2]
392
+ jacobian_y = (right_y - left_y) / 2.0
393
+ physical_y = left_y + jacobian_y * (quadrature_tensor + 1.0)
394
+ mesh_x, mesh_y = torch.meshgrid(physical_x, physical_y, indexing="ij")
395
+ coordinates = torch.stack(
396
+ (mesh_x.reshape(-1), mesh_y.reshape(-1)), dim=-1
397
+ )
398
+ coordinates.requires_grad_(True)
399
+ prediction = model(coordinates)
400
+ prediction_gradient = torch.autograd.grad(
401
+ prediction,
402
+ coordinates,
403
+ torch.ones_like(prediction),
404
+ create_graph=True,
405
+ )[0]
406
+ prediction_x = prediction_gradient[:, 0].reshape(n_quad, n_quad)
407
+ prediction_y = prediction_gradient[:, 1].reshape(n_quad, n_quad)
408
+ weighted_x = (
409
+ weight_tensor[:, None] * prediction_x * weight_tensor[None, :]
410
+ )
411
+ weighted_y = (
412
+ weight_tensor[:, None] * prediction_y * weight_tensor[None, :]
413
+ )
414
+ weak_prediction = -jacobian_y * (
415
+ derivative_tensor @ weighted_x @ test_tensor.T
416
+ ) - jacobian_x * (test_tensor @ weighted_y @ derivative_tensor.T)
417
+ loss = loss + torch.mean(
418
+ (weak_prediction - rhs_tensor[element_x, element_y]).square()
419
+ )
420
+ boundary_loss = torch.mean(
421
+ (model(boundary_tensor) - boundary_value_tensor).square()
422
+ )
423
+ return loss + float(config["boundary_weight"]) * boundary_loss
424
+
425
+ def relative_l2() -> float:
426
+ with torch.no_grad():
427
+ return (
428
+ torch.linalg.vector_norm(model(evaluation_points) - exact_values)
429
+ / torch.linalg.vector_norm(exact_values)
430
+ ).item()
431
+
432
+ optimizer = torch.optim.Adam(model.parameters(), lr=float(config["lr"]))
433
+ log_every = int(config["log_every"])
434
+ started = time.time()
435
+ print(
436
+ f"[2D] device={device} elements={n_el_x}x{n_el_y} n_test={n_test} "
437
+ f"n_quad={n_quad} epochs={epochs} lbfgs_iters={lbfgs_iters}"
438
+ )
439
+ for step in range(1, epochs + 1):
440
+ loss = variational_loss()
441
+ ensure_finite(loss, "2D Poisson", step)
442
+ optimizer.zero_grad(set_to_none=True)
443
+ loss.backward()
444
+ optimizer.step()
445
+ if step == 1 or step % log_every == 0 or step == epochs:
446
+ print(
447
+ f"[2D] step={step:5d} loss={loss.item():.3e} "
448
+ f"relative_l2={relative_l2():.3e}"
449
+ )
450
+
451
+ if lbfgs_iters > 0:
452
+ lbfgs = torch.optim.LBFGS(
453
+ model.parameters(),
454
+ lr=1.0,
455
+ max_iter=lbfgs_iters,
456
+ max_eval=max(1, 2 * lbfgs_iters),
457
+ history_size=50,
458
+ line_search_fn="strong_wolfe",
459
+ )
460
+
461
+ def closure() -> torch.Tensor:
462
+ lbfgs.zero_grad(set_to_none=True)
463
+ closure_loss = variational_loss()
464
+ ensure_finite(closure_loss, "2D Poisson L-BFGS", 0)
465
+ closure_loss.backward()
466
+ return closure_loss
467
+
468
+ lbfgs.step(closure)
469
+
470
+ error = relative_l2()
471
+ print(f"[2D] finished in {time.time() - started:.1f}s, relative L2={error:.6e}")
472
+ save_checkpoint(
473
+ weight_dir,
474
+ "hpvpinn_poisson2d.pt",
475
+ {
476
+ "case": "2d",
477
+ "architecture": "vpinn",
478
+ "model_state": model.state_dict(),
479
+ "layers": layers,
480
+ "n_el_x": n_el_x,
481
+ "n_el_y": n_el_y,
482
+ "n_test": n_test,
483
+ "n_quad": n_quad,
484
+ },
485
+ )
486
+
487
+
488
+ def main() -> None:
489
+ config_path = DEFAULT_CONFIG.resolve()
490
+ config = load_config(config_path)
491
+ common = config["common"]
492
+ device = resolve_device(str(common["device"]))
493
+ dtype = resolve_dtype(str(common["dtype"]))
494
+ weight_dir = project_path(common["weight_dir"])
495
+ seed_everything(int(common["seed"]))
496
+
497
+ print(f"Config: {config_path}")
498
+ selected_case = str(common["case"]).lower()
499
+ if selected_case not in (*CASES, "all"):
500
+ raise ValueError("common.case must be one of: 1d, 2d, all")
501
+ selected_cases = CASES if selected_case == "all" else (selected_case,)
502
+ if "1d" in selected_cases:
503
+ train_poisson1d(config["poisson1d"], device, dtype, weight_dir)
504
+ if "2d" in selected_cases:
505
+ train_poisson2d(config["poisson2d"], device, dtype, weight_dir)
506
+
507
+
508
+ if __name__ == "__main__":
509
+ main()
weight/hpvpinn_poisson1d.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:4f2fee7b078bedede0c8453c2fc925031f887d90c768956aac6c16f59892ea3a
3
+ size 11827