def _compute_energy_contribution(self, hybrid_thermodynamic_state,
initial_sampler_state,
final_sampler_state):
"""
Compute NCMC energy contribution to log probability.
See Eqs. 62 and 63 (two-stage) and Eq. 45 (hybrid) of reference document.
In both cases, the contribution is u(final_positions, final_lambda) - u(initial_positions, initial_lambda).
Parameters
----------
hybrid_thermodynamic_state : openmmtools.states.CompoundThermodynamicState
The thermodynamic state of the hybrid sampler.
initial_sampler_state : openmmtools.states.SamplerState
The sampler state of the nonalchemical system at the start of the NCMC protocol with box vectors
final_sampler_state : openmmtools.states.SamplerState
The sampler state of the nonalchemical system at the end of the NCMC protocol
Returns
-------
logP_energy : float
The NCMC energy contribution to log probability.
"""
hybrid_thermodynamic_state.set_alchemical_parameters(0.0)
initial_reduced_potential = compute_reduced_potential(
hybrid_thermodynamic_state, initial_sampler_state)
hybrid_thermodynamic_state.set_alchemical_parameters(1.0)
final_reduced_potential = compute_reduced_potential(
hybrid_thermodynamic_state, final_sampler_state)
return final_reduced_potential - initial_reduced_potential
def _compute_energy_contribution(self, hybrid_thermodynamic_state, initial_sampler_state, final_sampler_state):
"""
Compute NCMC energy contribution to log probability.
See Eqs. 62 and 63 (two-stage) and Eq. 45 (hybrid) of reference document.
In both cases, the contribution is u(final_positions, final_lambda) - u(initial_positions, initial_lambda).
Parameters
----------
hybrid_thermodynamic_state : openmmtools.states.CompoundThermodynamicState
The thermodynamic state of the hybrid sampler.
initial_sampler_state : openmmtools.states.SamplerState
The sampler state of the nonalchemical system at the start of the NCMC protocol with box vectors
final_sampler_state : openmmtools.states.SamplerState
The sampler state of the nonalchemical system at the end of the NCMC protocol
Returns
-------
logP_energy : float
The NCMC energy contribution to log probability.
"""
hybrid_thermodynamic_state.set_alchemical_parameters(0.0)
initial_reduced_potential = compute_reduced_potential(hybrid_thermodynamic_state, initial_sampler_state)
hybrid_thermodynamic_state.set_alchemical_parameters(1.0)
final_reduced_potential = compute_reduced_potential(hybrid_thermodynamic_state, final_sampler_state)
return final_reduced_potential - initial_reduced_potential
def run_rj_proposals(top_prop, configuration_traj, use_sterics, ncmc_nsteps, n_replicates, box_vectors, temperature=300.0*unit.kelvin):
ncmc_engine = NCMCEngine(nsteps=ncmc_nsteps, pressure=1.0*unit.atmosphere)
geometry_engine = FFAllAngleGeometryEngine(use_sterics=use_sterics)
initial_thermodynamic_state = states.ThermodynamicState(top_prop.old_system, temperature=temperature, pressure=1.0*unit.atmosphere)
final_thermodynamic_state = states.ThermodynamicState(top_prop.new_system, temperature=temperature, pressure=1.0*unit.atmosphere)
traj_indices = np.arange(0, configuration_traj.n_frames)
results = np.zeros([n_replicates, 7])
beta = 1.0 / (temperature * constants.kB)
for i in tqdm.trange(n_replicates):
frame_index = np.random.choice(traj_indices)
initial_sampler_state = traj_frame_to_sampler_state(configuration_traj, frame_index,box_vectors)
initial_logP = - compute_reduced_potential(initial_thermodynamic_state, initial_sampler_state)
proposed_geometry, logP_geometry_forward = geometry_engine.propose(top_prop, initial_sampler_state.positions, beta)
proposed_sampler_state = states.SamplerState(proposed_geometry, box_vectors=initial_sampler_state.box_vectors)
final_old_sampler_state, final_sampler_state, logP_work, initial_hybrid_logP, final_hybrid_logP = ncmc_engine.integrate(top_prop, initial_sampler_state, proposed_sampler_state)
final_logP = - compute_reduced_potential(final_thermodynamic_state, final_sampler_state)
logP_reverse = geometry_engine.logp_reverse(top_prop, final_sampler_state.positions, final_old_sampler_state.positions, beta)
results[i, 0] = initial_logP
results[i, 1] = logP_reverse
results[i, 2] = final_logP
results[i, 3] = logP_work
results[i, 4] = initial_hybrid_logP
results[i, 5] = final_hybrid_logP
results[i, 6] = logP_geometry_forward
return results
def run_rj_proposals(top_prop, configuration_traj, use_sterics, ncmc_nsteps, n_replicates, bond_softening_constant=1.0, angle_softening_constant=1.0):
ncmc_engine = NCMCEngine(nsteps=ncmc_nsteps, pressure=1.0*unit.atmosphere, bond_softening_constant=bond_softening_constant, angle_softening_constant=angle_softening_constant)
geometry_engine = FFAllAngleGeometryEngine(use_sterics=use_sterics, bond_softening_constant=bond_softening_constant, angle_softening_constant=angle_softening_constant)
initial_thermodynamic_state = states.ThermodynamicState(top_prop.old_system, temperature=temperature, pressure=1.0*unit.atmosphere)
final_thermodynamic_state = states.ThermodynamicState(top_prop.new_system, temperature=temperature, pressure=1.0*unit.atmosphere)
traj_indices = np.arange(0, configuration_traj.n_frames)
results = np.zeros([n_replicates, 4])
for i in tqdm.trange(n_replicates):
frame_index = np.random.choice(traj_indices)
initial_sampler_state = traj_frame_to_sampler_state(configuration_traj, frame_index)
initial_logP = - compute_reduced_potential(initial_thermodynamic_state, initial_sampler_state)
proposed_geometry, logP_geometry_forward = geometry_engine.propose(top_prop, initial_sampler_state.positions, beta)
proposed_sampler_state = states.SamplerState(proposed_geometry, box_vectors=initial_sampler_state.box_vectors)
final_old_sampler_state, final_sampler_state, logP_work, initial_hybrid_logP, final_hybrid_logP = ncmc_engine.integrate(top_prop, initial_sampler_state, proposed_sampler_state)
final_logP = - compute_reduced_potential(final_thermodynamic_state, final_sampler_state)
logP_reverse = geometry_engine.logp_reverse(top_prop, final_sampler_state.positions, final_old_sampler_state.positions, beta)
results[i, 0] = initial_hybrid_logP - initial_logP
results[i, 1] = logP_reverse - logP_geometry_forward
results[i, 2] = final_logP - final_hybrid_logP
results[i, 3] = logP_work
return results
def _geometry_ncmc_geometry(self, topology_proposal, sampler_state, old_log_weight, new_log_weight):
"""
Use a hybrid NCMC protocol to switch from the old system to new system
Will calculate new positions for the new system first, then give both
sets of positions to the hybrid NCMC integrator, and finally use the
final positions of the old and new systems to calculate the reverse
geometry probability
Parameters
----------
topology_proposal : TopologyProposal
Contains old/new Topology and System objects and atom mappings.
sampler_state : openmmtools.states.SamplerState
Configurational properties of old atoms at the beginning of the NCMC switching.
old_log_weight : float
Chemical state weight from SAMSSampler
new_log_weight : float
Chemical state weight from SAMSSampler
Returns
-------
logP_accept : float
Log of acceptance probability of entire Expanded Ensemble switch (Eq. 25 or 46)
ncmc_new_sampler_state : openmmtools.states.SamplerState
Configurational properties of new atoms at the end of the NCMC switching.
"""
if self.verbose: print("Updating chemical state with geometry-ncmc-geometry scheme...")
from perses.tests.utils import compute_potential
logP_chemical_proposal = topology_proposal.logp_proposal
old_thermodynamic_state = self.sampler.thermodynamic_state
new_thermodynamic_state = self._system_to_thermodynamic_state(topology_proposal.new_system)
initial_reduced_potential = feptasks.compute_reduced_potential(old_thermodynamic_state, sampler_state)
logP_initial_nonalchemical = - initial_reduced_potential
new_geometry_sampler_state, logP_geometry_forward = self._geometry_forward(topology_proposal, sampler_state)
#if we aren't doing any switching, then skip running the NCMC engine at all.
if self._switching_nsteps == 0:
ncmc_old_sampler_state = sampler_state
ncmc_new_sampler_state = new_geometry_sampler_state
logP_work = 0.0
logP_initial_hybrid = 0.0
logP_final_hybrid = 0.0
else:
ncmc_old_sampler_state, ncmc_new_sampler_state, logP_work, logP_initial_hybrid, logP_final_hybrid = self._ncmc_hybrid(topology_proposal, sampler_state, new_geometry_sampler_state)
if logP_work > -np.inf and logP_initial_hybrid > -np.inf and logP_final_hybrid > -np.inf:
logP_geometry_reverse = self._geometry_reverse(topology_proposal, ncmc_new_sampler_state, ncmc_old_sampler_state)
logP_to_hybrid = logP_initial_hybrid - logP_initial_nonalchemical
final_reduced_potential = feptasks.compute_reduced_potential(new_thermodynamic_state, ncmc_new_sampler_state)
logP_final_nonalchemical = -final_reduced_potential
logP_from_hybrid = logP_final_nonalchemical - logP_final_hybrid
logP_sams_weight = new_log_weight - old_log_weight
# Compute total log acceptance probability according to Eq. 46
logP_accept = logP_to_hybrid - logP_geometry_forward + logP_work + logP_from_hybrid + logP_geometry_reverse + logP_sams_weight
else:
logP_geometry_reverse = 0.0
logP_final = 0.0
logP_to_hybrid = 0.0
logP_from_hybrid = 0.0
logP_sams_weight = new_log_weight - old_log_weight
logP_accept = logP_to_hybrid - logP_geometry_forward + logP_work + logP_from_hybrid + logP_geometry_reverse + logP_sams_weight
#TODO: mark failed proposals as unproposable
if self.verbose:
print("logP_accept = %+10.4e [logP_to_hybrid = %+10.4e, logP_chemical_proposal = %10.4e, logP_reverse = %+10.4e, -logP_forward = %+10.4e, logP_work = %+10.4e, logP_from_hybrid = %+10.4e, logP_sams_weight = %+10.4e]"
% (logP_accept, logP_to_hybrid, logP_chemical_proposal, logP_geometry_reverse, -logP_geometry_forward, logP_work, logP_from_hybrid, logP_sams_weight))
# Write to storage.
if self.storage:
self.storage.write_quantity('logP_accept', logP_accept, iteration=self.iteration)
# Write components to storage
self.storage.write_quantity('logP_ncmc_work', logP_work, iteration=self.iteration)
self.storage.write_quantity('logP_from_hybrid', logP_from_hybrid, iteration=self.iteration)
self.storage.write_quantity('logP_to_hybrid', logP_to_hybrid, iteration=self.iteration)
self.storage.write_quantity('logP_chemical_proposal', logP_chemical_proposal, iteration=self.iteration)
self.storage.write_quantity('logP_reverse', logP_geometry_reverse, iteration=self.iteration)
self.storage.write_quantity('logP_forward', logP_geometry_forward, iteration=self.iteration)
self.storage.write_quantity('logP_sams_weight', logP_sams_weight, iteration=self.iteration)
# Write some aggregate statistics to storage to make contributions to acceptance probability easier to analyze
self.storage.write_quantity('logP_groups_chemical', logP_chemical_proposal, iteration=self.iteration)
self.storage.write_quantity('logP_groups_geometry', logP_geometry_reverse - logP_geometry_forward, iteration=self.iteration)
return logP_accept, ncmc_new_sampler_state
def _geometry_ncmc_geometry(self, topology_proposal, sampler_state, old_log_weight, new_log_weight):
"""
Use a hybrid NCMC protocol to switch from the old system to new system
Will calculate new positions for the new system first, then give both
sets of positions to the hybrid NCMC integrator, and finally use the
final positions of the old and new systems to calculate the reverse
geometry probability
Parameters
----------
topology_proposal : TopologyProposal
Contains old/new Topology and System objects and atom mappings.
sampler_state : openmmtools.states.SamplerState
Configurational properties of old atoms at the beginning of the NCMC switching.
old_log_weight : float
Chemical state weight from SAMSSampler
new_log_weight : float
Chemical state weight from SAMSSampler
Returns
-------
logP_accept : float
Log of acceptance probability of entire Expanded Ensemble switch (Eq. 25 or 46)
ncmc_new_sampler_state : openmmtools.states.SamplerState
Configurational properties of new atoms at the end of the NCMC switching.
"""
if self.verbose: print("Updating chemical state with geometry-ncmc-geometry scheme...")
from perses.tests.utils import compute_potential
logP_chemical_proposal = topology_proposal.logp_proposal
old_thermodynamic_state = self.sampler.thermodynamic_state
new_thermodynamic_state = self._system_to_thermodynamic_state(topology_proposal.new_system)
initial_reduced_potential = feptasks.compute_reduced_potential(old_thermodynamic_state, sampler_state)
logP_initial_nonalchemical = - initial_reduced_potential
new_geometry_sampler_state, logP_geometry_forward = self._geometry_forward(topology_proposal, sampler_state)
#if we aren't doing any switching, then skip running the NCMC engine at all.
if self._switching_nsteps == 0:
ncmc_old_sampler_state = sampler_state
ncmc_new_sampler_state = new_geometry_sampler_state
logP_work = 0.0
logP_initial_hybrid = 0.0
logP_final_hybrid = 0.0
else:
ncmc_old_sampler_state, ncmc_new_sampler_state, logP_work, logP_initial_hybrid, logP_final_hybrid = self._ncmc_hybrid(topology_proposal, sampler_state, new_geometry_sampler_state)
if logP_work > -np.inf and logP_initial_hybrid > -np.inf and logP_final_hybrid > -np.inf:
logP_geometry_reverse = self._geometry_reverse(topology_proposal, ncmc_new_sampler_state, ncmc_old_sampler_state)
logP_to_hybrid = logP_initial_hybrid - logP_initial_nonalchemical
final_reduced_potential = feptasks.compute_reduced_potential(new_thermodynamic_state, ncmc_new_sampler_state)
logP_final_nonalchemical = -final_reduced_potential
logP_from_hybrid = logP_final_nonalchemical - logP_final_hybrid
logP_sams_weight = new_log_weight - old_log_weight
# Compute total log acceptance probability according to Eq. 46
logP_accept = logP_to_hybrid - logP_geometry_forward + logP_work + logP_from_hybrid + logP_geometry_reverse + logP_sams_weight
else:
logP_geometry_reverse = 0.0
logP_final = 0.0
logP_to_hybrid = 0.0
logP_from_hybrid = 0.0
logP_sams_weight = new_log_weight - old_log_weight
logP_accept = logP_to_hybrid - logP_geometry_forward + logP_work + logP_from_hybrid + logP_geometry_reverse + logP_sams_weight
#TODO: mark failed proposals as unproposable
if self.verbose:
print("logP_accept = %+10.4e [logP_to_hybrid = %+10.4e, logP_chemical_proposal = %10.4e, logP_reverse = %+10.4e, -logP_forward = %+10.4e, logP_work = %+10.4e, logP_from_hybrid = %+10.4e, logP_sams_weight = %+10.4e]"
% (logP_accept, logP_to_hybrid, logP_chemical_proposal, logP_geometry_reverse, -logP_geometry_forward, logP_work, logP_from_hybrid, logP_sams_weight))
# Write to storage.
if self.storage:
self.storage.write_quantity('logP_accept', logP_accept, iteration=self.iteration)
# Write components to storage
self.storage.write_quantity('logP_ncmc_work', logP_work, iteration=self.iteration)
self.storage.write_quantity('logP_from_hybrid', logP_from_hybrid, iteration=self.iteration)
self.storage.write_quantity('logP_to_hybrid', logP_to_hybrid, iteration=self.iteration)
self.storage.write_quantity('logP_chemical_proposal', logP_chemical_proposal, iteration=self.iteration)
self.storage.write_quantity('logP_reverse', logP_geometry_reverse, iteration=self.iteration)
self.storage.write_quantity('logP_forward', logP_geometry_forward, iteration=self.iteration)
self.storage.write_quantity('logP_sams_weight', logP_sams_weight, iteration=self.iteration)
# Write some aggregate statistics to storage to make contributions to acceptance probability easier to analyze
self.storage.write_quantity('logP_groups_chemical', logP_chemical_proposal, iteration=self.iteration)
self.storage.write_quantity('logP_groups_geometry', logP_geometry_reverse - logP_geometry_forward, iteration=self.iteration)
return logP_accept, ncmc_new_sampler_state
