https://github.com/gregversteeg/esh_dynamics

Hamiltonian Dynamics with Non-Newtonian Momentum for Rapid Sampling
https://github.com/gregversteeg/esh_dynamics

Last synced: 3 months ago
JSON representation

Hamiltonian Dynamics with Non-Newtonian Momentum for Rapid Sampling

• Host: GitHub
• URL: https://github.com/gregversteeg/esh_dynamics
• Owner: gregversteeg
• License: mit
• Created: 2021-10-16T17:08:08.000Z (about 3 years ago)
• Default Branch: master
• Last Pushed: 2021-11-05T00:34:05.000Z (about 3 years ago)
• Last Synced: 2024-07-03T23:46:14.028Z (4 months ago)
• Language: Python
• Size: 23 MB
• Stars: 34
• Watchers: 4
• Forks: 1
• Open Issues: 0
• Metadata Files:
  • Readme: README.md
  • License: LICENSE.md

Awesome Lists containing this project

README

        
# Hamiltonian Dynamics with Non-Newtonian Momentum for Rapid Sampling

Code for the paper:

> Greg Ver Steeg and Aram Galstyan. "Hamiltonian Dynamics with Non-Newtonian Momentum for Rapid Sampling", NeurIPS 2021.

> [[arxiv]](http://arxiv.org/abs/2111.02434) [[bibtex]](#bibtex)







Non-Newtonian Momentum Animation:







This repo contains code for implementing **E**nergy **S**ampling **H**amiltonian Dynamics, 

so-called because the Hamiltonian dynamics with this special form of Non-Newtonian momentum 

ergodically samples from a target un-normalized density specified by an energy function. 

## Requirements

The core ESH dynamics sampler code (import esh) uses only PyTorch. 

```

python -m pip install git+https://github.com/gregversteeg/esh_dynamics

```

Use ``pip install -r requirements.txt`` to install requirements for all comparison code. 

## Usage

Here's a small example where we load a pytorch energy function, then sample Langevin versus ESH trajectories. 

```python

import torch as t

import esh  # ESH Dynamics integrator

from esh.datasets import ToyDataset  # Example energy models

from esh.samplers import hmc_integrate  # Sampling comparison methods, like Langevin

# Energy to sample - any pytorch function/module that outputs a scalar per batch item

energy = ToyDataset(toy_type='gmm').energy  # Gaussian mixture model

epsilon = 0.01  # Step size should be < 1

n_steps = 100  # Number of steps to take

x0 = t.tensor([[0., 0.5]])  # Initial state, size (batch_size, ...)

xs, vs, rs = esh.leap_integrate_chain(energy, x0, n_steps, epsilon, store=True)  # "Store" returns whole trajectory

xs_ula, vs_ula, _ = hmc_integrate(energy, x0, n_steps, epsilon=epsilon, k=1, mh_reject=False)  # Unadjusted Langevin Alg

```

To get just the last state instead of the whole trajectory, set store=False. 

To do ergodic reservoir sampling, set reservoir=True, store=False. 

## Generating figures

See the README in the ``generate_figures`` for scripts to generate each figure in the paper, 

and to see more example usage. 

## BibTeX

```markdown

@inproceedings{esh,

  title={Hamiltonian Dynamics with Non-Newtonian Momentum for Rapid Sampling},

  author={Greg {Ver Steeg} and Aram Galstyan},

  Booktitle={Advances in Neural Information Processing Systems},

  year={2021}

}

```