Download Notebook View in GitHub Open in Google Colab
Virtual Sites Showcase
import time
from typing import List, Tuple
import numpy as np
import openmm
import openmm.app
import openmm.unit
from openff.interchange import Interchange
from openff.units import unit
from openff.toolkit import ForceField, Molecule, Topology
def prepare_simulation(interchange: Interchange) -> openmm.app.Simulation:
"""Propagate an OpenMM System with Langevin dynamics."""
time_step = 2 * openmm.unit.femtoseconds
temperature = 300 * openmm.unit.kelvin
friction = 1 / openmm.unit.picosecond
integrator = openmm.LangevinIntegrator(temperature, friction, time_step)
trj_freq, data_freq = 100, 100
simulation = interchange.to_openmm_simulation(integrator=integrator)
# It's important to run energy minimization before computing velocities; otherwise the initial
# velocities may be too high as a result of high initial forces, causing a crash
# See https://github.com/openmm/openmm/issues/3736#issuecomment-1217250635
simulation.minimizeEnergy()
# Since we placed all virtual sites at [0.0, 0.0, 0.0], compute virtual site positions to avoid a crash
simulation.context.computeVirtualSites()
simulation.context.setVelocitiesToTemperature(temperature)
pdb_reporter = openmm.app.PDBReporter("trajectory.pdb", trj_freq)
state_data_reporter = openmm.app.StateDataReporter(
"data.csv",
data_freq,
step=True,
potentialEnergy=True,
temperature=True,
density=True,
)
simulation.reporters.append(pdb_reporter)
simulation.reporters.append(state_data_reporter)
return simulation
def run_simulation(simulation: openmm.app.Simulation, num_steps: int = 1000):
print("Starting simulation")
start = time.process_time()
simulation.step(num_steps)
end = time.process_time()
print("Elapsed time %.2f seconds" % (end - start))
print("Done!")
Part 1: Run a short simulation with a virtual sites on sulfur lone pair and chloride groups
vsite_offxml = """<?xml version="1.0" encoding="utf-8"?>
<SMIRNOFF version="0.3" aromaticity_model="OEAroModel_MDL">
<VirtualSites version="0.3">
<VirtualSite
type="DivalentLonePair"
name="EP"
smirks="[*:2]-[#16X2:1]-[*:3]"
distance="0.70 * angstrom"
charge_increment1="0.1205*elementary_charge"
charge_increment2="0.0*elementary_charge"
charge_increment3="0.1205*elementary_charge"
sigma="0.1*angstrom"
epsilon="0.0*kilocalories_per_mole"
outOfPlaneAngle="54.71384225*degree"
match="all_permutations" >
</VirtualSite>
<VirtualSite
type="BondCharge"
name="EP"
smirks="[*:2][Cl:1]"
distance="0.4*angstrom"
charge_increment1="0.2*elementary_charge"
charge_increment2="0.0*elementary_charge"
sigma="0.1*angstrom"
epsilon="0.05*kilocalories_per_mole"
match="all_permutations" >
</VirtualSite>
<VirtualSite
type="BondCharge"
name="EP"
smirks="[*:2][F:1]"
distance="0.4*angstrom"
charge_increment1="0.2*elementary_charge"
charge_increment2="0.0*elementary_charge"
sigma="0.1*angstrom"
epsilon="0.05*kilocalories_per_mole"
match="all_permutations" >
</VirtualSite>
</VirtualSites>
</SMIRNOFF>
"""
force_field = ForceField("openff-2.0.0.offxml", vsite_offxml)
molecule = Molecule.from_smiles("c1cc(Cl)ccc1C(=O)CS[C]1=CO[C](F)(F)CC1")
molecule.generate_conformers(n_conformers=1)
molecule.visualize()
# Create an Interchange object, which stores information needed for OpenMM (and other engines)
# to understand virtual sites as applied by a force field
interchange = force_field.create_interchange(topology=molecule.to_topology())
assert "VirtualSites" in interchange.collections.keys()
n_virtual_sites = len(interchange.collections["VirtualSites"].key_map)
print(f"There are {n_virtual_sites} virtual particles in this topology.")
There are 5 virtual particles in this topology.
At this point, each of our OpenFF objects have processed and stored all of the information needed to run a simulation in OpenMM. Different OpenMM objects can be created from an Interchange object:
openmm.System:Interchange.to_openmm()openmm.app.Topology:Interchange.to_openmm_topology()openmm.app.Simulation:Interchange.to_openmm_simulation()
For simplicity, we’ll just create a Simulation object with the function we defined above since that’s all we need to run a simulation in OpenMM.
simulation = prepare_simulation(interchange)
run_simulation(simulation, 10000)
Starting simulation
Elapsed time 1.57 seconds
Done!
import mdtraj
# Visualize the trajectory. There will probably be errant bonds drawn between atoms and virtual sites.
import nglview
nglview.show_mdtraj(mdtraj.load("trajectory.pdb"))
Part 2: Numerical comparison of OpenMM’s TIP5P and an equivalent SMIRNOFF implementation
Parameterizes a water box with OpenFF and OpenMM force fields. Currently set up to use a TIP5P definition. The code examines the geometry and energy between the two, and examines cases where minimization is performed. Specifically, the code compares the four possible combinations: - oFF and oMM geometry/energy, minimized separately and then compared - oFF and oMM geometry/energy using no minimization - Geometry minimized using oFF, then a single point is done with oMM - Geometry minimized using oMM, then a single point is done with oFF
The virtual site definitions give differences in geometry and energy mostly due to how they were defined from their parent atoms. OpenMM uses an OutOfPlaneSite definition, whereas OpenFF uses the LocalCoordinatesSite definition (both are OpenMM types). In the OutOfPlaneSite definition, both angle and distance are variable since the defined out-of-plane angle depends on a weighted vector cross. The cross is a function of the O-H vectors, so the virtual sites are sensitive to the molecular geometry. In the OpenFF version, the distance is fixed to a constant value, and the out-of-plane angle is explicitly required in the OpenFF spec.
In this example, the OpenFF parameter definition (the “offxml”) is a string
further below in the main function, and can be easily modified to explore
force field parameterization. The OpenMM definition is loaded from its internal
default location, and acts as a reference. One can change this this to a different
filename to compare other force fields.
This example is somewhat hardcoded to operate on water molecules, but can be easily modified to examine other cases as well. The only major assumption is that the system is homogenous, i.e., all of the molecules are same. The reason this is assumed is mostly due to the difference in how virtual sites are handled between OpenFF and OpenMM. Whereas OpenMM interleaves the virtual site particles between the atomic particles, OpenFF instead aggregates all virtual sites and places them last. The code below does assume that, barring this difference, the virtual sites are added in the same order.
The example begins in run_tests by defining a grid of water molecules, with
the default being a single water molecule (Nx=1, Ny=1, Nz=1). From this, the
calculations described in the first paragraph above are performed. The energy
difference and distance between the two geometries, per atom, is then reported.
There are commented lines that print the entire set of coordinates, and can be
uncommented if desired.
def _collate_virtual_site_positions(atom_positions: np.ndarray) -> np.ndarray:
"""Given an array of atomic positions of water, collate virtual particles between molecules."""
padded_positions = np.zeros(shape=(2, 3))
num_atoms_per_mol = 3
def mol_positions(i, atom_positions):
this_mol_atom_coordinates = atom_positions[
i * num_atoms_per_mol : (i + 1) * num_atoms_per_mol
]
return np.vstack([this_mol_atom_coordinates, padded_positions])
return np.vstack([mol_positions(i, atom_positions) for i in range(2)])
def _evaluate_positions_and_energy(
openmm_topology: openmm.app.Topology,
openmm_system: openmm.app.Simulation,
particle_positions: openmm.unit.Quantity,
minimize=False,
) -> Tuple[np.ndarray, unit.Quantity]:
"""
Calculate particle positions and potential energy of the a system.
Parameters
----------
openmm_topology: openmm.app.Topology, OpenMM Topology
openmm_system: openmm.app.System, OpenMM System
particle_positions: openff.units.unit.Quantity, (N, 3) array of positions in nanometers
minimize: bool, Whether or not to perform an energy minimization before calculating energy
Returns
-------
positions: numpy.ndarray, array of particle positions (as nanometers)
energy: openmm.unit.Quantity, The potential energy of the systems
"""
integrator = openmm.LangevinIntegrator(
300 * openmm.unit.kelvin,
1 / openmm.unit.picosecond,
0.002 * openmm.unit.picoseconds,
)
sim = openmm.app.Simulation(openmm_topology, openmm_system, integrator)
sim.context.setPositions(particle_positions)
sim.context.applyConstraints(1e-5)
if minimize:
sim.minimizeEnergy()
state = sim.context.getState(getEnergy=True, getPositions=True)
ene = state.getPotentialEnergy()
pos = openmm.unit.Quantity(
[
list(xyz)
for xyz in state.getPositions().value_in_unit(openmm.unit.nanometer)
],
openmm.unit.nanometer,
)
return pos, ene
def build_water_lattice(
num_duplicates: List[int] = [1, 1, 1],
spacing: List[float] = [2.0, 2.0, 2.0],
) -> List[Molecule]:
"""
Generate a box of water molecules as OpenFF Molecules
Parameters
----------
num_duplicates: List[int], The number of molecules in each dimension
spacing: List[float], The spacing between the molecules in each dimension, implicitly in Angstrom
Returns
-------
water_box: List[Molecule], A list of Molecule objcets with a 3D conformation
"""
Lx, Ly, Lz = (num_duplicates[i] * spacing[i] for i in range(3))
Z, Y, X = np.mgrid[
0 : Lz : spacing[0],
0 : Ly : spacing[1],
0 : Lx : spacing[2],
]
XYZ = [list(xyz) for xyz in zip(X.flat, Y.flat, Z.flat)]
water_box = list([None] * len(XYZ))
water_reference = Molecule.from_mapped_smiles("[O:1]([H:2])[H:3]")
water_reference.atoms[0].name = "O"
water_reference.atoms[1].name = "H1"
water_reference.atoms[2].name = "H2"
# Add ideal TIP5P geometry
water_reference.add_conformer(
[[0.0, 0.0, 0.0], [-0.7815, 0.5526, 0.0], [0.7815, 0.5526, 0.0]] * unit.angstrom
)
for i, xyz in enumerate(XYZ):
water_box[i] = Molecule(water_reference)
water_box[i].conformers[0] = water_box[i].conformers[0] + xyz * unit.angstrom
return water_box
def evaluate_openmm(
water: List[Molecule],
openmm_force_field: openmm.app.ForceField,
minimize: bool = False,
):
"""
Given a list of molecules and a force field definition, calculate the
positions and energy.
Parameters
----------
water: List[Molecule], each with a 3D conformation
openmm_force_field: openmm.app.ForceField, OpenMM ForceField object
minimize: bool, default = False, whether the structure should be minimized
Returns
-------
xyz: List The coordinates of all particles in the system (OpenMM ordering)
ene: float, The potential energy
"""
# First, get an OpenMM Topology and atom positions with no virtual sites
_topology: openmm.app.Topology = Topology.from_molecules(water).to_openmm()
atom_positions_unitless = np.vstack(
[mol.conformers[0].m_as(unit.nanometer) for mol in water]
)
atom_positions = openmm.unit.Quantity(
atom_positions_unitless, openmm.unit.nanometer
)
# Use OpenMM's Modeller to add virtual particles as perscribed by the force field
modeller = openmm.app.Modeller(_topology, atom_positions)
modeller.addExtraParticles(openmm_force_field)
# This topology includes virtual particles, so we can use it to create a System
topology = modeller.getTopology()
system: openmm.System = openmm_force_field.createSystem(
topology, nonbondedMethod=openmm.app.NoCutoff
)
# Add positions of virtual particles now that the topology includes them
particle_positions = openmm.unit.Quantity(
_collate_virtual_site_positions(atom_positions), openmm.unit.nanometer
)
return _evaluate_positions_and_energy(
topology, system, particle_positions, minimize=minimize
)
def evaluate_openff(
water: List[Molecule],
force_field: ForceField,
minimize: bool = False,
):
"""
Given a list of molecules and a force field definition, calculate the
positions and energy.
Parameters
----------
water: List[Molecule], each with a 3D conformation
force_field: ForceField, an OpenFF ForceField object
minimize: boolean, default= False, whether the structure should be minimized
Returns
-------
xyz: list The coordinates of all particles in the system (OpenMM ordering)
ene: float, The potential energy
"""
from openff.interchange.interop.openmm import to_openmm_positions
openff_topology = Topology.from_molecules(water)
interchange = Interchange.from_smirnoff(force_field, openff_topology)
return _evaluate_positions_and_energy(
interchange.to_openmm_topology(),
interchange.to_openmm(combine_nonbonded_forces=True),
to_openmm_positions(interchange),
)
def print_info(xyz, ene, name, crd_units=unit.angstrom) -> str:
print(f"Results for: {name}\nEnergy: {ene}\nCoordinates:{xyz * crd_units}\n")
tip5p_offxml = """<?xml version="1.0" encoding="utf-8"?>
<SMIRNOFF version="0.3" aromaticity_model="OEAroModel_MDL">
<LibraryCharges version="0.3">
<LibraryCharge name="tip5p" smirks="[#1:1]-[#8X2H2+0:2]-[#1:3]" charge1="0.*elementary_charge" charge2="0.*elementary_charge" charge3="0.*elementary_charge"/>
</LibraryCharges>
<vdW version="0.3" potential="Lennard-Jones-12-6" combining_rules="Lorentz-Berthelot" scale12="0.0" scale13="0.0" scale14="0.5" scale15="1.0" switch_width="0.0 * angstrom" cutoff="10.0 * angstrom" method="cutoff">
<Atom smirks="[#1:1]-[#8X2H2+0]-[#1]" epsilon="0. * mole**-1 * kilojoule" id="n35" sigma="1.0 * nanometer"/>
<Atom smirks="[#1]-[#8X2H2+0:1]-[#1]" epsilon="0.66944 * mole**-1 * kilojoule" id="n35" sigma="0.312 * nanometer"/>
</vdW>
<Bonds version="0.4" potential="harmonic" fractional_bondorder_method="AM1-Wiberg" fractional_bondorder_interpolation="linear">
<Bond smirks="[#1:1]-[#8X2H2+0:2]-[#1]" length="0.9572 * angstrom" k="462750.4 * nanometer**-2 * mole**-1 * kilojoule" id="b1" />
</Bonds>
<Angles version="0.3" potential="harmonic">
<Angle smirks="[#1:1]-[#8X2H2+0:2]-[#1:3]" angle="1.82421813418 * radian" k="836.8 * mole**-1 * radian**-2 * kilojoule" id="a1" />
</Angles>
<VirtualSites version="0.3">
<VirtualSite
type="DivalentLonePair"
name="EP"
smirks="[#1:2]-[#8X2H2+0:1]-[#1:3]"
distance="0.70 * angstrom"
charge_increment1="0.0*elementary_charge"
charge_increment2="0.1205*elementary_charge"
charge_increment3="0.1205*elementary_charge"
sigma="10.0*angstrom"
epsilon="0.0*kilocalories_per_mole"
outOfPlaneAngle="54.735*degree"
match="all_permutations" >
</VirtualSite>
</VirtualSites>
<Electrostatics version="0.3" method="PME" scale12="0.0" scale13="0.0" scale14="0.833333" scale15="1.0" switch_width="0.0 * angstrom" cutoff="9.0 * angstrom"/>
</SMIRNOFF>
"""
constraints = """
<Constraints version="0.3">
<Constraint smirks="[#1:1]-[#8X2H2+0:2]-[#1]" id="c1" distance="0.9572 * angstrom"/>
<Constraint smirks="[#1:1]-[#8X2H2+0]-[#1:2]" id="c2" distance="1.5139006545247014 * angstrom"/>
</Constraints>
"""
# The TIP5P force field in SMIRNOFF format
openff_force_field = ForceField(tip5p_offxml)
# The standard OpenMM definition of tip5p
openmm_force_field = openmm.app.ForceField("tip5p.xml")
minimize = False
num_duplicates = (2, 1, 1) # 2x2x2 = 8 water molecules
spacing = (3.0, 3.0, 3.0) # water spaced 3A apart in each direction
np.set_printoptions(formatter={"float_kind": "{:13.10f}".format})
waters = build_water_lattice(num_duplicates, spacing)
off_crds, off_ene = evaluate_openff(waters, openff_force_field, minimize=minimize)
off_crds = np.array(off_crds.value_in_unit(openmm.unit.angstrom))
omm_crds, omm_ene = evaluate_openmm(waters, openmm_force_field, minimize=minimize)
omm_crds = np.array(omm_crds.value_in_unit(openmm.unit.angstrom))
print_info(
np.linalg.norm(off_crds - omm_crds, axis=1),
off_ene - omm_ene,
"OpenFF - OpenMM norm",
)
coordinate_difference = np.linalg.norm(off_crds - omm_crds) / np.prod(num_duplicates)
energy_difference = abs(off_ene - omm_ene)
# For some reason, there is a slight coordinate/energy difference between the OpenFF
# TIP5P virtual sites and the OpenMM TIP5P virtual sites. The energy difference appears
# to be entirely due to slightly different geometry.
assert (
coordinate_difference < 0.04
), f"Coordinates differ by a norm of {coordinate_difference}"
assert (
energy_difference < 1.0 * openmm.unit.kilojoule_per_mole
), f"Energies differ by {energy_difference}"
Results for: OpenFF - OpenMM norm
Energy: 0.6343612670898438 kJ/mol
Coordinates:[0.0 0.0 0.0 0.02741583295825021 0.02741583295825021 0.0 0.0 0.0 0.02741583295825033 0.02741583295825033] angstrom