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