def test_is_file_and_not_empty(tmpdir):
path = "testfile"
with tempfile.TemporaryDirectory():
with open(path, "w") as f:
pass
assert os.path.isfile(path)
assert os.path.getsize(path) == 0
assert not is_file_and_not_empty(path)
path = "testfile2"
with tempfile.TemporaryDirectory():
with open(path, "w") as f:
f.write("ubiquitous mendacious polyglottal")
assert os.path.isfile(path)
assert os.path.getsize(path) != 0
assert is_file_and_not_empty(path)
def _simulate(self, context, integrator):
"""Performs the simulation using a given context
and integrator.
Parameters
----------
context: simtk.openmm.Context
The OpenMM context to run with.
integrator: simtk.openmm.Integrator
The integrator to evolve the simulation with.
"""
from simtk.openmm import app
# Define how many steps should be taken.
total_number_of_steps = (self.total_number_of_iterations *
self.steps_per_iteration)
# Try to load the current state from any available checkpoint information
current_step = self._resume_from_checkpoint(context)
if current_step == total_number_of_steps:
return
# Build the reporters which we will use to report the state
# of the simulation.
append_trajectory = is_file_and_not_empty(self._local_trajectory_path)
dcd_reporter = app.DCDReporter(self._local_trajectory_path, 0,
append_trajectory)
statistics_file = open(self._local_statistics_path, "a+")
statistics_reporter = app.StateDataReporter(
statistics_file,
0,
step=True,
potentialEnergy=True,
kineticEnergy=True,
totalEnergy=True,
temperature=True,
volume=True,
density=True,
speed=True,
)
# Create the object which will transfer simulation output to the
# reporters.
topology = app.PDBFile(self.input_coordinate_file).topology
system = self.parameterized_system.system
simulation = self._Simulation(integrator, topology, system,
current_step)
# Perform the simulation.
checkpoint_counter = 0
while current_step < total_number_of_steps:
steps_to_take = min(self.output_frequency,
total_number_of_steps - current_step)
integrator.step(steps_to_take)
current_step += steps_to_take
state = context.getState(
getPositions=True,
getEnergy=True,
getVelocities=False,
getForces=False,
getParameters=False,
enforcePeriodicBox=self.enable_pbc,
)
simulation.currentStep = current_step
# Write out the current state using the reporters.
dcd_reporter.report(simulation, state)
statistics_reporter.report(simulation, state)
if checkpoint_counter >= self.checkpoint_frequency:
# Save to the checkpoint file if needed.
self._write_checkpoint_file(current_step, context)
checkpoint_counter = 0
checkpoint_counter += 1
# Save out the final positions.
self._write_checkpoint_file(current_step, context)
final_state = context.getState(getPositions=True)
positions = final_state.getPositions()
topology.setPeriodicBoxVectors(final_state.getPeriodicBoxVectors())
with open(self.output_coordinate_file, "w+") as configuration_file:
app.PDBFile.writeFile(topology, positions, configuration_file)
def _resume_from_checkpoint(self, context):
"""Resumes the simulation from a checkpoint file.
Parameters
----------
context: simtk.openmm.Context
The current OpenMM context.
Returns
-------
int
The current step number.
"""
from simtk import openmm
current_step_number = 0
# Check whether the checkpoint files actually exists.
if not is_file_and_not_empty(
self._checkpoint_path) or not is_file_and_not_empty(
self._state_path):
logger.info("No checkpoint files were found.")
return current_step_number
if not is_file_and_not_empty(
self._local_statistics_path) or not is_file_and_not_empty(
self._local_trajectory_path):
raise ValueError(
"Checkpoint files were correctly found, but the trajectory "
"or statistics files seem to be missing. This should not happen."
)
logger.info("Restoring the system state from checkpoint files.")
# If they do, load the current state from disk.
with open(self._state_path, "r") as file:
current_state = openmm.XmlSerializer.deserialize(file.read())
with open(self._checkpoint_path, "r") as file:
checkpoint = json.load(file, cls=TypedJSONDecoder)
if (self.output_frequency != checkpoint.output_frequency
or self.checkpoint_frequency != checkpoint.checkpoint_frequency
or self.steps_per_iteration != checkpoint.steps_per_iteration):
raise ValueError("Neither the output frequency, the checkpoint "
"frequency, nor the steps per iteration can "
"currently be changed during the course of the "
"simulation. Only the number of iterations is "
"allowed to change.")
# Make sure this simulation hasn't already finished.
total_expected_number_of_steps = (self.total_number_of_iterations *
self.steps_per_iteration)
if checkpoint.current_step_number == total_expected_number_of_steps:
return checkpoint.current_step_number
context.setState(current_state)
# Make sure that the number of frames in the trajectory /
# statistics file correspond to the recorded number of steps.
# This is to handle possible cases where only some of the files
# have been written from the current step (i.e only the trajectory may
# have been written to before this protocol gets unexpectedly killed.
expected_number_of_frames = int(checkpoint.current_step_number /
self.output_frequency)
# Handle the truncation of the statistics file.
self._truncate_statistics_file(expected_number_of_frames)
# Handle the truncation of the trajectory file.
self._truncate_trajectory_file(expected_number_of_frames)
logger.info("System state restored from checkpoint files.")
return checkpoint.current_step_number
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 _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)