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
