def _run_yank(directory, available_resources, setup_only):
"""Runs YANK within the specified directory which contains a `yank.yaml`
input file.
Parameters
----------
directory: str
The directory within which to run yank.
available_resources: ComputeResources
The compute resources available to yank.
setup_only: bool
If true, YANK will only create and validate the setup files,
but not actually run any simulations. This argument is mainly
only to be used for testing purposes.
Returns
-------
simtk.pint.Quantity
The free energy returned by yank.
simtk.pint.Quantity
The uncertainty in the free energy returned by yank.
"""
from yank.experiment import ExperimentBuilder
from yank.analyze import ExperimentAnalyzer
from simtk import unit as simtk_unit
with temporarily_change_directory(directory):
# Set the default properties on the desired platform
# before calling into yank.
setup_platform_with_resources(available_resources)
exp_builder = ExperimentBuilder("yank.yaml")
if setup_only is True:
return (
0.0 * simtk_unit.kilojoule_per_mole,
0.0 * simtk_unit.kilojoule_per_mole,
)
exp_builder.run_experiments()
analyzer = ExperimentAnalyzer("experiments")
output = analyzer.auto_analyze()
free_energy = output["free_energy"]["free_energy_diff_unit"]
free_energy_uncertainty = output["free_energy"][
"free_energy_diff_error_unit"]
return free_energy, free_energy_uncertainty
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 _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