Upload folder using huggingface_hub
Browse files- .gitattributes +24 -10
- README.md +125 -0
- conf/config.yaml +31 -0
- configuration.json +1 -0
- data/.gitkeep +0 -0
- model/__pycache__/vpinn.cpython-311.pyc +0 -0
- model/vpinn.py +203 -0
- scripts/inference.py +186 -0
- scripts/train.py +509 -0
- weight/hpvpinn_poisson1d.pt +3 -0
.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 |
-
|
| 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
|