This repository is the official implementation of a NeurIPS 2022 paper: "Physics-Embedded Neural Networks: Graph Neural PDE Solvers with Mixed Boundary Conditions."
Please cite us as:
@inproceedings{
horie2022physicsembedded,
title={Physics-Embedded Neural Networks: Graph Neural {PDE} Solvers with Mixed Boundary Conditions},
author={Masanobu Horie and Naoto Mitsume},
booktitle={Advances in Neural Information Processing Systems},
editor={Alice H. Oh and Alekh Agarwal and Danielle Belgrave and Kyunghyun Cho},
year={2022},
url={https://openreview.net/forum?id=B3TOg-YCtzo}
}
This code is based on the original IsoGCN's code: https://github.com/yellowshippo/isogcn-iclr2021 (including FEMIO (https://github.com/ricosjp/femio) and SiML (https://github.com/ricosjp/siml) as submodules with some modifications), except for mp-neural-pde-solvers based on https://github.com/brandstetter-johannes/MP-Neural-PDE-Solvers.
NOTE: This project requires Python3.9 and CUDA 11.1.
This repository uses submodules, so to clone:
git clone git@github.com:yellowshippo/penn_neurips2022.git --recurse-submodules
To install requirements:
make install
Or, if you have no GPUs:
make install_cpu
To download the gradient dataset (1.7 GB):
make grad_data
Or to generate dataset:
make generate_grad
NOTE: We do not plan to add a persistent identifier, metadata, and license to this dataset because these data can be generated from the code easily.
To download the advection-diffusion dataset (350 MB):
make ad_data
NOTE: We do not plan to add a persistent identifier, metadata, and license to this dataset because these data can be generated from the code easily.
To download the incompressible flow dataset (24 GB):
make fluid_data
NOTE: We plan to add a persistent identifier, metadata, and license to this dataset upon publication.
Raw data (OpenFOAM analysis results) are available in (5 GB for each, 303 GB in total):
- https://savanna.ritc.jp/~horiem/penn_neurips2022/data/fluid/raw.tar.gz.partaa
- https://savanna.ritc.jp/~horiem/penn_neurips2022/data/fluid/raw.tar.gz.partab
- ...
- https://savanna.ritc.jp/~horiem/penn_neurips2022/data/fluid/raw.tar.gz.partci
You can download pretrained models using:
make pretrained_models
The downloaded models look as follows:
pretrained
│ # Models for the gradient dataset
├── grad
│ ├── isogcn # The original IsoGCN model
│ └── neumannisogcn # The proposed NeumannIsoGCN (NIsoGCN) model
│
│ # Models for the advection-diffusion dataset
├── ad
│ │ # The proposed model (PENN)
│ ├── penn
│ │
│ │ # Ablation study models
│ ├── wo_boundary_condition_in_nns
│ ├── wo_boundary_condition_input
│ ├── wo_dirichlet_layer
│ ├── wo_encoded_boundary
│ ├── wo_neural_nonlinear_solver
│ ├── wo_pseudoinverse_decoder
│ └── wo_pseudoinverse_decoder_w_dirichlet_layer_after_decoding
│
│ # Models for the incompressible dataset
└── fluid
│
│ # The proposed model (PENN) (n: hidden feature size, rep: iteration count)
├── penn_n16_rep8 # Reference
├── penn_n16_rep4
├── penn_n8_rep8
├── penn_n8_rep4
├── penn_n4_rep8
├── penn_n4_rep4
│
│ # MP-PDE models (tw: time window size, n: hidden feature size)
├── mp-pde_tw20_n128
├── mp-pde_tw20_n64
├── mp-pde_tw20_n32
├── mp-pde_tw10_n128
├── mp-pde_tw10_n64
├── mp-pde_tw10_n32
├── mp-pde_tw4_n128
├── mp-pde_tw4_n64
├── mp-pde_tw4_n32
├── mp-pde_tw2_n128
├── mp-pde_tw2_n64
├── mp-pde_tw2_n32
│
│ # Ablation study models
├── wo_boundary_condition_in_nns
├── wo_boundary_condition_input
├── wo_dirichlet_layer
├── wo_encoded_boundary
├── wo_neural_nonlinear_solver
├── wo_pseudoinverse_decoder
└── wo_pseudoinverse_decoder_w_dirichlet_layer_after_decodingTo train the NIsoGCN model in the paper, run:
make grad_train
To train the original IsoGCN model (baseline), run:
make grad_train GRAD_MODEL=isogcn
The corresponding model architectures are written in the following YAML files:
data/grad
├── neumannisogcn.yml # NIsoGCN
└── isogcn.yml # IsoGCNTo train the proposed model in the paper (PENN), run:
make fluid_train
For the ablation study run in the paper, run:
make fluid_train FLUID_MODEL=model_name
The possible options of model_name are the following (see Appendix C.7 for more details):
-
A:
wo_encoded_boundary -
B:
wo_boundary_condition_in_nns -
C:
wo_neural_nonlinear_solver -
D:
wo_boundary_condition_input -
E:
wo_dirichlet_layer -
F:
wo_pseudoinverse_decoder -
G:
wo_pseudoinverse_decoder_w_dirichlet_layer_after_decoding
The corresponding model architectures are written in the following YAML files:
data/fluid
│
│ # The proposed model (PENN) (n: hidden feature size, rep: iteration count)
├── penn_n16_rep8.yml
├── penn_n16_rep4.yml
├── penn_n4_rep8.yml
├── penn_n4_rep4.yml
├── penn_n8_rep8.yml
├── penn_n8_rep4.yml
│
│ # Ablation study models
├── wo_boundary_condition_in_nns.yml
├── wo_boundary_condition_input.yml
├── wo_dirichlet_layer.yml
├── wo_encoded_boundary.yml
├── wo_neural_nonlinear_solver.yml
├── wo_pseudoinverse_decoder_w_dirichlet_layer_after_decoding.yml
└── wo_pseudoinverse_decoder.ymlTo train the MP-PDE model (baseline), run:
make fluid_train_mppde TW=time_window
, where time_window is the time window size, which can be 2, 4, 10, or 20.
To train the proposed model in the paper (PENN), run:
make ad_train
For the ablation study run in the paper, run:
make ad_train AD_MODEL=model_name
The possible options of model_name are the same as these in the Incompressible flow dataset.
The corresponding model architectures are written in the following YAML files:
data/ad
│
│ # The proposed model (PENN)
├── penn.yml
│
│ # Ablation study models
├── wo_boundary_condition_in_nns.yml
├── wo_boundary_condition_input.yml
├── wo_dirichlet_layer.yml
├── wo_encoded_boundary.yml
├── wo_neural_nonlinear_solver.yml
├── wo_pseudoinverse_decoder_w_dirichlet_layer_after_decoding.yml
└── wo_pseudoinverse_decoder.ymlTo evaluate the NIsoGCN model or original IsoGCN model, run:
make grad_eval PRETRAINED_GRAD_MODEL=path/to/model/directory
To evaluate the proposed model (PENN) or its ablation, run:
make fluid_eval PRETRAINED_PENN_MODEL=path/to/model/directory
To evaluate MP-PDE model (baseline), run:
make fluid_eval_mppde PRETRAINED_MPPDE_MODEL=path/to/model/directory TW=time_window
To evaluate the proposed model (PENN) or its ablation on the transformed incompressible fluid dataset, run:
make transformed_fluid_eval PRETRAINED_PENN_MODEL=path/to/model/directory
To evaluate MP-PDE model (baseline) on the transformed incompressible fluid dataset, run:
make transformed_fluid_eval_mppde PRETRAINED_MPPDE_MODEL=path/to/model/directory TW=time_window
To evaluate the proposed model (PENN) or its ablation, run:
make ad_eval PRETRAINED_AD_MODEL=path/to/model/directory
| Model name | Command to reproduce | Loss (1e-3) | Neumann loss (1e-3) |
|---|---|---|---|
| IsoGCN | make grad_eval PRETRAINED_GRAD_MODEL=pretrained/grad/isogcn |
192.72 | 1390.95 |
| NIsoGCN (Ours) | make grad_eval |
6.70 | 3.52 |
| Model name | Command to reproduce | Loss T (1e-4) | Dirichlet loss T (1e-4) |
|---|---|---|---|
| PENN (Ours) | make ad_eval |
1.795 | 0.00 |
| Model name | Command to reproduce | Loss u (1e-4) | Loss p (1e-3) | Dirichlet loss u (1e-4) | Dirichlet loss p (1e-3) |
|---|---|---|---|---|---|
| MP-PDE TW=20 | make fluid_eval_mppde |
1.30 | 1.32 | 0.45 | 0.28 |
| PENN (Ours) | make fluid_eval |
4.36 | 1.17 | 0.00 | 0.00 |
| Model name | Command to reproduce | Loss u (1e-4) | Loss p (1e-3) | Dirichlet loss u (1e-4) | Dirichlet loss p (1e-3) |
|---|---|---|---|---|---|
| MP-PDE TW=20 | make transformed_fluid_eval_mppde |
1953.62 | 281.86 | 924.73 | 202.97 |
| PENN (Ours) | make transformed_fluid_eval |
4.36 | 1.17 | 0.00 | 0.00 |
Please use a browser and input URLs written in the Makefile if downloads fail.