def compute(self, input_data: "AtomicInput", config: "TaskConfig") -> "AtomicResult":
"""
Runs OpenMM on given structure, inputs, in vacuum.
"""
self.found(raise_error=True)
from simtk import openmm
from simtk import unit
import openforcefield.topology as offtop
# Failure flag
ret_data = {"success": False}
# generate basis, not given
if not input_data.model.basis:
basis = self._generate_basis(input_data)
ret_data["basis"] = basis
# get number of threads to use from `TaskConfig.ncores`; otherwise, try environment variable
nthreads = config.ncores
if nthreads is None:
nthreads = os.environ.get("OPENMM_CPU_THREADS")
# Set workdir to scratch
# Location resolution order config.scratch_dir, /tmp
parent = config.scratch_directory
with temporary_directory(parent=parent, suffix="_openmm_scratch") as tmpdir:
# Grab molecule, forcefield
jmol = input_data.molecule
# TODO: If urls are supported by
# `openforcefield.typing.engines.smirnoff.ForceField` already, we
# can eliminate the `offxml` and `url` distinction
# URL processing can happen there instead
if getattr(input_data.model, "offxml", None):
# we were given a file path or relative path
offxml = input_data.model.offxml
# Load an Open Force Field `ForceField`
off_forcefield = self._get_off_forcefield(offxml, offxml)
elif getattr(input_data.model, "url", None):
# we were given a url
with urllib.request.urlopen(input_data.model.url) as req:
xml = req.read()
# Load an Open Force Field `ForceField`
off_forcefield = self._get_off_forcefield(xml.decode(), xml)
else:
raise InputError("OpenMM requires either `model.offxml` or `model.url` to be set")
# Process molecule with RDKit
rdkit_mol = RDKitHarness._process_molecule_rdkit(jmol)
# Create an Open Force Field `Molecule` from the RDKit Molecule
off_mol = offtop.Molecule(rdkit_mol)
# Create OpenMM system in vacuum from forcefield, molecule
off_top = off_mol.to_topology()
openmm_system = self._get_openmm_system(off_forcefield, off_top)
# 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
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 simtk.unit.Quantity, put into units of hartree
q = state.getPotentialEnergy() / unit.hartree
ret_data["properties"] = {"return_energy": q.value_in_unit(q.unit)}
# Execute driver
if input_data.driver == "energy":
ret_data["return_result"] = ret_data["properties"]["return_energy"]
elif input_data.driver == "gradient":
# Get number of atoms
n_atoms = len(jmol.symbols)
# Compute the forces
state = context.getState(getForces=True)
# Get the gradient as a simtk.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([n_atoms * 3])
ret_data["return_result"] = q.value_in_unit(q.unit)
else:
raise InputError(
f"OpenMM can only compute energy and gradient driver methods. Found {input_data.driver}."
)
ret_data["success"] = True
# Move several pieces up a level
ret_data["provenance"] = Provenance(creator="openmm", version=openmm.__version__, nthreads=nthreads)
return AtomicResult(**{**input_data.dict(), **ret_data})
def compute(self, input_data: "AtomicInput",
config: "TaskConfig") -> "AtomicResult":
"""
Runs OpenMM on given structure, inputs, in vacuum.
"""
self.found(raise_error=True)
from simtk import openmm
from simtk import unit
with capture_stdout():
import openforcefield.topology as offtop
# Failure flag
ret_data = {"success": False}
# generate basis, not given
if not input_data.model.basis:
raise InputError("Method must contain a basis set.")
# Make sure we are using smirnoff or antechamber
basis = input_data.model.basis.lower()
if basis in ["smirnoff", "antechamber"]:
with capture_stdout():
# try and make the molecule from the cmiles
cmiles = None
if input_data.molecule.extras:
cmiles = input_data.molecule.extras.get(
"canonical_isomeric_explicit_hydrogen_mapped_smiles",
None)
if cmiles is None:
cmiles = input_data.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_data.molecule.geometry),
unit=unit.bohr)
off_mol.add_conformer(conformer)
else:
# Process molecule with RDKit
rdkit_mol = RDKitHarness._process_molecule_rdkit(
input_data.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_data.model.method,
keywords=input_data.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 simtk.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_data.driver == "energy":
ret_data["return_result"] = ret_data["properties"]["return_energy"]
elif input_data.driver == "gradient":
# Compute the forces
state = context.getState(getForces=True)
# Get the gradient as a simtk.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"OpenMM can only compute energy and gradient driver methods. Found {input_data.driver}."
)
ret_data["success"] = True
ret_data["extras"] = input_data.extras
# Move several pieces up a level
ret_data["provenance"] = Provenance(creator="openmm",
version=openmm.__version__,
nthreads=nthreads)
return AtomicResult(**{**input_data.dict(), **ret_data})