Simplifying Hamiltonian and Lagrangian Neural Networks via Explicit Constraints

This repo contains the implementation and the experiments for the paper

Simplifying Hamiltonian and Lagrangian Neural Networks via Explicit Constraints by Marc Finzi, Alex Wang, and Andrew Gordon Wilson.

Some example systems

Each of these animations were produced by running

from IPython.display import HTML
from biases.systems import ChainPendulum, CoupledPendulum, MagnetPendulum
HTML(CoupledPendulum(3).animate())

Code

Our code in the biases directory relies on some publically available codebases which we package together as a conda environment.

Installation instructions

Install PyTorch>=1.0.0

(Optional) Create a wandb account for experiment tracking

Pip

git clone https://github.com/mfinzi/constrained-hamiltonian-neural-networks.git
cd constrained-hamiltonian-neural-networks
pip install -e .

Conda

git clone https://github.com/mfinzi/constrained-hamiltonian-neural-networks.git
cd constrained-hamiltonian-neural-networks
conda env create -f conda_env.yml
pip install ./

Train Models

We have implemented a variety of challenging benchmarks for modeling physical dynamical systems such as ChainPendulum, CoupledPendulum,MagnetPendulum,Gyroscope,Rotor which can be selected with the --body-class argument.

You can run our models CHNN and CLNN as well as the baseline NN (NeuralODE), DeLaN, and HNN models with the network-class argument as shown below.

python pl_trainer.py --network-class CHNN --body-class Gyroscope --wandb-project "YOUR WANDB PROJECT"
python pl_trainer.py --network-class CLNN --body-class Gyroscope --wandb-project "YOUR WANDB PROJECT"
python pl_trainer.py --network-class HNN --body-class Gyroscope --wandb-project "YOUR WANDB PROJECT"
python pl_trainer.py --network-class DeLaN --body-class Gyroscope --wandb-project "YOUR WANDB PROJECT"
python pl_trainer.py --network-class NN --body-class Gyroscope --wandb-project "YOUR WANDB PROJECT"

Our explicitly constrained CHNN and CLNN outperform the competing methods by several orders of magnitude across the different benchmarks as shown below.

If you find our work helpful, please cite it with

@article{finzi2020simplifying,
  title={Simplifying Hamiltonian and Lagrangian Neural Networks via Explicit Constraints},
  author={Finzi, Marc and Wang, Alex and Wilson, Andrew Gordon},
  journal={NeurIPS},
  year={2020}
}