Download Notebook View in GitHub Open in Google Colab

Alchemical free energy calculations

Computing the free energy of inserting a Lennard-Jones particle in a Lennard-Jones fluid.

This tutorial is described in OpenMM 7 publication.

Basic concept

OpenMM’s custom forces—which allow the programmer to express a potential algebraically, potentially with multiple parameters that can be adjusted on the fly—allow a great deal of flexibility and simplicity in encoding potentials while still achieving high performance on GPUs. One common use of this facility is to convert standard interactions (such as Lennard-Jones potentials) into alchemically-modified potentials for the purposes of computing free energy differences. The alchemical free energy code YANK, for example, uses a variety of custom forces to represent alchemically-modified potentials for protein-ligand alchemical binding free energy calculations.

Defining alchemically-modified potentials

As a simple example of how this is facilitated by custom forces, consider computing the chemical potential of liquid argon by estimating the free energy of alchemically annihilating a Lennard-Jones particle. First, we create a simple Lennard-Jones fluid to represent liquid argon at 120 K and 80 atm, which can be conveniently done using the testsystems module of the conda-installable openmmtools package:

# openmmtools is also needed for this notebook
!mamba install -y -c conda-forge openmmtools

from openmm.app import *
from openmm import *
from openmm.unit import *
from openmmtools.testsystems import LennardJonesFluid


# Simulation settings
pressure = 80*atmospheres
temperature = 120*kelvin
collision_rate = 5/picoseconds
timestep = 2.5*femtoseconds


# Create a Lennard Jones test fluid
sigma = 3.4*angstrom
epsilon = 0.238 * kilocalories_per_mole
fluid = LennardJonesFluid(sigma=sigma, epsilon=epsilon)
[topology, system, positions] = [fluid.topology, fluid.system, fluid.positions]

# Add a barostat
barostat = MonteCarloBarostat(pressure, temperature)
system.addForce(barostat)

To allow one of the Lennard-Jones particles to be alchemically eliminated, we create a CustomNonbondedForce that will compute the interactions between the alchemical particle and the remaining chemical particles using a softcore potential. The alchemically-modified particle has its Lennard-Jones well depth (epsilon parameter) set to zero in the original NonbondedForce, while the CustomNonbondedForce is set to evaluate only the interactions between the alchemically-modified particle and the remaining particles using addInteractionGroup() to specify only interactions between these groups are to be computed. A global context parameter lambda is created to control the coupling of the alchemically-modified particle with the rest of the system during the simulation. The Lennard-Jones parameters sigma and epsilon are implemented as per-particle parameters, though this is not strictly necessary in this case since all particles are equivalent.

# Retrieve the NonbondedForce
forces = { force.__class__.__name__ : force for force in system.getForces() }
nbforce = forces['NonbondedForce']

# Add a CustomNonbondedForce to handle only alchemically-modified interactions

# Make two sets of particles, one that contains just the particle we will alchemically annihilate
# and the other which contains all the other particles.
alchemical_particles = set([0])
chemical_particles = set(range(system.getNumParticles())) - alchemical_particles


# Define the energy function for the CustomNonbondedForce
# when lambda is 1.0 it is a normal LJ potential, when lambda is 0.0 the interaction vanishes
energy_function = 'lambda*4*epsilon*x*(x-1.0); x = (sigma/reff_sterics)^6;'
energy_function += 'reff_sterics = sigma*(0.5*(1.0-lambda) + (r/sigma)^6)^(1/6);'
energy_function += 'sigma = 0.5*(sigma1+sigma2); epsilon = sqrt(epsilon1*epsilon2);'
custom_force = CustomNonbondedForce(energy_function)

# Add lambda as a parameter we can change during the simulation
custom_force.addGlobalParameter('lambda', 1.0)

# set the values of sigma and epsilon by copying them from the existing NonBondedForce
custom_force.addPerParticleParameter('sigma')
custom_force.addPerParticleParameter('epsilon')
for index in range(system.getNumParticles()):
    [charge, sigma, epsilon] = nbforce.getParticleParameters(index)
    custom_force.addParticle([sigma, epsilon])
    if index in alchemical_particles:
        # remove the alchemical particle from the existing NonBondedForce
        nbforce.setParticleParameters(index, charge*0, sigma, epsilon*0)

# Set the custom force to occur between just the alchemical particle and the other particles
custom_force.addInteractionGroup(alchemical_particles, chemical_particles)
system.addForce(custom_force)

Simulating alchemically-modified systems

We then create a LangevinIntegrator and Simulation to run the simulation, and run a series of simulations at different values of lambda by using simulation.context.setParameter() to update the alchemical parameter on the fly. For each configuration sample that is collected, we can easily scan through the energy at different lambda values by simply alternating between simulation.context.setParameter() to update lambda and simulation.context.getState() to retrieve potential energies at the new alchemical state.

# Create an integrator
integrator = LangevinIntegrator(temperature, collision_rate, timestep)

# Create a simulation
simulation = Simulation(topology, system, integrator)
simulation.context.setPositions(positions)

# Minimize energy
print('Minimizing energy...')
#LocalEnergyMinimizer.minimize(context)
simulation.minimizeEnergy()


# Collect data

# number of steps per sample
nsteps = 2500

# number of samples to collect per alchemical state
niterations = 500

import numpy as np
lambdas = np.linspace(1.0, 0.0, 10) # alchemical lambda schedule
nstates = len(lambdas)
u_kln = np.zeros([nstates,nstates,niterations], np.float64)
kT = AVOGADRO_CONSTANT_NA * BOLTZMANN_CONSTANT_kB * integrator.getTemperature()
for k in range(nstates):
    for iteration in range(niterations):
        print('state %5d iteration %5d / %5d' % (k, iteration, niterations))
        # Set alchemical state
        simulation.context.setParameter('lambda', lambdas[k])
        # Run some dynamics
        simulation.step(nsteps)
        # Compute energies at all alchemical states
        for l in range(nstates):
            simulation.context.setParameter('lambda', lambdas[l])
            u_kln[k,l,iteration] = simulation.context.getState(getEnergy=True).getPotentialEnergy() / kT

Analyzing the data with MBAR

Finally, the multistate Bennett acceptance ratio (MBAR) is used to estimate the free energy of annihilating the particle using the conda-installable pymbar package. In order to estimate how much data must be discarded to equilibration, we use a scheme for automated equilibration detection and subsequent extraction of decorrelated samples found in the pymbar.timeseries module.

# Estimate free energy of Lennard-Jones particle insertion
from pymbar import MBAR, timeseries

## Subsample data to extract uncorrelated equilibrium timeseries
N_k = np.zeros([nstates], np.int32) # number of uncorrelated samples
for k in range(nstates):
    [nequil, g, Neff_max] = timeseries.detectEquilibration(u_kln[k,k,:])
    indices = timeseries.subsampleCorrelatedData(u_kln[k,k,:], g=g)
    N_k[k] = len(indices)
    u_kln[k,:,0:N_k[k]] = u_kln[k,:,indices].T

# Compute free energy differences
mbar = MBAR(u_kln, N_k)

# dont compute uncertainties here, if you do it may fail with an error for
# pymbar versions > 3.0.3. See this issue: https://github.com/choderalab/pymbar/issues/419
[DeltaF_ij] = mbar.getFreeEnergyDifferences(compute_uncertainty=False)

print("Free energy change to insert a particle = ",DeltaF_ij[nstates-1][0])