def test_pint_to_openmm(pint_unit, value):
pint_quantity = value * pint_unit
pint_raw_value = pint_quantity.magnitude
openmm_quantity = pint_quantity_to_openmm(pint_quantity)
openmm_raw_value = openmm_quantity.value_in_unit(openmm_quantity.unit)
assert np.allclose(openmm_raw_value, pint_raw_value)
def _execute(self, directory, available_resources):
platform = setup_platform_with_resources(available_resources)
input_pdb_file = app.PDBFile(self.input_coordinate_file)
with open(self.system_path, "rb") as file:
system = openmm.XmlSerializer.deserialize(file.read().decode())
if not self.enable_pbc:
for force_index in range(system.getNumForces()):
force = system.getForce(force_index)
if not isinstance(force, openmm.NonbondedForce):
continue
force.setNonbondedMethod(
0) # NoCutoff = 0, NonbondedMethod.CutoffNonPeriodic = 1
# TODO: Expose the constraint tolerance
integrator = openmm.VerletIntegrator(0.002 * simtk_unit.picoseconds)
simulation = app.Simulation(input_pdb_file.topology, system,
integrator, platform)
box_vectors = input_pdb_file.topology.getPeriodicBoxVectors()
if box_vectors is None:
box_vectors = simulation.system.getDefaultPeriodicBoxVectors()
simulation.context.setPeriodicBoxVectors(*box_vectors)
simulation.context.setPositions(input_pdb_file.positions)
simulation.minimizeEnergy(pint_quantity_to_openmm(self.tolerance),
self.max_iterations)
positions = simulation.context.getState(
getPositions=True).getPositions()
self.output_coordinate_file = os.path.join(directory, "minimised.pdb")
with open(self.output_coordinate_file, "w+") as minimised_file:
app.PDBFile.writeFile(simulation.topology, positions,
minimised_file)
def test_combinatorial_pint_to_openmm():
all_pint_units = _get_all_pint_units()
for i in range(len(all_pint_units)):
for j in range(i, len(all_pint_units)):
pint_unit = all_pint_units[i] * all_pint_units[j]
pint_quantity = random() * pint_unit
pint_raw_value = pint_quantity.magnitude
openmm_quantity = pint_quantity_to_openmm(pint_quantity)
openmm_raw_value = openmm_quantity.value_in_unit(
openmm_quantity.unit)
assert np.isclose(openmm_raw_value, pint_raw_value)
def test_pint_unit_conversions(pint_unit, value):
pint_quantity = value * pint_unit
pint_base_quantity = pint_quantity.to_base_units()
openmm_base_quantity = pint_quantity_to_openmm(pint_base_quantity)
openmm_unit = pint_unit_to_openmm(pint_unit)
openmm_quantity = openmm_base_quantity.in_units_of(openmm_unit)
pint_raw_value = pint_quantity.magnitude
if pint_unit == unit.dimensionless and (
isinstance(openmm_quantity, float) or isinstance(
openmm_quantity, int) or isinstance(openmm_quantity, list)
or isinstance(openmm_quantity, np.ndarray)):
openmm_raw_value = openmm_quantity
else:
openmm_raw_value = openmm_quantity.value_in_unit(openmm_unit)
assert np.allclose(openmm_raw_value, pint_raw_value)
def _get_options_dictionary(self, available_resources):
"""Returns a dictionary of options which will be serialized
to a yaml file and passed to YANK.
Parameters
----------
available_resources: ComputeResources
The resources available to execute on.
Returns
-------
dict of str and Any
A yaml compatible dictionary of YANK options.
"""
from openforcefield.utils import quantity_to_string
platform_name = "CPU"
if available_resources.number_of_gpus > 0:
# A platform which runs on GPUs has been requested.
from evaluator.backends import ComputeResources
toolkit_enum = ComputeResources.GPUToolkit(
available_resources.preferred_gpu_toolkit)
# A platform which runs on GPUs has been requested.
platform_name = ("CUDA" if toolkit_enum
== ComputeResources.GPUToolkit.CUDA else
ComputeResources.GPUToolkit.OpenCL)
return {
"verbose":
self.verbose,
"output_dir":
".",
"temperature":
quantity_to_string(
pint_quantity_to_openmm(self.thermodynamic_state.temperature)),
"pressure":
quantity_to_string(
pint_quantity_to_openmm(self.thermodynamic_state.pressure)),
"minimize":
True,
"number_of_equilibration_iterations":
self.number_of_equilibration_iterations,
"default_number_of_iterations":
self.number_of_iterations,
"default_nsteps_per_iteration":
self.steps_per_iteration,
"checkpoint_interval":
self.checkpoint_interval,
"default_timestep":
quantity_to_string(pint_quantity_to_openmm(self.timestep)),
"annihilate_electrostatics":
True,
"annihilate_sterics":
False,
"platform":
platform_name,
}
def _execute(self, directory, available_resources):
from openforcefield.topology import Molecule, Topology
force_field_source = ForceFieldSource.from_json(self.force_field_path)
cutoff = pint_quantity_to_openmm(force_field_source.cutoff)
# Load in the systems topology
openmm_pdb_file = app.PDBFile(self.coordinate_file_path)
# Create an OFF topology for better insight into the layout of the system
# topology.
unique_molecules = {}
for component in self.substance:
unique_molecule = Molecule.from_smiles(component.smiles)
unique_molecules[unique_molecule.to_smiles()] = unique_molecule
# Parameterize each component in the system.
system_templates = {}
for index, (smiles,
unique_molecule) in enumerate(unique_molecules.items()):
if smiles in ["O", "[H]O[H]", "[H][O][H]"]:
component_system = self._build_tip3p_system(
cutoff,
openmm_pdb_file.topology.getUnitCellDimensions(),
)
else:
component_directory = os.path.join(directory, str(index))
os.makedirs(component_directory, exist_ok=True)
with temporarily_change_directory(component_directory):
component_system = self._parameterize_molecule(
unique_molecule, force_field_source, cutoff)
system_templates[smiles] = component_system
# Apply the parameters to the topology.
topology = Topology.from_openmm(openmm_pdb_file.topology,
unique_molecules.values())
# Create the full system object from the component templates.
system = self._create_empty_system(cutoff)
for topology_molecule in topology.topology_molecules:
smiles = topology_molecule.reference_molecule.to_smiles()
system_template = system_templates[smiles]
index_map = {}
for index, topology_atom in enumerate(topology_molecule.atoms):
index_map[topology_atom.atom.molecule_particle_index] = index
# Append the component template to the full system.
self._append_system(system, system_template, index_map)
if openmm_pdb_file.topology.getPeriodicBoxVectors() is not None:
system.setDefaultPeriodicBoxVectors(
*openmm_pdb_file.topology.getPeriodicBoxVectors())
# Serialize the system object.
self.system_path = os.path.join(directory, "system.xml")
with open(self.system_path, "w") as file:
file.write(openmm.XmlSerializer.serialize(system))
def _evaluate_reduced_potential(
self,
system,
trajectory,
file_path,
compute_resources,
subset_energy_corrections=None,
):
"""Computes the reduced potential of each frame in a trajectory
using the provided system.
Parameters
----------
system: simtk.openmm.System
The system which encodes the interaction forces for the
specified parameter.
trajectory: mdtraj.Trajectory
A trajectory of configurations to evaluate.
file_path: str
The path to save the reduced potentials to.
compute_resources: ComputeResources
The compute resources available to execute on.
subset_energy_corrections: pint.Quantity, optional
A pint.Quantity wrapped numpy.ndarray which contains a set
of energies to add to the re-evaluated potential energies.
This is mainly used to correct the potential energies evaluated
using a subset of the force field back to energies as if evaluated
using the full thing.
Returns
---------
StatisticsArray
The array containing the reduced potentials.
"""
integrator = openmm.VerletIntegrator(0.1 * simtk_unit.femtoseconds)
platform = setup_platform_with_resources(compute_resources, True)
openmm_context = openmm.Context(system, integrator, platform)
potentials = np.zeros(trajectory.n_frames, dtype=np.float64)
reduced_potentials = np.zeros(trajectory.n_frames, dtype=np.float64)
temperature = pint_quantity_to_openmm(
self.thermodynamic_state.temperature)
beta = 1.0 / (simtk_unit.BOLTZMANN_CONSTANT_kB * temperature)
pressure = pint_quantity_to_openmm(self.thermodynamic_state.pressure)
if subset_energy_corrections is None:
subset_energy_corrections = (
np.zeros(trajectory.n_frames, dtype=np.float64) *
simtk_unit.kilojoules_per_mole)
else:
subset_energy_corrections = pint_quantity_to_openmm(
subset_energy_corrections)
for frame_index in range(trajectory.n_frames):
positions = trajectory.xyz[frame_index]
box_vectors = trajectory.openmm_boxes(frame_index)
if self.enable_pbc:
openmm_context.setPeriodicBoxVectors(*box_vectors)
openmm_context.setPositions(positions)
state = openmm_context.getState(getEnergy=True)
potential_energy = (state.getPotentialEnergy() +
subset_energy_corrections[frame_index])
unreduced_potential = potential_energy / simtk_unit.AVOGADRO_CONSTANT_NA
if pressure is not None and self.enable_pbc:
unreduced_potential += pressure * state.getPeriodicBoxVolume()
potentials[frame_index] = potential_energy.value_in_unit(
simtk_unit.kilojoule_per_mole)
reduced_potentials[frame_index] = unreduced_potential * beta
potentials *= unit.kilojoule / unit.mole
reduced_potentials *= unit.dimensionless
statistics_array = StatisticsArray()
statistics_array[ObservableType.ReducedPotential] = reduced_potentials
statistics_array[ObservableType.PotentialEnergy] = potentials
statistics_array.to_pandas_csv(file_path)
return statistics_array
def _execute(self, directory, available_resources):
import openmmtools
import mdtraj
trajectory = mdtraj.load_dcd(self.trajectory_file_path,
self.coordinate_file_path)
with open(self.system_path, "r") as file:
system = openmm.XmlSerializer.deserialize(file.read())
temperature = pint_quantity_to_openmm(
self.thermodynamic_state.temperature)
pressure = pint_quantity_to_openmm(self.thermodynamic_state.pressure)
if self.enable_pbc:
system.setDefaultPeriodicBoxVectors(*trajectory.openmm_boxes(0))
else:
pressure = None
openmm_state = openmmtools.states.ThermodynamicState(
system=system, temperature=temperature, pressure=pressure)
integrator = openmmtools.integrators.VelocityVerletIntegrator(
0.01 * simtk_unit.femtoseconds)
# Setup the requested platform:
platform = setup_platform_with_resources(available_resources,
self.high_precision)
openmm_system = openmm_state.get_system(True, True)
if not self.enable_pbc:
disable_pbc(openmm_system)
openmm_context = openmm.Context(openmm_system, integrator, platform)
potential_energies = np.zeros(trajectory.n_frames)
reduced_potentials = np.zeros(trajectory.n_frames)
for frame_index in range(trajectory.n_frames):
if self.enable_pbc:
box_vectors = trajectory.openmm_boxes(frame_index)
openmm_context.setPeriodicBoxVectors(*box_vectors)
positions = trajectory.xyz[frame_index]
openmm_context.setPositions(positions)
potential_energy = openmm_context.getState(
getEnergy=True).getPotentialEnergy()
potential_energies[frame_index] = potential_energy.value_in_unit(
simtk_unit.kilojoule_per_mole)
reduced_potentials[frame_index] = openmm_state.reduced_potential(
openmm_context)
kinetic_energies = StatisticsArray.from_pandas_csv(
self.kinetic_energies_path)[ObservableType.KineticEnergy]
statistics_array = StatisticsArray()
statistics_array[ObservableType.PotentialEnergy] = (
potential_energies * unit.kilojoule / unit.mole)
statistics_array[ObservableType.KineticEnergy] = kinetic_energies
statistics_array[ObservableType.ReducedPotential] = (
reduced_potentials * unit.dimensionless)
statistics_array[ObservableType.TotalEnergy] = (
statistics_array[ObservableType.PotentialEnergy] +
statistics_array[ObservableType.KineticEnergy])
statistics_array[ObservableType.Enthalpy] = (
statistics_array[ObservableType.ReducedPotential] *
self.thermodynamic_state.inverse_beta + kinetic_energies)
if self.use_internal_energy:
statistics_array[ObservableType.ReducedPotential] += (
kinetic_energies * self.thermodynamic_state.beta)
self.statistics_file_path = os.path.join(directory, "statistics.csv")
statistics_array.to_pandas_csv(self.statistics_file_path)
def _setup_simulation_objects(self, temperature, pressure,
available_resources):
"""Initializes the objects needed to perform the simulation.
This comprises of a context, and an integrator.
Parameters
----------
temperature: simtk.pint.Quantity
The temperature to run the simulation at.
pressure: simtk.pint.Quantity
The pressure to run the simulation at.
available_resources: ComputeResources
The resources available to run on.
Returns
-------
simtk.openmm.Context
The created openmm context which takes advantage
of the available compute resources.
openmmtools.integrators.LangevinIntegrator
The Langevin integrator which will propogate
the simulation.
"""
import openmmtools
from simtk.openmm import XmlSerializer
# Create a platform with the correct resources.
if not self.allow_gpu_platforms:
from evaluator.backends import ComputeResources
available_resources = ComputeResources(
available_resources.number_of_threads)
platform = setup_platform_with_resources(available_resources,
self.high_precision)
# Load in the system object from the provided xml file.
with open(self.system_path, "r") as file:
system = XmlSerializer.deserialize(file.read())
# Disable the periodic boundary conditions if requested.
if not self.enable_pbc:
disable_pbc(system)
pressure = None
# Use the openmmtools ThermodynamicState object to help
# set up a system which contains the correct barostat if
# one should be present.
openmm_state = openmmtools.states.ThermodynamicState(
system=system, temperature=temperature, pressure=pressure)
system = openmm_state.get_system(remove_thermostat=True)
# Set up the integrator.
thermostat_friction = pint_quantity_to_openmm(self.thermostat_friction)
timestep = pint_quantity_to_openmm(self.timestep)
integrator = openmmtools.integrators.LangevinIntegrator(
temperature=temperature,
collision_rate=thermostat_friction,
timestep=timestep,
)
# Create the simulation context.
context = openmm.Context(system, integrator, platform)
# Initialize the context with the correct positions etc.
input_pdb_file = app.PDBFile(self.input_coordinate_file)
if self.enable_pbc:
# Optionally set up the box vectors.
box_vectors = input_pdb_file.topology.getPeriodicBoxVectors()
if box_vectors is None:
raise ValueError(
"The input file must contain box vectors when running with PBC."
)
context.setPeriodicBoxVectors(*box_vectors)
context.setPositions(input_pdb_file.positions)
context.setVelocitiesToTemperature(temperature)
return context, integrator
def _execute(self, directory, available_resources):
# We handle most things in OMM units here.
temperature = self.thermodynamic_state.temperature
openmm_temperature = pint_quantity_to_openmm(temperature)
pressure = (None if self.ensemble == Ensemble.NVT else
self.thermodynamic_state.pressure)
openmm_pressure = pint_quantity_to_openmm(pressure)
if openmm_temperature is None:
raise ValueError(
"A temperature must be set to perform a simulation in any ensemble"
)
if Ensemble(self.ensemble) == Ensemble.NPT and openmm_pressure is None:
raise ValueError(
"A pressure must be set to perform an NPT simulation")
if Ensemble(
self.ensemble) == Ensemble.NPT and self.enable_pbc is False:
raise ValueError(
"PBC must be enabled when running in the NPT ensemble.")
# Set up the internal file paths
self._checkpoint_path = os.path.join(directory, "checkpoint.json")
self._state_path = os.path.join(directory, "checkpoint_state.xml")
self._local_trajectory_path = os.path.join(directory, "trajectory.dcd")
self._local_statistics_path = os.path.join(directory,
"openmm_statistics.csv")
# Set up the simulation objects.
if self._context is None or self._integrator is None:
self._context, self._integrator = self._setup_simulation_objects(
openmm_temperature, openmm_pressure, available_resources)
# Save a copy of the starting configuration if it doesn't already exist
local_input_coordinate_path = os.path.join(directory, "input.pdb")
if not os.path.isfile(local_input_coordinate_path):
input_pdb_file = app.PDBFile(self.input_coordinate_file)
with open(local_input_coordinate_path, "w+") as configuration_file:
app.PDBFile.writeFile(
input_pdb_file.topology,
input_pdb_file.positions,
configuration_file,
)
# Run the simulation.
self._simulate(directory, self._context, self._integrator)
# Set the output paths.
self.trajectory_file_path = self._local_trajectory_path
self.statistics_file_path = os.path.join(directory, "statistics.csv")
# Save out the final statistics in the evaluator format
self._save_final_statistics(self.statistics_file_path, temperature,
pressure)
