def test_readme():
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)
# In[13]:
bonds, angles, torsions, cartesian, dlogp = coordinate_transform.forward(
training_data[:3])
bonds.shape, angles.shape, torsions.shape, cartesian.shape, dlogp.shape
# ## Prior Distribution
#
# The next step is to define a prior distribution that we can easily sample from. The normalizing flow will be trained to transform such latent samples into molecular coordinates. Here, we just take a normal distribution, which is a rather naive choice for reasons that will be discussed in other notebooks.
# In[14]:
dim_ics = dim_bonds + dim_angles + dim_torsions + dim_cartesian
mean = torch.zeros(dim_ics).to(ctx)
# passing the mean explicitly to create samples on the correct device
prior = bg.NormalDistribution(dim_ics, mean=mean)
# ## Normalizing Flow
#
# Next, we set up the normalizing flow by stacking together different neural networks. For now, we will do this in a rather naive way, not distinguishing between bonds, angles, and torsions. Therefore, we will first define a flow that splits the output from the prior into the different IC terms.
#
# ### Split Layer
# In[15]:
split_into_ics_flow = bg.SplitFlow(dim_bonds, dim_angles, dim_torsions,
dim_cartesian)
# In[16]:
# test
"""
Minimal Boltzmann Generator Example
====================================
In this example a simple Boltzmann Generator is created using coupling layers.
"""
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)))
def test_bg_basic_multiple(device, dtype):
dim = 4
mean = torch.zeros(dim // 2, dtype=dtype, device=device)
import bgflow as bg
prior = bg.ProductDistribution([
bg.NormalDistribution(dim // 2, mean),
bg.NormalDistribution(dim // 2, mean)
])
# RealNVP
flow = bg.SequentialFlow([
bg.CouplingFlow(
bg.AffineTransformer(bg.DenseNet([dim // 2, dim, dim // 2]),
bg.DenseNet([dim // 2, dim, dim // 2]))),
bg.SwapFlow(),
bg.CouplingFlow(
bg.AffineTransformer(bg.DenseNet([dim // 2, dim, dim // 2]),
bg.DenseNet([dim // 2, dim, dim // 2]))),
bg.SwapFlow(),
]).to(mean)
target = bg.ProductDistribution([
bg.NormalDistribution(dim // 2, mean),
bg.NormalDistribution(dim // 2, mean)
])
generator = bg.BoltzmannGenerator(prior, flow, target)
# set parameters to 0 -> flow = id
for p in generator.parameters():
p.data.zero_()
z = prior.sample(10)
*x, dlogp = flow.forward(*z)
for zi, xi in zip(z, x):
assert torch.allclose(zi, xi)
assert torch.allclose(dlogp, torch.zeros_like(dlogp))
# Test losses
generator.zero_grad()
kll = generator.kldiv(100000)
kll.mean().backward()
# gradients should be small, as the network is already optimal
for p in generator.parameters():
assert torch.allclose(p.grad,
torch.zeros_like(p.grad),
rtol=0.0,
atol=5e-2)
generator.zero_grad()
samples = target.sample(100000)
nll = generator.energy(*samples)
nll.mean().backward()
# gradients should be small, as the network is already optimal
for p in generator.parameters():
assert torch.allclose(p.grad,
torch.zeros_like(p.grad),
rtol=0.0,
atol=5e-2)
# just testing the API for the following:
generator.log_weights(*samples)
*z, dlogp = flow.forward(*samples, inverse=True)
generator.log_weights_given_latent(samples, z, dlogp)
generator.sample(10000)
generator.force(*z)
# test trainers
trainer = bg.KLTrainer(generator)
sampler = bg.ProductSampler(
[bg.DataSetSampler(samples[0]),
bg.DataSetSampler(samples[1])]).to(device=device, dtype=dtype)
trainer.train(100, sampler)