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 minimize(thermodynamic_state: states.ThermodynamicState,
sampler_state: states.SamplerState,
max_iterations: int = 100) -> states.SamplerState:
"""
Minimize the given system and state, up to a maximum number of steps.
This does not return a copy of the samplerstate; it is an update-in-place.
Parameters
----------
thermodynamic_state : openmmtools.states.ThermodynamicState
The state at which the system could be minimized
sampler_state : openmmtools.states.SamplerState
The starting state at which to minimize the system.
max_iterations : int, optional, default 100
The maximum number of minimization steps. Default is 100.
Returns
-------
sampler_state : openmmtools.states.SamplerState
The posititions and accompanying state following minimization
"""
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)
_logger.debug(f"using dummy context cache")
else:
_logger.debug(f"using global context cache")
context, integrator = cache.global_context_cache.get_context(
thermodynamic_state)
sampler_state.apply_to_context(context, ignore_velocities=True)
openmm.LocalEnergyMinimizer.minimize(context, maxIterations=max_iterations)
sampler_state.update_from_context(context)
def create_langevin_integrator(htf, constraint_tol):
"""
create lambda alchemical states, thermodynamic states, sampler states, integrator, and return context, thermostate, sampler_state, integrator
"""
fast_lambda_alchemical_state = RelativeAlchemicalState.from_system(
htf.hybrid_system)
fast_lambda_alchemical_state.set_alchemical_parameters(
0.0, LambdaProtocol(functions='default'))
fast_thermodynamic_state = CompoundThermodynamicState(
ThermodynamicState(htf.hybrid_system, temperature=temperature),
composable_states=[fast_lambda_alchemical_state])
fast_sampler_state = SamplerState(
positions=htf._hybrid_positions,
box_vectors=htf.hybrid_system.getDefaultPeriodicBoxVectors())
integrator_1 = integrators.LangevinIntegrator(
temperature=temperature,
timestep=4.0 * unit.femtoseconds,
splitting='V R O R V',
measure_shadow_work=False,
measure_heat=False,
constraint_tolerance=constraint_tol,
collision_rate=5.0 / unit.picoseconds)
# mcmc_moves=mcmc.LangevinSplittingDynamicsMove(timestep = 4.0 * unit.femtoseconds,
# collision_rate=5.0 / unit.picosecond,
# n_steps=1,
# reassign_velocities=False,
# n_restart_attempts=20,
# splitting="V R R R O R R R V",
# constraint_tolerance=constraint_tol)
#print(integrator_1.getConstraintTolerance())
fast_context, fast_integrator = cache.global_context_cache.get_context(
fast_thermodynamic_state, integrator_1)
fast_sampler_state.apply_to_context(fast_context)
return fast_context, fast_thermodynamic_state, fast_sampler_state, fast_integrator
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)
class Propagator(OMMBIP):
"""
Propagator pseudocode:
Step 1: initialization--
set iteration = 0, n_iterations = n_iterations, lambda = 0 (i.e. iteration / n_iterations); work_accumulated = 0.0
generate sample x_0 ~ e^(-p(x))
evaluate work_incremental = 0 (i.e. u_mm(x_0) - g(x_0), but we presume that g = u_mm(.))
work_accumulated <- work_accumulated + work_incremental
x' = x_0
Step 2: sampling
for increment in range(n_iterations):
x = x'
ante_perturbation_potential = (1 - lambda) * u_mm(x) + lambda * u_ani_mm_mix(x)
set iteration <- iteration + 1.0; lambda <- iteration / n_iterations
evaluate work_incremental = [(1 - lambda) * u_mm(x) + lambda * u_ani_mm_mix(x)] - ante_perturbation_potential
work_accumulated <- work_accumulated + work_incremental
create a modified force: modified_f = (1 - lambda) * f_mm + lambda * f_ani_mm_mix (where f_. = -grad(u_.) )
x' = V R O R (where V deterministic update is according to modified_f defined above) w.r.t x
NOTE: in this regime, the last x' is propagated w.r.t. a propagator whose invariant distribution respects u_ani_mm_mix;
this should _not_ be the case. There should be an exception in the Step 2 for loop that breaks once the final work_incremental is computed and updated
to the work_accumulated. Regardless, the distribution of accumulated works is unaffected by this 'bug'; only expectations (as a function of x) w.r.t. these
weights may be affected.
See: 3.1.1. of https://www.stats.ox.ac.uk/~doucet/delmoral_doucet_jasra_sequentialmontecarlosamplersJRSSB.pdf (esp. Remark 1.)
"""
def __init__(self,
openmm_pdf_state,
openmm_pdf_state_subset,
subset_indices_map,
integrator,
ani_handler,
context_cache=None,
reassign_velocities=True,
n_restart_attempts=0,
reporter=None,
write_trajectory_interval = 1,
**kwargs):
"""
arguments
openmm_pdf_state : openmmtools.states.ThermodynamicState
the pdf state of the propagator
openmm_pdf_state_subset : openmmtools.states.ThermodynamicState
the pdf state of the atom subset
subset_indices_map : dict
dict of {openmm_pdf_state atom_index : openmm_pdf_state_subset atom index}
integrator : openmm.Integrator
integrator of dynamics
ani_handler : ANI1_force_and_energy
handler for ani forces and potential energy
context_cache : openmmtools.cache.ContextCache, optional default:None
The ContextCache to use for Context creation. If None, the global cache
openmmtools.cache.global_context_cache is used.
reassign_velocities : bool, optional default:False
If True, the velocities will be reassigned from the Maxwell-Boltzmann
distribution at the beginning of the move.
n_restart_attempts : int, optional default:0
When greater than 0, if after the integration there are NaNs in energies,
the move will restart. When the integrator has a random component, this
may help recovering. On the last attempt, the ``Context`` is
re-initialized in a slower process, but better than the simulation
crashing. An IntegratorMoveError is raised after the given number of
attempts if there are still NaNs.
reporter : coddiwomple.openmm.reporter.OpenMMReporter, default None
a reporter object to write trajectories
write_trajectory_interval : int
frequency of writing trajectory
"""
super().__init__(openmm_pdf_state,
integrator,
context_cache,
reassign_velocities,
n_restart_attempts)
#create a pdf state for the subset indices (usually a vacuum system)
self.pdf_state_subset = openmm_pdf_state_subset
assert self.pdf_state_subset.temperature == self.pdf_state.temperature, f"the temperatures of the pdf states do not match"
#create a dictionary for subset indices
self._subset_indices_map = subset_indices_map
#create an ani handler attribute that can be referenced
self.ani_handler = ani_handler
#create a context for the subset atoms that can be referenced
self.context_subset, _ = cache.global_context_cache.get_context(self.pdf_state_subset)
#create a reporter for the accumulated works
self._state_works = {}
self._state_works_counter = 0
#create a reporter
self._write_trajectory = False if reporter is None else True
self.reporter=reporter
if self._write_trajectory:
from coddiwomple.particles import Particle
self.particle = Particle(0)
self.write_trajectory_interval=write_trajectory_interval
else:
self.particle = None
self.write_trajectory_interval=None
def _initialize_state_works(self):
"""
initialize an empty list and add 0.0 to it (state works)
"""
self._current_state_works = [] #define an interim (auxiliary) list that will track the thermodynamic work of the current application
self._current_state_works.append(0.0) #the first incremental work is always 0 since the importance function is identical to the first target distribution (i.e. fully interacting MM)
def _initialize_iterations(self, n_iterations):
"""
initialize the iteration counter
"""
self._iteration = 0.0 #define the first iteration as 0
self._n_iterations = n_iterations #the number of iterations in the protocol is equal to the number of steps in the application
def _update_particle_state_substate(self, particle_state, new_state_subset=False):
"""
update the particle state from the context, create a particle substate and update from context
"""
#update the particle state and the particle state subset
particle_state.update_from_context(self.context, ignore_velocities=True) #update the particle state from the context
if new_state_subset:
self.particle_state_subset = SamplerState(positions = particle_state.positions[list(self._subset_indices_map.keys())]) #create a particle state from the subset context
else:
self.particle_state_subset.positions = particle_state.positions[list(self._subset_indices_map.keys())] #update the particle subset positions appropriately
self.particle_state_subset.apply_to_context(self.context_subset, ignore_velocities=True) #apply the subset particle state to its context
self.particle_state_subset.update_from_context(self.context_subset, ignore_velocities=True) #update the subset particle state from its context to updated the potential energy
def _update_current_state_works(self, particle_state):
"""
update the current state and associated works
"""
#get the reduced potential
reduced_potential = self._compute_hybrid_potential(_lambda = self._iteration / self._n_iterations, particle_state = particle_state)
perturbed_reduced_potential = self._compute_hybrid_potential(_lambda = (self._iteration + 1.0) / self._n_iterations, particle_state = particle_state)
self._current_state_works.append(self._current_state_works[-1] + (perturbed_reduced_potential - reduced_potential))
def _update_force(self, particle_state):
"""
update the force
"""
mm_force_matrix = self._compute_hybrid_forces(_lambda = (self._iteration + 1.0) / self._n_iterations, particle_state = particle_state).value_in_unit_system(unit.md_unit_system)
self.integrator.setPerDofVariableByName('modified_force', mm_force_matrix)
def _before_integration(self, *args, **kwargs):
particle_state = args[0] #define the particle state
n_iterations = args[1] #define the number of iterations
self._initialize_state_works()
self._initialize_iterations(n_iterations)
#update the particle state and the particle state subset
self._update_particle_state_substate(particle_state, new_state_subset=True)
self._update_current_state_works(particle_state)
self._update_force(particle_state)
#report
if self._write_trajectory: # the first state is always saved for processing purposes
self.particle.update_state(particle_state)
self.reporter.record([self.particle])
def _during_integration(self, *args, **kwargs):
particle_state = args[0]
self._iteration += 1.0
self._update_particle_state_substate(particle_state)
#get the reduced potential
if self._iteration < self._n_iterations:
self._update_current_state_works(particle_state)
self._update_force(particle_state)
else:
#we are done
pass
if self._write_trajectory and int(self._iteration) % self.write_trajectory_interval == 0:
self.particle.update_state(particle_state)
if self._iteration == self._n_iterations:
self.reporter.record([self.particle], save_to_disk=True)
else:
self.reporter.record([self.particle], save_to_disk=False)
def _after_integration(self, *args, **kwargs):
self._state_works[self._state_works_counter] = deepcopy(self._current_state_works)
self._state_works_counter += 1
if self._write_trajectory:
self.reporter.reset()
#self._log_context_parameters()
def _compute_hybrid_potential(self,_lambda, particle_state):
"""
function to compute the hybrid reduced potential defined as follows:
U(x_rec, x_lig) = u_mm,rec(x_rec) - lambda*u_mm,lig(x_lig) + lambda*u_ani,lig(x_lig)
"""
reduced_potential = (self.pdf_state.reduced_potential(particle_state)
- _lambda * self.pdf_state_subset.reduced_potential(self.particle_state_subset)
+ _lambda * self.ani_handler.calculate_energy(self.particle_state_subset.positions) * self.pdf_state.beta)
return reduced_potential
def _compute_hybrid_forces(self, _lambda, particle_state):
"""
function to compute a hybrid force matrix of shape num_particles x 3
in the spirit of the _compute_hybrid_potential, we compute the forces in the following way
F(x_rec, x_lig) = F_mm(x_rec, x_lig) - lambda * F_mm(x_lig) + lambda * F_ani(x_lig)
"""
# get the complex mm forces
state = self.context.getState(getForces=True)
mm_force_matrix = state.getForces(asNumpy=True) # returns forces in kJ/(nm mol)
# get the ligand mm forces
subset_state = self.context_subset.getState(getForces=True)
mm_force_matrix_subset = subset_state.getForces(asNumpy=True)
# get the ligand ani forces
coords = self.particle_state_subset.positions
subset_ani_force_matrix, energie = self.ani_handler.calculate_force(coords) # returns force in kJ/(A mol)
#print(f"ani force matrix head: ",subset_ani_force_matrix[0])
# now combine the ligand forces
subset_force_matrix = _lambda * (subset_ani_force_matrix - mm_force_matrix_subset) #we are adding two Quantities with different units, but they are compatible
#print(f"mm subset force matrix head", mm_force_matrix_subset[0])
# and append to the complex forces...
#print(f"mm force matrix head", mm_force_matrix[0])
mm_force_matrix[list(self._subset_indices_map.keys()), :] += subset_force_matrix #and same, here...
#print(f"mm force matrix head (after ani modification)", mm_force_matrix[0])
return mm_force_matrix
def _get_context_subset_parameters(self):
"""
return a dictionary of the self.context_subset's parameters
returns
context_parameters : dict
{parameter name <str> : parameter value value <float>}
"""
swig_parameters = self.context_subset.getParameters()
context_parameters = {q: swig_parameters[q] for q in swig_parameters}
return context_parameters
def _log_context_parameters(self):
"""
log the context and context subset parameters
"""
context_parameters = self._get_context_parameters()
context_subset_parameters = self._get_context_subset_parameters()
_logger.debug(f"\tcontext_parameters during integration:")
for key, val in context_parameters.items():
_logger.debug(f"\t\t{key}: {val}")
_logger.debug(f"\tcontext subset parameters during integration:")
for key, val in context_subset_parameters:
_logger.debug(f"\t\t{key}: {val}")
@property
def state_works(self):
return self._state_works
def HybridTopologyFactory_energies(
current_mol='toluene',
proposed_mol='1,2-bis(trifluoromethyl) benzene'):
"""
Test whether the difference in the nonalchemical zero and alchemical zero states is the forward valence energy. Also test for the one states.
"""
from perses.tests.utils import generate_solvated_hybrid_test_topology, generate_endpoint_thermodynamic_states
import openmmtools.cache as cache
#Just test the solvated system
top_proposal, old_positions, _ = generate_solvated_hybrid_test_topology(
current_mol_name=current_mol, proposed_mol_name=proposed_mol)
#remove the dispersion correction
top_proposal._old_system.getForce(3).setUseDispersionCorrection(False)
top_proposal._new_system.getForce(3).setUseDispersionCorrection(False)
# run geometry engine to generate old and new positions
_geometry_engine = FFAllAngleGeometryEngine(metadata=None,
use_sterics=False,
n_bond_divisions=100,
n_angle_divisions=180,
n_torsion_divisions=360,
verbose=True,
storage=None,
bond_softening_constant=1.0,
angle_softening_constant=1.0,
neglect_angles=False)
_new_positions, _lp = _geometry_engine.propose(top_proposal, old_positions,
beta)
_lp_rev = _geometry_engine.logp_reverse(top_proposal, _new_positions,
old_positions, beta)
# make the hybrid system, reset the CustomNonbondedForce cutoff
HTF = HybridTopologyFactory(top_proposal, old_positions, _new_positions)
hybrid_system = HTF.hybrid_system
nonalch_zero, nonalch_one, alch_zero, alch_one = generate_endpoint_thermodynamic_states(
hybrid_system, top_proposal)
# compute reduced energies
#for the nonalchemical systems...
attrib_list = [(nonalch_zero, old_positions,
top_proposal._old_system.getDefaultPeriodicBoxVectors()),
(alch_zero, HTF._hybrid_positions,
hybrid_system.getDefaultPeriodicBoxVectors()),
(alch_one, HTF._hybrid_positions,
hybrid_system.getDefaultPeriodicBoxVectors()),
(nonalch_one, _new_positions,
top_proposal._new_system.getDefaultPeriodicBoxVectors())]
rp_list = []
for (state, pos, box_vectors) in attrib_list:
context, integrator = cache.global_context_cache.get_context(state)
samplerstate = SamplerState(positions=pos, box_vectors=box_vectors)
samplerstate.apply_to_context(context)
rp = state.reduced_potential(context)
rp_list.append(rp)
#valence energy definitions
forward_added_valence_energy = _geometry_engine.forward_final_context_reduced_potential - _geometry_engine.forward_atoms_with_positions_reduced_potential
reverse_subtracted_valence_energy = _geometry_engine.reverse_final_context_reduced_potential - _geometry_engine.reverse_atoms_with_positions_reduced_potential
nonalch_zero_rp, alch_zero_rp, alch_one_rp, nonalch_one_rp = rp_list[
0], rp_list[1], rp_list[2], rp_list[3]
# print(f"Difference between zeros: {nonalch_zero_rp - alch_zero_rp}; forward added: {forward_added_valence_energy}")
# print(f"Difference between ones: {nonalch_zero_rp - alch_zero_rp}; forward added: {forward_added_valence_energy}")
assert abs(
nonalch_zero_rp - alch_zero_rp + forward_added_valence_energy
) < ENERGY_THRESHOLD, f"The zero state alchemical and nonalchemical energy absolute difference {abs(nonalch_zero_rp - alch_zero_rp + forward_added_valence_energy)} is greater than the threshold of {ENERGY_THRESHOLD}."
assert abs(
nonalch_one_rp - alch_one_rp + reverse_subtracted_valence_energy
) < ENERGY_THRESHOLD, f"The one state alchemical and nonalchemical energy absolute difference {abs(nonalch_one_rp - alch_one_rp + reverse_subtracted_valence_energy)} is greater than the threshold of {ENERGY_THRESHOLD}."
print(
f"Abs difference in zero alchemical vs nonalchemical systems: {abs(nonalch_zero_rp - alch_zero_rp + forward_added_valence_energy)}"
)
print(
f"Abs difference in one alchemical vs nonalchemical systems: {abs(nonalch_one_rp - alch_one_rp + reverse_subtracted_valence_energy)}"
)
def run_endpoint_perturbation(lambda_thermodynamic_state,
nonalchemical_thermodynamic_state,
initial_hybrid_sampler_state,
mc_move,
n_iterations,
factory,
lambda_index=0,
print_work=False,
write_system=False,
write_state=False,
write_trajectories=False):
"""
Parameters
----------
lambda_thermodynamic_state : ThermodynamicState
The thermodynamic state corresponding to the hybrid system at a lambda endpoint
nonalchemical_thermodynamic_state : ThermodynamicState
The nonalchemical thermodynamic state for the relevant endpoint
initial_hybrid_sampler_state : SamplerState
Starting positions for the sampler. Must be compatible with lambda_thermodynamic_state
mc_move : MCMCMove
The MCMove that will be used for sampling at the lambda endpoint
n_iterations : int
The number of iterations
factory : HybridTopologyFactory
The hybrid topology factory
lambda_index : int, optional, default=0
The index, 0 or 1, at which to retrieve nonalchemical positions
print_work : bool, optional, default=False
If True, will print work values
write_system : bool, optional, default=False
If True, will write alchemical and nonalchemical System XML files
write_state : bool, optional, default=False
If True, write alchemical (hybrid) State XML files each iteration
write_trajectories : bool, optional, default=False
If True, will write trajectories
Returns
-------
df : float
Free energy difference between alchemical and nonalchemical systems, estimated with EXP
ddf : float
Standard deviation of estimate, corrected for correlation, from EXP estimator.
"""
import mdtraj as md
#run an initial minimization:
mcmc_sampler = mcmc.MCMCSampler(lambda_thermodynamic_state,
initial_hybrid_sampler_state, mc_move)
mcmc_sampler.minimize(max_iterations=20)
new_sampler_state = mcmc_sampler.sampler_state
if write_system:
with open(f'hybrid{lambda_index}-system.xml', 'w') as outfile:
outfile.write(
openmm.XmlSerializer.serialize(
lambda_thermodynamic_state.system))
with open(f'nonalchemical{lambda_index}-system.xml', 'w') as outfile:
outfile.write(
openmm.XmlSerializer.serialize(
nonalchemical_thermodynamic_state.system))
#initialize work array
w = np.zeros([n_iterations])
non_potential = np.zeros([n_iterations])
hybrid_potential = np.zeros([n_iterations])
#run n_iterations of the endpoint perturbation:
hybrid_trajectory = unit.Quantity(
np.zeros([
n_iterations,
lambda_thermodynamic_state.system.getNumParticles(), 3
]), unit.nanometers) # DEBUG
nonalchemical_trajectory = unit.Quantity(
np.zeros([
n_iterations,
nonalchemical_thermodynamic_state.system.getNumParticles(), 3
]), unit.nanometers) # DEBUG
for iteration in range(n_iterations):
# Generate a new sampler state for the hybrid system
mc_move.apply(lambda_thermodynamic_state, new_sampler_state)
# Compute the hybrid reduced potential at the new sampler state
hybrid_context, integrator = cache.global_context_cache.get_context(
lambda_thermodynamic_state)
new_sampler_state.apply_to_context(hybrid_context,
ignore_velocities=True)
hybrid_reduced_potential = lambda_thermodynamic_state.reduced_potential(
hybrid_context)
if write_state:
state = hybrid_context.getState(getPositions=True,
getParameters=True)
state_xml = openmm.XmlSerializer.serialize(state)
with open(f'state{iteration}_l{lambda_index}.xml', 'w') as outfile:
outfile.write(state_xml)
# Construct a sampler state for the nonalchemical system
if lambda_index == 0:
nonalchemical_positions = factory.old_positions(
new_sampler_state.positions)
elif lambda_index == 1:
nonalchemical_positions = factory.new_positions(
new_sampler_state.positions)
else:
raise ValueError(
"The lambda index needs to be either one or zero for this to be meaningful"
)
nonalchemical_sampler_state = SamplerState(
nonalchemical_positions, box_vectors=new_sampler_state.box_vectors)
if write_trajectories:
state = hybrid_context.getState(getPositions=True)
hybrid_trajectory[iteration, :, :] = state.getPositions(
asNumpy=True)
nonalchemical_trajectory[iteration, :, :] = nonalchemical_positions
# Compute the nonalchemical reduced potential
nonalchemical_context, integrator = cache.global_context_cache.get_context(
nonalchemical_thermodynamic_state)
nonalchemical_sampler_state.apply_to_context(nonalchemical_context,
ignore_velocities=True)
nonalchemical_reduced_potential = nonalchemical_thermodynamic_state.reduced_potential(
nonalchemical_context)
# Compute and store the work
w[iteration] = nonalchemical_reduced_potential - hybrid_reduced_potential
non_potential[iteration] = nonalchemical_reduced_potential
hybrid_potential[iteration] = hybrid_reduced_potential
if print_work:
print(
f'{iteration:8d} {hybrid_reduced_potential:8.3f} {nonalchemical_reduced_potential:8.3f} => {w[iteration]:8.3f}'
)
if write_trajectories:
if lambda_index == 0:
nonalchemical_mdtraj_topology = md.Topology.from_openmm(
factory._topology_proposal.old_topology)
elif lambda_index == 1:
nonalchemical_mdtraj_topology = md.Topology.from_openmm(
factory._topology_proposal.new_topology)
md.Trajectory(
hybrid_trajectory / unit.nanometers,
factory.hybrid_topology).save(f'hybrid{lambda_index}.pdb')
md.Trajectory(nonalchemical_trajectory / unit.nanometers,
nonalchemical_mdtraj_topology).save(
f'nonalchemical{lambda_index}.pdb')
# Analyze data and return results
[t0, g, Neff_max] = timeseries.detectEquilibration(w)
w_burned_in = w[t0:]
[df, ddf] = pymbar.EXP(w_burned_in)
ddf_corrected = ddf * np.sqrt(g)
results = [df, ddf_corrected, t0, Neff_max]
return results, non_potential, hybrid_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 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_endpoint_perturbation(lambda_thermodynamic_state,
nonalchemical_thermodynamic_state,
initial_hybrid_sampler_state,
mc_move,
n_iterations,
factory,
lambda_index=0):
"""
Parameters
----------
lambda_thermodynamic_state : ThermodynamicState
The thermodynamic state corresponding to the hybrid system at a lambda endpoint
nonalchemical_thermodynamic_state : ThermodynamicState
The nonalchemical thermodynamic state for the relevant endpoint
initial_hybrid_sampler_state : SamplerState
Starting positions for the sampler. Must be compatible with lambda_thermodynamic_state
mc_move : MCMCMove
The MCMove that will be used for sampling at the lambda endpoint
n_iterations : int
The number of iterations
factory : HybridTopologyFactory
The hybrid topology factory
lambda_index : int, optional default 0
The index, 0 or 1, at which to retrieve nonalchemical positions
Returns
-------
df : float
Free energy difference between alchemical and nonalchemical systems, estimated with EXP
ddf : float
Standard deviation of estimate, corrected for correlation, from EXP estimator.
"""
#run an initial minimization:
mcmc_sampler = mcmc.MCMCSampler(lambda_thermodynamic_state,
initial_hybrid_sampler_state, mc_move)
mcmc_sampler.minimize(max_iterations=20)
new_sampler_state = mcmc_sampler.sampler_state
#initialize work array
w = np.zeros([n_iterations])
#run n_iterations of the endpoint perturbation:
for iteration in range(n_iterations):
mc_move.apply(lambda_thermodynamic_state, new_sampler_state)
#compute the reduced potential at the new state
hybrid_context, integrator = cache.global_context_cache.get_context(
lambda_thermodynamic_state)
new_sampler_state.apply_to_context(hybrid_context,
ignore_velocities=True)
hybrid_reduced_potential = lambda_thermodynamic_state.reduced_potential(
hybrid_context)
#generate a sampler state for the nonalchemical system
if lambda_index == 0:
nonalchemical_positions = factory.old_positions(
new_sampler_state.positions)
elif lambda_index == 1:
nonalchemical_positions = factory.new_positions(
new_sampler_state.positions)
else:
raise ValueError(
"The lambda index needs to be either one or zero for this to be meaningful"
)
nonalchemical_sampler_state = SamplerState(
nonalchemical_positions, box_vectors=new_sampler_state.box_vectors)
#compute the reduced potential at the nonalchemical system as well:
nonalchemical_context, integrator = cache.global_context_cache.get_context(
nonalchemical_thermodynamic_state)
nonalchemical_sampler_state.apply_to_context(nonalchemical_context,
ignore_velocities=True)
nonalchemical_reduced_potential = nonalchemical_thermodynamic_state.reduced_potential(
nonalchemical_context)
w[iteration] = nonalchemical_reduced_potential - hybrid_reduced_potential
[t0, g, Neff_max] = timeseries.detectEquilibration(w)
print(Neff_max)
w_burned_in = w[t0:]
[df, ddf] = pymbar.EXP(w_burned_in)
ddf_corrected = ddf * np.sqrt(g)
return [df, ddf_corrected, Neff_max]
def run_endpoint_perturbation(lambda_thermodynamic_state, nonalchemical_thermodynamic_state, initial_hybrid_sampler_state, mc_move, n_iterations, factory,
lambda_index=0, print_work=False, write_system=False, write_state=False, write_trajectories=False):
"""
Parameters
----------
lambda_thermodynamic_state : ThermodynamicState
The thermodynamic state corresponding to the hybrid system at a lambda endpoint
nonalchemical_thermodynamic_state : ThermodynamicState
The nonalchemical thermodynamic state for the relevant endpoint
initial_hybrid_sampler_state : SamplerState
Starting positions for the sampler. Must be compatible with lambda_thermodynamic_state
mc_move : MCMCMove
The MCMove that will be used for sampling at the lambda endpoint
n_iterations : int
The number of iterations
factory : HybridTopologyFactory
The hybrid topology factory
lambda_index : int, optional, default=0
The index, 0 or 1, at which to retrieve nonalchemical positions
print_work : bool, optional, default=False
If True, will print work values
write_system : bool, optional, default=False
If True, will write alchemical and nonalchemical System XML files
write_state : bool, optional, default=False
If True, write alchemical (hybrid) State XML files each iteration
write_trajectories : bool, optional, default=False
If True, will write trajectories
Returns
-------
df : float
Free energy difference between alchemical and nonalchemical systems, estimated with EXP
ddf : float
Standard deviation of estimate, corrected for correlation, from EXP estimator.
"""
import mdtraj as md
#run an initial minimization:
mcmc_sampler = mcmc.MCMCSampler(lambda_thermodynamic_state, initial_hybrid_sampler_state, mc_move)
mcmc_sampler.minimize(max_iterations=20)
new_sampler_state = mcmc_sampler.sampler_state
if write_system:
with open(f'hybrid{lambda_index}-system.xml', 'w') as outfile:
outfile.write(openmm.XmlSerializer.serialize(lambda_thermodynamic_state.system))
with open(f'nonalchemical{lambda_index}-system.xml', 'w') as outfile:
outfile.write(openmm.XmlSerializer.serialize(nonalchemical_thermodynamic_state.system))
#initialize work array
w = np.zeros([n_iterations])
non_potential = np.zeros([n_iterations])
hybrid_potential = np.zeros([n_iterations])
#run n_iterations of the endpoint perturbation:
hybrid_trajectory = unit.Quantity(np.zeros([n_iterations, lambda_thermodynamic_state.system.getNumParticles(), 3]), unit.nanometers) # DEBUG
nonalchemical_trajectory = unit.Quantity(np.zeros([n_iterations, nonalchemical_thermodynamic_state.system.getNumParticles(), 3]), unit.nanometers) # DEBUG
for iteration in range(n_iterations):
# Generate a new sampler state for the hybrid system
mc_move.apply(lambda_thermodynamic_state, new_sampler_state)
# Compute the hybrid reduced potential at the new sampler state
hybrid_context, integrator = cache.global_context_cache.get_context(lambda_thermodynamic_state)
new_sampler_state.apply_to_context(hybrid_context, ignore_velocities=True)
hybrid_reduced_potential = lambda_thermodynamic_state.reduced_potential(hybrid_context)
if write_state:
state = hybrid_context.getState(getPositions=True, getParameters=True)
state_xml = openmm.XmlSerializer.serialize(state)
with open(f'state{iteration}_l{lambda_index}.xml', 'w') as outfile:
outfile.write(state_xml)
# Construct a sampler state for the nonalchemical system
if lambda_index == 0:
nonalchemical_positions = factory.old_positions(new_sampler_state.positions)
elif lambda_index == 1:
nonalchemical_positions = factory.new_positions(new_sampler_state.positions)
else:
raise ValueError("The lambda index needs to be either one or zero for this to be meaningful")
nonalchemical_sampler_state = SamplerState(nonalchemical_positions, box_vectors=new_sampler_state.box_vectors)
if write_trajectories:
state = hybrid_context.getState(getPositions=True)
hybrid_trajectory[iteration,:,:] = state.getPositions(asNumpy=True)
nonalchemical_trajectory[iteration,:,:] = nonalchemical_positions
# Compute the nonalchemical reduced potential
nonalchemical_context, integrator = cache.global_context_cache.get_context(nonalchemical_thermodynamic_state)
nonalchemical_sampler_state.apply_to_context(nonalchemical_context, ignore_velocities=True)
nonalchemical_reduced_potential = nonalchemical_thermodynamic_state.reduced_potential(nonalchemical_context)
# Compute and store the work
w[iteration] = nonalchemical_reduced_potential - hybrid_reduced_potential
non_potential[iteration] = nonalchemical_reduced_potential
hybrid_potential[iteration] = hybrid_reduced_potential
if print_work:
print(f'{iteration:8d} {hybrid_reduced_potential:8.3f} {nonalchemical_reduced_potential:8.3f} => {w[iteration]:8.3f}')
if write_trajectories:
if lambda_index == 0:
nonalchemical_mdtraj_topology = md.Topology.from_openmm(factory._topology_proposal.old_topology)
elif lambda_index == 1:
nonalchemical_mdtraj_topology = md.Topology.from_openmm(factory._topology_proposal.new_topology)
md.Trajectory(hybrid_trajectory / unit.nanometers, factory.hybrid_topology).save(f'hybrid{lambda_index}.pdb')
md.Trajectory(nonalchemical_trajectory / unit.nanometers, nonalchemical_mdtraj_topology).save(f'nonalchemical{lambda_index}.pdb')
# Analyze data and return results
[t0, g, Neff_max] = timeseries.detectEquilibration(w)
w_burned_in = w[t0:]
[df, ddf] = pymbar.EXP(w_burned_in)
ddf_corrected = ddf * np.sqrt(g)
results = [df, ddf_corrected, t0, Neff_max]
return results, non_potential, hybrid_potential