/bgflow

Boltzmann Generators and Normalizing Flows in PyTorch

Primary LanguageJupyter NotebookMIT LicenseMIT

bgflow

Language grade: Python

Bgflow is a pytorch framework for Boltzmann Generators (BG) and other sampling methods.

Boltzmann Generators use neural networks to learn a coordinate transformation of the complex configurational equilibrium distribution to a distribution that can be easily sampled. Accurate computation of free-energy differences and discovery of new configurations are demonstrated, providing a statistical mechanics tool that can avoid rare events during sampling without prior knowledge of reaction coordinates. -- Noe et al. (2019)

This framework provides:


This library is alpha software under active development. Certain elements of its API are going to change in the future and some current implementations may not be tested.

If you are interested in contributing to the library, please feel free to open a pull request or report an issue.

When using bgflow in your research, please cite our preprint (coming soon).


Implementation of a BG with a single Real NVP coupling block as the invertible transformation. The two-dimensional target potential is given by a double well potential in one dimension and a harmonic potential in the other. The prior distribution is a two-dimensional standard normal distribution. Note that the training procedure is not included in this example.

import torch
import matplotlib.pyplot as plt
import bgflow as bg

# define prior and target
dim = 2
prior = bg.NormalDistribution(dim)
target = bg.DoubleWellEnergy(dim)

# here we aggregate all layers of the flow
layers = []
layers.append(bg.SplitFlow(dim // 2))
layers.append(bg.CouplingFlow(
        # we use a affine transformation to transform 
        # the RHS conditioned on the LHS
        bg.AffineTransformer(
            # use simple dense nets for the affine shift/scale
            shift_transformation=bg.DenseNet(
                [dim // 2, 4, dim // 2], 
                activation=torch.nn.ReLU()
            ), 
            scale_transformation=bg.DenseNet(
                [dim // 2, 4, dim // 2], 
                activation=torch.nn.Tanh()
            )
        )
))
layers.append(bg.InverseFlow(bg.SplitFlow(dim // 2)))
    
# now define the flow as a sequence of all operations stored in layers
flow = bg.SequentialFlow(layers)

# The BG is defined by a prior, target and a flow
generator = bg.BoltzmannGenerator(prior, flow, target)

# sample from the BG
samples = generator.sample(1000)
plt.hist2d(
    samples[:, 0].detach().numpy(), 
    samples[:, 1].detach().numpy(), bins=100
)


  • Clone this repository from github
  • Navigate to the cloned repository
  • Run the installation scrip
python setup.py install
  • Install all required dependencies
  • Validate your installation by running all tests in the repository with the command
pytest
  • Depending on the optional installations some tests might be skipped.

  • Mandatory
  • Optional
    • matplotlib
    • pytest (for testing)
    • nflows (for Neural Spline Flows)
    • torchdiffeq (for neural ODEs)
    • ANODE (for neural ODEs)
    • OpenMM (for molecular mechanics energies)
    • ase (for quantum and molecular mechanics energies through the atomic simulation environment)
    • xtb-python (for semi-empirical GFN quantum energies)
    • netCDF4 (for the ReplayBufferReporter)
    • jax (for smooth flows / implicit backprop)
    • jax2torch (for smooth flows / implicit backprop)
    • allegro (for Graph Neural Networks)
    • nequip (for Graph Neural Networks)
    • bgmol (for some example notebooks)

MIT License