class ExpandedEnsembleSampler(object):
"""
Method of expanded ensembles sampling engine.
The acceptance criteria is given in the reference document. Roughly, the proposal scheme is:
* Draw a proposed chemical state k', and calculate reverse proposal probability
* Conditioned on k' and the current positions x, generate new positions with the GeometryEngine
* With new positions, jump to a hybrid system at lambda=0
* Anneal from lambda=0 to lambda=1, accumulating work
* Jump from the hybrid system at lambda=1 to the k' system, and compute reverse GeometryEngine proposal
* Add weight of chemical states k and k' to acceptance probabilities
Properties
----------
sampler : MCMCSampler
The MCMC sampler used for updating positions.
proposal_engine : ProposalEngine
The ProposalEngine to use for proposing new sampler states and topologies.
system_generator : SystemGenerator
The SystemGenerator to use for creating System objects following proposals.
state : hashable object
The current sampler state. Can be any hashable object.
states : set of hashable object
All known states.
iteration : int
Iterations completed.
naccepted : int
Number of accepted thermodynamic/chemical state changes.
nrejected : int
Number of rejected thermodynamic/chemical state changes.
number_of_state_visits : dict of state_key
Cumulative counts of visited states.
verbose : bool
If True, verbose output is printed.
References
----------
[1] Lyubartsev AP, Martsinovski AA, Shevkunov SV, and Vorontsov-Velyaminov PN. New approach to Monte Carlo calculation of the free energy: Method of expanded ensembles. JCP 96:1776, 1992
http://dx.doi.org/10.1063/1.462133
Examples
--------
>>> # Create a test system
>>> test = testsystems.AlanineDipeptideVacuum()
>>> # Create a SystemGenerator and rebuild the System.
>>> from perses.rjmc.topology_proposal import SystemGenerator
>>> system_generator = SystemGenerator(['amber99sbildn.xml'], forcefield_kwargs={'implicitSolvent' : None, 'constraints' : None }, nonperiodic_forcefield_kwargs={'nonbondedMethod' : app.NoCutoff})
>>> test.system = system_generator.build_system(test.topology)
>>> # Create a sampler state.
>>> sampler_state = SamplerState(system=test.system, positions=test.positions)
>>> # Create a thermodynamic state.
>>> thermodynamic_state = ThermodynamicState(system=test.system, temperature=298.0*unit.kelvin)
>>> # Create an MCMC sampler
>>> mcmc_sampler = MCMCSampler(thermodynamic_state, sampler_state)
>>> # Turn off verbosity
>>> mcmc_sampler.verbose = False
>>> # Create an Expanded Ensemble sampler
>>> from perses.rjmc.topology_proposal import PointMutationEngine
>>> from perses.rjmc.geometry import FFAllAngleGeometryEngine
>>> geometry_engine = FFAllAngleGeometryEngine(metadata={})
>>> allowed_mutations = [[('2','ALA')],[('2','VAL'),('2','LEU')]]
>>> proposal_engine = PointMutationEngine(test.topology, system_generator, max_point_mutants=1, chain_id='1', proposal_metadata=None, allowed_mutations=allowed_mutations)
>>> exen_sampler = ExpandedEnsembleSampler(mcmc_sampler, test.topology, 'ACE-ALA-NME', proposal_engine, geometry_engine)
>>> # Run the sampler
>>> exen_sampler.run()
"""
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)
@property
def state_keys(self):
return self.log_weights.keys()
def get_log_weight(self, state_key):
"""
Get the log weight of the specified state.
Parameters
----------
state_key : hashable object
The state key (e.g. chemical state key) to look up.
Returns
-------
log_weight : float
The log weight of the provided state key.
Notes
-----
This adds the key to the self.log_weights dict.
"""
if state_key not in self.log_weights:
self.log_weights[state_key] = 0.0
return self.log_weights[state_key]
def _system_to_thermodynamic_state(self, system):
"""
Given an OpenMM system object, create a corresponding ThermodynamicState that has the same
temperature and pressure as the current thermodynamic state.
Parameters
----------
system : openmm.System
The OpenMM system for which to create the thermodynamic state
Returns
-------
new_thermodynamic_state : openmmtools.states.ThermodynamicState
The thermodynamic state object representing the given system
"""
return ThermodynamicState(system, temperature=self._temperature, pressure=self._pressure)
def _geometry_forward(self, topology_proposal, old_sampler_state):
"""
Run geometry engine to propose new positions and compute logP
Parameters
----------
topology_proposal : TopologyProposal
Contains old/new Topology and System objects and atom mappings.
old_sampler_state : openmmtools.states.SamplerState
Configurational properties of the old system atoms.
Returns
-------
new_sampler_state : openmmtools.states.SamplerState
Configurational properties of new atoms proposed by geometry engine calculation.
geometry_logp_propose : float
The log probability of the forward-only proposal
"""
if self.verbose: print("Geometry engine proposal...")
# Generate coordinates for new atoms and compute probability ratio of old and new probabilities.
initial_time = time.time()
new_positions, geometry_logp_propose = self.geometry_engine.propose(topology_proposal, old_sampler_state.positions, self.sampler.thermodynamic_state.beta)
if self.verbose: print('proposal took %.3f s' % (time.time() - initial_time))
if self.geometry_pdbfile is not None:
print("Writing proposed geometry...")
from simtk.openmm.app import PDBFile
PDBFile.writeFile(topology_proposal.new_topology, new_positions, file=self.geometry_pdbfile)
self.geometry_pdbfile.flush()
new_sampler_state = SamplerState(new_positions, box_vectors=old_sampler_state.box_vectors)
return new_sampler_state, geometry_logp_propose
def _geometry_reverse(self, topology_proposal, new_sampler_state, old_sampler_state):
"""
Run geometry engine reverse calculation to determine logP
of proposing the old positions based on the new positions
Parameters
----------
topology_proposal : TopologyProposal
Contains old/new Topology and System objects and atom mappings.
new_sampler_state : openmmtools.states.SamplerState
Configurational properties of the new atoms.
old_sampler_state : openmmtools.states.SamplerState
Configurational properties of the old atoms.
Returns
-------
geometry_logp_reverse : float
The log probability of the proposal for the given transformation
"""
if self.verbose: print("Geometry engine logP_reverse calculation...")
initial_time = time.time()
geometry_logp_reverse = self.geometry_engine.logp_reverse(topology_proposal, new_sampler_state.positions, old_sampler_state.positions, self.sampler.thermodynamic_state.beta)
if self.verbose: print('calculation took %.3f s' % (time.time() - initial_time))
return geometry_logp_reverse
def _ncmc_hybrid(self, topology_proposal, old_sampler_state, new_sampler_state):
"""
Run a hybrid NCMC protocol from lambda = 0 to lambda = 1
Parameters
----------
topology_proposal : TopologyProposal
Contains old/new Topology and System objects and atom mappings.
old_sampler_State : openmmtools.states.SamplerState
SamplerState of old system at the beginning of NCMCSwitching
new_sampler_state : openmmtools.states.SamplerState
SamplerState of new system at the beginning of NCMCSwitching
Returns
-------
old_final_sampler_state : openmmtools.states.SamplerState
SamplerState of old system at the end of switching
new_final_sampler_state : openmmtools.states.SamplerState
SamplerState of new system at the end of switching
logP_work : float
The NCMC work contribution to the log acceptance probability (Eq. 44)
logP_energy : float
The contribution of switching to and from the hybrid system to the acceptance probability (Eq. 45)
"""
if self.verbose: print("Performing NCMC switching")
initial_time = time.time()
[ncmc_old_sampler_state, ncmc_new_sampler_state, logP_work, logP_initial_hybrid, logP_final_hybrid] = self.ncmc_engine.integrate(topology_proposal, old_sampler_state, new_sampler_state, iteration=self.iteration)
if self.verbose: print('NCMC took %.3f s' % (time.time() - initial_time))
# Check that positions are not NaN
if new_sampler_state.has_nan():
raise Exception("Positions are NaN after NCMC insert with %d steps" % self._switching_nsteps)
return ncmc_old_sampler_state, ncmc_new_sampler_state, logP_work, logP_initial_hybrid, logP_final_hybrid
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 update_positions(self, n_iterations=1):
"""
Sample new positions.
"""
self.sampler.run(n_iterations=n_iterations)
def update_state(self):
"""
Sample the thermodynamic state.
"""
initial_time = time.time()
# Propose new chemical state.
if self.verbose: print("Proposing new topology...")
[system, positions] = [self.sampler.thermodynamic_state.get_system(remove_thermostat=True), self.sampler.sampler_state.positions]
#omm_topology = topology.to_openmm() #convert to OpenMM topology for proposal engine
self._omm_topology.setPeriodicBoxVectors(self.sampler.sampler_state.box_vectors) #set the box vectors because in OpenMM topology has these...
topology_proposal = self.proposal_engine.propose(system, self._omm_topology)
if self.verbose: print("Proposed transformation: %s => %s" % (topology_proposal.old_chemical_state_key, topology_proposal.new_chemical_state_key))
# Determine state keys
old_state_key = self.state_key
new_state_key = topology_proposal.new_chemical_state_key
# Determine log weight
old_log_weight = self.get_log_weight(old_state_key)
new_log_weight = self.get_log_weight(new_state_key)
logp_accept, ncmc_new_sampler_state = self._geometry_ncmc_geometry(topology_proposal, self.sampler.sampler_state, old_log_weight, new_log_weight)
# Accept or reject.
if np.isnan(logp_accept):
accept = False
print('logp_accept = NaN')
else:
accept = ((logp_accept>=0.0) or (np.random.uniform() < np.exp(logp_accept)))
if self.accept_everything:
print('accept_everything option is turned on; accepting')
accept = True
if accept:
self.sampler.thermodynamic_state.set_system(topology_proposal.new_system, fix_state=True)
self.sampler.sampler_state.system = topology_proposal.new_system
self.topology = md.Topology.from_openmm(topology_proposal.new_topology)
self.sampler.sampler_state = ncmc_new_sampler_state
self.sampler.topology = self.topology
self.state_key = topology_proposal.new_chemical_state_key
self.naccepted += 1
if self.verbose: print(" accepted")
else:
self.nrejected += 1
if self.verbose: print(" rejected")
if self.storage:
self.storage.write_configuration('positions', self.sampler.sampler_state.positions, self.topology, iteration=self.iteration)
self.storage.write_object('state_key', self.state_key, iteration=self.iteration)
self.storage.write_object('proposed_state_key', topology_proposal.new_chemical_state_key, iteration=self.iteration)
self.storage.write_quantity('naccepted', self.naccepted, iteration=self.iteration)
self.storage.write_quantity('nrejected', self.nrejected, iteration=self.iteration)
self.storage.write_quantity('logp_accept', logp_accept, iteration=self.iteration)
self.storage.write_quantity('logp_topology_proposal', topology_proposal.logp_proposal, iteration=self.iteration)
# Update statistics.
self.update_statistics()
def update(self):
"""
Update the sampler with one step of sampling.
"""
if self.verbose:
print("-" * 80)
print("Expanded Ensemble sampler iteration %8d" % self.iteration)
self.update_positions(n_iterations=self._n_iterations_per_update)
self.update_state()
self.iteration += 1
if self.verbose:
print("-" * 80)
if self.pdbfile is not None:
print("Writing frame...")
from simtk.openmm.app import PDBFile
PDBFile.writeModel(self.topology.to_openmm(), self.sampler.sampler_state.positions, self.pdbfile, self.iteration)
self.pdbfile.flush()
if self.storage:
self.storage.sync()
def run(self, niterations=1):
"""
Run the sampler for the specified number of iterations
Parameters
----------
niterations : int, optional, default=1
Number of iterations to run the sampler for.
"""
for iteration in range(niterations):
self.update()
def update_statistics(self):
"""
Update sampler statistics.
"""
if self.state_key not in self.number_of_state_visits:
self.number_of_state_visits[self.state_key] = 0
self.number_of_state_visits[self.state_key] += 1
class NCMCEngine(object):
"""
NCMC switching engine
Examples
--------
Create a transformation for an alanine dipeptide test system where the N-methyl group is eliminated.
>>> from openmmtools import testsystems
>>> testsystem = testsystems.AlanineDipeptideVacuum()
>>> from perses.rjmc.topology_proposal import TopologyProposal
>>> new_to_old_atom_map = { index : index for index in range(testsystem.system.getNumParticles()) if (index > 3) } # all atoms but N-methyl
>>> topology_proposal = TopologyProposal(old_system=testsystem.system, old_topology=testsystem.topology, old_chemical_state_key='AA', new_chemical_state_key='AA', new_system=testsystem.system, new_topology=testsystem.topology, logp_proposal=0.0, new_to_old_atom_map=new_to_old_atom_map, metadata=dict())
>>> ncmc_engine = NCMCEngine(temperature=300.0*unit.kelvin, functions=default_functions, nsteps=50, timestep=1.0*unit.femtoseconds)
>>> positions = testsystem.positions
>>> [positions, logP_delete, potential_delete] = ncmc_engine.integrate(topology_proposal, positions, direction='delete')
>>> [positions, logP_insert, potential_insert] = ncmc_engine.integrate(topology_proposal, positions, direction='insert')
"""
def __init__(self, temperature=default_temperature, functions=None, nsteps=default_nsteps,
steps_per_propagation=default_steps_per_propagation, timestep=default_timestep,
constraint_tolerance=None, platform=None, write_ncmc_interval=1, measure_shadow_work=False,
integrator_splitting='V R O H R V', storage=None, verbose=False, LRUCapacity=10, pressure=None, bond_softening_constant=1.0, angle_softening_constant=1.0):
"""
This is the base class for NCMC switching between two different systems.
Arguments
---------
temperature : simtk.unit.Quantity with units compatible with kelvin
The temperature at which switching is to be run
functions : dict of str:str, optional, default=default_functions
functions[parameter] is the function (parameterized by 't' which switched from 0 to 1) that
controls how alchemical context parameter 'parameter' is switched
nsteps : int, optional, default=1
The number of steps to use for switching.
steps_per_propagation : int, optional, default=1
The number of intermediate propagation steps taken at each switching step
timestep : simtk.unit.Quantity with units compatible with femtoseconds, optional, default=1*femtosecond
The timestep to use for integration of switching velocity Verlet steps.
constraint_tolerance : float, optional, default=None
If not None, this relative constraint tolerance is used for position and velocity constraints.
platform : simtk.openmm.Platform, optional, default=None
If specified, the platform to use for OpenMM simulations.
write_ncmc_interval : int, optional, default=None
If a positive integer is specified, a snapshot frame will be written to storage with the specified interval on NCMC switching.
'storage' must also be specified.
measure_shadow_work : bool, optional, default False
Whether to measure shadow work
integrator_splitting : str, optional, default='V R O H R V'
NCMC internal integrator splitting based on OpenMMTools Langevin splittings
storage : NetCDFStorageView, optional, default=None
If specified, write data using this class.
verbose : bool, optional, default=False
If True, print debug information.
LRUCapacity : int, default 10
Capacity of LRU cache for hybrid systems
pressure : float, default None
The pressure to use for the simulation. If None, no barostat
"""
# Handle some defaults.
if functions == None:
functions = python_hybrid_functions
if nsteps == None:
nsteps = default_nsteps
if timestep == None:
timestep = default_timestep
if temperature == None:
temperature = default_temperature
self._temperature = temperature
self._functions = copy.deepcopy(functions)
self._nsteps = nsteps
self._timestep = timestep
self._constraint_tolerance = constraint_tolerance
self._platform = platform
self._integrator_splitting = integrator_splitting
self._steps_per_propagation = steps_per_propagation
self._verbose = verbose
self._pressure = pressure
self._bond_softening_constant = bond_softening_constant
self._angle_softening_constant = angle_softening_constant
self._disable_barostat = False
self._hybrid_cache = LRUCache(capacity=LRUCapacity)
self._measure_shadow_work = measure_shadow_work
self._nattempted = 0
self._storage = None
if storage is not None:
self._storage = NetCDFStorageView(storage, modname=self.__class__.__name__)
self._save_configuration = True
else:
self._save_configuration = False
if write_ncmc_interval is not None:
self._write_ncmc_interval = write_ncmc_interval
else:
self._write_ncmc_interval = 1
self._work_save_interval = write_ncmc_interval
@property
def beta(self):
kT = kB * self._temperature
beta = 1.0 / kT
return beta
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 _topology_proposal_to_thermodynamic_states(self, topology_proposal):
"""
Convert a topology proposal to thermodynamic states for the end systems. This will be used to compute the
"logP_energy" quantity.
Arguments
---------
topology_proposal : perses.rjmc.TopologyProposal
topology proposal for whose endpoint systems we want ThermodynamicStates
Returns
-------
old_thermodynamic_state : openmmtools.states.ThermodynamicState
The old system (nonalchemical) thermodynamic state
new_thermodynamic_state : openmmtools.states.ThermodynamicState
The new system (nonalchemical) thermodynamic state
"""
systems = [topology_proposal.old_system, topology_proposal.new_system]
thermostates = []
for system in systems:
thermodynamic_state = ThermodynamicState(system, temperature=self._temperature, pressure=self._pressure)
thermostates.append(thermodynamic_state)
return thermostates[0], thermostates[1]
def make_alchemical_system(self, topology_proposal, current_positions, new_positions):
"""
Generate an alchemically-modified system at the correct atoms
based on the topology proposal. This method generates a hybrid system using the new
HybridTopologyFactory. It memoizes so that calling multiple times (within a recent time period)
will immediately return a cached object.
Arguments
---------
topology_proposal : perses.rjmc.TopologyProposal
Unmodified real system corresponding to appropriate leg of transformation.
current_positions : np.ndarray of float
Positions of "old" system
new_positions : np.ndarray of float
Positions of "new" system atoms
Returns
-------
hybrid_factory : perses.annihilation.new_relative.HybridTopologyFactory
a factory object containing the hybrid system
"""
try:
hybrid_factory = self._hybrid_cache[topology_proposal]
#If we've retrieved the factory from the cache, update it to include the relevant positions
hybrid_factory._old_positions = current_positions
hybrid_factory._new_positions = new_positions
hybrid_factory._compute_hybrid_positions()
except KeyError:
try:
hybrid_factory = HybridTopologyFactory(topology_proposal, current_positions, new_positions, bond_softening_constant=self._bond_softening_constant, angle_softening_constant=self._angle_softening_constant)
self._hybrid_cache[topology_proposal] = hybrid_factory
except:
hybrid_factory = None
return hybrid_factory
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]
class NCMCEngine(object):
"""
NCMC switching engine
Examples
--------
Create a transformation for an alanine dipeptide test system where the N-methyl group is eliminated.
>>> from openmmtools import testsystems
>>> testsystem = testsystems.AlanineDipeptideVacuum()
>>> from perses.rjmc.topology_proposal import TopologyProposal
>>> new_to_old_atom_map = { index : index for index in range(testsystem.system.getNumParticles()) if (index > 3) } # all atoms but N-methyl
>>> topology_proposal = TopologyProposal(old_system=testsystem.system, old_topology=testsystem.topology, old_chemical_state_key='AA', new_chemical_state_key='AA', new_system=testsystem.system, new_topology=testsystem.topology, logp_proposal=0.0, new_to_old_atom_map=new_to_old_atom_map, metadata=dict())
>>> ncmc_engine = NCMCEngine(temperature=300.0*unit.kelvin, functions=default_functions, nsteps=50, timestep=1.0*unit.femtoseconds)
>>> positions = testsystem.positions
>>> [positions, logP_delete, potential_delete] = ncmc_engine.integrate(topology_proposal, positions, direction='delete')
>>> [positions, logP_insert, potential_insert] = ncmc_engine.integrate(topology_proposal, positions, direction='insert')
"""
def __init__(self,
temperature=default_temperature,
functions=None,
nsteps=default_nsteps,
steps_per_propagation=default_steps_per_propagation,
timestep=default_timestep,
constraint_tolerance=None,
platform=None,
write_ncmc_interval=1,
measure_shadow_work=False,
integrator_splitting='V R O H R V',
storage=None,
verbose=False,
LRUCapacity=10,
pressure=None,
bond_softening_constant=1.0,
angle_softening_constant=1.0):
"""
This is the base class for NCMC switching between two different systems.
Parameters
----------
temperature : simtk.unit.Quantity with units compatible with kelvin
The temperature at which switching is to be run
functions : dict of str:str, optional, default=default_functions
functions[parameter] is the function (parameterized by 't' which switched from 0 to 1) that
controls how alchemical context parameter 'parameter' is switched
nsteps : int, optional, default=1
The number of steps to use for switching.
steps_per_propagation : int, optional, default=1
The number of intermediate propagation steps taken at each switching step
timestep : simtk.unit.Quantity with units compatible with femtoseconds, optional, default=1*femtosecond
The timestep to use for integration of switching velocity Verlet steps.
constraint_tolerance : float, optional, default=None
If not None, this relative constraint tolerance is used for position and velocity constraints.
platform : simtk.openmm.Platform, optional, default=None
If specified, the platform to use for OpenMM simulations.
write_ncmc_interval : int, optional, default=None
If a positive integer is specified, a snapshot frame will be written to storage with the specified interval on NCMC switching.
'storage' must also be specified.
measure_shadow_work : bool, optional, default False
Whether to measure shadow work
integrator_splitting : str, optional, default='V R O H R V'
NCMC internal integrator splitting based on OpenMMTools Langevin splittings
storage : NetCDFStorageView, optional, default=None
If specified, write data using this class.
verbose : bool, optional, default=False
If True, print debug information.
LRUCapacity : int, default 10
Capacity of LRU cache for hybrid systems
pressure : float, default None
The pressure to use for the simulation. If None, no barostat
"""
# Handle some defaults.
if functions == None:
functions = LambdaProtocol.default_functions
if nsteps == None:
nsteps = default_nsteps
if timestep == None:
timestep = default_timestep
if temperature == None:
temperature = default_temperature
self._temperature = temperature
self._functions = copy.deepcopy(functions)
self._nsteps = nsteps
self._timestep = timestep
self._constraint_tolerance = constraint_tolerance
self._platform = platform
self._integrator_splitting = integrator_splitting
self._steps_per_propagation = steps_per_propagation
self._verbose = verbose
self._pressure = pressure
self._bond_softening_constant = bond_softening_constant
self._angle_softening_constant = angle_softening_constant
self._disable_barostat = False
self._hybrid_cache = LRUCache(capacity=LRUCapacity)
self._measure_shadow_work = measure_shadow_work
self._nattempted = 0
self._storage = None
if storage is not None:
self._storage = NetCDFStorageView(storage,
modname=self.__class__.__name__)
self._save_configuration = True
else:
self._save_configuration = False
if write_ncmc_interval is not None:
self._write_ncmc_interval = write_ncmc_interval
else:
self._write_ncmc_interval = 1
self._work_save_interval = write_ncmc_interval
@property
def beta(self):
kT = kB * self._temperature
beta = 1.0 / kT
return beta
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 _topology_proposal_to_thermodynamic_states(self, topology_proposal):
"""
Convert a topology proposal to thermodynamic states for the end systems. This will be used to compute the
"logP_energy" quantity.
Parameters
----------
topology_proposal : perses.rjmc.TopologyProposal
topology proposal for whose endpoint systems we want ThermodynamicStates
Returns
-------
old_thermodynamic_state : openmmtools.states.ThermodynamicState
The old system (nonalchemical) thermodynamic state
new_thermodynamic_state : openmmtools.states.ThermodynamicState
The new system (nonalchemical) thermodynamic state
"""
systems = [topology_proposal.old_system, topology_proposal.new_system]
thermostates = []
for system in systems:
thermodynamic_state = ThermodynamicState(
system, temperature=self._temperature, pressure=self._pressure)
thermostates.append(thermodynamic_state)
return thermostates[0], thermostates[1]
def make_alchemical_system(self, topology_proposal, current_positions,
new_positions):
"""
Generate an alchemically-modified system at the correct atoms
based on the topology proposal. This method generates a hybrid system using the new
HybridTopologyFactory. It memoizes so that calling multiple times (within a recent time period)
will immediately return a cached object.
Parameters
----------
topology_proposal : perses.rjmc.TopologyProposal
Unmodified real system corresponding to appropriate leg of transformation.
current_positions : np.ndarray of float
Positions of "old" system
new_positions : np.ndarray of float
Positions of "new" system atoms
Returns
-------
hybrid_factory : perses.annihilation.relative.HybridTopologyFactory
a factory object containing the hybrid system
"""
try:
hybrid_factory = self._hybrid_cache[topology_proposal]
#If we've retrieved the factory from the cache, update it to include the relevant positions
hybrid_factory._old_positions = current_positions
hybrid_factory._new_positions = new_positions
hybrid_factory._compute_hybrid_positions()
except KeyError:
try:
hybrid_factory = HybridTopologyFactory(
topology_proposal,
current_positions,
new_positions,
bond_softening_constant=self._bond_softening_constant,
angle_softening_constant=self._angle_softening_constant)
self._hybrid_cache[topology_proposal] = hybrid_factory
except:
hybrid_factory = None
return hybrid_factory
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
]
class ExpandedEnsembleSampler(object):
"""
Method of expanded ensembles sampling engine.
The acceptance criteria is given in the reference document. Roughly, the proposal scheme is:
* Draw a proposed chemical state k', and calculate reverse proposal probability
* Conditioned on k' and the current positions x, generate new positions with the GeometryEngine
* With new positions, jump to a hybrid system at lambda=0
* Anneal from lambda=0 to lambda=1, accumulating work
* Jump from the hybrid system at lambda=1 to the k' system, and compute reverse GeometryEngine proposal
* Add weight of chemical states k and k' to acceptance probabilities
Properties
----------
sampler : MCMCSampler
The MCMC sampler used for updating positions.
proposal_engine : ProposalEngine
The ProposalEngine to use for proposing new sampler states and topologies.
system_generator : SystemGenerator
The SystemGenerator to use for creating System objects following proposals.
state : hashable object
The current sampler state. Can be any hashable object.
states : set of hashable object
All known states.
iteration : int
Iterations completed.
naccepted : int
Number of accepted thermodynamic/chemical state changes.
nrejected : int
Number of rejected thermodynamic/chemical state changes.
number_of_state_visits : dict of state_key
Cumulative counts of visited states.
verbose : bool
If True, verbose output is printed.
References
----------
[1] Lyubartsev AP, Martsinovski AA, Shevkunov SV, and Vorontsov-Velyaminov PN. New approach to Monte Carlo calculation of the free energy: Method of expanded ensembles. JCP 96:1776, 1992
http://dx.doi.org/10.1063/1.462133
Examples
--------
>>> # Create a test system
>>> test = testsystems.AlanineDipeptideVacuum()
>>> # Create a SystemGenerator and rebuild the System.
>>> from perses.rjmc.topology_proposal import SystemGenerator
>>> system_generator = SystemGenerator(['amber99sbildn.xml'], forcefield_kwargs={ 'nonbondedMethod' : app.NoCutoff, 'implicitSolvent' : None, 'constraints' : None })
>>> test.system = system_generator.build_system(test.topology)
>>> # Create a sampler state.
>>> sampler_state = SamplerState(system=test.system, positions=test.positions)
>>> # Create a thermodynamic state.
>>> thermodynamic_state = ThermodynamicState(system=test.system, temperature=298.0*unit.kelvin)
>>> # Create an MCMC sampler
>>> mcmc_sampler = MCMCSampler(thermodynamic_state, sampler_state)
>>> # Turn off verbosity
>>> mcmc_sampler.verbose = False
>>> # Create an Expanded Ensemble sampler
>>> from perses.rjmc.topology_proposal import PointMutationEngine
>>> from perses.rjmc.geometry import FFAllAngleGeometryEngine
>>> geometry_engine = FFAllAngleGeometryEngine(metadata={})
>>> allowed_mutations = [[('2','ALA')],[('2','VAL'),('2','LEU')]]
>>> proposal_engine = PointMutationEngine(test.topology, system_generator, max_point_mutants=1, chain_id='1', proposal_metadata=None, allowed_mutations=allowed_mutations)
>>> exen_sampler = ExpandedEnsembleSampler(mcmc_sampler, test.topology, 'ACE-ALA-NME', proposal_engine, geometry_engine)
>>> # Run the sampler
>>> exen_sampler.run()
"""
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.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)
@property
def state_keys(self):
return self.log_weights.keys()
def get_log_weight(self, state_key):
"""
Get the log weight of the specified state.
Parameters
----------
state_key : hashable object
The state key (e.g. chemical state key) to look up.
Returns
-------
log_weight : float
The log weight of the provided state key.
Note
----
This adds the key to the self.log_weights dict.
"""
if state_key not in self.log_weights:
self.log_weights[state_key] = 0.0
return self.log_weights[state_key]
def _system_to_thermodynamic_state(self, system):
"""
Given an OpenMM system object, create a corresponding ThermodynamicState that has the same
temperature and pressure as the current thermodynamic state.
Arguments
---------
system : openmm.System
The OpenMM system for which to create the thermodynamic state
Returns
-------
new_thermodynamic_state : openmmtools.states.ThermodynamicState
The thermodynamic state object representing the given system
"""
return ThermodynamicState(system, temperature=self._temperature, pressure=self._pressure)
def _geometry_forward(self, topology_proposal, old_sampler_state):
"""
Run geometry engine to propose new positions and compute logP
Parameters
----------
topology_proposal : TopologyProposal
Contains old/new Topology and System objects and atom mappings.
old_sampler_state : openmmtools.states.SamplerState
Configurational properties of the old system atoms.
Returns
-------
new_sampler_state : openmmtools.states.SamplerState
Configurational properties of new atoms proposed by geometry engine calculation.
geometry_logp_propose : float
The log probability of the forward-only proposal
"""
if self.verbose: print("Geometry engine proposal...")
# Generate coordinates for new atoms and compute probability ratio of old and new probabilities.
initial_time = time.time()
new_positions, geometry_logp_propose = self.geometry_engine.propose(topology_proposal, old_sampler_state.positions, self.sampler.thermodynamic_state.beta)
if self.verbose: print('proposal took %.3f s' % (time.time() - initial_time))
if self.geometry_pdbfile is not None:
print("Writing proposed geometry...")
from simtk.openmm.app import PDBFile
PDBFile.writeFile(topology_proposal.new_topology, new_positions, file=self.geometry_pdbfile)
self.geometry_pdbfile.flush()
new_sampler_state = SamplerState(new_positions, box_vectors=old_sampler_state.box_vectors)
return new_sampler_state, geometry_logp_propose
def _geometry_reverse(self, topology_proposal, new_sampler_state, old_sampler_state):
"""
Run geometry engine reverse calculation to determine logP
of proposing the old positions based on the new positions
Parameters
----------
topology_proposal : TopologyProposal
Contains old/new Topology and System objects and atom mappings.
new_sampler_state : openmmtools.states.SamplerState
Configurational properties of the new atoms.
old_sampler_state : openmmtools.states.SamplerState
Configurational properties of the old atoms.
Returns
-------
geometry_logp_reverse : float
The log probability of the proposal for the given transformation
"""
if self.verbose: print("Geometry engine logP_reverse calculation...")
initial_time = time.time()
geometry_logp_reverse = self.geometry_engine.logp_reverse(topology_proposal, new_sampler_state.positions, old_sampler_state.positions, self.sampler.thermodynamic_state.beta)
if self.verbose: print('calculation took %.3f s' % (time.time() - initial_time))
return geometry_logp_reverse
def _ncmc_hybrid(self, topology_proposal, old_sampler_state, new_sampler_state):
"""
Run a hybrid NCMC protocol from lambda = 0 to lambda = 1
Parameters
----------
topology_proposal : TopologyProposal
Contains old/new Topology and System objects and atom mappings.
old_sampler_State : openmmtools.states.SamplerState
SamplerState of old system at the beginning of NCMCSwitching
new_sampler_state : openmmtools.states.SamplerState
SamplerState of new system at the beginning of NCMCSwitching
Returns
-------
old_final_sampler_state : openmmtools.states.SamplerState
SamplerState of old system at the end of switching
new_final_sampler_state : openmmtools.states.SamplerState
SamplerState of new system at the end of switching
logP_work : float
The NCMC work contribution to the log acceptance probability (Eq. 44)
logP_energy : float
The contribution of switching to and from the hybrid system to the acceptance probability (Eq. 45)
"""
if self.verbose: print("Performing NCMC switching")
initial_time = time.time()
[ncmc_old_sampler_state, ncmc_new_sampler_state, logP_work, logP_initial_hybrid, logP_final_hybrid] = self.ncmc_engine.integrate(topology_proposal, old_sampler_state, new_sampler_state, iteration=self.iteration)
if self.verbose: print('NCMC took %.3f s' % (time.time() - initial_time))
# Check that positions are not NaN
if new_sampler_state.has_nan():
raise Exception("Positions are NaN after NCMC insert with %d steps" % self._switching_nsteps)
return ncmc_old_sampler_state, ncmc_new_sampler_state, logP_work, logP_initial_hybrid, logP_final_hybrid
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 update_positions(self, n_iterations=1):
"""
Sample new positions.
"""
self.sampler.run(n_iterations=n_iterations)
def update_state(self):
"""
Sample the thermodynamic state.
"""
initial_time = time.time()
# Propose new chemical state.
if self.verbose: print("Proposing new topology...")
[system, topology, positions] = [self.sampler.thermodynamic_state.get_system(remove_thermostat=True), self.topology, self.sampler.sampler_state.positions]
omm_topology = topology.to_openmm() #convert to OpenMM topology for proposal engine
omm_topology.setPeriodicBoxVectors(self.sampler.sampler_state.box_vectors) #set the box vectors because in OpenMM topology has these...
topology_proposal = self.proposal_engine.propose(system, omm_topology)
if self.verbose: print("Proposed transformation: %s => %s" % (topology_proposal.old_chemical_state_key, topology_proposal.new_chemical_state_key))
# Determine state keys
old_state_key = self.state_key
new_state_key = topology_proposal.new_chemical_state_key
# Determine log weight
old_log_weight = self.get_log_weight(old_state_key)
new_log_weight = self.get_log_weight(new_state_key)
logp_accept, ncmc_new_sampler_state = self._geometry_ncmc_geometry(topology_proposal, self.sampler.sampler_state, old_log_weight, new_log_weight)
# Accept or reject.
if np.isnan(logp_accept):
accept = False
print('logp_accept = NaN')
else:
accept = ((logp_accept>=0.0) or (np.random.uniform() < np.exp(logp_accept)))
if self.accept_everything:
print('accept_everything option is turned on; accepting')
accept = True
if accept:
self.sampler.thermodynamic_state.set_system(topology_proposal.new_system, fix_state=True)
self.sampler.sampler_state.system = topology_proposal.new_system
self.topology = md.Topology.from_openmm(topology_proposal.new_topology)
self.sampler.sampler_state = ncmc_new_sampler_state
self.sampler.topology = self.topology
self.state_key = topology_proposal.new_chemical_state_key
self.naccepted += 1
if self.verbose: print(" accepted")
else:
self.nrejected += 1
if self.verbose: print(" rejected")
if self.storage:
self.storage.write_configuration('positions', self.sampler.sampler_state.positions, self.topology, iteration=self.iteration)
self.storage.write_object('state_key', self.state_key, iteration=self.iteration)
self.storage.write_object('proposed_state_key', topology_proposal.new_chemical_state_key, iteration=self.iteration)
self.storage.write_quantity('naccepted', self.naccepted, iteration=self.iteration)
self.storage.write_quantity('nrejected', self.nrejected, iteration=self.iteration)
self.storage.write_quantity('logp_accept', logp_accept, iteration=self.iteration)
self.storage.write_quantity('logp_topology_proposal', topology_proposal.logp_proposal, iteration=self.iteration)
# Update statistics.
self.update_statistics()
def update(self):
"""
Update the sampler with one step of sampling.
"""
if self.verbose:
print("-" * 80)
print("Expanded Ensemble sampler iteration %8d" % self.iteration)
self.update_positions(n_iterations=self._n_iterations_per_update)
self.update_state()
self.iteration += 1
if self.verbose:
print("-" * 80)
if self.pdbfile is not None:
print("Writing frame...")
from simtk.openmm.app import PDBFile
PDBFile.writeModel(self.topology.to_openmm(), self.sampler.sampler_state.positions, self.pdbfile, self.iteration)
self.pdbfile.flush()
if self.storage:
self.storage.sync()
def run(self, niterations=1):
"""
Run the sampler for the specified number of iterations
Parameters
----------
niterations : int, optional, default=1
Number of iterations to run the sampler for.
"""
for iteration in range(niterations):
self.update()
def update_statistics(self):
"""
Update sampler statistics.
"""
if self.state_key not in self.number_of_state_visits:
self.number_of_state_visits[self.state_key] = 0
self.number_of_state_visits[self.state_key] += 1
