Project Page | Video | Paper
Official PyTorch implementation.
Learning to Solve PDE-constrained Inverse Problems with Graph Networks
Qingqing Zhao*,
David B. Lindell,
Gordon Wetzstein
Stanford University
To setup a conda environment use these commands
conda env create -f environment.yml
conda activate gnn
We also need to install pytorch and pytorch-geometric with following commands:
pip install torch==1.10.0+cu113 torchvision==0.11.1+cu113 torchaudio==0.10.0+cu113 -f https://download.pytorch.org/whl/cu113/torch_stable.html
wget https://data.pyg.org/whl/torch-1.10.0%2Bcu113/torch_cluster-1.5.9-cp37-cp37m-linux_x86_64.whl
pip install torch_cluster-1.5.9-cp37-cp37m-linux_x86_64.whl
wget https://data.pyg.org/whl/torch-1.10.0%2Bcu113/torch_scatter-2.0.9-cp37-cp37m-linux_x86_64.whl
pip install torch_scatter-2.0.9-cp37-cp37m-linux_x86_64.whl
wget https://data.pyg.org/whl/torch-1.10.0%2Bcu113/torch_sparse-0.6.12-cp37-cp37m-linux_x86_64.whl
pip install torch_sparse-0.6.12-cp37-cp37m-linux_x86_64.whl
wget https://data.pyg.org/whl/torch-1.10.0%2Bcu113/torch_spline_conv-1.2.1-cp37-cp37m-linux_x86_64.whl
pip install torch_spline_conv-1.2.1-cp37-cp37m-linux_x86_64.whl
pip install torch-geometric
Dataset and pretrained model and validation samples can be download here. Unzip the data.zip folder in the root directory.
Now you can solve invere problem with 2D wave equation with the following commands.
# with prior
python InverseProblem/experiment_scripts/run_gnn.py --config InverseProblem/config/density_gnn_p.ini
python InverseProblem/experiment_scripts/run_gnn.py --config InverseProblem/config/init_state_gnn_p.ini
# without prior
python InverseProblem/experiment_scripts/run_gnn.py --config InverseProblem/config/density_gnn_np.ini
python InverseProblem/experiment_scripts/run_gnn.py --config InverseProblem/config/init_state_gnn_np.ini
You may also run the notebooks for a quick demo and visualization.
| File | Description |
|---|---|
| notebook/inverse_wave_equation_density.ipynb | Full Waveform Inversion |
| notebook/inverse_wave_equation_init.ipynb | Initial State Recovery |
We also provide sample training script for both GNN and prior network. Training dataset for both can be downloaded from here and unzip the data.zip folder in the root directory.
# train GNN forward model
python GNN/train_2d_wave_equation.py --file ./data/training --diffML --normalize --log --lr_schedule
# train generative prior
python Prior/autodecoder.py --num_pe_fns 3 --use_pe --dataset_size 10000 --batch_size 32 --gpu 1 --regularize --irregular_mesh --jitter --prior init_state
python Prior/autodecoder.py --num_pe_fns 3 --use_pe --dataset_size 10000 --batch_size 32 --gpu 1 --regularize --irregular_mesh --jitter --prior density
@inproceedings{qzhao2022graphpde,
title={Learning to Solve PDE-constrained Inverse Problems with Graph Networks},
author={Qingqing Zhao and David B. Lindell and Gordon Wetzstein}
journal={ICML},
year={2022}
}