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 openff.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 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):
from simtk import openmm
from simtk import unit as simtk_unit
from simtk.openmm import app
platform = setup_platform_with_resources(available_resources)
input_pdb_file = app.PDBFile(self.input_coordinate_file)
system = self.parameterized_system.system
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 _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: openff.evaluator.unit.Quantity
The temperature to run the simulation at.
pressure: openff.evaluator.unit.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 import openmm
from simtk.openmm import app
# Create a platform with the correct resources.
if not self.allow_gpu_platforms:
from openff.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.
system = self.parameterized_system.system
# 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 _evaluate_energies(
thermodynamic_state: ThermodynamicState,
system: "openmm.System",
trajectory: "Trajectory",
compute_resources: ComputeResources,
enable_pbc: bool = True,
high_precision: bool = True,
) -> ObservableFrame:
"""Evaluates the reduced and potential energies of each frame in a trajectory
using the specified system and at a particular state.
Parameters
----------
thermodynamic_state
The thermodynamic state to evaluate the reduced potentials at.
system
The system to evaluate the energies using.
trajectory
A trajectory of configurations to evaluate.
compute_resources: ComputeResources
The compute resources available to execute on.
enable_pbc
Whether PBC are enabled. This controls whether box vectors
are set or not.
high_precision
Whether to compute the energies using double precision.
Returns
-------
The array containing the evaluated potentials.
"""
from simtk import openmm
from simtk import unit as simtk_unit
integrator = openmm.VerletIntegrator(0.1 * simtk_unit.femtoseconds)
platform = setup_platform_with_resources(compute_resources, high_precision)
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(thermodynamic_state.temperature)
beta = 1.0 / (simtk_unit.BOLTZMANN_CONSTANT_kB * temperature)
pressure = pint_quantity_to_openmm(thermodynamic_state.pressure)
for frame_index in range(trajectory.n_frames):
positions = trajectory.xyz[frame_index]
if enable_pbc:
box_vectors = trajectory.openmm_boxes(frame_index)
openmm_context.setPeriodicBoxVectors(*box_vectors)
openmm_context.setPositions(positions)
state = openmm_context.getState(getEnergy=True)
potential_energy = state.getPotentialEnergy()
unreduced_potential = potential_energy / simtk_unit.AVOGADRO_CONSTANT_NA
if pressure is not None and 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
observables_frame = ObservableFrame({
ObservableType.PotentialEnergy:
ObservableArray(potentials),
ObservableType.ReducedPotential:
ObservableArray(reduced_potentials),
})
return observables_frame
def _execute(self, directory, available_resources):
import mdtraj
from simtk.openmm import app
# 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 is_file_and_not_empty(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,
)
self.output_coordinate_file = os.path.join(directory, "output.pdb")
# Run the simulation.
self._simulate(self._context, self._integrator)
# Set the output paths.
self.trajectory_file_path = self._local_trajectory_path
self.observables_file_path = os.path.join(directory,
"observables.json")
# Compute the final observables
self.observables = self._compute_final_observables(
temperature, pressure)
# Optionally compute any gradients.
if len(self.gradient_parameters) == 0:
return
if not isinstance(self.parameterized_system.force_field,
SmirnoffForceFieldSource):
raise ValueError(
"Derivates can only be computed for systems parameterized with SMIRNOFF "
"force fields.")
_compute_gradients(
self.gradient_parameters,
self.observables,
self.parameterized_system.force_field.to_force_field(),
self.thermodynamic_state,
self.parameterized_system.topology,
mdtraj.load_dcd(self.trajectory_file_path,
top=self.input_coordinate_file),
available_resources,
self.enable_pbc,
)
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 mdtraj
import openmmtools
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 _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 is_file_and_not_empty(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 openff-evaluator format
self._save_final_statistics(self.statistics_file_path, temperature,
pressure)
def _execute(self, directory, available_resources):
from openff.toolkit.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.
system_path = os.path.join(directory, "system.xml")
with open(system_path, "w") as file:
file.write(openmm.XmlSerializer.serialize(system))
self.parameterized_system = ParameterizedSystem(
substance=self.substance,
force_field=force_field_source,
topology_path=self.coordinate_file_path,
system_path=system_path,
)