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 check_alchemical_elimination(ncmc_nsteps=50): """ Test alchemical elimination engine on alanine dipeptide null transformation. """ NSIGMA_MAX = 6.0 # number of standard errors away from analytical solution tolerated before Exception is thrown # Create an alanine dipeptide null transformation, where N-methyl group is deleted and then inserted. 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_positions=testsystem.positions, new_system=testsystem.system, new_topology=testsystem.topology, logp_proposal=0.0, new_to_old_atom_map=new_to_old_atom_map, metadata=dict()) # Initialize engine from perses.annihilation.ncmc_switching import NCMCEngine ncmc_engine = NCMCEngine(nsteps=ncmc_nsteps) niterations = 20 # number of round-trip switching trials positions = testsystem.positions logP_insert_n = np.zeros([niterations], np.float64) logP_delete_n = np.zeros([niterations], np.float64) for iteration in range(niterations): # Equilibrate positions = simulate(testsystem.system, positions) # Delete atoms [positions, logP_delete] = ncmc_engine.integrate(topology_proposal, positions, direction='delete') # Insert atoms [positions, logP_insert] = ncmc_engine.integrate(topology_proposal, positions, direction='insert') # Compute total probability logP_delete_n[iteration] = logP_delete logP_insert_n[iteration] = logP_insert # Check free energy difference is withing NSIGMA_MAX standard errors of zero. logP_n = logP_delete_n + logP_insert_n from pymbar import EXP [df, ddf] = EXP(logP_n) if (abs(df) > NSIGMA_MAX * ddf): msg = 'Delta F (%d steps switching) = %f +- %f kT; should be within %f sigma of 0' % (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)
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 __init__(self, sampler, topology, state_key, proposal_engine, log_weights=None, scheme='ncmc-geometry-ncmc', options=dict(), platform=None): """ 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 log_weights : dict of object : float Log weights to use for expanded ensemble biases. scheme : str, optional, default='ncmc-geometry-ncmc' Update scheme. One of ['ncmc-geometry-ncmc', 'geometry-ncmc-geometry', 'geometry-ncmc'] 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. """ # Keep copies of initializing arguments. # TODO: Make deep copies? self.sampler = sampler self.topology = topology self.state_key = state_key self.proposal_engine = proposal_engine self.log_weights = log_weights self.scheme = scheme if self.log_weights is None: self.log_weights = dict() # Initialize self.iteration = 0 option_names = ['timestep', 'nsteps', 'functions'] for option_name in option_names: if option_name not in options: options[option_name] = None from perses.annihilation.ncmc_switching import NCMCEngine self.ncmc_engine = NCMCEngine(temperature=self.sampler.thermodynamic_state.temperature, timestep=options['timestep'], nsteps=options['nsteps'], functions=options['functions'], platform=platform) from perses.rjmc.geometry import FFAllAngleGeometryEngine self.geometry_engine = FFAllAngleGeometryEngine({'data': 0}) 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
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)
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
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)
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
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)
class ExpandedEnsembleSampler(object): """ Method of expanded ensembles sampling engine. 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 >>> allowed_mutations = [[('2','ALA')],[('2','VAL'),('2','LEU')]] >>> proposal_engine = PointMutationEngine(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) >>> # Run the sampler >>> exen_sampler.run() """ def __init__(self, sampler, topology, state_key, proposal_engine, log_weights=None, scheme='ncmc-geometry-ncmc', options=dict(), platform=None): """ 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 log_weights : dict of object : float Log weights to use for expanded ensemble biases. scheme : str, optional, default='ncmc-geometry-ncmc' Update scheme. One of ['ncmc-geometry-ncmc', 'geometry-ncmc-geometry', 'geometry-ncmc'] 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. """ # Keep copies of initializing arguments. # TODO: Make deep copies? self.sampler = sampler self.topology = topology self.state_key = state_key self.proposal_engine = proposal_engine self.log_weights = log_weights self.scheme = scheme if self.log_weights is None: self.log_weights = dict() # Initialize self.iteration = 0 option_names = ['timestep', 'nsteps', 'functions'] for option_name in option_names: if option_name not in options: options[option_name] = None from perses.annihilation.ncmc_switching import NCMCEngine self.ncmc_engine = NCMCEngine(temperature=self.sampler.thermodynamic_state.temperature, timestep=options['timestep'], nsteps=options['nsteps'], functions=options['functions'], platform=platform) from perses.rjmc.geometry import FFAllAngleGeometryEngine self.geometry_engine = FFAllAngleGeometryEngine({'data': 0}) 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 @property def state_keys(self): return 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 update_positions(self): """ Sample new positions. """ self.sampler.update() def update_state(self): """ Sample the thermodynamic state. """ # Check that system and topology have same number of atoms. old_system = self.sampler.sampler_state.system old_topology = self.topology old_topology_natoms = sum([1 for atom in old_topology.atoms()]) # number of topology atoms old_system_natoms = old_system.getNumParticles() if old_topology_natoms != old_system_natoms: msg = 'ExpandedEnsembleSampler: topology has %d atoms, while system has %d atoms' % (old_topology_natoms, old_system_natoms) raise Exception(msg) if self.scheme == 'ncmc-geometry-ncmc': if self.verbose: print("Updating chemical state with ncmc-geometry-ncmc scheme...") # DEBUG: Check current topology can be built. try: self.proposal_engine._system_generator.build_system(self.topology) except Exception as e: msg = str(e) msg += '\n' msg += 'ExpandedEnsembleSampler.update_sampler: self.topology before ProposalEngine call cannot be built into a system' raise Exception(msg) # Propose new chemical state. if self.verbose: print("Proposing new topology...") [system, topology, positions] = [self.sampler.thermodynamic_state.system, self.topology, self.sampler.sampler_state.positions] topology_proposal = self.proposal_engine.propose(system, topology) if self.verbose: print("Proposed transformation: %s => %s" % (topology_proposal.old_chemical_state_key, topology_proposal.new_chemical_state_key)) # DEBUG: Check current topology can be built. if self.verbose: print("Generating new system...") try: self.proposal_engine._system_generator.build_system(topology_proposal.new_topology) except Exception as e: msg = str(e) msg += '\n' msg += 'ExpandedEnsembleSampler.update_sampler: toology_proposal.new_topology before ProposalEngine call cannot be built into a system' raise Exception(msg) # Check to make sure no out-of-bounds atoms are present in new_to_old_atom_map natoms_old = topology_proposal.old_system.getNumParticles() natoms_new = topology_proposal.new_system.getNumParticles() if not set(topology_proposal.new_to_old_atom_map.values()).issubset(range(natoms_old)): msg = "Some old atoms in TopologyProposal.new_to_old_atom_map are not in span of old atoms (1..%d):\n" % natoms_old msg += str(topology_proposal.new_to_old_atom_map) raise Exception(msg) if not set(topology_proposal.new_to_old_atom_map.keys()).issubset(range(natoms_new)): msg = "Some new atoms in TopologyProposal.new_to_old_atom_map are not in span of old atoms (1..%d):\n" % natoms_new msg += str(topology_proposal.new_to_old_atom_map) raise Exception(msg) # 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) if self.verbose: print("Performing NCMC annihilation") # Alchemically eliminate atoms being removed. [ncmc_old_positions, ncmc_elimination_logp, potential_delete] = self.ncmc_engine.integrate(topology_proposal, positions, direction='delete') # Check that positions are not NaN if np.any(np.isnan(ncmc_old_positions)): raise Exception("Positions are NaN after NCMC delete with %d steps" % switching_nsteps) if self.verbose: print("Geometry engine proposal...") # Generate coordinates for new atoms and compute probability ratio of old and new probabilities. geometry_old_positions = ncmc_old_positions geometry_new_positions, geometry_logp_propose = self.geometry_engine.propose(topology_proposal, geometry_old_positions, self.sampler.thermodynamic_state.beta) if self.geometry_pdbfile is not None: print("Writing proposed geometry...") #self.geometry_pdbfile.write('MODEL %4d\n' % (self.iteration+1)) # PyMOL doesn't render connectivity correctly this way from simtk.openmm.app import PDBFile PDBFile.writeFile(topology_proposal.new_topology, geometry_new_positions, file=self.geometry_pdbfile) #self.geometry_pdbfile.write('ENDMDL\n') self.geometry_pdbfile.flush() geometry_logp_reverse = self.geometry_engine.logp_reverse(topology_proposal, geometry_new_positions, geometry_old_positions, self.sampler.thermodynamic_state.beta) geometry_logp = geometry_logp_reverse - geometry_logp_propose if self.verbose: print("Performing NCMC insertion") # Alchemically introduce new atoms. [ncmc_new_positions, ncmc_introduction_logp, potential_insert] = self.ncmc_engine.integrate(topology_proposal, geometry_new_positions, direction='insert') # Check that positions are not NaN if np.any(np.isnan(ncmc_new_positions)): raise Exception("Positions are NaN after NCMC insert with %d steps" % switching_nsteps) # Compute change in eliminated potential contribution. switch_logp = - (potential_insert - potential_delete) if self.verbose: print('potential before geometry : %12.3f kT' % potential_delete) print('potential after geometry : %12.3f kT' % potential_insert) print('---------------------------------------------------------') print('switch_logp : %12.3f' % switch_logp) print('geometry_logp_propose : %12.3f' % geometry_logp_propose) print('geometry_logp_reverse : %12.3f' % geometry_logp_reverse) # Compute total log acceptance probability, including all components. logp_accept = topology_proposal.logp_proposal + geometry_logp + switch_logp + ncmc_elimination_logp + ncmc_introduction_logp + new_log_weight - old_log_weight if self.verbose: print("logp_accept = %+10.4e [logp_proposal %+10.4e geometry_logp %+10.4e switch_logp %+10.4e ncmc_elimination_logp %+10.4e ncmc_introduction_logp %+10.4e old_log_weight %+10.4e new_log_weight %+10.4e]" % (logp_accept, topology_proposal.logp_proposal, geometry_logp, switch_logp, ncmc_elimination_logp, ncmc_introduction_logp, 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.system = topology_proposal.new_system self.sampler.sampler_state.system = topology_proposal.new_system self.topology = topology_proposal.new_topology self.sampler.sampler_state.positions = ncmc_new_positions 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") else: raise Exception("Expanded ensemble state proposal scheme '%s' unsupported" % self.scheme) # 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() 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, self.sampler.sampler_state.positions, self.pdbfile, self.iteration) self.pdbfile.flush() 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