amazon-science/probconserv

Datasets and code for results presented in the ProbConserv paper

ProbConserv: Probabilistic Framework to Enforce Conservation Laws

Derek Hansen, Danielle C. Maddix, Shima Alizadeh, Gaurav Gupta, Michael W. Mahoney
Learning Physical Models that Can Respect Conservation Laws
Proceedings of the 40th International Conference on Machine Learning (ICML), PMLR. 202:12469-12510, 2023.

Installation

This project uses poetry to manage dependencies.

From the root directory:

poetry install

Some of the plots require certain LaTeX packages are present. On Ubuntu, these are

sudo apt install cm-super dvipng texlive-latex-extra texlive-fonts-recommended

You can then use poetry run followed by a command, or poetry shell to open a shell with the correct virtual environment.

To run the tests:

poetry run pytest

The code for this project is located in the deep_pdes folder. It consists of two libraries, attentive_neural_process and datasets, that comprise the models and datasets respectively. These libraries are imported by the scripts in experiments that configure and run the specific case studies explored in the ProbConserv paper.

Running experiments

The experiment code in deep_pdes/experiment uses Hydra to manage configuration and run experiments. The different stages of the experiments are broken into distinct commands for easier reproduceability

• generate.py: Generate synthetic datasets for training
• train.py: Train ProbConserv-ANP, ANP, and other baseline methods such as Physics-Informed Neural Networks (PINNs)
• analyze.py: Evaluate the trained models on test datasets and create tables/plots from the results.
• plots.py: Generate all plots used in the submission. Does not use the Hydra CLI but uses the compose API internally.

Each script is run by passing an +experiments=* flag. The available experiments can be found in deep_pdes/experiments/conf/experiments. For example, to recreate the results on the Stefan GPME setting:

EXPERIMENT=2b_stefan_var_p

python generate.py +experiments=$EXPERIMENT
python train.py +experiments=$EXPERIMENT +train=${EXPERIMENT}_anp
python train.py +experiments=$EXPERIMENT +train=${EXPERIMENT}_pinp
python analyze.py +experiments=$EXPERIMENT

These commands are also available in convenience scripts; for example, the above is in deep_pdes/experiments/2b_stefan_var_p.sh.

Solution Profiles and UQ for the Stefan Equation

Downstream Task: Shock location detection

For the diffusion equation with constant diffusivity, see deep_pdes/experiments/3b_heat_var_c.sh.
Conservation of mass can be violated by the black-box deep learning models, even with applying the PDE as a soft-constraint to the loss function a la Physics informed Neural Networks (PINNs). The true mass for this diffusion equation is zero over time since there is zero net flux from the domain boundaries and mass cannot be created or destroyed on the interior.

Sources

This repo contains modified versions of the code found in the following repos:

• https://github.com/a1k12/characterizing-pinns-failure-modes: For diffusion/heat equation analytical solution (MIT license)
• https://github.com/soobinseo/Attentive-Neural-Process: For implementation of the Attentive Neural Process (Apache 2.0 license)

Citation

If you use this code, or our work, please cite:

@inproceedings{hansen2023learning,
  title={Learning Physical Models that Can Respect Conservation Laws},
  author={Hansen, Derek and Maddix, Danielle C. and Alizadeh, Shima and Gupta, Gaurav and Mahoney, Michael W},
  booktitle={International Conference on Machine Learning},
  year={2023},
  volume = {202},
  pages={12469-12510},
  organization={PMLR}      
}

probconserv's Issues

Unable to run sh script inside packages

Dear All,

I tried to set up the ProbConserv package without any good result. I have a Ubuntu 22.04 LTS machine with a virtual environment dedicated. The Python version in 3.10.13.
I installed poetry, I created the poetry lock file commenting the line wemake-python-styleguide = "^0.16.1" (otherwise I had a warning concerning the version of python that is compliant with the warning) and then I installed using pip the version of wemake-python-styleguide required.
When I try to run, for example, the 2b_stefan_var_p.sh script this is the output:
`
Traceback (most recent call last):
File "/home/munerato/PhysicalLearning/probconserv/deep_pdes/experiments/generate.py", line 7, in
from deep_pdes.datasets.base import ANPDataset
ModuleNotFoundError: No module named 'deep_pdes'
Error executing job with overrides: ['+experiments=2b_stefan_var_p', '+train=2b_stefan_var_p_anp']
Error locating target 'deep_pdes.attentive_neural_process.anp.ANP', see chained exception above.
full_key: methods.anp.model

Set the environment variable HYDRA_FULL_ERROR=1 for a complete stack trace.
Error executing job with overrides: ['+experiments=2b_stefan_var_p', '+train=2b_stefan_var_p_pinp']
Error locating target 'deep_pdes.attentive_neural_process.softc.PINP', see chained exception above.
full_key: methods.pinp.model

Set the environment variable HYDRA_FULL_ERROR=1 for a complete stack trace.
Traceback (most recent call last):
File "/home/munerato/PhysicalLearning/probconserv/deep_pdes/experiments/analyze.py", line 16, in
from deep_pdes.attentive_neural_process.anp import ANP
ModuleNotFoundError: No module named 'deep_pdes'
`

Do you have any ideas how to resolve it?

Thanks in advance,
Mauro

Git LFS files are missing

In the original repo, we stored some training datasets in Git LFS, not GIt itself. It seems that these weren't copied over properly when we created the public repo.

Here's a list of the files:

00f25b170c - deep_pdes/experiments/output/paper/1b_pme_var_m/datasets/test.pt
5238cd6b0d - deep_pdes/experiments/output/paper/1b_pme_var_m/datasets/train.pt
58ca6bdfb3 - deep_pdes/experiments/output/paper/1b_pme_var_m/datasets/valid.pt
6c63edb64b - deep_pdes/experiments/output/paper/1b_pme_var_m/train/anp.pt
cbeca801e3 - deep_pdes/experiments/output/paper/1b_pme_var_m/train/physnp.pt
18307bf894 - deep_pdes/experiments/output/paper/1b_pme_var_m/train/pinp.pt
2444d833ee - deep_pdes/experiments/output/paper/2b_stefan_var_p/datasets/test.pt
2fb4a83aeb - deep_pdes/experiments/output/paper/2b_stefan_var_p/datasets/train.pt
6da1c5ae30 - deep_pdes/experiments/output/paper/2b_stefan_var_p/datasets/valid.pt
ae90105256 - deep_pdes/experiments/output/paper/2b_stefan_var_p/train/anp.pt
88f21da3bc - deep_pdes/experiments/output/paper/2b_stefan_var_p/train/physnp.pt
b2e0c79c55 - deep_pdes/experiments/output/paper/2b_stefan_var_p/train/physnp_second_deriv.pt

We need to get these files from the old repo and upload them as Git LFS files to this repo. Otherwise, you'll get an error when trying to reproduce the results.