def energy_decomposition_system(structure,
system,
platform=None,
nrg=u.kilocalories_per_mole):
"""
This function computes the energy contribution for all of the different
force groups.
Parameters
----------
structure : Structure
The Structure with the coordinates for this system
system : mm.System
The OpenMM System object to get decomposed energies from
platform : str
The platform to use. Options are None (default), 'CPU', 'Reference',
'CUDA', and 'OpenCL'. None will pick default OpenMM platform for this
installation and computer
nrg : energy unit, optional
The unit to convert all energies into. Default is kcal/mol
Returns
-------
energies : list of tuple
Each entry is a tuple with the name of the force followed by its
contribution
"""
import openmm as mm
# First get all of the old force groups so we can restore them
old_groups = [f.getForceGroup() for f in system.getForces()]
old_recip_group = []
def _ene(context, grp):
st = context.getState(getEnergy=True, groups=1 << grp)
return (type(system.getForce(grp)).__name__,
st.getPotentialEnergy().value_in_unit(nrg))
try:
for i, f in enumerate(system.getForces()):
if isinstance(f, mm.NonbondedForce):
old_recip_group.append(f.getReciprocalSpaceForceGroup())
f.setReciprocalSpaceForceGroup(i)
f.setForceGroup(i)
if platform is None:
con = mm.Context(system, mm.VerletIntegrator(0.001))
else:
con = mm.Context(system, mm.VerletIntegrator(0.001),
mm.Platform.getPlatformByName(platform))
con.setPositions(structure.positions)
if structure.box is not None:
con.setPeriodicBoxVectors(*structure.box_vectors)
return list(map(lambda x: _ene(con, x), range(system.getNumForces())))
finally:
idx = 0
for grp, force in zip(old_groups, system.getForces()):
if isinstance(force, mm.NonbondedForce):
force.setReciprocalSpaceForceGroup(old_recip_group[idx])
idx += 1
force.setForceGroup(grp)
def __init__(self, topology, system, integrator, platform=None, platformProperties=None, state=None):
"""Create a Simulation.
Parameters
----------
topology : Topology
A Topology describing the the system to simulate
system : System or XML file name
The OpenMM System object to simulate (or the name of an XML file
with a serialized System)
integrator : Integrator or XML file name
The OpenMM Integrator to use for simulating the System (or the name
of an XML file with a serialized System)
platform : Platform=None
If not None, the OpenMM Platform to use
platformProperties : map=None
If not None, a set of platform-specific properties to pass to the
Context's constructor
state : XML file name=None
The name of an XML file containing a serialized State. If not None,
the information stored in state will be transferred to the generated
Simulation object.
"""
self.topology = topology
## The System being simulated
if isinstance(system, string_types):
with open(system, 'r') as f:
self.system = mm.XmlSerializer.deserialize(f.read())
else:
self.system = system
## The Integrator used to advance the simulation
if isinstance(integrator, string_types):
with open(integrator, 'r') as f:
self.integrator = mm.XmlSerializer.deserialize(f.read())
else:
self.integrator = integrator
## A list of reporters to invoke during the simulation
self.reporters = []
if platform is None:
## The Context containing the current state of the simulation
self.context = mm.Context(self.system, self.integrator)
elif platformProperties is None:
self.context = mm.Context(self.system, self.integrator, platform)
else:
self.context = mm.Context(self.system, self.integrator, platform, platformProperties)
if state is not None:
with open(state, 'r') as f:
self.context.setState(mm.XmlSerializer.deserialize(f.read()))
## Determines whether or not we are using PBC. Try from the System first,
## fall back to Topology if that doesn't work
try:
self._usesPBC = self.system.usesPeriodicBoundaryConditions()
except Exception: # OpenMM just raises Exception if it's not implemented everywhere
self._usesPBC = topology.getUnitCellDimensions() is not None
def testForce(model_file, output_forces):
# Create a random cloud of particles.
numParticles = 10
system = mm.System()
positions = np.random.rand(numParticles, 3)
for i in range(numParticles):
system.addParticle(1.0)
# Create a force
force = ot.TorchForce(model_file)
assert not force.getOutputsForces() # Check the default
force.setOutputsForces(output_forces)
assert force.getOutputsForces() == output_forces
system.addForce(force)
# Compute the forces and energy.
integ = mm.VerletIntegrator(1.0)
context = mm.Context(system, integ,
mm.Platform.getPlatformByName('Reference'))
context.setPositions(positions)
state = context.getState(getEnergy=True, getForces=True)
# See if the energy and forces are correct. The network defines a potential of the form E(r) = |r|^2
expectedEnergy = np.sum(positions * positions)
assert np.allclose(
expectedEnergy,
state.getPotentialEnergy().value_in_unit(unit.kilojoules_per_mole))
assert np.allclose(-2 * positions, state.getForces(asNumpy=True))
def setUp(self):
system = mm.System()
system.addParticle(1.0)
for i in range(32):
force = mm.CustomExternalForce(str(i))
force.addParticle(0, [])
force.setForceGroup(i)
system.addForce(force)
platform = mm.Platform.getPlatformByName('Reference')
context = mm.Context(system, mm.VerletIntegrator(0), platform)
context.setPositions([(0, 0, 0)])
self.context = context
def testModuleArguments(deviceString, precision):
if pt.cuda.device_count() < 1 and deviceString == 'cuda:0':
pytest.skip('A CUDA device is not available')
if pt.cuda.device_count() < 2 and deviceString == 'cuda:1':
pytest.skip('Two CUDA devices are not available')
class TestModule(pt.nn.Module):
def __init__(self, device, dtype, positions):
super().__init__()
self.device = device
self.dtype = dtype
self.register_buffer('positions', pt.tensor(positions).to(dtype))
def forward(self, positions):
assert self.positions.device == self.device
assert positions.device == self.device
assert positions.dtype == self.dtype
assert pt.all(positions == self.positions)
return pt.sum(positions)
with NamedTemporaryFile() as fd:
numParticles = 10
system = mm.System()
positions = np.random.rand(numParticles, 3)
for _ in range(numParticles):
system.addParticle(1.0)
device = pt.device(deviceString)
if device.type == 'cpu' or precision == 'double':
dtype = pt.float64
else:
dtype = pt.float32
module = TestModule(device, dtype, positions)
pt.jit.script(module).save(fd.name)
force = ot.TorchForce(fd.name)
system.addForce(force)
integrator = mm.VerletIntegrator(1.0)
platform = mm.Platform.getPlatformByName(device.type.upper())
properties = {}
if device.type == 'cuda':
properties['DeviceIndex'] = str(device.index)
properties['Precision'] = precision
context = mm.Context(system, integrator, platform, properties)
context.setPositions(positions)
context.getState(getEnergy=True, getForces=True)
def test_createCheckpoint(self):
system = mm.System()
system.addParticle(1.0)
refPositions = [(0,0,0)]
platform = mm.Platform.getPlatformByName('Reference')
context = mm.Context(system, mm.VerletIntegrator(0), platform)
context.setPositions(refPositions)
chk = context.createCheckpoint()
# check that the return value of createCheckpoint is of type bytes (non-unicode)
assert isinstance(chk, bytes)
# set the positions to something random then reload the checkpoint, and
# make sure that the positions get restored correctly
context.setPositions([(12345, 12345, 123451)])
context.loadCheckpoint(chk)
newPositions = context.getState(getPositions=True).getPositions()._value
assert newPositions == refPositions
def runOneTest(testName, options):
"""Perform a single benchmarking simulation."""
explicit = (testName not in ('gbsa', 'amoebagk'))
amoeba = (testName in ('amoebagk', 'amoebapme'))
apoa1 = testName.startswith('apoa1')
amber = (testName.startswith('amber'))
hydrogenMass = None
print()
if amoeba:
print('Test: %s (epsilon=%g)' % (testName, options.epsilon))
elif testName == 'pme':
print('Test: pme (cutoff=%g)' % options.cutoff)
else:
print('Test: %s' % testName)
print('Ensemble: %s' % options.ensemble)
platform = mm.Platform.getPlatformByName(options.platform)
# Create the System.
temperature = 300 * unit.kelvin
if explicit:
friction = 1 * (1 / unit.picoseconds)
else:
friction = 91 * (1 / unit.picoseconds)
if amoeba:
constraints = None
epsilon = float(options.epsilon)
if explicit:
ff = app.ForceField('amoeba2009.xml')
pdb = app.PDBFile('5dfr_solv-cube_equil.pdb')
cutoff = 0.7 * unit.nanometers
vdwCutoff = 0.9 * unit.nanometers
system = ff.createSystem(pdb.topology,
nonbondedMethod=app.PME,
nonbondedCutoff=cutoff,
vdwCutoff=vdwCutoff,
constraints=constraints,
ewaldErrorTolerance=0.00075,
mutualInducedTargetEpsilon=epsilon,
polarization=options.polarization)
else:
ff = app.ForceField('amoeba2009.xml', 'amoeba2009_gk.xml')
pdb = app.PDBFile('5dfr_minimized.pdb')
cutoff = 2.0 * unit.nanometers
vdwCutoff = 1.2 * unit.nanometers
system = ff.createSystem(pdb.topology,
nonbondedMethod=app.NoCutoff,
constraints=constraints,
mutualInducedTargetEpsilon=epsilon,
polarization=options.polarization)
for f in system.getForces():
if isinstance(f, mm.AmoebaMultipoleForce) or isinstance(
f, mm.AmoebaVdwForce) or isinstance(
f, mm.AmoebaGeneralizedKirkwoodForce) or isinstance(
f, mm.AmoebaWcaDispersionForce):
f.setForceGroup(1)
dt = 0.002 * unit.picoseconds
if options.ensemble == 'NVE':
integ = mm.MTSIntegrator(dt, [(0, 2), (1, 1)])
else:
integ = mm.MTSLangevinIntegrator(temperature, friction, dt,
[(0, 2), (1, 1)])
positions = pdb.positions
elif amber:
dirname = downloadAmberSuite()
names = {
'amber20-dhfr': 'JAC',
'amber20-factorix': 'FactorIX',
'amber20-cellulose': 'Cellulose',
'amber20-stmv': 'STMV'
}
fileName = names[testName]
prmtop = app.AmberPrmtopFile(
os.path.join(dirname, f'PME/Topologies/{fileName}.prmtop'))
inpcrd = app.AmberInpcrdFile(
os.path.join(dirname, f'PME/Coordinates/{fileName}.inpcrd'))
topology = prmtop.topology
positions = inpcrd.positions
dt = 0.004 * unit.picoseconds
method = app.PME
cutoff = options.cutoff
constraints = app.HBonds
system = prmtop.createSystem(nonbondedMethod=method,
nonbondedCutoff=cutoff,
constraints=constraints)
if options.ensemble == 'NVE':
integ = mm.VerletIntegrator(dt)
else:
integ = mm.LangevinMiddleIntegrator(temperature, friction, dt)
else:
if apoa1:
ff = app.ForceField('amber14/protein.ff14SB.xml',
'amber14/lipid17.xml', 'amber14/tip3p.xml')
pdb = app.PDBFile('apoa1.pdb')
if testName == 'apoa1pme':
method = app.PME
cutoff = options.cutoff
elif testName == 'apoa1ljpme':
method = app.LJPME
cutoff = options.cutoff
else:
method = app.CutoffPeriodic
cutoff = 1 * unit.nanometers
hydrogenMass = 1.5 * unit.amu
elif explicit:
ff = app.ForceField('amber99sb.xml', 'tip3p.xml')
pdb = app.PDBFile('5dfr_solv-cube_equil.pdb')
if testName == 'pme':
method = app.PME
cutoff = options.cutoff
else:
method = app.CutoffPeriodic
cutoff = 1 * unit.nanometers
else:
ff = app.ForceField('amber99sb.xml', 'amber99_obc.xml')
pdb = app.PDBFile('5dfr_minimized.pdb')
method = app.CutoffNonPeriodic
cutoff = 2 * unit.nanometers
if options.heavy:
dt = 0.005 * unit.picoseconds
constraints = app.AllBonds
hydrogenMass = 4 * unit.amu
if options.ensemble == 'NVE':
integ = mm.VerletIntegrator(dt)
else:
integ = mm.LangevinIntegrator(temperature, friction, dt)
else:
dt = 0.004 * unit.picoseconds
constraints = app.HBonds
if options.ensemble == 'NVE':
integ = mm.VerletIntegrator(dt)
else:
integ = mm.LangevinMiddleIntegrator(temperature, friction, dt)
positions = pdb.positions
system = ff.createSystem(pdb.topology,
nonbondedMethod=method,
nonbondedCutoff=cutoff,
constraints=constraints,
hydrogenMass=hydrogenMass)
if options.ensemble == 'NPT':
system.addForce(mm.MonteCarloBarostat(1 * unit.bar, temperature, 100))
print('Step Size: %g fs' % dt.value_in_unit(unit.femtoseconds))
properties = {}
initialSteps = 5
if options.device is not None and platform.getName() in ('CUDA', 'OpenCL'):
properties['DeviceIndex'] = options.device
if ',' in options.device or ' ' in options.device:
initialSteps = 250
if options.precision is not None and platform.getName() in ('CUDA',
'OpenCL'):
properties['Precision'] = options.precision
# Run the simulation.
integ.setConstraintTolerance(1e-5)
if len(properties) > 0:
context = mm.Context(system, integ, platform, properties)
else:
context = mm.Context(system, integ, platform)
context.setPositions(positions)
if amber:
if inpcrd.boxVectors is not None:
context.setPeriodicBoxVectors(*inpcrd.boxVectors)
mm.LocalEnergyMinimizer.minimize(
context, 100 * unit.kilojoules_per_mole / unit.nanometer)
context.setVelocitiesToTemperature(temperature)
steps = 20
while True:
time = timeIntegration(context, steps, initialSteps)
if time >= 0.5 * options.seconds:
break
if time < 0.5:
steps = int(
steps * 1.0 / time
) # Integrate enough steps to get a reasonable estimate for how many we'll need.
else:
steps = int(steps * options.seconds / time)
print('Integrated %d steps in %g seconds' % (steps, time))
print('%g ns/day' %
(dt * steps * 86400 / time).value_in_unit(unit.nanoseconds))
def compute(self, input_model: "AtomicInput",
config: "TaskConfig") -> "AtomicResult":
"""
Runs OpenMM on given structure, inputs, in vacuum.
"""
self.found(raise_error=True)
try:
import openmm
from openmm import unit
except ImportError:
from simtk import openmm, unit
with capture_stdout():
from openff.toolkit import topology as offtop
# Failure flag
ret_data = {"success": False}
# generate basis, not given
if not input_model.model.basis:
raise InputError("Method must contain a basis set.")
if isinstance(input_model.model.basis, BasisSet):
raise InputError(
"QCSchema BasisSet for model.basis not implemented since not suitable for OpenMM."
)
# Make sure we are using smirnoff or antechamber
basis = input_model.model.basis.lower()
if basis in ["smirnoff", "antechamber"]:
with capture_stdout():
# try and make the molecule from the cmiles
cmiles = None
if input_model.molecule.extras:
cmiles = input_model.molecule.extras.get(
"canonical_isomeric_explicit_hydrogen_mapped_smiles",
None)
if cmiles is None:
cmiles = input_model.molecule.extras.get(
"cmiles", {}
).get(
"canonical_isomeric_explicit_hydrogen_mapped_smiles",
None)
if cmiles is not None:
off_mol = offtop.Molecule.from_mapped_smiles(
mapped_smiles=cmiles)
# add the conformer
conformer = unit.Quantity(value=np.array(
input_model.molecule.geometry),
unit=unit.bohr)
off_mol.add_conformer(conformer)
else:
# Process molecule with RDKit
rdkit_mol = RDKitHarness._process_molecule_rdkit(
input_model.molecule)
# Create an Open Force Field `Molecule` from the RDKit Molecule
off_mol = offtop.Molecule(rdkit_mol)
# now we need to create the system
openmm_system = self._generate_openmm_system(
molecule=off_mol,
method=input_model.model.method,
keywords=input_model.keywords)
else:
raise InputError(
"Accepted bases are: {'smirnoff', 'antechamber', }")
# Need an integrator for simulation even if we don't end up using it really
integrator = openmm.VerletIntegrator(1.0 * unit.femtoseconds)
# Set platform to CPU explicitly
platform = openmm.Platform.getPlatformByName("CPU")
# Set number of threads to use
# if `nthreads` is `None`, OpenMM default of all logical cores on
# processor will be used
nthreads = config.ncores
if nthreads is None:
nthreads = os.environ.get("OPENMM_CPU_THREADS")
if nthreads:
properties = {"Threads": str(nthreads)}
else:
properties = {}
# Initialize context
context = openmm.Context(openmm_system, integrator, platform,
properties)
# Set positions from our Open Force Field `Molecule`
context.setPositions(off_mol.conformers[0])
# Compute the energy of the configuration
state = context.getState(getEnergy=True)
# Get the potential as a unit.Quantity, put into units of hartree
q = state.getPotentialEnergy(
) / unit.hartree / unit.AVOGADRO_CONSTANT_NA
ret_data["properties"] = {"return_energy": q}
# Execute driver
if input_model.driver == "energy":
ret_data["return_result"] = ret_data["properties"]["return_energy"]
elif input_model.driver == "gradient":
# Compute the forces
state = context.getState(getForces=True)
# Get the gradient as a unit.Quantity with shape (n_atoms, 3)
gradient = state.getForces(asNumpy=True)
# Convert to hartree/bohr and reformat as 1D array
q = (gradient / (unit.hartree / unit.bohr)
).reshape(-1) / unit.AVOGADRO_CONSTANT_NA
# Force to gradient
ret_data["return_result"] = -1 * q
else:
raise InputError(
f"Driver {input_model.driver} not implemented for OpenMM.")
ret_data["success"] = True
ret_data["extras"] = input_model.extras
# Move several pieces up a level
ret_data["provenance"] = Provenance(
creator="openmm",
version=openmm.version.short_version,
nthreads=nthreads)
return AtomicResult(**{**input_model.dict(), **ret_data})
def openmm_energy(prm,
structure,
coords,
box=None,
cutoff=None,
switch_dist=None):
from openmm import unit
from openmm import app
import openmm
from parmed.amber import AmberParm
if box is not None and not np.all(box == 0):
if cutoff is None:
raise RuntimeError(
"You need to provide a cutoff when passing a box")
a = unit.Quantity(
(box[0] * unit.angstrom, 0 * unit.angstrom, 0 * unit.angstrom))
b = unit.Quantity(
(0 * unit.angstrom, box[1] * unit.angstrom, 0 * unit.angstrom))
c = unit.Quantity(
(0 * unit.angstrom, 0 * unit.angstrom, box[2] * unit.angstrom))
structure.box_vectors = (a, b, c)
if isinstance(structure, AmberParm):
system = structure.createSystem(
nonbondedMethod=app.CutoffPeriodic,
nonbondedCutoff=0 if cutoff is None else cutoff *
unit.angstrom,
switchDistance=0 if switch_dist is None else switch_dist *
unit.angstrom,
)
else:
system = structure.createSystem(
prm,
nonbondedMethod=app.CutoffPeriodic,
nonbondedCutoff=0 if cutoff is None else cutoff *
unit.angstrom,
switchDistance=0 if switch_dist is None else switch_dist *
unit.angstrom,
)
system.setDefaultPeriodicBoxVectors(a, b, c)
else:
if isinstance(structure, AmberParm):
system = structure.createSystem()
else:
system = structure.createSystem(prm)
disableDispersionCorrection(system)
integrator = openmm.LangevinIntegrator(300 * unit.kelvin,
1 / unit.picoseconds,
2 * unit.femtoseconds)
try:
platform = openmm.Platform.getPlatformByName("CUDA")
properties = {"CudaPrecision": "double"}
except Exception:
platform = openmm.Platform.getPlatformByName("CPU")
properties = {}
context = openmm.Context(system, integrator, platform, properties)
# Run OpenMM with given coordinates
context.setPositions(coords * unit.angstrom)
energies = parmed.openmm.energy_decomposition(structure, context)
state = context.getState(getForces=True)
forces = state.getForces(asNumpy=True).value_in_unit(
unit.kilocalories_per_mole / unit.angstrom)
if "angle" not in energies:
energies["angle"] = 0
if "dihedral" not in energies:
energies["dihedral"] = 0
if "improper" not in energies:
energies["improper"] = 0
return energies, forces
system.setDefaultPeriodicBoxVectors([3.0, 0.0, 0.0] * unit.nanometers,
[0.0, 3.0, 0.0] * unit.nanometers,
[0.0, 0.0, 3.0] * unit.nanometers)
_ = system.addForce(non_bonded_force)
step_size = 2 * unit.femtoseconds
temperature = 300 * unit.kelvin
friction = 1.0 / unit.picosecond
integrator = mm.LangevinIntegrator(temperature, friction, step_size)
platform_name = 'CUDA'
platform = mm.Platform.getPlatformByName(platform_name)
context = mm.Context(system, integrator, platform)
initial_positions = np.zeros([2, 3], np.float32) * unit.angstroms
initial_velocities = np.zeros([2, 3],
np.float32) * unit.angstroms / unit.picoseconds
initial_positions[1, 0] = 1.0 * unit.nanometers
context.setPositions(initial_positions)
context.setVelocities(initial_velocities)
simulation_time = 20.0 * unit.nanosecond
saving_time = 1.0 * unit.picoseconds
n_steps_per_saving_period = int(saving_time / step_size)
n_saving_periods = int(simulation_time / saving_time)
