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)