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:

• A general API for Boltzmann Generators
• Different invertible normalizing flow structures to build BGs
  • Coupling flows
  • Neural Ordinary Differential Equations
  • Stochastic Normalizing Flows
  • Equivariant flows
  • Temperature-steerable flows
  • P4Inv Linear Flows
  • Neural Spline Flows
  • Augmented Normalizing Flows
• Other sampling methods
  • Markov-Chain Monte Carlo
  • Molecular dynamics
  • ... replica exchange and more coming soon ...
• API to combine BGs with other sampling methods
• OpenMM bridge
• Internal Coordinate Transformations to map between Z-matrices and Cartesian coordinates

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
)

• Basic Boltzmann Generator example
• Training a Boltzmann Generator for Alanine Dipeptide
• Equivariant kernel flow example
• Equivariant Real NVP flow example

• 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
  • einops
  • pytorch
  • numpy
• 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