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 __init__(self, sampler, topology, state_key, proposal_engine, geometry_engine, log_weights=None, options=None, platform=None, envname=None, storage=None, ncmc_write_interval=1):
"""
Create an expanded ensemble sampler.
p(x,k) \propto \exp[-u_k(x) + g_k]
where g_k is the log weight.
Parameters
----------
sampler : MCMCSampler
MCMCSampler initialized with current SamplerState
topology : simtk.openmm.app.Topology
Current topology
state : hashable object
Current chemical state
proposal_engine : ProposalEngine
ProposalEngine to use for proposing new chemical states
geometry_engine : GeometryEngine
GeometryEngine to use for dimension matching
log_weights : dict of object : float
Log weights to use for expanded ensemble biases.
options : dict, optional, default=dict()
Options for initializing switching scheme, such as 'timestep', 'nsteps', 'functions' for NCMC
platform : simtk.openmm.Platform, optional, default=None
Platform to use for NCMC switching. If `None`, default (fastest) platform is used.
storage : NetCDFStorageView, optional, default=None
If specified, use this storage layer.
ncmc_write_interval : int, default 1
How frequently to write out NCMC protocol steps.
"""
# Keep copies of initializing arguments.
# TODO: Make deep copies?
self.sampler = sampler
self._pressure = sampler.thermodynamic_state.pressure
self._temperature = sampler.thermodynamic_state.temperature
self._omm_topology = topology
self.topology = md.Topology.from_openmm(topology)
self.state_key = state_key
self.proposal_engine = proposal_engine
self.log_weights = log_weights
if self.log_weights is None: self.log_weights = dict()
self.storage = None
if storage is not None:
self.storage = NetCDFStorageView(storage, modname=self.__class__.__name__)
# Initialize
self.iteration = 0
option_names = ['timestep', 'nsteps', 'functions', 'nsteps_mcmc', 'splitting']
if options is None:
options = dict()
for option_name in option_names:
if option_name not in options:
options[option_name] = None
if options['splitting']:
self._ncmc_splitting = options['splitting']
else:
self._ncmc_splitting = "V R O H R V"
if options['nsteps']:
self._switching_nsteps = options['nsteps']
self.ncmc_engine = NCMCEngine(temperature=self.sampler.thermodynamic_state.temperature,
timestep=options['timestep'], nsteps=options['nsteps'],
functions=options['functions'], integrator_splitting=self._ncmc_splitting,
platform=platform, storage=self.storage,
write_ncmc_interval=ncmc_write_interval)
else:
self._switching_nsteps = 0
if options['nsteps_mcmc']:
self._n_iterations_per_update = options['nsteps_mcmc']
else:
self._n_iterations_per_update = 100
self.geometry_engine = geometry_engine
self.naccepted = 0
self.nrejected = 0
self.number_of_state_visits = dict()
self.verbose = False
self.pdbfile = None # if not None, write PDB file
self.geometry_pdbfile = None # if not None, write PDB file of geometry proposals
self.accept_everything = False # if True, will accept anything that doesn't lead to NaNs
self.logPs = list()
self.sampler.minimize(max_iterations=40)
def check_alchemical_null_elimination(topology_proposal,
positions,
ncmc_nsteps=50,
NSIGMA_MAX=6.0,
geometry=False):
"""
Test alchemical elimination engine on null transformations, where some atoms are deleted and then reinserted in a cycle.
Parameters
----------
topology_proposal : TopologyProposal
The topology proposal to test.
This must be a null transformation, where topology_proposal.old_system == topology_proposal.new_system
ncmc_steps : int, optional, default=50
Number of NCMC switching steps, or 0 for instantaneous switching.
NSIGMA_MAX : float, optional, default=6.0
Number of standard errors away from analytical solution tolerated before Exception is thrown
geometry : bool, optional, default=None
If True, will also use geometry engine in the middle of the null transformation.
"""
# Initialize engine
from perses.annihilation.ncmc_switching import NCMCEngine
ncmc_engine = NCMCEngine(temperature=temperature, nsteps=ncmc_nsteps)
# Make sure that old system and new system are identical.
if not (topology_proposal.old_system == topology_proposal.new_system):
raise Exception(
"topology_proposal must be a null transformation for this test (old_system == new_system)"
)
for (k, v) in topology_proposal.new_to_old_atom_map.items():
if k != v:
raise Exception(
"topology_proposal must be a null transformation for this test (retailed atoms must map onto themselves)"
)
nequil = 5 # number of equilibration iterations
niterations = 50 # number of round-trip switching trials
logP_insert_n = np.zeros([niterations], np.float64)
logP_delete_n = np.zeros([niterations], np.float64)
logP_switch_n = np.zeros([niterations], np.float64)
for iteration in range(nequil):
[positions, velocities] = simulate(topology_proposal.old_system,
positions)
for iteration in range(niterations):
# Equilibrate
[positions, velocities] = simulate(topology_proposal.old_system,
positions)
# Check that positions are not NaN
if (np.any(np.isnan(positions / unit.angstroms))):
raise Exception("Positions became NaN during equilibration")
# Delete atoms
[positions, logP_delete,
potential_delete] = ncmc_engine.integrate(topology_proposal,
positions,
direction='delete')
# Check that positions are not NaN
if (np.any(np.isnan(positions / unit.angstroms))):
raise Exception("Positions became NaN on NCMC deletion")
# Insert atoms
[positions, logP_insert,
potential_insert] = ncmc_engine.integrate(topology_proposal,
positions,
direction='insert')
# Check that positions are not NaN
if (np.any(np.isnan(positions / unit.angstroms))):
raise Exception("Positions became NaN on NCMC insertion")
# Compute probability of switching geometries.
logP_switch = -(potential_insert - potential_delete)
# Compute total probability
logP_delete_n[iteration] = logP_delete
logP_insert_n[iteration] = logP_insert
logP_switch_n[iteration] = logP_switch
#print("Iteration %5d : delete %16.8f kT | insert %16.8f kT | geometry switch %16.8f" % (iteration, logP_delete, logP_insert, logP_switch))
# Check free energy difference is withing NSIGMA_MAX standard errors of zero.
logP_n = logP_delete_n + logP_insert_n + logP_switch_n
work_n = -logP_n
from pymbar import EXP
[df, ddf] = EXP(work_n)
#print("df = %12.6f +- %12.5f kT" % (df, ddf))
if (abs(df) > NSIGMA_MAX * ddf):
msg = 'Delta F (%d steps switching) = %f +- %f kT; should be within %f sigma of 0\n' % (
ncmc_nsteps, df, ddf, NSIGMA_MAX)
msg += 'delete logP:\n'
msg += str(logP_delete_n) + '\n'
msg += 'insert logP:\n'
msg += str(logP_insert_n) + '\n'
msg += 'logP:\n'
msg += str(logP_n) + '\n'
raise Exception(msg)
