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