Upload folder using huggingface_hub
Browse files- .gitattributes +25 -10
- README.md +128 -0
- conf/config.yaml +45 -0
- configuration.json +1 -0
- data/XPINN_2D_PoissonEqn.mat +3 -0
- download.sh +4 -0
- model/__pycache__/xpinn.cpython-311.pyc +0 -0
- model/xpinn.py +131 -0
- scripts/__pycache__/data_utils.cpython-311.pyc +0 -0
- scripts/data_utils.py +212 -0
- scripts/inference.py +188 -0
- scripts/train.py +208 -0
- weight/.gitkeep +0 -0
- weight/xpinn_poisson_2d.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 |
+
*.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
|
| 50 |
+
data/XPINN_2D_PoissonEqn.mat filter=lfs diff=lfs merge=lfs -text
|
README.md
ADDED
|
@@ -0,0 +1,128 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
---
|
| 2 |
+
frameworks:
|
| 3 |
+
- pytorch
|
| 4 |
+
language:
|
| 5 |
+
- en
|
| 6 |
+
license: apache-2.0
|
| 7 |
+
tags:
|
| 8 |
+
- OneScience
|
| 9 |
+
- XPINNs
|
| 10 |
+
- physics-informed-neural-networks
|
| 11 |
+
- domain-decomposition
|
| 12 |
+
- partial-differential-equations
|
| 13 |
+
- Poisson
|
| 14 |
+
tasks:
|
| 15 |
+
- pde-solving
|
| 16 |
+
---
|
| 17 |
+
<p align="center">
|
| 18 |
+
<strong>
|
| 19 |
+
<span style="font-size: 30px;">XPINNs</span>
|
| 20 |
+
</strong>
|
| 21 |
+
</p>
|
| 22 |
+
|
| 23 |
+
# Model Overview
|
| 24 |
+
|
| 25 |
+
XPINNs (Extended Physics-Informed Neural Networks) decompose a complex computational domain into multiple subdomains and assign an independent neural network to each. In addition to enforcing PDE residuals and external boundary conditions, training constrains solution continuity and residual consistency across subdomain interfaces.
|
| 26 |
+
|
| 27 |
+
This model package reproduces the two-dimensional Poisson benchmark from the XPINNs paper. An irregular X-shaped domain is decomposed into three subdomains governed by:
|
| 28 |
+
|
| 29 |
+
```text
|
| 30 |
+
u_xx + u_yy = exp(x) + exp(y)
|
| 31 |
+
```
|
| 32 |
+
|
| 33 |
+
Paper: Extended Physics-Informed Neural Networks (XPINNs): A Generalized Space-Time Domain Decomposition Based Deep Learning Framework for Nonlinear Partial Differential Equations
|
| 34 |
+
https://doi.org/10.4208/cicp.OA-2020-0164
|
| 35 |
+
|
| 36 |
+
# Model Description
|
| 37 |
+
|
| 38 |
+
XPINNs assign a separate neural network to each of the three subdomains, using `tanh`, `sin`, and `cos` activation functions by default. The model solves the equation on the irregular domain by jointly optimizing boundary loss, PDE residual loss, interface-value loss, and interface-residual loss.
|
| 39 |
+
|
| 40 |
+
# Use Cases
|
| 41 |
+
|
| 42 |
+
| Use Case | Description |
|
| 43 |
+
| :---: | :--- |
|
| 44 |
+
| PDE solving on complex domains | Solve a two-dimensional Poisson equation on an irregular X-shaped domain |
|
| 45 |
+
| Domain-decomposition research | Configure independent network architectures and activation functions for different subdomains |
|
| 46 |
+
| Interface-constraint research | Compare losses for interface solution continuity and residual consistency |
|
| 47 |
+
| Pipeline validation | Validate training and inference using the bundled data and a small-scale configuration |
|
| 48 |
+
|
| 49 |
+
# Usage
|
| 50 |
+
|
| 51 |
+
## 1. OneCode
|
| 52 |
+
|
| 53 |
+
Use the online OneCode environment for an intelligent, one-click AI for Science (AI4S) programming experience:
|
| 54 |
+
|
| 55 |
+
[Launch OneCode for one-click AI4S programming](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)
|
| 56 |
+
|
| 57 |
+
## 2. Manual Setup
|
| 58 |
+
|
| 59 |
+
**Hardware Requirements**
|
| 60 |
+
|
| 61 |
+
- A GPU or DCU is recommended for training.
|
| 62 |
+
- A CPU can be used for small-scale pipeline validation, but training with the full sample set will be slow.
|
| 63 |
+
- DCU users must install DTK and a PyTorch environment compatible with the target cluster.
|
| 64 |
+
|
| 65 |
+
### Download the Model Package
|
| 66 |
+
|
| 67 |
+
```bash
|
| 68 |
+
modelscope download --model OneScience/XPINNs --local_dir ./XPINNs
|
| 69 |
+
cd XPINNs
|
| 70 |
+
```
|
| 71 |
+
|
| 72 |
+
### Set Up the Runtime Environment
|
| 73 |
+
|
| 74 |
+
**DCU Environment**
|
| 75 |
+
|
| 76 |
+
```bash
|
| 77 |
+
# Activate DTK and Conda first
|
| 78 |
+
conda create -n onescience311 python=3.11 -y
|
| 79 |
+
conda activate onescience311
|
| 80 |
+
pip install onescience[cfd-dcu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
|
| 81 |
+
```
|
| 82 |
+
|
| 83 |
+
**GPU Environment**
|
| 84 |
+
|
| 85 |
+
```bash
|
| 86 |
+
# Activate Conda first
|
| 87 |
+
conda create -n onescience311 python=3.11 -y libstdcxx-ng=12 libgcc-ng=12 gcc_linux-64=12 gxx_linux-64=12
|
| 88 |
+
conda activate onescience311
|
| 89 |
+
pip install onescience[cfd-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
|
| 90 |
+
```
|
| 91 |
+
|
| 92 |
+
### Training Data
|
| 93 |
+
|
| 94 |
+
The model package includes the two-dimensional Poisson data file `data/XPINN_2D_PoissonEqn.mat`, which contains interior, boundary, and interface points for the three subdomains together with the exact solution. The number of training samples of each type can be adjusted in `conf/config.yaml`.
|
| 95 |
+
|
| 96 |
+
### Training
|
| 97 |
+
|
| 98 |
+
```bash
|
| 99 |
+
python scripts/train.py
|
| 100 |
+
```
|
| 101 |
+
|
| 102 |
+
The default weights are saved to `weight/xpinn_poisson_2d.pt`.
|
| 103 |
+
|
| 104 |
+
### Model Weights
|
| 105 |
+
|
| 106 |
+
This repository provides weights trained on the two-dimensional Poisson dataset in the `weight/` directory.
|
| 107 |
+
|
| 108 |
+
### Inference, Evaluation, and Visualization
|
| 109 |
+
|
| 110 |
+
After training, run:
|
| 111 |
+
|
| 112 |
+
```bash
|
| 113 |
+
python scripts/inference.py
|
| 114 |
+
```
|
| 115 |
+
|
| 116 |
+
The script reports the overall relative L2 error and saves plots of the exact solution, prediction, and absolute error to `result/xpinn_poisson_2d.png`. Model, data, loss, and inference parameters can all be modified in `conf/config.yaml`.
|
| 117 |
+
|
| 118 |
+
# Official OneScience Resources
|
| 119 |
+
|
| 120 |
+
| Platform | OneScience Repository | Skills Repository |
|
| 121 |
+
| --- | --- | --- |
|
| 122 |
+
| Gitee | https://gitee.com/onescience-ai/onescience | https://gitee.com/onescience-ai/oneskills |
|
| 123 |
+
| GitHub | https://github.com/onescience-ai/OneScience | https://github.com/onescience-ai/oneskills |
|
| 124 |
+
|
| 125 |
+
# Citations and License
|
| 126 |
+
|
| 127 |
+
- Jagtap, A. D., Kharazmi, E., and Karniadakis, G. E. Extended Physics-Informed Neural Networks (XPINNs): A Generalized Space-Time Domain Decomposition Based Deep Learning Framework for Nonlinear Partial Differential Equations. Communications in Computational Physics, 28(5), 2002-2041, 2020.
|
| 128 |
+
- This model package is released under the Apache-2.0 license and retains attribution to the original paper and data sources.
|
conf/config.yaml
ADDED
|
@@ -0,0 +1,45 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
root:
|
| 2 |
+
common:
|
| 3 |
+
device: "auto"
|
| 4 |
+
dtype: "float64"
|
| 5 |
+
seed: 1234
|
| 6 |
+
weight_dir: "weight"
|
| 7 |
+
result_dir: "result"
|
| 8 |
+
|
| 9 |
+
data:
|
| 10 |
+
mat_file: "data/XPINN_2D_PoissonEqn.mat"
|
| 11 |
+
samples:
|
| 12 |
+
residual_1: 5000
|
| 13 |
+
residual_2: 1800
|
| 14 |
+
residual_3: 1200
|
| 15 |
+
boundary: 200
|
| 16 |
+
interface_1: 100
|
| 17 |
+
interface_2: 100
|
| 18 |
+
|
| 19 |
+
model:
|
| 20 |
+
name: "XPINNPoisson2D"
|
| 21 |
+
subnetworks:
|
| 22 |
+
domain1:
|
| 23 |
+
layers: [2, 30, 30, 1]
|
| 24 |
+
activation: "tanh"
|
| 25 |
+
domain2:
|
| 26 |
+
layers: [2, 20, 20, 20, 20, 1]
|
| 27 |
+
activation: "sin"
|
| 28 |
+
domain3:
|
| 29 |
+
layers: [2, 25, 25, 25, 1]
|
| 30 |
+
activation: "cos"
|
| 31 |
+
|
| 32 |
+
loss:
|
| 33 |
+
boundary: 20.0
|
| 34 |
+
pde: 1.0
|
| 35 |
+
interface_residual: 1.0
|
| 36 |
+
interface_value: 20.0
|
| 37 |
+
|
| 38 |
+
training:
|
| 39 |
+
steps: 501
|
| 40 |
+
lr: 0.0008
|
| 41 |
+
log_interval: 20
|
| 42 |
+
checkpoint_name: "xpinn_poisson_2d.pt"
|
| 43 |
+
|
| 44 |
+
inference:
|
| 45 |
+
figure_name: "xpinn_poisson_2d.png"
|
configuration.json
ADDED
|
@@ -0,0 +1 @@
|
|
|
|
|
|
|
| 1 |
+
{"framework":"Pytorch","task":"other"}
|
data/XPINN_2D_PoissonEqn.mat
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:13283f7494b976f65329bf160b66da916199d0092f8d943bdb8c5c04487fde1a
|
| 3 |
+
size 1353695
|
download.sh
ADDED
|
@@ -0,0 +1,4 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#!/bin/bash
|
| 2 |
+
# Download files larger than 1MB, excluding .sh .py .md .yaml .yml
|
| 3 |
+
# Total large files: 1
|
| 4 |
+
modelscope download --model OneScience/XPINNs data/XPINN_2D_PoissonEqn.mat --local_dir ./
|
model/__pycache__/xpinn.cpython-311.pyc
ADDED
|
Binary file (9.31 kB). View file
|
|
|
model/xpinn.py
ADDED
|
@@ -0,0 +1,131 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from __future__ import annotations
|
| 2 |
+
|
| 3 |
+
from collections.abc import Mapping, Sequence
|
| 4 |
+
|
| 5 |
+
import torch
|
| 6 |
+
import torch.nn as nn
|
| 7 |
+
|
| 8 |
+
|
| 9 |
+
class SubNet(nn.Module):
|
| 10 |
+
"""Independent neural network assigned to one XPINN subdomain."""
|
| 11 |
+
|
| 12 |
+
def __init__(
|
| 13 |
+
self,
|
| 14 |
+
layers: Sequence[int],
|
| 15 |
+
activation: str = "tanh",
|
| 16 |
+
dtype: torch.dtype = torch.float64,
|
| 17 |
+
) -> None:
|
| 18 |
+
super().__init__()
|
| 19 |
+
if len(layers) < 2:
|
| 20 |
+
raise ValueError("layers must contain at least an input and an output size")
|
| 21 |
+
if layers[0] != 2 or layers[-1] != 1:
|
| 22 |
+
raise ValueError("XPINN subnetworks must have two inputs and one output")
|
| 23 |
+
if activation not in {"tanh", "sin", "cos"}:
|
| 24 |
+
raise ValueError(f"unsupported activation: {activation}")
|
| 25 |
+
|
| 26 |
+
self.activation = activation
|
| 27 |
+
self.linears = nn.ModuleList(
|
| 28 |
+
nn.Linear(layers[index], layers[index + 1], dtype=dtype)
|
| 29 |
+
for index in range(len(layers) - 1)
|
| 30 |
+
)
|
| 31 |
+
# Keep one value per layer for compatibility with the original implementation.
|
| 32 |
+
self.a = nn.ParameterList(
|
| 33 |
+
nn.Parameter(torch.tensor(0.05, dtype=dtype))
|
| 34 |
+
for _ in range(len(layers) - 1)
|
| 35 |
+
)
|
| 36 |
+
for linear in self.linears:
|
| 37 |
+
nn.init.xavier_normal_(linear.weight)
|
| 38 |
+
nn.init.zeros_(linear.bias)
|
| 39 |
+
|
| 40 |
+
def forward(self, coordinates: torch.Tensor) -> torch.Tensor:
|
| 41 |
+
hidden = coordinates
|
| 42 |
+
activation = getattr(torch, self.activation)
|
| 43 |
+
for index, linear in enumerate(self.linears[:-1]):
|
| 44 |
+
hidden = activation(20.0 * self.a[index] * linear(hidden))
|
| 45 |
+
return self.linears[-1](hidden)
|
| 46 |
+
|
| 47 |
+
|
| 48 |
+
class XPINNPoisson2D(nn.Module):
|
| 49 |
+
"""Three-subdomain XPINN for the two-dimensional Poisson benchmark."""
|
| 50 |
+
|
| 51 |
+
def __init__(
|
| 52 |
+
self, config: Mapping, dtype: torch.dtype = torch.float64
|
| 53 |
+
) -> None:
|
| 54 |
+
super().__init__()
|
| 55 |
+
try:
|
| 56 |
+
subnetworks = config["subnetworks"]
|
| 57 |
+
domain1 = subnetworks["domain1"]
|
| 58 |
+
domain2 = subnetworks["domain2"]
|
| 59 |
+
domain3 = subnetworks["domain3"]
|
| 60 |
+
except (KeyError, TypeError) as error:
|
| 61 |
+
raise ValueError("model config must define three subnetworks") from error
|
| 62 |
+
|
| 63 |
+
self.n1 = SubNet(domain1["layers"], domain1["activation"], dtype=dtype)
|
| 64 |
+
self.n2 = SubNet(domain2["layers"], domain2["activation"], dtype=dtype)
|
| 65 |
+
self.n3 = SubNet(domain3["layers"], domain3["activation"], dtype=dtype)
|
| 66 |
+
|
| 67 |
+
@staticmethod
|
| 68 |
+
def _gradient(output: torch.Tensor, inputs: torch.Tensor) -> torch.Tensor:
|
| 69 |
+
return torch.autograd.grad(
|
| 70 |
+
output,
|
| 71 |
+
inputs,
|
| 72 |
+
torch.ones_like(output),
|
| 73 |
+
create_graph=True,
|
| 74 |
+
)[0]
|
| 75 |
+
|
| 76 |
+
@staticmethod
|
| 77 |
+
def _source(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
|
| 78 |
+
return torch.exp(x) + torch.exp(y)
|
| 79 |
+
|
| 80 |
+
def _residual(
|
| 81 |
+
self, network: nn.Module, x: torch.Tensor, y: torch.Tensor
|
| 82 |
+
) -> tuple[torch.Tensor, torch.Tensor]:
|
| 83 |
+
prediction = network(torch.cat((x, y), dim=1))
|
| 84 |
+
prediction_x = self._gradient(prediction, x)
|
| 85 |
+
prediction_y = self._gradient(prediction, y)
|
| 86 |
+
prediction_xx = self._gradient(prediction_x, x)
|
| 87 |
+
prediction_yy = self._gradient(prediction_y, y)
|
| 88 |
+
residual = prediction_xx + prediction_yy - self._source(x, y)
|
| 89 |
+
return prediction, residual
|
| 90 |
+
|
| 91 |
+
def training_outputs(self, batch: Mapping[str, torch.Tensor]) -> dict[str, torch.Tensor]:
|
| 92 |
+
boundary_prediction = self.n1(torch.cat((batch["xb"], batch["yb"]), dim=1))
|
| 93 |
+
|
| 94 |
+
_, residual1 = self._residual(self.n1, batch["x1"], batch["y1"])
|
| 95 |
+
_, residual2 = self._residual(self.n2, batch["x2"], batch["y2"])
|
| 96 |
+
_, residual3 = self._residual(self.n3, batch["x3"], batch["y3"])
|
| 97 |
+
|
| 98 |
+
interface1_domain1, interface1_residual1 = self._residual(
|
| 99 |
+
self.n1, batch["xi1"], batch["yi1"]
|
| 100 |
+
)
|
| 101 |
+
interface1_domain2, interface1_residual2 = self._residual(
|
| 102 |
+
self.n2, batch["xi1"], batch["yi1"]
|
| 103 |
+
)
|
| 104 |
+
interface2_domain1, interface2_residual1 = self._residual(
|
| 105 |
+
self.n1, batch["xi2"], batch["yi2"]
|
| 106 |
+
)
|
| 107 |
+
interface2_domain3, interface2_residual3 = self._residual(
|
| 108 |
+
self.n3, batch["xi2"], batch["yi2"]
|
| 109 |
+
)
|
| 110 |
+
|
| 111 |
+
interface1_average = 0.5 * (interface1_domain1 + interface1_domain2)
|
| 112 |
+
interface2_average = 0.5 * (interface2_domain1 + interface2_domain3)
|
| 113 |
+
return {
|
| 114 |
+
"boundary_prediction": boundary_prediction,
|
| 115 |
+
"residual1": residual1,
|
| 116 |
+
"residual2": residual2,
|
| 117 |
+
"residual3": residual3,
|
| 118 |
+
"interface1_residual": interface1_residual1 - interface1_residual2,
|
| 119 |
+
"interface2_residual": interface2_residual1 - interface2_residual3,
|
| 120 |
+
"interface1_average": interface1_average,
|
| 121 |
+
"interface2_average": interface2_average,
|
| 122 |
+
"interface1_domain1": interface1_domain1,
|
| 123 |
+
"interface1_domain2": interface1_domain2,
|
| 124 |
+
"interface2_domain1": interface2_domain1,
|
| 125 |
+
"interface2_domain3": interface2_domain3,
|
| 126 |
+
}
|
| 127 |
+
|
| 128 |
+
def predict(
|
| 129 |
+
self, domain1: torch.Tensor, domain2: torch.Tensor, domain3: torch.Tensor
|
| 130 |
+
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
| 131 |
+
return self.n1(domain1), self.n2(domain2), self.n3(domain3)
|
scripts/__pycache__/data_utils.cpython-311.pyc
ADDED
|
Binary file (9.23 kB). View file
|
|
|
scripts/data_utils.py
ADDED
|
@@ -0,0 +1,212 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from __future__ import annotations
|
| 2 |
+
|
| 3 |
+
from collections.abc import Mapping
|
| 4 |
+
from pathlib import Path
|
| 5 |
+
|
| 6 |
+
import numpy as np
|
| 7 |
+
import scipy.io
|
| 8 |
+
import torch
|
| 9 |
+
|
| 10 |
+
|
| 11 |
+
REQUIRED_FIELDS = {
|
| 12 |
+
"x_f1",
|
| 13 |
+
"y_f1",
|
| 14 |
+
"x_f2",
|
| 15 |
+
"y_f2",
|
| 16 |
+
"x_f3",
|
| 17 |
+
"y_f3",
|
| 18 |
+
"xi1",
|
| 19 |
+
"yi1",
|
| 20 |
+
"xi2",
|
| 21 |
+
"yi2",
|
| 22 |
+
"xb",
|
| 23 |
+
"yb",
|
| 24 |
+
"ub",
|
| 25 |
+
"u_exact",
|
| 26 |
+
"u_exact1",
|
| 27 |
+
"u_exact2",
|
| 28 |
+
"u_exact3",
|
| 29 |
+
}
|
| 30 |
+
|
| 31 |
+
|
| 32 |
+
def load_mat_data(path: Path) -> dict:
|
| 33 |
+
if not path.is_file():
|
| 34 |
+
raise FileNotFoundError(f"XPINN MATLAB data not found: {path}")
|
| 35 |
+
data = scipy.io.loadmat(path)
|
| 36 |
+
missing = REQUIRED_FIELDS.difference(data)
|
| 37 |
+
if missing:
|
| 38 |
+
raise ValueError(f"MATLAB data is missing fields: {sorted(missing)}")
|
| 39 |
+
return data
|
| 40 |
+
|
| 41 |
+
|
| 42 |
+
def column(data: Mapping, key: str) -> np.ndarray:
|
| 43 |
+
return np.asarray(data[key], dtype=np.float64).reshape(-1, 1)
|
| 44 |
+
|
| 45 |
+
|
| 46 |
+
def sample_indices(
|
| 47 |
+
generator: np.random.Generator, total_size: int, sample_size: int, name: str
|
| 48 |
+
) -> np.ndarray:
|
| 49 |
+
if sample_size <= 0:
|
| 50 |
+
raise ValueError(f"{name} sample size must be positive")
|
| 51 |
+
if sample_size > total_size:
|
| 52 |
+
raise ValueError(
|
| 53 |
+
f"{name} sample size {sample_size} exceeds available points {total_size}"
|
| 54 |
+
)
|
| 55 |
+
return generator.choice(total_size, sample_size, replace=False)
|
| 56 |
+
|
| 57 |
+
|
| 58 |
+
def tensor(
|
| 59 |
+
values: np.ndarray,
|
| 60 |
+
device: torch.device,
|
| 61 |
+
dtype: torch.dtype,
|
| 62 |
+
requires_grad: bool = False,
|
| 63 |
+
) -> torch.Tensor:
|
| 64 |
+
return torch.as_tensor(values, dtype=dtype, device=device).clone().requires_grad_(
|
| 65 |
+
requires_grad
|
| 66 |
+
)
|
| 67 |
+
|
| 68 |
+
|
| 69 |
+
def paired_sample(
|
| 70 |
+
data: Mapping,
|
| 71 |
+
x_key: str,
|
| 72 |
+
y_key: str,
|
| 73 |
+
sample_size: int,
|
| 74 |
+
generator: np.random.Generator,
|
| 75 |
+
device: torch.device,
|
| 76 |
+
dtype: torch.dtype,
|
| 77 |
+
name: str,
|
| 78 |
+
) -> tuple[torch.Tensor, torch.Tensor]:
|
| 79 |
+
x = column(data, x_key)
|
| 80 |
+
y = column(data, y_key)
|
| 81 |
+
if x.shape != y.shape:
|
| 82 |
+
raise ValueError(f"coordinate shape mismatch for {name}: {x.shape} and {y.shape}")
|
| 83 |
+
indices = sample_indices(generator, x.shape[0], sample_size, name)
|
| 84 |
+
return (
|
| 85 |
+
tensor(x[indices], device, dtype, requires_grad=True),
|
| 86 |
+
tensor(y[indices], device, dtype, requires_grad=True),
|
| 87 |
+
)
|
| 88 |
+
|
| 89 |
+
|
| 90 |
+
def build_training_batch(
|
| 91 |
+
data: Mapping,
|
| 92 |
+
sample_counts: Mapping[str, int],
|
| 93 |
+
seed: int,
|
| 94 |
+
device: torch.device,
|
| 95 |
+
dtype: torch.dtype,
|
| 96 |
+
) -> dict[str, torch.Tensor]:
|
| 97 |
+
generator = np.random.default_rng(seed)
|
| 98 |
+
x1, y1 = paired_sample(
|
| 99 |
+
data,
|
| 100 |
+
"x_f1",
|
| 101 |
+
"y_f1",
|
| 102 |
+
int(sample_counts["residual_1"]),
|
| 103 |
+
generator,
|
| 104 |
+
device,
|
| 105 |
+
dtype,
|
| 106 |
+
"residual_1",
|
| 107 |
+
)
|
| 108 |
+
x2, y2 = paired_sample(
|
| 109 |
+
data,
|
| 110 |
+
"x_f2",
|
| 111 |
+
"y_f2",
|
| 112 |
+
int(sample_counts["residual_2"]),
|
| 113 |
+
generator,
|
| 114 |
+
device,
|
| 115 |
+
dtype,
|
| 116 |
+
"residual_2",
|
| 117 |
+
)
|
| 118 |
+
x3, y3 = paired_sample(
|
| 119 |
+
data,
|
| 120 |
+
"x_f3",
|
| 121 |
+
"y_f3",
|
| 122 |
+
int(sample_counts["residual_3"]),
|
| 123 |
+
generator,
|
| 124 |
+
device,
|
| 125 |
+
dtype,
|
| 126 |
+
"residual_3",
|
| 127 |
+
)
|
| 128 |
+
xi1, yi1 = paired_sample(
|
| 129 |
+
data,
|
| 130 |
+
"xi1",
|
| 131 |
+
"yi1",
|
| 132 |
+
int(sample_counts["interface_1"]),
|
| 133 |
+
generator,
|
| 134 |
+
device,
|
| 135 |
+
dtype,
|
| 136 |
+
"interface_1",
|
| 137 |
+
)
|
| 138 |
+
xi2, yi2 = paired_sample(
|
| 139 |
+
data,
|
| 140 |
+
"xi2",
|
| 141 |
+
"yi2",
|
| 142 |
+
int(sample_counts["interface_2"]),
|
| 143 |
+
generator,
|
| 144 |
+
device,
|
| 145 |
+
dtype,
|
| 146 |
+
"interface_2",
|
| 147 |
+
)
|
| 148 |
+
|
| 149 |
+
boundary_x = column(data, "xb")
|
| 150 |
+
boundary_y = column(data, "yb")
|
| 151 |
+
boundary_values = column(data, "ub")
|
| 152 |
+
if boundary_x.shape != boundary_y.shape or boundary_x.shape != boundary_values.shape:
|
| 153 |
+
raise ValueError("boundary coordinate and value shapes do not match")
|
| 154 |
+
boundary_indices = sample_indices(
|
| 155 |
+
generator,
|
| 156 |
+
boundary_x.shape[0],
|
| 157 |
+
int(sample_counts["boundary"]),
|
| 158 |
+
"boundary",
|
| 159 |
+
)
|
| 160 |
+
return {
|
| 161 |
+
"xb": tensor(boundary_x[boundary_indices], device, dtype),
|
| 162 |
+
"yb": tensor(boundary_y[boundary_indices], device, dtype),
|
| 163 |
+
"ub": tensor(boundary_values[boundary_indices], device, dtype),
|
| 164 |
+
"x1": x1,
|
| 165 |
+
"y1": y1,
|
| 166 |
+
"x2": x2,
|
| 167 |
+
"y2": y2,
|
| 168 |
+
"x3": x3,
|
| 169 |
+
"y3": y3,
|
| 170 |
+
"xi1": xi1,
|
| 171 |
+
"yi1": yi1,
|
| 172 |
+
"xi2": xi2,
|
| 173 |
+
"yi2": yi2,
|
| 174 |
+
}
|
| 175 |
+
|
| 176 |
+
|
| 177 |
+
def build_evaluation_points(
|
| 178 |
+
data: Mapping, device: torch.device, dtype: torch.dtype
|
| 179 |
+
) -> dict[str, torch.Tensor]:
|
| 180 |
+
points = {}
|
| 181 |
+
for domain in (1, 2, 3):
|
| 182 |
+
x = column(data, f"x_f{domain}")
|
| 183 |
+
y = column(data, f"y_f{domain}")
|
| 184 |
+
if x.shape != y.shape:
|
| 185 |
+
raise ValueError(f"evaluation coordinate mismatch in domain {domain}")
|
| 186 |
+
points[f"xy{domain}"] = tensor(np.hstack((x, y)), device, dtype)
|
| 187 |
+
return points
|
| 188 |
+
|
| 189 |
+
|
| 190 |
+
def exact_subdomain_values(
|
| 191 |
+
data: Mapping, device: torch.device, dtype: torch.dtype
|
| 192 |
+
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
| 193 |
+
return tuple(
|
| 194 |
+
tensor(column(data, f"u_exact{domain}"), device, dtype)
|
| 195 |
+
for domain in (1, 2, 3)
|
| 196 |
+
)
|
| 197 |
+
|
| 198 |
+
|
| 199 |
+
def combined_coordinates(data: Mapping) -> tuple[np.ndarray, np.ndarray]:
|
| 200 |
+
x = np.concatenate([column(data, f"x_f{domain}").reshape(-1) for domain in (1, 2, 3)])
|
| 201 |
+
y = np.concatenate([column(data, f"y_f{domain}").reshape(-1) for domain in (1, 2, 3)])
|
| 202 |
+
return x, y
|
| 203 |
+
|
| 204 |
+
|
| 205 |
+
def combined_exact_solution(data: Mapping) -> np.ndarray:
|
| 206 |
+
exact = column(data, "u_exact").reshape(-1)
|
| 207 |
+
expected_size = sum(column(data, f"x_f{domain}").size for domain in (1, 2, 3))
|
| 208 |
+
if exact.size != expected_size:
|
| 209 |
+
raise ValueError(
|
| 210 |
+
f"combined exact solution has {exact.size} values, expected {expected_size}"
|
| 211 |
+
)
|
| 212 |
+
return exact
|
scripts/inference.py
ADDED
|
@@ -0,0 +1,188 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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 matplotlib.tri as tri
|
| 12 |
+
import numpy as np
|
| 13 |
+
import torch
|
| 14 |
+
import yaml
|
| 15 |
+
from matplotlib.patches import Polygon
|
| 16 |
+
|
| 17 |
+
|
| 18 |
+
PROJECT_ROOT = Path(__file__).resolve().parents[1]
|
| 19 |
+
sys.path.insert(0, str(PROJECT_ROOT))
|
| 20 |
+
sys.path.insert(0, str(Path(__file__).resolve().parent))
|
| 21 |
+
|
| 22 |
+
from data_utils import ( # noqa: E402
|
| 23 |
+
build_evaluation_points,
|
| 24 |
+
combined_coordinates,
|
| 25 |
+
combined_exact_solution,
|
| 26 |
+
load_mat_data,
|
| 27 |
+
)
|
| 28 |
+
from model.xpinn import XPINNPoisson2D # noqa: E402
|
| 29 |
+
|
| 30 |
+
|
| 31 |
+
DEFAULT_CONFIG = PROJECT_ROOT / "conf" / "config.yaml"
|
| 32 |
+
|
| 33 |
+
|
| 34 |
+
def load_config(path: Path) -> dict:
|
| 35 |
+
with path.open("r", encoding="utf-8") as stream:
|
| 36 |
+
config = yaml.safe_load(stream)
|
| 37 |
+
if not isinstance(config, dict) or "root" not in config:
|
| 38 |
+
raise ValueError(f"config must contain a 'root' mapping: {path}")
|
| 39 |
+
return config["root"]
|
| 40 |
+
|
| 41 |
+
|
| 42 |
+
def project_path(value: str | Path) -> Path:
|
| 43 |
+
path = Path(value).expanduser()
|
| 44 |
+
return path if path.is_absolute() else PROJECT_ROOT / path
|
| 45 |
+
|
| 46 |
+
|
| 47 |
+
def resolve_device(requested: str) -> torch.device:
|
| 48 |
+
if requested == "auto":
|
| 49 |
+
return torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
| 50 |
+
device = torch.device(requested)
|
| 51 |
+
if device.type == "cuda" and not torch.cuda.is_available():
|
| 52 |
+
raise RuntimeError("CUDA/DCU was requested but torch.cuda.is_available() is false")
|
| 53 |
+
return device
|
| 54 |
+
|
| 55 |
+
|
| 56 |
+
def resolve_dtype(name: str) -> torch.dtype:
|
| 57 |
+
try:
|
| 58 |
+
return {"float32": torch.float32, "float64": torch.float64}[name]
|
| 59 |
+
except KeyError as error:
|
| 60 |
+
raise ValueError(f"unsupported dtype: {name}") from error
|
| 61 |
+
|
| 62 |
+
|
| 63 |
+
def load_model(
|
| 64 |
+
checkpoint_path: Path,
|
| 65 |
+
fallback_model_config: dict,
|
| 66 |
+
device: torch.device,
|
| 67 |
+
dtype: torch.dtype,
|
| 68 |
+
) -> XPINNPoisson2D:
|
| 69 |
+
if not checkpoint_path.is_file():
|
| 70 |
+
raise FileNotFoundError(
|
| 71 |
+
f"checkpoint not found: {checkpoint_path}. Run scripts/train.py first."
|
| 72 |
+
)
|
| 73 |
+
checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=True)
|
| 74 |
+
if not isinstance(checkpoint, Mapping):
|
| 75 |
+
raise ValueError(f"invalid XPINN checkpoint: {checkpoint_path}")
|
| 76 |
+
if "model_state" in checkpoint:
|
| 77 |
+
state = checkpoint["model_state"]
|
| 78 |
+
model_config = checkpoint.get("model_config", fallback_model_config)
|
| 79 |
+
elif checkpoint and all(torch.is_tensor(value) for value in checkpoint.values()):
|
| 80 |
+
state = checkpoint
|
| 81 |
+
model_config = fallback_model_config
|
| 82 |
+
else:
|
| 83 |
+
raise ValueError(f"checkpoint contains no model state: {checkpoint_path}")
|
| 84 |
+
model = XPINNPoisson2D(model_config, dtype=dtype).to(device=device, dtype=dtype)
|
| 85 |
+
model.load_state_dict(state, strict=True)
|
| 86 |
+
model.eval()
|
| 87 |
+
return model
|
| 88 |
+
|
| 89 |
+
|
| 90 |
+
def predict(
|
| 91 |
+
model: XPINNPoisson2D, evaluation_points: dict[str, torch.Tensor]
|
| 92 |
+
) -> np.ndarray:
|
| 93 |
+
with torch.no_grad():
|
| 94 |
+
predictions = model.predict(
|
| 95 |
+
evaluation_points["xy1"],
|
| 96 |
+
evaluation_points["xy2"],
|
| 97 |
+
evaluation_points["xy3"],
|
| 98 |
+
)
|
| 99 |
+
return torch.cat(predictions).cpu().numpy().reshape(-1)
|
| 100 |
+
|
| 101 |
+
|
| 102 |
+
def mask_polygon(data: Mapping) -> np.ndarray:
|
| 103 |
+
boundary_x = np.asarray(data["xb"]).reshape(-1)
|
| 104 |
+
boundary_y = np.asarray(data["yb"]).reshape(-1)
|
| 105 |
+
return np.vstack(
|
| 106 |
+
(
|
| 107 |
+
np.column_stack((boundary_x, boundary_y)),
|
| 108 |
+
np.array(
|
| 109 |
+
[
|
| 110 |
+
[1.8, boundary_y[-1]],
|
| 111 |
+
[1.8, -1.7],
|
| 112 |
+
[-1.6, -1.7],
|
| 113 |
+
[-1.6, 1.55],
|
| 114 |
+
[1.8, 1.55],
|
| 115 |
+
[1.8, boundary_y[-1]],
|
| 116 |
+
]
|
| 117 |
+
),
|
| 118 |
+
np.array([[boundary_x[-1], boundary_y[-1]]]),
|
| 119 |
+
)
|
| 120 |
+
)
|
| 121 |
+
|
| 122 |
+
|
| 123 |
+
def save_figure(
|
| 124 |
+
data: Mapping,
|
| 125 |
+
exact: np.ndarray,
|
| 126 |
+
prediction: np.ndarray,
|
| 127 |
+
output_path: Path,
|
| 128 |
+
) -> None:
|
| 129 |
+
x, y = combined_coordinates(data)
|
| 130 |
+
if exact.size != x.size or prediction.size != x.size:
|
| 131 |
+
raise ValueError("plot coordinates, exact values, and predictions must have equal size")
|
| 132 |
+
triangulation = tri.Triangulation(x, y)
|
| 133 |
+
polygon = mask_polygon(data)
|
| 134 |
+
fields = (
|
| 135 |
+
("Exact", exact),
|
| 136 |
+
("XPINN", prediction),
|
| 137 |
+
("Absolute error", np.abs(exact - prediction)),
|
| 138 |
+
)
|
| 139 |
+
figure, axes = plt.subplots(1, 3, figsize=(18, 5))
|
| 140 |
+
for axis, (title, field) in zip(axes, fields, strict=True):
|
| 141 |
+
contour = axis.tricontourf(triangulation, field, 100, cmap="jet")
|
| 142 |
+
axis.add_patch(
|
| 143 |
+
Polygon(polygon, closed=True, facecolor="white", edgecolor="white")
|
| 144 |
+
)
|
| 145 |
+
axis.plot(data["xi1"].reshape(-1), data["yi1"].reshape(-1), "w-", linewidth=0.5)
|
| 146 |
+
axis.plot(data["xi2"].reshape(-1), data["yi2"].reshape(-1), "w-", linewidth=0.5)
|
| 147 |
+
axis.set_title(title)
|
| 148 |
+
axis.set_xlabel("x")
|
| 149 |
+
axis.set_ylabel("y")
|
| 150 |
+
axis.set_aspect("equal")
|
| 151 |
+
figure.colorbar(contour, ax=axis)
|
| 152 |
+
figure.tight_layout()
|
| 153 |
+
figure.savefig(output_path, dpi=150)
|
| 154 |
+
plt.close(figure)
|
| 155 |
+
|
| 156 |
+
|
| 157 |
+
def main() -> None:
|
| 158 |
+
config_path = DEFAULT_CONFIG.resolve()
|
| 159 |
+
config = load_config(config_path)
|
| 160 |
+
common = config["common"]
|
| 161 |
+
device = resolve_device(str(common["device"]))
|
| 162 |
+
dtype = resolve_dtype(str(common["dtype"]))
|
| 163 |
+
data_path = project_path(config["data"]["mat_file"])
|
| 164 |
+
weight_dir = project_path(common["weight_dir"])
|
| 165 |
+
result_dir = project_path(common["result_dir"])
|
| 166 |
+
checkpoint_path = weight_dir / config["training"]["checkpoint_name"]
|
| 167 |
+
output_path = result_dir / config["inference"]["figure_name"]
|
| 168 |
+
result_dir.mkdir(parents=True, exist_ok=True)
|
| 169 |
+
|
| 170 |
+
print(f"Config: {config_path}")
|
| 171 |
+
print(f"Data: {data_path}")
|
| 172 |
+
print(f"Checkpoint: {checkpoint_path}")
|
| 173 |
+
print(f"Device: {device}")
|
| 174 |
+
data = load_mat_data(data_path)
|
| 175 |
+
model = load_model(checkpoint_path, config["model"], device, dtype)
|
| 176 |
+
evaluation_points = build_evaluation_points(data, device, dtype)
|
| 177 |
+
prediction = predict(model, evaluation_points)
|
| 178 |
+
exact = combined_exact_solution(data)
|
| 179 |
+
relative_l2 = float(
|
| 180 |
+
np.linalg.norm(exact - prediction) / np.linalg.norm(exact)
|
| 181 |
+
)
|
| 182 |
+
save_figure(data, exact, prediction, output_path)
|
| 183 |
+
print(f"Relative L2={relative_l2:.6e}")
|
| 184 |
+
print(f"Plot: {output_path}")
|
| 185 |
+
|
| 186 |
+
|
| 187 |
+
if __name__ == "__main__":
|
| 188 |
+
main()
|
scripts/train.py
ADDED
|
@@ -0,0 +1,208 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from __future__ import annotations
|
| 2 |
+
|
| 3 |
+
import random
|
| 4 |
+
import sys
|
| 5 |
+
import time
|
| 6 |
+
from pathlib import Path
|
| 7 |
+
|
| 8 |
+
import numpy as np
|
| 9 |
+
import torch
|
| 10 |
+
import yaml
|
| 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 data_utils import ( # noqa: E402
|
| 18 |
+
build_evaluation_points,
|
| 19 |
+
build_training_batch,
|
| 20 |
+
exact_subdomain_values,
|
| 21 |
+
load_mat_data,
|
| 22 |
+
)
|
| 23 |
+
from model.xpinn import XPINNPoisson2D # noqa: E402
|
| 24 |
+
|
| 25 |
+
|
| 26 |
+
DEFAULT_CONFIG = PROJECT_ROOT / "conf" / "config.yaml"
|
| 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 |
+
random.seed(seed)
|
| 60 |
+
np.random.seed(seed)
|
| 61 |
+
torch.manual_seed(seed)
|
| 62 |
+
if torch.cuda.is_available():
|
| 63 |
+
torch.cuda.manual_seed_all(seed)
|
| 64 |
+
|
| 65 |
+
|
| 66 |
+
def compute_loss(
|
| 67 |
+
outputs: dict[str, torch.Tensor],
|
| 68 |
+
batch: dict[str, torch.Tensor],
|
| 69 |
+
weights: dict,
|
| 70 |
+
) -> tuple[torch.Tensor, dict[str, float]]:
|
| 71 |
+
boundary_loss = torch.mean(
|
| 72 |
+
(batch["ub"] - outputs["boundary_prediction"]).square()
|
| 73 |
+
)
|
| 74 |
+
pde_loss = sum(
|
| 75 |
+
torch.mean(outputs[f"residual{domain}"].square()) for domain in (1, 2, 3)
|
| 76 |
+
)
|
| 77 |
+
interface_residual_loss = torch.mean(
|
| 78 |
+
outputs["interface1_residual"].square()
|
| 79 |
+
) + torch.mean(outputs["interface2_residual"].square())
|
| 80 |
+
interface_value_loss = (
|
| 81 |
+
torch.mean(
|
| 82 |
+
(outputs["interface1_domain1"] - outputs["interface1_average"]).square()
|
| 83 |
+
)
|
| 84 |
+
+ torch.mean(
|
| 85 |
+
(outputs["interface1_domain2"] - outputs["interface1_average"]).square()
|
| 86 |
+
)
|
| 87 |
+
+ torch.mean(
|
| 88 |
+
(outputs["interface2_domain1"] - outputs["interface2_average"]).square()
|
| 89 |
+
)
|
| 90 |
+
+ torch.mean(
|
| 91 |
+
(outputs["interface2_domain3"] - outputs["interface2_average"]).square()
|
| 92 |
+
)
|
| 93 |
+
)
|
| 94 |
+
total = (
|
| 95 |
+
float(weights["boundary"]) * boundary_loss
|
| 96 |
+
+ float(weights["pde"]) * pde_loss
|
| 97 |
+
+ float(weights["interface_residual"]) * interface_residual_loss
|
| 98 |
+
+ float(weights["interface_value"]) * interface_value_loss
|
| 99 |
+
)
|
| 100 |
+
return total, {
|
| 101 |
+
"boundary": boundary_loss.item(),
|
| 102 |
+
"pde": pde_loss.item(),
|
| 103 |
+
"interface_residual": interface_residual_loss.item(),
|
| 104 |
+
"interface_value": interface_value_loss.item(),
|
| 105 |
+
}
|
| 106 |
+
|
| 107 |
+
|
| 108 |
+
def clear_coordinate_grads(batch: dict[str, torch.Tensor]) -> None:
|
| 109 |
+
for value in batch.values():
|
| 110 |
+
if value.requires_grad and value.grad is not None:
|
| 111 |
+
value.grad = None
|
| 112 |
+
|
| 113 |
+
|
| 114 |
+
def evaluate(
|
| 115 |
+
model: XPINNPoisson2D,
|
| 116 |
+
points: dict[str, torch.Tensor],
|
| 117 |
+
references: tuple[torch.Tensor, torch.Tensor, torch.Tensor],
|
| 118 |
+
) -> tuple[float, float, float, float]:
|
| 119 |
+
with torch.no_grad():
|
| 120 |
+
predictions = model.predict(points["xy1"], points["xy2"], points["xy3"])
|
| 121 |
+
errors = tuple(
|
| 122 |
+
(
|
| 123 |
+
torch.linalg.vector_norm(prediction - reference)
|
| 124 |
+
/ torch.linalg.vector_norm(reference)
|
| 125 |
+
).item()
|
| 126 |
+
for prediction, reference in zip(predictions, references, strict=True)
|
| 127 |
+
)
|
| 128 |
+
combined_prediction = torch.cat(predictions)
|
| 129 |
+
combined_reference = torch.cat(references)
|
| 130 |
+
combined_error = (
|
| 131 |
+
torch.linalg.vector_norm(combined_prediction - combined_reference)
|
| 132 |
+
/ torch.linalg.vector_norm(combined_reference)
|
| 133 |
+
).item()
|
| 134 |
+
return errors[0], errors[1], errors[2], combined_error
|
| 135 |
+
|
| 136 |
+
|
| 137 |
+
def main() -> None:
|
| 138 |
+
config_path = DEFAULT_CONFIG.resolve()
|
| 139 |
+
config = load_config(config_path)
|
| 140 |
+
common = config["common"]
|
| 141 |
+
device = resolve_device(str(common["device"]))
|
| 142 |
+
dtype = resolve_dtype(str(common["dtype"]))
|
| 143 |
+
seed = int(common["seed"])
|
| 144 |
+
data_path = project_path(config["data"]["mat_file"])
|
| 145 |
+
weight_dir = project_path(common["weight_dir"])
|
| 146 |
+
checkpoint_path = weight_dir / config["training"]["checkpoint_name"]
|
| 147 |
+
steps = int(config["training"]["steps"])
|
| 148 |
+
if steps <= 0:
|
| 149 |
+
raise ValueError("training steps must be positive")
|
| 150 |
+
counts = {key: int(value) for key, value in config["data"]["samples"].items()}
|
| 151 |
+
seed_everything(seed)
|
| 152 |
+
|
| 153 |
+
print(f"Config: {config_path}")
|
| 154 |
+
print(f"Data: {data_path}")
|
| 155 |
+
print(f"Device: {device}")
|
| 156 |
+
print(f"Samples: {counts}")
|
| 157 |
+
data = load_mat_data(data_path)
|
| 158 |
+
batch = build_training_batch(data, counts, seed, device, dtype)
|
| 159 |
+
evaluation_points = build_evaluation_points(data, device, dtype)
|
| 160 |
+
references = exact_subdomain_values(data, device, dtype)
|
| 161 |
+
|
| 162 |
+
model = XPINNPoisson2D(config["model"], dtype=dtype).to(
|
| 163 |
+
device=device, dtype=dtype
|
| 164 |
+
)
|
| 165 |
+
optimizer = torch.optim.Adam(
|
| 166 |
+
model.parameters(), lr=float(config["training"]["lr"])
|
| 167 |
+
)
|
| 168 |
+
log_interval = int(config["training"]["log_interval"])
|
| 169 |
+
started = time.time()
|
| 170 |
+
for step in range(1, steps + 1):
|
| 171 |
+
clear_coordinate_grads(batch)
|
| 172 |
+
outputs = model.training_outputs(batch)
|
| 173 |
+
loss, parts = compute_loss(outputs, batch, config["loss"])
|
| 174 |
+
if not torch.isfinite(loss):
|
| 175 |
+
raise FloatingPointError(f"XPINN loss became non-finite at step {step}")
|
| 176 |
+
optimizer.zero_grad(set_to_none=True)
|
| 177 |
+
loss.backward()
|
| 178 |
+
optimizer.step()
|
| 179 |
+
|
| 180 |
+
if step == 1 or step % log_interval == 0 or step == steps:
|
| 181 |
+
error1, error2, error3, combined_error = evaluate(
|
| 182 |
+
model, evaluation_points, references
|
| 183 |
+
)
|
| 184 |
+
print(
|
| 185 |
+
f"step={step:4d} loss={loss.item():.3e} "
|
| 186 |
+
f"boundary={parts['boundary']:.3e} pde={parts['pde']:.3e} "
|
| 187 |
+
f"interface_r={parts['interface_residual']:.3e} "
|
| 188 |
+
f"interface_u={parts['interface_value']:.3e} "
|
| 189 |
+
f"l2=({error1:.3e}, {error2:.3e}, {error3:.3e}) "
|
| 190 |
+
f"combined={combined_error:.3e}"
|
| 191 |
+
)
|
| 192 |
+
|
| 193 |
+
weight_dir.mkdir(parents=True, exist_ok=True)
|
| 194 |
+
checkpoint = {
|
| 195 |
+
"case": "poisson2d",
|
| 196 |
+
"architecture": "xpinn_poisson2d",
|
| 197 |
+
"model_state": model.state_dict(),
|
| 198 |
+
"model_config": config["model"],
|
| 199 |
+
"sample_counts": counts,
|
| 200 |
+
"step": steps,
|
| 201 |
+
}
|
| 202 |
+
torch.save(checkpoint, checkpoint_path)
|
| 203 |
+
print(f"Training finished in {time.time() - started:.1f}s")
|
| 204 |
+
print(f"Saved checkpoint: {checkpoint_path}")
|
| 205 |
+
|
| 206 |
+
|
| 207 |
+
if __name__ == "__main__":
|
| 208 |
+
main()
|
weight/.gitkeep
ADDED
|
File without changes
|
weight/xpinn_poisson_2d.pt
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:8619759612f1cc0ce3a5a7f14db1de2e7da86179a0c193901130cfddda29626e
|
| 3 |
+
size 42450
|