def compute_reduced_potential(thermodynamic_state: states.ThermodynamicState, sampler_state: states.SamplerState) -> float: """ Compute the reduced potential of the given SamplerState under the given ThermodynamicState. Arguments ---------- thermodynamic_state : openmmtools.states.ThermodynamicState The thermodynamic state under which to compute the reduced potential sampler_state : openmmtools.states.SamplerState The sampler state for which to compute the reduced potential Returns ------- reduced_potential : float unitless reduced potential (kT) """ if type(cache.global_context_cache) == cache.DummyContextCache: integrator = openmm.VerletIntegrator( 1.0) #we won't take any steps, so use a simple integrator context, integrator = cache.global_context_cache.get_context( thermodynamic_state, integrator) else: context, integrator = cache.global_context_cache.get_context( thermodynamic_state) sampler_state.apply_to_context(context, ignore_velocities=True) return thermodynamic_state.reduced_potential(context)
def compute_reduced_potential(thermodynamic_state: states.ThermodynamicState, sampler_state: states.SamplerState) -> float: """ Compute the reduced potential of the given SamplerState under the given ThermodynamicState. Parameters ---------- thermodynamic_state : openmmtools.states.ThermodynamicState The thermodynamic state under which to compute the reduced potential sampler_state : openmmtools.states.SamplerState The sampler state for which to compute the reduced potential Returns ------- reduced_potential : float unitless reduced potential (kT) """ context, integrator = cache.global_context_cache.get_context( thermodynamic_state) sampler_state.apply_to_context(context, ignore_velocities=True) return thermodynamic_state.reduced_potential(context)
def compute_reduced_potential(thermodynamic_state: states.ThermodynamicState, sampler_state: states.SamplerState) -> float: """ Compute the reduced potential of the given SamplerState under the given ThermodynamicState. Parameters ---------- thermodynamic_state : openmmtools.states.ThermodynamicState The thermodynamic state under which to compute the reduced potential sampler_state : openmmtools.states.SamplerState The sampler state for which to compute the reduced potential Returns ------- reduced_potential : float unitless reduced potential (kT) """ if type(cache.global_context_cache) == cache.DummyContextCache: integrator = openmm.VerletIntegrator(1.0) #we won't take any steps, so use a simple integrator context, integrator = cache.global_context_cache.get_context(thermodynamic_state, integrator) else: context, integrator = cache.global_context_cache.get_context(thermodynamic_state) sampler_state.apply_to_context(context, ignore_velocities=True) return thermodynamic_state.reduced_potential(context)
def run_equilibrium( equilibrium_result: EquilibriumResult, thermodynamic_state: states.ThermodynamicState, nsteps_equil: int, topology: md.Topology, n_iterations: int, atom_indices_to_save: List[int] = None, trajectory_filename: str = None, splitting: str = "V R O R V", timestep: unit.Quantity = 1.0 * unit.femtoseconds ) -> EquilibriumResult: """ Run nsteps of equilibrium sampling at the specified thermodynamic state and return the final sampler state as well as a trajectory of the positions after each application of an MCMove. This means that if the MCMove is configured to run 1000 steps of dynamics, and n_iterations is 100, there will be 100 frames in the resulting trajectory; these are the result of 100,000 steps (1000*100) of dynamics. Parameters ---------- equilibrium_result : EquilibriumResult EquilibriumResult namedtuple containing the information necessary to resume thermodynamic_state : openmmtools.states.ThermodynamicState The thermodynamic state (including context parameters) that should be used nsteps_equil : int The number of equilibrium steps that a move should make when apply is called topology : mdtraj.Topology an MDTraj topology object used to construct the trajectory n_iterations : int The number of times to apply the move. Note that this is not the number of steps of dynamics; it is n_iterations*n_steps (which is set in the MCMove). splitting: str, default "V R O H R V" The splitting string for the dynamics atom_indices_to_save : list of int, default None list of indices to save (when excluding waters, for instance). If None, all indices are saved. trajectory_filename : str, optional, default None Full filepath of trajectory files. If none, trajectory files are not written. splitting: str, default "V R O H R V" The splitting string for the dynamics Returns ------- equilibrium_result : EquilibriumResult Container namedtuple that has the SamplerState for resuming, an MDTraj trajectory, and the reduced potential of the final frame. """ sampler_state = equilibrium_result.sampler_state #get the atom indices we need to subset the topology and positions if atom_indices_to_save is None: atom_indices = list(range(topology.n_atoms)) subset_topology = topology else: subset_topology = topology.subset(atom_indices_to_save) atom_indices = atom_indices_to_save n_atoms = subset_topology.n_atoms #construct the MCMove: mc_move = mcmc.LangevinSplittingDynamicsMove(n_steps=nsteps_equil, splitting=splitting) mc_move.n_restart_attempts = 10 #create a numpy array for the trajectory trajectory_positions = np.zeros([n_iterations, n_atoms, 3]) trajectory_box_lengths = np.zeros([n_iterations, 3]) trajectory_box_angles = np.zeros([n_iterations, 3]) #loop through iterations and apply MCMove, then collect positions into numpy array for iteration in range(n_iterations): mc_move.apply(thermodynamic_state, sampler_state) trajectory_positions[iteration, :] = sampler_state.positions[ atom_indices, :].value_in_unit_system(unit.md_unit_system) #get the box lengths and angles a, b, c, alpha, beta, gamma = mdtrajutils.unitcell.box_vectors_to_lengths_and_angles( *sampler_state.box_vectors) trajectory_box_lengths[iteration, :] = [a, b, c] trajectory_box_angles[iteration, :] = [alpha, beta, gamma] #construct trajectory object: trajectory = md.Trajectory(trajectory_positions, subset_topology, unitcell_lengths=trajectory_box_lengths, unitcell_angles=trajectory_box_angles) #get the reduced potential from the final frame for endpoint perturbations reduced_potential_final_frame = thermodynamic_state.reduced_potential( sampler_state) #construct equilibrium result object equilibrium_result = EquilibriumResult(sampler_state, reduced_potential_final_frame) #If there is a trajectory filename passed, write out the results here: if trajectory_filename is not None: write_equilibrium_trajectory(equilibrium_result, trajectory, trajectory_filename) return equilibrium_result
def compare_energies(REST_system, other_system, positions, rest_atoms, T_min, T): # Create thermodynamic state lambda_zero_alchemical_state = RESTState.from_system(REST_system) thermostate = ThermodynamicState(REST_system, temperature=T_min) compound_thermodynamic_state = CompoundThermodynamicState(thermostate, composable_states=[lambda_zero_alchemical_state]) # Set alchemical parameters beta_0 = 1 / (kB * T_min) beta_m = 1 / (kB * T) compound_thermodynamic_state.set_alchemical_parameters(beta_0, beta_m) # Minimize and save energy integrator = openmm.VerletIntegrator(1.0 * unit.femtosecond) context = compound_thermodynamic_state.create_context(integrator) context.setPositions(positions) sampler_state = SamplerState.from_context(context) REST_energy = compound_thermodynamic_state.reduced_potential(sampler_state) # Compute energy for non-RESTified system # Determine regions and scaling factors solute = rest_atoms solvent = [i for i in range(other_system.getNumParticles()) if i not in solute] solute_scaling = beta_m / beta_0 inter_scaling = np.sqrt(beta_m / beta_0) # Scale the terms in the bond force appropriately bond_force = other_system.getForce(0) for bond in range(bond_force.getNumBonds()): p1, p2, length, k = bond_force.getBondParameters(bond) if p1 in solute and p2 in solute: bond_force.setBondParameters(bond, p1, p2, length, k * solute_scaling) elif (p1 in solute and p2 in solvent) or (p1 in solvent and p2 in solute): bond_force.setBondParameters(bond, p1, p2, length, k * inter_scaling) # Scale the terms in the angle force appropriately angle_force = other_system.getForce(1) for angle_index in range(angle_force.getNumAngles()): p1, p2, p3, angle, k = angle_force.getAngleParameters(angle_index) if p1 in solute and p2 in solute and p3 in solute: angle_force.setAngleParameters(angle_index, p1, p2, p3, angle, k * solute_scaling) elif set([p1, p2, p3]).intersection(set(solute)) != set() and set([p1, p2, p3]).intersection( set(solvent)) != set(): angle_force.setAngleParameters(angle_index, p1, p2, p3, angle, k * inter_scaling) # Scale the terms in the torsion force appropriately torsion_force = other_system.getForce(2) for torsion_index in range(torsion_force.getNumTorsions()): p1, p2, p3, p4, periodicity, phase, k = torsion_force.getTorsionParameters(torsion_index) if p1 in solute and p2 in solute and p3 in solute and p4 in solute: torsion_force.setTorsionParameters(torsion_index, p1, p2, p3, p4, periodicity, phase, k * solute_scaling) elif set([p1, p2, p3, p4]).intersection(set(solute)) != set() and set([p1, p2, p3, p4]).intersection( set(solvent)) != set(): torsion_force.setTorsionParameters(torsion_index, p1, p2, p3, p4, periodicity, phase, k * inter_scaling) # Scale the exceptions in the nonbonded force appropriately nb_force = other_system.getForce(3) for nb_index in range(nb_force.getNumExceptions()): p1, p2, chargeProd, sigma, epsilon = nb_force.getExceptionParameters(nb_index) if p1 in solute and p2 in solute: nb_force.setExceptionParameters(nb_index, p1, p2, solute_scaling * chargeProd, sigma, solute_scaling * epsilon) elif (p1 in solute and p2 in solvent) or (p1 in solvent and p2 in solute): nb_force.setExceptionParameters(nb_index, p1, p2, inter_scaling * chargeProd, sigma, inter_scaling * epsilon) # Scale nonbonded interactions for solute-solute region by adding exceptions for all pairs of atoms exception_pairs = [tuple(sorted([nb_force.getExceptionParameters(nb_index)[0], nb_force.getExceptionParameters(nb_index)[1]])) for nb_index in range(nb_force.getNumExceptions())] solute_pairs = set([tuple(sorted(pair)) for pair in list(itertools.product(solute, solute))]) for pair in list(solute_pairs): p1 = pair[0] p2 = pair[1] p1_charge, p1_sigma, p1_epsilon = nb_force.getParticleParameters(p1) p2_charge, p2_sigma, p2_epsilon = nb_force.getParticleParameters(p2) if p1 != p2: if pair not in exception_pairs: nb_force.addException(p1, p2, p1_charge * p2_charge * solute_scaling, 0.5 * (p1_sigma + p2_sigma), np.sqrt(p1_epsilon * p2_epsilon) * solute_scaling) # Scale nonbonded interactions for inter region by adding exceptions for all pairs of atoms for pair in list(itertools.product(solute, solvent)): p1 = pair[0] p2 = int(pair[1]) # otherwise, will be a numpy int p1_charge, p1_sigma, p1_epsilon = nb_force.getParticleParameters(p1) p2_charge, p2_sigma, p2_epsilon = nb_force.getParticleParameters(p2) if tuple(sorted(pair)) not in exception_pairs: nb_force.addException(p1, p2, p1_charge * p2_charge * inter_scaling, 0.5 * (p1_sigma + p2_sigma), np.sqrt(p1_epsilon * p2_epsilon) * inter_scaling) # Get energy thermostate = ThermodynamicState(other_system, temperature=T_min) integrator = openmm.VerletIntegrator(1.0 * unit.femtosecond) context = thermostate.create_context(integrator) context.setPositions(positions) sampler_state = SamplerState.from_context(context) nonREST_energy = thermostate.reduced_potential(sampler_state) assert REST_energy - nonREST_energy < 1, f"The energy of the REST system ({REST_energy}) does not match " \ f"that of the non-REST system with terms manually scaled according to REST2({nonREST_energy})."
def integrate(self, topology_proposal, initial_sampler_state, proposed_sampler_state, iteration=None): """ Performs NCMC switching to either delete or insert atoms according to the provided `topology_proposal`. For `delete`, the system is first modified from fully interacting to alchemically modified, and then NCMC switching is used to eliminate atoms. For `insert`, the system begins with eliminated atoms in an alchemically noninteracting form and NCMC switching is used to turn atoms on, followed by making system real. Parameters ---------- topology_proposal : TopologyProposal Contains old/new Topology and System objects and atom mappings. initial_sampler_state : openmmtools.states.SamplerState representing the initial (old) system Configurational properties of the atoms at the beginning of the NCMC switching. proposed_sampler_state : openmmtools.states.SamplerState representing the proposed (post-geometry new) system Configurational properties new system atoms at beginning of NCMC switching iteration : int, optional, default=None Iteration number, for storage purposes. Returns ------- final_old_sampler_state : openmmtools.State.SamplerState The final configurational properties of the old system after hybrid alchemical switching final_sampler_state : openmmtools.states.SamplerState The final configurational properties after `nsteps` steps of alchemical switching, and reversion to the nonalchemical system logP_work : float The NCMC work contribution to the log acceptance probability (Eqs. 62 and 63) logP_initial : float The initial logP of the hybrid configuration logP_final : float The final logP of the hybrid configuration """ assert not initial_sampler_state.has_nan( ) and not proposed_sampler_state.has_nan() #generate or retrieve the hybrid topology factory: hybrid_factory = self.make_alchemical_system( topology_proposal, initial_sampler_state.positions, proposed_sampler_state.positions) if hybrid_factory is None: _logger.warning( "Unable to construct hybrid system for {} -> {}".format( topology_proposal.old_chemical_state_key, topology_proposal.new_chemical_state_key)) return initial_sampler_state, proposed_sampler_state, -np.inf, 0.0, 0.0 topology = hybrid_factory.hybrid_topology #generate the corresponding thermodynamic and sampler states so that we can use the NonequilibriumSwitchingMove: #First generate the thermodynamic state: hybrid_system = hybrid_factory.hybrid_system hybrid_thermodynamic_state = ThermodynamicState( hybrid_system, temperature=self._temperature, pressure=self._pressure) #Now create an RelativeAlchemicalState from the hybrid system: alchemical_state = RelativeAlchemicalState.from_system(hybrid_system) alchemical_state.set_alchemical_parameters(0.0) #Now create a compound thermodynamic state that combines the hybrid thermodynamic state with the alchemical state: compound_thermodynamic_state = CompoundThermodynamicState( hybrid_thermodynamic_state, composable_states=[alchemical_state]) #construct a sampler state from the hybrid positions and the box vectors of the initial sampler state: initial_hybrid_positions = hybrid_factory.hybrid_positions initial_hybrid_box_vectors = initial_sampler_state.box_vectors initial_hybrid_sampler_state = SamplerState( initial_hybrid_positions, box_vectors=initial_hybrid_box_vectors) final_hybrid_sampler_state = copy.deepcopy( initial_hybrid_sampler_state) #create the nonequilibrium move: #ne_move = NonequilibriumSwitchingMove(self._functions, self._integrator_splitting, self._temperature, self._nsteps, self._timestep, # work_save_interval=self._write_ncmc_interval, top=topology,subset_atoms=None, # save_configuration=self._save_configuration, measure_shadow_work=self._measure_shadow_work) ne_move = ExternalNonequilibriumSwitchingMove( self._functions, nsteps_neq=self._nsteps, timestep=self._timestep, temperature=self._temperature, work_configuration_save_interval=self._work_save_interval, splitting="V R O R V") #run the NCMC protocol try: ne_move.apply(compound_thermodynamic_state, final_hybrid_sampler_state) except Exception as e: _logger.warn("NCMC failed because {}; rejecting.".format(str(e))) logP_work = -np.inf return [ initial_sampler_state, proposed_sampler_state, -np.inf, 0.0, 0.0 ] #get the total work: logP_work = -ne_move.cumulative_work[-1] # Compute contribution of transforming to and from the hybrid system: context, integrator = global_context_cache.get_context( hybrid_thermodynamic_state) #set all alchemical parameters to zero: for parameter in self._functions.keys(): context.setParameter(parameter, 0.0) initial_hybrid_sampler_state.apply_to_context(context, ignore_velocities=True) initial_reduced_potential = hybrid_thermodynamic_state.reduced_potential( context) #set all alchemical parameters to one: for parameter in self._functions.keys(): context.setParameter(parameter, 1.0) final_hybrid_sampler_state.apply_to_context(context, ignore_velocities=True) final_reduced_potential = hybrid_thermodynamic_state.reduced_potential( context) #reset the parameters back to zero just in case for parameter in self._functions.keys(): context.setParameter(parameter, 0.0) #compute the output SamplerState, which has the atoms only for the new system post-NCMC: new_positions = hybrid_factory.new_positions( final_hybrid_sampler_state.positions) new_box_vectors = final_hybrid_sampler_state.box_vectors final_sampler_state = SamplerState(new_positions, box_vectors=new_box_vectors) #compute the output SamplerState for the atoms only in the old system (required for geometry_logP_reverse) old_positions = hybrid_factory.old_positions( final_hybrid_sampler_state.positions) old_box_vectors = copy.deepcopy( new_box_vectors) #these are the same as the new system final_old_sampler_state = SamplerState(old_positions, box_vectors=old_box_vectors) #extract the trajectory and box vectors from the move: trajectory = ne_move.trajectory[::-self. _write_ncmc_interval, :, :][::-1] topology = hybrid_factory.hybrid_topology position_varname = "ncmcpositions" nframes = np.shape(trajectory)[0] #extract box vectors: box_vec_varname = "ncmcboxvectors" box_lengths = ne_move.box_lengths[::-self. _write_ncmc_interval, :][::-1] box_angles = ne_move.box_angles[::-self._write_ncmc_interval, :][::-1] box_lengths_and_angles = np.stack([box_lengths, box_angles]) #write out the positions of the topology if self._storage: for frame in range(nframes): self._storage.write_configuration(position_varname, trajectory[frame, :, :], topology, iteration=iteration, frame=frame, nframes=nframes) #write out the periodict box vectors: if self._storage: self._storage.write_array(box_vec_varname, box_lengths_and_angles, iteration=iteration) #retrieve the protocol work and write that out too: protocol_work = ne_move.cumulative_work if self._storage: self._storage.write_array("protocolwork", protocol_work, iteration=iteration) # Return return [ final_old_sampler_state, final_sampler_state, logP_work, -initial_reduced_potential, -final_reduced_potential ]
def integrate(self, topology_proposal, initial_sampler_state, proposed_sampler_state, iteration=None): """ Performs NCMC switching to either delete or insert atoms according to the provided `topology_proposal`. For `delete`, the system is first modified from fully interacting to alchemically modified, and then NCMC switching is used to eliminate atoms. For `insert`, the system begins with eliminated atoms in an alchemically noninteracting form and NCMC switching is used to turn atoms on, followed by making system real. Parameters ---------- topology_proposal : TopologyProposal Contains old/new Topology and System objects and atom mappings. initial_sampler_state : openmmtools.states.SamplerState representing the initial (old) system Configurational properties of the atoms at the beginning of the NCMC switching. proposed_sampler_state : openmmtools.states.SamplerState representing the proposed (post-geometry new) system Configurational properties new system atoms at beginning of NCMC switching iteration : int, optional, default=None Iteration number, for storage purposes. Returns ------- final_old_sampler_state : openmmtools.State.SamplerState The final configurational properties of the old system after hybrid alchemical switching final_sampler_state : openmmtools.states.SamplerState The final configurational properties after `nsteps` steps of alchemical switching, and reversion to the nonalchemical system logP_work : float The NCMC work contribution to the log acceptance probability (Eqs. 62 and 63) logP_initial : float The initial logP of the hybrid configuration logP_final : float The final logP of the hybrid configuration """ assert not initial_sampler_state.has_nan() and not proposed_sampler_state.has_nan() #generate or retrieve the hybrid topology factory: hybrid_factory = self.make_alchemical_system(topology_proposal, initial_sampler_state.positions, proposed_sampler_state.positions) if hybrid_factory is None: _logger.warning("Unable to construct hybrid system for {} -> {}".format(topology_proposal.old_chemical_state_key, topology_proposal.new_chemical_state_key)) return initial_sampler_state, proposed_sampler_state, -np.inf, 0.0, 0.0 topology = hybrid_factory.hybrid_topology #generate the corresponding thermodynamic and sampler states so that we can use the NonequilibriumSwitchingMove: #First generate the thermodynamic state: hybrid_system = hybrid_factory.hybrid_system hybrid_thermodynamic_state = ThermodynamicState(hybrid_system, temperature=self._temperature, pressure=self._pressure) #Now create an RelativeAlchemicalState from the hybrid system: alchemical_state = RelativeAlchemicalState.from_system(hybrid_system) alchemical_state.set_alchemical_parameters(0.0) #Now create a compound thermodynamic state that combines the hybrid thermodynamic state with the alchemical state: compound_thermodynamic_state = CompoundThermodynamicState(hybrid_thermodynamic_state, composable_states=[alchemical_state]) #construct a sampler state from the hybrid positions and the box vectors of the initial sampler state: initial_hybrid_positions = hybrid_factory.hybrid_positions initial_hybrid_box_vectors = initial_sampler_state.box_vectors initial_hybrid_sampler_state = SamplerState(initial_hybrid_positions, box_vectors=initial_hybrid_box_vectors) final_hybrid_sampler_state = copy.deepcopy(initial_hybrid_sampler_state) #create the nonequilibrium move: #ne_move = NonequilibriumSwitchingMove(self._functions, self._integrator_splitting, self._temperature, self._nsteps, self._timestep, # work_save_interval=self._write_ncmc_interval, top=topology,subset_atoms=None, # save_configuration=self._save_configuration, measure_shadow_work=self._measure_shadow_work) ne_move = ExternalNonequilibriumSwitchingMove(self._functions, nsteps_neq=self._nsteps, timestep=self._timestep, temperature=self._temperature, work_configuration_save_interval=self._work_save_interval, splitting="V R O R V") #run the NCMC protocol try: ne_move.apply(compound_thermodynamic_state, final_hybrid_sampler_state) except Exception as e: _logger.warn("NCMC failed because {}; rejecting.".format(str(e))) logP_work = -np.inf return [initial_sampler_state, proposed_sampler_state, -np.inf, 0.0, 0.0] #get the total work: logP_work = - ne_move.cumulative_work[-1] # Compute contribution of transforming to and from the hybrid system: context, integrator = global_context_cache.get_context(hybrid_thermodynamic_state) #set all alchemical parameters to zero: for parameter in self._functions.keys(): context.setParameter(parameter, 0.0) initial_hybrid_sampler_state.apply_to_context(context, ignore_velocities=True) initial_reduced_potential = hybrid_thermodynamic_state.reduced_potential(context) #set all alchemical parameters to one: for parameter in self._functions.keys(): context.setParameter(parameter, 1.0) final_hybrid_sampler_state.apply_to_context(context, ignore_velocities=True) final_reduced_potential = hybrid_thermodynamic_state.reduced_potential(context) #reset the parameters back to zero just in case for parameter in self._functions.keys(): context.setParameter(parameter, 0.0) #compute the output SamplerState, which has the atoms only for the new system post-NCMC: new_positions = hybrid_factory.new_positions(final_hybrid_sampler_state.positions) new_box_vectors = final_hybrid_sampler_state.box_vectors final_sampler_state = SamplerState(new_positions, box_vectors=new_box_vectors) #compute the output SamplerState for the atoms only in the old system (required for geometry_logP_reverse) old_positions = hybrid_factory.old_positions(final_hybrid_sampler_state.positions) old_box_vectors = copy.deepcopy(new_box_vectors) #these are the same as the new system final_old_sampler_state = SamplerState(old_positions, box_vectors=old_box_vectors) #extract the trajectory and box vectors from the move: trajectory = ne_move.trajectory[::-self._write_ncmc_interval, :, :][::-1] topology = hybrid_factory.hybrid_topology position_varname = "ncmcpositions" nframes = np.shape(trajectory)[0] #extract box vectors: box_vec_varname = "ncmcboxvectors" box_lengths = ne_move.box_lengths[::-self._write_ncmc_interval, :][::-1] box_angles = ne_move.box_angles[::-self._write_ncmc_interval, :][::-1] box_lengths_and_angles = np.stack([box_lengths, box_angles]) #write out the positions of the topology if self._storage: for frame in range(nframes): self._storage.write_configuration(position_varname, trajectory[frame, :, :], topology, iteration=iteration, frame=frame, nframes=nframes) #write out the periodict box vectors: if self._storage: self._storage.write_array(box_vec_varname, box_lengths_and_angles, iteration=iteration) #retrieve the protocol work and write that out too: protocol_work = ne_move.cumulative_work if self._storage: self._storage.write_array("protocolwork", protocol_work, iteration=iteration) # Return return [final_old_sampler_state, final_sampler_state, logP_work, -initial_reduced_potential, -final_reduced_potential]
def run_equilibrium(equilibrium_result: EquilibriumResult, thermodynamic_state: states.ThermodynamicState, nsteps_equil: int, topology: md.Topology, n_iterations : int, atom_indices_to_save: List[int] = None, trajectory_filename: str = None, splitting: str="V R O R V", timestep: unit.Quantity=1.0*unit.femtoseconds) -> EquilibriumResult: """ Run nsteps of equilibrium sampling at the specified thermodynamic state and return the final sampler state as well as a trajectory of the positions after each application of an MCMove. This means that if the MCMove is configured to run 1000 steps of dynamics, and n_iterations is 100, there will be 100 frames in the resulting trajectory; these are the result of 100,000 steps (1000*100) of dynamics. Parameters ---------- equilibrium_result : EquilibriumResult EquilibriumResult namedtuple containing the information necessary to resume thermodynamic_state : openmmtools.states.ThermodynamicState The thermodynamic state (including context parameters) that should be used nsteps_equil : int The number of equilibrium steps that a move should make when apply is called topology : mdtraj.Topology an MDTraj topology object used to construct the trajectory n_iterations : int The number of times to apply the move. Note that this is not the number of steps of dynamics; it is n_iterations*n_steps (which is set in the MCMove). splitting: str, default "V R O H R V" The splitting string for the dynamics atom_indices_to_save : list of int, default None list of indices to save (when excluding waters, for instance). If None, all indices are saved. trajectory_filename : str, optional, default None Full filepath of trajectory files. If none, trajectory files are not written. splitting: str, default "V R O H R V" The splitting string for the dynamics Returns ------- equilibrium_result : EquilibriumResult Container namedtuple that has the SamplerState for resuming, an MDTraj trajectory, and the reduced potential of the final frame. """ sampler_state = equilibrium_result.sampler_state #get the atom indices we need to subset the topology and positions if atom_indices_to_save is None: atom_indices = list(range(topology.n_atoms)) subset_topology = topology else: subset_topology = topology.subset(atom_indices_to_save) atom_indices = atom_indices_to_save n_atoms = subset_topology.n_atoms #construct the MCMove: mc_move = mcmc.LangevinSplittingDynamicsMove(n_steps=nsteps_equil, splitting=splitting) mc_move.n_restart_attempts = 10 #create a numpy array for the trajectory trajectory_positions = np.zeros([n_iterations, n_atoms, 3]) trajectory_box_lengths = np.zeros([n_iterations, 3]) trajectory_box_angles = np.zeros([n_iterations, 3]) #loop through iterations and apply MCMove, then collect positions into numpy array for iteration in range(n_iterations): mc_move.apply(thermodynamic_state, sampler_state) trajectory_positions[iteration, :] = sampler_state.positions[atom_indices, :].value_in_unit_system(unit.md_unit_system) #get the box lengths and angles a, b, c, alpha, beta, gamma = mdtrajutils.unitcell.box_vectors_to_lengths_and_angles(*sampler_state.box_vectors) trajectory_box_lengths[iteration, :] = [a, b, c] trajectory_box_angles[iteration, :] = [alpha, beta, gamma] #construct trajectory object: trajectory = md.Trajectory(trajectory_positions, subset_topology, unitcell_lengths=trajectory_box_lengths, unitcell_angles=trajectory_box_angles) #get the reduced potential from the final frame for endpoint perturbations reduced_potential_final_frame = thermodynamic_state.reduced_potential(sampler_state) #construct equilibrium result object equilibrium_result = EquilibriumResult(sampler_state, reduced_potential_final_frame) #If there is a trajectory filename passed, write out the results here: if trajectory_filename is not None: write_equilibrium_trajectory(equilibrium_result, trajectory, trajectory_filename) return equilibrium_result