jhbastek / PhysicsInformedDiffusionModels

Implementation of Physics-Informed Diffusion Models

Physics-Informed Diffusion Models

Jan-Hendrik Bastek¹, WaiChing Sun² and Dennis M. Kochmann¹

¹ETH Zurich, ²Columbia University

$~$

Introduction & Setup

We introduce a framework to inform diffusion models of constraints generated samples must adhere to during model training, as presented in Physics-Informed Diffusion Models. To conduct similar studies as those presented in the preprint, start by cloning this repository via

git clone https://github.com/jhbastek/PhysicsInformedDiffusionModels.git

We provide three scripts:

main_toy.py reproduces the toy study presented in Appendix F.1. It is helpful to understand the implications of the PIDM loss and several variants. Simply change the config file and run the script to reproduce the results or experiment with different parameters.

To reproduce the results for the Darcy flow and topology optimization study, you will first have to download the data and pretrained models from the ETHZ Research Collection and place them (unzipped) as follows:

.
├── data
│   ├── darcy
│   │   └── ...
│   └── mechanics
│       └── ...
└── trained_models
    ├── darcy
    │   └── ...
    └── mechanics
        └── ...

After this, you can run the following scripts:

main.py reproduces the Darcy flow and topology optimization study presented in Section 4. Simply adjust the parameters and governing equations in model.yaml and run the script to train the models. Note that the name of the run and logging parameters can be directly adjusted in main.py, if necessary.

sample.py evaluates trained models. Provide the directory_path, name, and load_model_step of the model to evaluate and run the script. Note that the full evaluation of the in- and out-of-distribution test sets for the topology optimization study may take some time.

Dependencies

The framework was developed and tested on Python 3.11 using CUDA 12.0. To run the toy model, the following packages are required:

Package	Version (>=)
pytorch	2.0.1
tqdm	4.65.0
matplotlib	3.7.2
imageio	2.28.1
einops	0.6.1
wandb (optional)	0.15.2

To run the Darcy flow and topology optimization study, the following additional packages are required:

Package	Version (>=)
findiff	0.10.0
solidspy	1.0.16
pandas	2.1.3
einops-exts	0.0.4
rotary_embedding_torch	0.2.3
torchvision	0.15.2
opencv	4.9.0.80

Citation

If this code is useful for your research, please cite our preprint:

@article{Bastek2024,
      title={Physics-Informed Diffusion Models}, 
      author={Jan-Hendrik Bastek and WaiChing Sun and Dennis M. Kochmann},
      year={2024},
      eprint={2403.14404},
      archivePrefix={arXiv},
      primaryClass={cs.LG}
}