def run(self, system, topology, positions, box, resume=False): """ Run the sampler for a specified number of iterations Parameters resume : bool, default=False If True, resume the simulation """ self.system = system self.box = box if resume: # Resume the simulation print('Storage {} exists; resuming...'.format( self.output_filename)) from yank.multistate import SAMSSampler sampler = SAMSSampler.from_storage(self.output_filename) # Run the remainder of the simulation if sampler._iteration < self.n_iterations: # Resume sampler.run() else: # Extend sampler.extend(n_iterations=self.n_iterations) else: # Set up the simulation if it has not yet been set up if not self._setup_complete: self._setup(topology, positions, box) # Run the simulation #print(self.simulation.sampler_states[0].collective_variables) #print(self.simulation.sampler_states[0].potential_energy) self.simulation.run()
def run(self, platform_name=None, precision='auto', max_n_contexts=None, resume=False): """ Run the sampler for a specified number of iterations Parameters ---------- platform_name : str, optional, default=None Name of platform, or 'fastest' if fastest platform should be automatically selected precision : str, optional, default='auto' Precision to use, or None to automatically select ('mixed' is used if supported) max_n_contexts : int, optional, default=None Maximum number of contexts to use resume : bool, default=False If True, resume the simulation """ # Configure ContextCache, platform and precision from yank.experiment import ExperimentBuilder platform = ExperimentBuilder._configure_platform( platform_name, precision) try: openmmtools.cache.global_context_cache.platform = platform except RuntimeError: # The cache has been already used. Empty it before switching platform. openmmtools.cache.global_context_cache.empty() openmmtools.cache.global_context_cache.platform = platform if max_n_contexts is not None: openmmtools.cache.global_context_cache.capacity = max_n_contexts if resume: # Resume the simulation print('Storage {} exists; resuming...'.format( self.output_filename)) from yank.multistate import SAMSSampler, MultiStateReporter sampler = SAMSSampler.from_storage(self.output_filename) # Run the remainder of the simulation if sampler._iteration < self.n_iterations: # Resume sampler.run() else: # Extend sampler.extend(n_iterations=self.n_iterations) else: # Set up the simulation if it has not yet been set up if not self._setup_complete: self._setup() # Run the simulation self.simulation.run()
class SimulatePermeation(object): """ Properties ---------- anneal_ligand : bool, optional, default=True If True, will relax the ligand by alchemically annealing it first. Recommended, but can be slow on CPUs. """ def __init__(self, gromacs_input_path=None, ligand_resseq=None, output_filename=None, verbose=False): """Set up a SAMS permeation PMF simulation. Parameters ---------- gromacs_input_path : str, optional, default=None If specified, use gromacs files in this directory ligand_resseq : str, optional, default='MOL' Resdiue sequence id for ligand in reference PDB file output_filename : str, optional, default=None NetCDF output filename verbose : bool, optional, default=False If True, print verbose output """ # Setup general logging logging.root.setLevel(logging.DEBUG) logging.basicConfig(level=logging.DEBUG) yank.utils.config_root_logger(verbose=verbose, log_file_path=None) self._setup_complete = False # Set default parameters self.temperature = 310.0 * unit.kelvin self.pressure = 1.0 * unit.atmospheres self.collision_rate = 1.0 / unit.picoseconds self.timestep = 4.0 * unit.femtoseconds self.n_steps_per_iteration = 1250 self.n_iterations = 10000 self.checkpoint_interval = 50 self.gamma0 = 10.0 self.flatness_threshold = 10.0 self.anneal_ligand = True # Check input if gromacs_input_path is None: raise ValueError('gromacs_input_path must be specified') if ligand_resseq is None: raise ValueError('ligand_resseq must be specified') if output_filename is None: raise ValueError('output_filename must be specified') # Discover contents of the input path by suffix contents = { pathlib.Path(filename).suffix: filename for filename in os.listdir(gromacs_input_path) } gro_filename = os.path.join(gromacs_input_path, contents['.gro']) top_filename = os.path.join(gromacs_input_path, contents['.top']) pdb_filename = os.path.join(gromacs_input_path, contents['.pdb']) # Load system files print('Reading system from path: {}'.format(gromacs_input_path)) self.grofile = app.GromacsGroFile(gro_filename) self.topfile = app.GromacsTopFile( top_filename, periodicBoxVectors=self.grofile.getPeriodicBoxVectors()) self.pdbfile = app.PDBFile(pdb_filename) # Create MDTraj Trajectory for reference PDB file for use in atom selections and slicing self.mdtraj_refpdb = md.load(pdb_filename) self.mdtraj_topology = self.mdtraj_refpdb.topology self.analysis_particle_indices = self.mdtraj_topology.select( 'not water') # Store output filename self.output_filename = output_filename # Store ligand resseq self.ligand_resseq = ligand_resseq def _setup(self): """ Set up calculation. """ # Signal that system has now been set up if self._setup_complete: raise Exception( "System has already been set up---cannot run again.") self._setup_complete = True # Compute thermal energy and inverse temperature self.kT = kB * self.temperature self.beta = 1.0 / self.kT # Create the system self.system = self._create_system() # Add a barostat # TODO: Is this necessary, sicne ThermodynamicState handles this automatically? It may not correctly handle MonteCarloAnisotropicBarostat. self._add_barostat() # Restrain protein atoms in space # TODO: Allow protein atom selection to be configured selection = '((residue 342 and resname GLY) or (residue 97 and resname ASP) or (residue 184 and resname SER)) and (name CA)' protein_atoms_to_restrain = self.mdtraj_topology.select(selection) #self._restrain_protein(protein_atoms_to_restrain) # Create SamplerState for initial conditions self.sampler_state = states.SamplerState( positions=self.grofile.positions, box_vectors=self.grofile.getPeriodicBoxVectors()) # Create reference thermodynamic state self.reference_thermodynamic_state = states.ThermodynamicState( system=self.system, temperature=self.temperature, pressure=self.pressure) # Anneal ligand into binding site if self.anneal_ligand: self._anneal_ligand() # Create ThermodynamicStates for umbrella sampling along pore self.thermodynamic_states = self._create_thermodynamic_states( self.reference_thermodynamic_state) # Minimize initial thermodynamic state # TODO: Select initial thermodynamic state based on which state has minimum energy initial_state_index = 0 self.sampler_state = self._minimize_sampler_state( self.thermodynamic_states[initial_state_index], self.sampler_state) # Set up simulation from yank.multistate import SAMSSampler, MultiStateReporter # TODO: Change this to LangevinSplittingDynamicsMove move = mcmc.LangevinDynamicsMove(timestep=self.timestep, collision_rate=self.collision_rate, n_steps=self.n_steps_per_iteration, reassign_velocities=False) self.simulation = SAMSSampler( mcmc_moves=move, number_of_iterations=self.n_iterations, online_analysis_interval=None, gamma0=self.gamma0, flatness_threshold=self.flatness_threshold) self.reporter = MultiStateReporter( self.output_filename, checkpoint_interval=self.checkpoint_interval, analysis_particle_indices=self.analysis_particle_indices) self.simulation.create( thermodynamic_states=self.thermodynamic_states, unsampled_thermodynamic_states=[ self.reference_thermodynamic_state ], sampler_states=[self.sampler_state], initial_thermodynamic_states=[initial_state_index], storage=self.reporter) def run(self, platform_name=None, precision='auto', max_n_contexts=None, resume=False): """ Run the sampler for a specified number of iterations Parameters ---------- platform_name : str, optional, default=None Name of platform, or 'fastest' if fastest platform should be automatically selected precision : str, optional, default='auto' Precision to use, or None to automatically select ('mixed' is used if supported) max_n_contexts : int, optional, default=None Maximum number of contexts to use resume : bool, default=False If True, resume the simulation """ # Configure ContextCache, platform and precision from yank.experiment import ExperimentBuilder platform = ExperimentBuilder._configure_platform( platform_name, precision) try: openmmtools.cache.global_context_cache.platform = platform except RuntimeError: # The cache has been already used. Empty it before switching platform. openmmtools.cache.global_context_cache.empty() openmmtools.cache.global_context_cache.platform = platform if max_n_contexts is not None: openmmtools.cache.global_context_cache.capacity = max_n_contexts if resume: # Resume the simulation print('Storage {} exists; resuming...'.format( self.output_filename)) from yank.multistate import SAMSSampler, MultiStateReporter sampler = SAMSSampler.from_storage(self.output_filename) # Run the remainder of the simulation if sampler._iteration < self.n_iterations: # Resume sampler.run() else: # Extend sampler.extend(n_iterations=self.n_iterations) else: # Set up the simulation if it has not yet been set up if not self._setup_complete: self._setup() # Run the simulation self.simulation.run() def _create_thermodynamic_states(self, reference_thermodynamic_state, spacing=0.25 * unit.angstroms): """ Create thermodynamic states for sampling along pore axis. Here, the porin is restrained in the (x,z) plane and the pore axis is oriented along the y-axis. Parameters ---------- reference_thermodynamic_state : openmmtools.states.ThermodynamicState Reference ThermodynamicState containing system, temperature, and pressure spacing : simtk.unit.Quantity with units compatible with angstroms Spacing between umbrellas spanning pore axis """ # Determine relevant coordinates # TODO: Allow selections to be specified as class methods axis_center_atom = self.mdtraj_topology.select( 'residue 128 and resname ALA and name CA')[0] axis_center = self._get_reference_coordinates(axis_center_atom) print('axis center (x,z) : {} : {}'.format(axis_center_atom, axis_center)) pore_top_atom = self.mdtraj_topology.select( 'residue 226 and resname SER and name CA')[0] pore_top = self._get_reference_coordinates(pore_top_atom) print('pore top (y-max): {} : {}'.format(pore_top_atom, pore_top)) pore_bottom_atom = self.mdtraj_topology.select( 'residue 88 and resname VAL and name CA')[0] pore_bottom = self._get_reference_coordinates(pore_bottom_atom) print('pore bottom (y-min): {} : {}'.format(pore_bottom_atom, pore_bottom)) # Determine ligand atoms selection = '(residue {}) and (mass > 1.5)'.format(self.ligand_resseq) print('Determining ligand atoms using "{}"...'.format(selection)) ligand_atoms = self.mdtraj_topology.select(selection) if len(ligand_atoms) == 0: raise ValueError('Ligand residue name {} not found'.format( self.ligand_resseq)) print('ligand heavy atoms: {}'.format(ligand_atoms)) # Determine protein atoms bottom_selection = '((residue 342 and resname GLY) or (residue 97 and resname ASP) or (residue 184 and resname SER)) and (name CA)' bottom_protein_atoms = self.mdtraj_topology.select(bottom_selection) top_selection = '((residue 226 and resname SER) or (residue 421 and resname LEU) or (residue 147 and resname GLU)) and (name CA)' top_protein_atoms = self.mdtraj_topology.select(top_selection) # Compute pore bottom (lambda=0) and top (lambda=1) axis_bottom = (axis_center[0], pore_bottom[1], axis_center[2]) axis_top = (axis_center[0], pore_top[1], axis_center[2]) axis_distance = pore_top[1] - pore_bottom[1] print('axis_distance: {}'.format(axis_distance)) # Compute spacing and spring constant expansion_factor = 1.3 nstates = int(expansion_factor * axis_distance / spacing) + 1 print('nstates: {}'.format(nstates)) sigma_y = axis_distance / float( nstates) # stddev of force-free fluctuations in y-axis K_y = self.kT / (sigma_y**2) # spring constant print('vertical sigma_y = {:.3f} A'.format(sigma_y / unit.angstroms)) # Compute restraint width # TODO: Come up with a better way to define pore width? scale_factor = 0.25 sigma_xz = scale_factor * unit.sqrt( (axis_center[0] - pore_bottom[0])**2 + (axis_center[2] - pore_bottom[2])** 2) # stddev of force-free fluctuations in xz-plane K_xz = self.kT / (sigma_xz**2) # spring constant print('in-plane sigma_xz = {:.3f} A'.format(sigma_xz / unit.angstroms)) dr = axis_distance * (expansion_factor - 1.0) / 1.0 rmax = axis_distance + dr rmin = -1.0 * axis_distance # Create restraint state that encodes this axis # TODO: Rework this as CustomCVForce with in-plane and along-axis deviations as separate forces so we can store them each iteration print('Creating restraint...') from yank.restraints import RestraintState energy_expression = '(K_parallel/2)*(r_parallel-r0)^2 + (K_orthogonal/2)*r_orthogonal^2;' energy_expression += 'r_parallel = r*cos(theta);' energy_expression += 'r_orthogonal = r*sin(theta);' energy_expression += 'r = distance(g1,g2);' energy_expression += 'theta = angle(g1,g2,g3);' energy_expression += 'r0 = lambda_restraints * (rmax - rmin) + rmin;' force = openmm.CustomCentroidBondForce(3, energy_expression) force.addGlobalParameter('lambda_restraints', 1.0) force.addGlobalParameter('K_parallel', K_y) force.addGlobalParameter('K_orthogonal', K_xz) force.addGlobalParameter('rmax', rmax) force.addGlobalParameter('rmin', rmin) force.addGroup([int(index) for index in ligand_atoms]) force.addGroup([int(index) for index in bottom_protein_atoms]) force.addGroup([int(index) for index in top_protein_atoms]) force.addBond([0, 1, 2], []) self.system.addForce(force) # Update reference thermodynamic state print('Updating system in reference thermodynamic state...') self.reference_thermodynamic_state.set_system(self.system, fix_state=True) # Create alchemical state #from openmmtools.alchemy import AlchemicalState #alchemical_state = AlchemicalState.from_system(self.reference_thermodynamic_state.system) # Create restraint state restraint_state = RestraintState(lambda_restraints=1.0) # Create thermodynamic states to be sampled # TODO: Should we include an unbiased state? initial_time = time.time() thermodynamic_states = list() #compound_state = states.CompoundThermodynamicState(self.reference_thermodynamic_state, composable_states=[alchemical_state, restraint_state]) compound_state = states.CompoundThermodynamicState( self.reference_thermodynamic_state, composable_states=[restraint_state]) for lambda_restraints in np.linspace(0, 1, nstates): thermodynamic_state = copy.deepcopy(compound_state) thermodynamic_state.lambda_restraints = lambda_restraints #thermodynamic_state.lambda_sterics = 1.0 #thermodynamic_state.lambda_electrostatics = 1.0 thermodynamic_states.append(thermodynamic_state) elapsed_time = time.time() - initial_time print('Creating thermodynamic states took %.3f s' % elapsed_time) return thermodynamic_states def _get_reference_coordinates(self, atom): return self.mdtraj_refpdb.xyz[0, atom, :] * unit.nanometers def _create_system(self): """Create the System. """ # Create System print('Creating System...') # TODO: Allow these to be user-specified parameters if self.pressure is None: # We are simulating in a vacuum nonbonded_method = app.NoCutoff else: # We are simulating in solvent nonbonded_method = app.PME kwargs = { 'nonbondedMethod': nonbonded_method, 'constraints': app.HBonds, 'rigidWater': True, 'ewaldErrorTolerance': 1.0e-4, 'removeCMMotion': False, 'hydrogenMass': 3.0 * unit.amu } system = self.topfile.createSystem(**kwargs) # Fix particles with zero LJ sigma for force in system.getForces(): if force.__class__.__name__ == 'NonbondedForce': for index in range(system.getNumParticles()): [charge, sigma, epsilon] = force.getParticleParameters(index) if sigma / unit.nanometers == 0.0: force.setParticleParameters(index, charge, 1.0 * unit.angstroms, epsilon) print('System has {} particles'.format(system.getNumParticles())) return system def _alchemically_modify_ligand(self, reference_system): """ Alchemically soften ligand. """ # Determine ligand atoms ligand_atoms = self.mdtraj_topology.select('residue {}'.format( self.ligand_resseq)) if len(ligand_atoms) == 0: raise ValueError('Ligand residue name {} not found'.format( self.ligand_resseq)) print('ligand heavy atoms: {}'.format(ligand_atoms)) # Create alchemically-modified system # TODO: Use exact PME if PME is selected from openmmtools.alchemy import AbsoluteAlchemicalFactory, AlchemicalRegion, AlchemicalState #factory = AbsoluteAlchemicalFactory(consistent_exceptions=False, alchemical_pme_treatment='exact') factory = AbsoluteAlchemicalFactory( consistent_exceptions=False, disable_alchemical_dispersion_correction=True) alchemical_region = AlchemicalRegion(alchemical_atoms=ligand_atoms) alchemical_system = factory.create_alchemical_system( reference_system, alchemical_region) # Soften ligand #alchemical_state = AlchemicalState.from_system(alchemical_system) #alchemical_state.lambda_sterics = 1.0 #alchemical_state.lambda_electrostatics = 1.0 #alchemical_state.apply_to_system(alchemical_system) # Store alchemical system return alchemical_system def _anneal_ligand(self): """Anneal ligand interactions to clean up clashes. """ alchemical_system = self._alchemically_modify_ligand(self.system) from openmmtools.alchemy import AlchemicalState alchemical_state = AlchemicalState.from_system(alchemical_system) thermodynamic_state = states.ThermodynamicState( system=alchemical_system, temperature=self.temperature, pressure=self.pressure) alchemical_thermodynamic_state = states.CompoundThermodynamicState( thermodynamic_state=thermodynamic_state, composable_states=[alchemical_state]) # Initial softened minimization print('Minimizing softened ligand...') alchemical_thermodynamic_state.lambda_sterics = 0.01 alchemical_thermodynamic_state.lambda_electrostatics = 0.0 self.sampler_state = self._minimize_sampler_state( alchemical_thermodynamic_state, self.sampler_state) # Anneal n_annealing_steps = 1000 integrator = openmm.LangevinIntegrator(self.temperature, 90.0 / unit.picoseconds, 1.0 * unit.femtoseconds) context, integrator = openmmtools.cache.global_context_cache.get_context( alchemical_thermodynamic_state, integrator) self.sampler_state.apply_to_context(context) print('Annealing sterics...') for step in progressbar.progressbar(range(n_annealing_steps)): alchemical_state.lambda_sterics = float(step) / float( n_annealing_steps) alchemical_state.lambda_electrostatics = 0.0 alchemical_state.apply_to_context(context) integrator.step(1) print('Annealing electrostatics...') for step in progressbar.progressbar(range(n_annealing_steps)): alchemical_state.lambda_sterics = 1.0 alchemical_state.lambda_electrostatics = float(step) / float( n_annealing_steps) alchemical_state.apply_to_context(context) integrator.step(1) self.sampler_state.update_from_context(context) # Compute the final energy of the system for logging. final_energy = thermodynamic_state.reduced_potential(context) logger.debug('final alchemical energy {:8.3f}kT'.format(final_energy)) # Initial softened minimization print('Minimizing real ligand...') self.sampler_state = self._minimize_sampler_state( self.reference_thermodynamic_state, self.sampler_state) def _add_barostat(self): """Add a barostat to the system if one doesn't already exist. """ # TODO: For membrane protein simulations, we should add an anisotropic barostat has_barostat = False for (index, force) in enumerate(self.system.getForces()): if force.__class__.__name__ in [ 'MonteCarloBarostat', 'MonteCarloAnisotropicBarostat' ]: has_barostat = True force_index = index if self.pressure is not None: if not has_barostat: print('Adding a barostat...') barostat = openmm.MonteCarloBarostat(self.pressure, self.temperature) self.system.addForce(barostat) else: if has_barostat: # Remove barostat self.system.removeForce(force_index) def _restrain_protein(self, protein_atoms_to_restrain): """ Restrain protein atoms in system to keep protein from moving from reference position. Parameters ---------- protein_atoms_to_restrain : list of int List of atom indices to restraint """ # TODO: Allow sigma_protein to be configured print('Adding restraints to protein...') sigma_protein = 2.0 * unit.angstroms # stddev of fluctuations of protein atoms K_protein = self.kT / (sigma_protein**2) # spring constant energy_expression = '(K_protein/2)*((x-x0)^2 + (y-y0)^2 + (z-z0)^2);' force = openmm.CustomExternalForce(energy_expression) force.addGlobalParameter('K_protein', K_protein) force.addPerParticleParameter('x0') force.addPerParticleParameter('y0') force.addPerParticleParameter('z0') for particle_index in protein_atoms_to_restrain: x0, y0, z0 = self._get_reference_coordinates( particle_index) / unit.nanometers force.addParticle(int(particle_index), [x0, y0, z0]) self.system.addForce(force) @staticmethod def _minimize_sampler_state(thermodynamic_state, sampler_state, tolerance=None, max_iterations=100): """Minimize the specified sampler state at the given thermodynamic state. Parameters ---------- thermodynamic_state : openmmtools.states.ThermodynamicState The thermodynamic state for the simulation. sampler_state : openmmtools.states.SamplerState Initial sampler state tolerance : unit.Quantity compatible with kilojoules_per_mole/nanometer, optional, default=1.0*unit.kilojoules_per_mole/unit.angstroms Minimization will be terminated when RMS force reaches this tolerance max_iterations : int, optional, default=100 Maximum number of iterations for minimization. If 0, minimization continues until converged. Returns ------- minimized_sampler_state : openmmtools.states.SamplerState Minimized sampler state """ if tolerance is None: tolerance = 1.0 * unit.kilojoules_per_mole / unit.angstroms # Use the FIRE minimizer from yank.fire import FIREMinimizationIntegrator reference_integrator = FIREMinimizationIntegrator(tolerance=tolerance) # Create context #context = thermodynamic_state.create_context(integrator) context, integrator = openmmtools.cache.global_context_cache.get_context( thermodynamic_state, reference_integrator) # DEBUG: Reset FIRE minimizer integrator print('Resetting FIRE integrator...') for index in range(reference_integrator.getNumGlobalVariables()): value = reference_integrator.getGlobalVariable(index) integrator.setGlobalVariable(index, value) # Set initial positions and box vectors. sampler_state.apply_to_context(context) # Compute the initial energy of the system for logging. initial_energy = thermodynamic_state.reduced_potential(context) logger.debug('initial energy {:8.3f}kT'.format(initial_energy)) # Minimize energy. try: if max_iterations == 0: logger.debug( 'Using FIRE: tolerance {} minimizing to convergence'. format(tolerance)) while integrator.getGlobalVariableByName('converged') < 1: integrator.step(50) else: logger.debug( 'Using FIRE: tolerance {} max_iterations {}'.format( tolerance, max_iterations)) integrator.step(max_iterations) except Exception as e: if str(e) == 'Particle coordinate is nan': logger.debug( 'NaN encountered in FIRE minimizer; falling back to L-BFGS after resetting positions' ) sampler_state.apply_to_context(context) openmm.LocalEnergyMinimizer.minimize(context, tolerance, max_iterations) else: raise e # Get the minimized positions. sampler_state.update_from_context(context) # Compute the final energy of the system for logging. final_energy = thermodynamic_state.reduced_potential(context) logger.debug('final energy {:8.3f}kT'.format(final_energy)) if (final_energy >= initial_energy): logger.debug('minimizing again since no progress was made...') #sampler_state.apply_to_context(context) initial_energy = final_energy logger.debug('initial energy {:8.3f}kT'.format(initial_energy)) openmm.LocalEnergyMinimizer.minimize(context, tolerance, max_iterations) final_energy = thermodynamic_state.reduced_potential(context) logger.debug('final energy {:8.3f}kT'.format(final_energy)) sampler_state.update_from_context(context) # Clean up the integrator #del context, integrator return sampler_state
def _setup(self): """ Set up calculation. """ # Signal that system has now been set up if self._setup_complete: raise Exception( "System has already been set up---cannot run again.") self._setup_complete = True # Compute thermal energy and inverse temperature self.kT = kB * self.temperature self.beta = 1.0 / self.kT # Create the system self.system = self._create_system() # Add a barostat # TODO: Is this necessary, sicne ThermodynamicState handles this automatically? It may not correctly handle MonteCarloAnisotropicBarostat. self._add_barostat() # Restrain protein atoms in space # TODO: Allow protein atom selection to be configured selection = '((residue 342 and resname GLY) or (residue 97 and resname ASP) or (residue 184 and resname SER)) and (name CA)' protein_atoms_to_restrain = self.mdtraj_topology.select(selection) #self._restrain_protein(protein_atoms_to_restrain) # Create SamplerState for initial conditions self.sampler_state = states.SamplerState( positions=self.grofile.positions, box_vectors=self.grofile.getPeriodicBoxVectors()) # Create reference thermodynamic state self.reference_thermodynamic_state = states.ThermodynamicState( system=self.system, temperature=self.temperature, pressure=self.pressure) # Anneal ligand into binding site if self.anneal_ligand: self._anneal_ligand() # Create ThermodynamicStates for umbrella sampling along pore self.thermodynamic_states = self._create_thermodynamic_states( self.reference_thermodynamic_state) # Minimize initial thermodynamic state # TODO: Select initial thermodynamic state based on which state has minimum energy initial_state_index = 0 self.sampler_state = self._minimize_sampler_state( self.thermodynamic_states[initial_state_index], self.sampler_state) # Set up simulation from yank.multistate import SAMSSampler, MultiStateReporter # TODO: Change this to LangevinSplittingDynamicsMove move = mcmc.LangevinDynamicsMove(timestep=self.timestep, collision_rate=self.collision_rate, n_steps=self.n_steps_per_iteration, reassign_velocities=False) self.simulation = SAMSSampler( mcmc_moves=move, number_of_iterations=self.n_iterations, online_analysis_interval=None, gamma0=self.gamma0, flatness_threshold=self.flatness_threshold) self.reporter = MultiStateReporter( self.output_filename, checkpoint_interval=self.checkpoint_interval, analysis_particle_indices=self.analysis_particle_indices) self.simulation.create( thermodynamic_states=self.thermodynamic_states, unsampled_thermodynamic_states=[ self.reference_thermodynamic_state ], sampler_states=[self.sampler_state], initial_thermodynamic_states=[initial_state_index], storage=self.reporter)
protocol=protocol, constants=constants, composable_states=[composable_state]) # assign sampler_state sampler_state = states.SamplerState( positions=pdb.positions, box_vectors=system.getDefaultPeriodicBoxVectors()) # Set up the context for mtd simulation # at this step the CV and the system are separately passed to Metadynamics from yank.multistate import SAMSSampler, MultiStateReporter # TODO: You can use heavy hydrogens and 4 fs timesteps move = mcmc.LangevinDynamicsMove(timestep=2.0 * unit.femtoseconds, collision_rate=1.0 / unit.picosecond, n_steps=1000, reassign_velocities=False) print("Done specifying integrator for simulation.") simulation = SAMSSampler(mcmc_moves=move, number_of_iterations=iteration, online_analysis_interval=None, gamma0=1.0, flatness_threshold=0.2) storage_path = os.path.join(work_dir, 'traj.nc') reporter = MultiStateReporter(storage_path, checkpoint_interval=500) # We should also add unsampled_states with the fully-interacting system simulation.create(thermodynamic_states=thermo_states, sampler_states=[sampler_state], storage=reporter) print("Done specifying simulation.") simulation.run() print(f"Done with {iteration} iterations.")
def _setup(self, topology, positions, box): """ Set up calculation. """ # Signal that system has now been set up if self._setup_complete: raise Exception( "System has already been set up---cannot run again.") self._setup_complete = True # Compute thermal energy and inverse temperature self.kT = kB * self.temperature self.beta = 1.0 / self.kT # Add a barostat # TODO: Is this necessary, since ThermodynamicState handles this automatically? It may not correctly handle MonteCarloAnisotropicBarostat. self._add_barostat() # Create SamplerState for initial conditions self.sampler_state = states.SamplerState(positions=positions, box_vectors=self.box) # Create reference thermodynamic state self.reference_thermodynamic_state = states.ThermodynamicState( system=self.system, temperature=self.temperature, pressure=self.pressure) # Anneal ligand into binding site if self.anneal_ligand: self._anneal_ligand() positions = self.sampler_state.positions structure = pmd.openmm.load_topology(topology, system=self.system, xyz=positions) # Create ThermodynamicStates for umbrella sampling along pore self.thermodynamic_states = self._auto_create_thermodynamic_states( structure, topology, self.reference_thermodynamic_state) # Minimize initial thermodynamic state # TODO: Select initial thermodynamic state based on which state has minimum energy initial_state_index = 0 #self.sampler_state = self._minimize_sampler_state(self.thermodynamic_states[initial_state_index], self.sampler_state) #pickle.dump(self.sampler_state, open('sampler_state.obj', 'wb')) #pickle.dump(self.reference_thermodynamic_state, open('reference_thermodynamic_state.obj', 'wb')) #pickle.dump(self.thermodynamic_states, open('thermodynamic_states.obj', 'wb')) #pickle.dump(self.system, open('system.obj', 'wb')) # Set up simulation from yank.multistate import SAMSSampler, MultiStateReporter # TODO: Change this to LangevinSplittingDynamicsMove move = mcmc.LangevinDynamicsMove(timestep=self.timestep, collision_rate=self.collision_rate, n_steps=self.n_steps_per_iteration, reassign_velocities=False) self.simulation = SAMSSampler( mcmc_moves=move, number_of_iterations=self.n_iterations, online_analysis_interval=None, gamma0=self.gamma0, flatness_threshold=self.flatness_threshold) self.reporter = MultiStateReporter( self.output_filename, checkpoint_interval=self.checkpoint_interval, analysis_particle_indices=self.analysis_particle_indices) self.simulation.create( thermodynamic_states=self.thermodynamic_states, unsampled_thermodynamic_states=[ self.reference_thermodynamic_state ], sampler_states=[self.sampler_state], initial_thermodynamic_states=[initial_state_index], storage=self.reporter)
class SimulatePermeation(object): """ Properties ---------- anneal_ligand : bool, optional, default=True If True, will relax the ligand by alchemically annealing it first. Recommended, but can be slow on CPUs. """ def __init__(self, topology, ligand_resseq=None, membrane=None, output_filename=None): """Set up a SAMS permeation PMF simulation. Parameters ---------- topology membrane : str, optional, default=None output_filename : str, optional, default=None NetCDF output filename """ self._setup_complete = False # Set default parameters self.temperature = 310.0 * unit.kelvin self.pressure = 1.0 * unit.atmospheres self.collision_rate = 1.0 / unit.picoseconds self.timestep = 4.0 * unit.femtoseconds self.n_steps_per_iteration = 1250 self.n_iterations = 60000 self.checkpoint_interval = 50 self.gamma0 = 10.0 self.flatness_threshold = 10.0 self.anneal_ligand = True # Check input if output_filename is None: raise ValueError('output_filename must be specified') self.ligand_resseq = ligand_resseq # Create MDTraj Trajectory for reference PDB file for use in atom selections and slicings self.mdtraj_topology = md.Topology.from_openmm(topology) if membrane is not None: self.analysis_particle_indices = self.mdtraj_topology.select( 'not water and not resname ' + membrane) else: self.analysis_particle_indices = self.mdtraj_topology.select( 'not water') # Store output filename self.output_filename = output_filename def _setup(self, topology, positions, box): """ Set up calculation. """ # Signal that system has now been set up if self._setup_complete: raise Exception( "System has already been set up---cannot run again.") self._setup_complete = True # Compute thermal energy and inverse temperature self.kT = kB * self.temperature self.beta = 1.0 / self.kT # Add a barostat # TODO: Is this necessary, since ThermodynamicState handles this automatically? It may not correctly handle MonteCarloAnisotropicBarostat. self._add_barostat() # Create SamplerState for initial conditions self.sampler_state = states.SamplerState(positions=positions, box_vectors=self.box) # Create reference thermodynamic state self.reference_thermodynamic_state = states.ThermodynamicState( system=self.system, temperature=self.temperature, pressure=self.pressure) # Anneal ligand into binding site if self.anneal_ligand: self._anneal_ligand() positions = self.sampler_state.positions structure = pmd.openmm.load_topology(topology, system=self.system, xyz=positions) # Create ThermodynamicStates for umbrella sampling along pore self.thermodynamic_states = self._auto_create_thermodynamic_states( structure, topology, self.reference_thermodynamic_state) # Minimize initial thermodynamic state # TODO: Select initial thermodynamic state based on which state has minimum energy initial_state_index = 0 #self.sampler_state = self._minimize_sampler_state(self.thermodynamic_states[initial_state_index], self.sampler_state) #pickle.dump(self.sampler_state, open('sampler_state.obj', 'wb')) #pickle.dump(self.reference_thermodynamic_state, open('reference_thermodynamic_state.obj', 'wb')) #pickle.dump(self.thermodynamic_states, open('thermodynamic_states.obj', 'wb')) #pickle.dump(self.system, open('system.obj', 'wb')) # Set up simulation from yank.multistate import SAMSSampler, MultiStateReporter # TODO: Change this to LangevinSplittingDynamicsMove move = mcmc.LangevinDynamicsMove(timestep=self.timestep, collision_rate=self.collision_rate, n_steps=self.n_steps_per_iteration, reassign_velocities=False) self.simulation = SAMSSampler( mcmc_moves=move, number_of_iterations=self.n_iterations, online_analysis_interval=None, gamma0=self.gamma0, flatness_threshold=self.flatness_threshold) self.reporter = MultiStateReporter( self.output_filename, checkpoint_interval=self.checkpoint_interval, analysis_particle_indices=self.analysis_particle_indices) self.simulation.create( thermodynamic_states=self.thermodynamic_states, unsampled_thermodynamic_states=[ self.reference_thermodynamic_state ], sampler_states=[self.sampler_state], initial_thermodynamic_states=[initial_state_index], storage=self.reporter) def run(self, system, topology, positions, box, resume=False): """ Run the sampler for a specified number of iterations Parameters resume : bool, default=False If True, resume the simulation """ self.system = system self.box = box if resume: # Resume the simulation print('Storage {} exists; resuming...'.format( self.output_filename)) from yank.multistate import SAMSSampler sampler = SAMSSampler.from_storage(self.output_filename) # Run the remainder of the simulation if sampler._iteration < self.n_iterations: # Resume sampler.run() else: # Extend sampler.extend(n_iterations=self.n_iterations) else: # Set up the simulation if it has not yet been set up if not self._setup_complete: self._setup(topology, positions, box) # Run the simulation #print(self.simulation.sampler_states[0].collective_variables) #print(self.simulation.sampler_states[0].potential_energy) self.simulation.run() #ts = self.simulation._thermodynamic_states[self.simulation._replica_thermodynamic_states[0]] #context, _ = openmmtools.cache.global_context_cache.get_context(ts) #self.simulation.sampler_states[0].update_from_context(context, ignore_positions=True, ignore_velocities=True) #print(self.simulation.sampler_states[0].collective_variables) #print(self.simulation.sampler_states[0].potential_energy) #for i in range(10): # self.simulation.extend(n_iterations=1) # self.simulation.sampler_states[0].update_from_context(context, ignore_positions=True, ignore_velocities=True) # print(self.simulation.sampler_states[0].collective_variables) # print(self.simulation.sampler_states[0].potential_energy) def _auto_create_thermodynamic_states(self, structure, topology, reference_thermodynamic_state, spacing=0.25 * unit.angstroms): """ Create thermodynamic states for sampling along pore axis. """ topo = md.Topology.from_openmm(topology) data = DataPoints(structure, topo) data.writeCoordinates(open('cylinder.xyz', 'w')) cylinder = CylinderFitting(data.coordinates) print(cylinder.vmdCommands(), file=open('vmd_commands.txt', 'w')) b_index, t_index = cylinder.atomsInExtremes(data.coordinates, n=5) cylinder.writeExtremesCoords(data.coordinates, b_index, t_index, open('extremes.xyz', 'w')) bottom_atoms = [] top_atoms = [] for idx in b_index: bottom_atoms.append(data.index[idx]) for idx in t_index: top_atoms.append(data.index[idx]) axis_distance = cylinder.height * unit.angstroms selection = '(residue {}) and (mass > 1.5)'.format(self.ligand_resseq) print('Determining ligand atoms using "{}"...'.format(selection)) ligand_atoms = self.mdtraj_topology.select(selection) expansion_factor = 1.3 nstates = int(expansion_factor * axis_distance / spacing) + 1 print('nstates: {}'.format(nstates)) sigma_y = axis_distance / float( nstates) # stddev of force-free fluctuations in y-axis K_y = self.kT / (sigma_y**2) # spring constant print('vertical sigma_y = {:.3f} A'.format(sigma_y / unit.angstroms)) # Compute restraint width scale_factor = 0.25 sigma_xz = scale_factor * cylinder.r * unit.angstroms # stddev of force-free fluctuations in xz-plane K_xz = self.kT / (sigma_xz**2) # spring constant Kmax = K_y Kmin = K_xz dr = axis_distance * (expansion_factor - 1.0) / 2.0 rmax = axis_distance + dr rmin = -dr # Create restraint state that encodes this axis print('Creating restraint...') from yank.restraints import RestraintState # Collective variable definitions common = 'r = distance(g1,g2);' common += 'theta = angle(g1,g2,g3);' r_parallel = openmm.CustomCentroidBondForce(3, 'r*cos(theta);' + common) r_orthogonal = openmm.CustomCentroidBondForce(3, 'r*sin(theta);' + common) for cv in [r_parallel, r_orthogonal]: cv.addGroup([int(index) for index in ligand_atoms]) cv.addGroup([int(index) for index in bottom_atoms]) cv.addGroup([int(index) for index in top_atoms]) cv.addBond([0, 1, 2], []) energy_parallel = '(K_parallel/2)*(r_parallel-r0)^2;' energy_parallel += 'r0 = lambda_restraints * (rmax - rmin) + rmin;' self.cvforce_parallel = openmm.CustomCVForce(energy_parallel) self.cvforce_parallel.addCollectiveVariable('r_parallel', r_parallel) energy_orthogonal = '(1/2)*(Kmin + S*(Kmax - Kmin))*(r_orthogonal^2);' energy_orthogonal += 'S = 1 - step(z_ext-z)*(1 + u^3*(15*u - 6*u^2 - 10));' energy_orthogonal += 'u = step(z-z_int)*(z - z_int)/(z_ext - z_int);' energy_orthogonal += 'z = abs(r0-z_c);' energy_orthogonal += 'r0 = lambda_restraints * (rmax - rmin) + rmin;' self.cvforce_orthogonal = openmm.CustomCVForce(energy_orthogonal) self.cvforce_orthogonal.addCollectiveVariable('r_orthogonal', r_orthogonal) for force in [self.cvforce_parallel, self.cvforce_orthogonal]: force.addGlobalParameter('rmax', rmax) force.addGlobalParameter('rmin', rmin) force.addGlobalParameter('lambda_restraints', 1.0) self.cvforce_parallel.addGlobalParameter('K_parallel', K_y) self.cvforce_orthogonal.addGlobalParameter('Kmax', Kmax) self.cvforce_orthogonal.addGlobalParameter('Kmin', Kmin) self.cvforce_orthogonal.addGlobalParameter('z_c', axis_distance / 2.0) self.cvforce_orthogonal.addGlobalParameter('z_int', 0.8 * (axis_distance / 2.0)) self.cvforce_orthogonal.addGlobalParameter('z_ext', 1.3 * (axis_distance / 2.0)) self.system.addForce(self.cvforce_parallel) self.system.addForce(self.cvforce_orthogonal) # Update reference thermodynamic state print('Updating system in reference thermodynamic state...') self.reference_thermodynamic_state.set_system(self.system, fix_state=True) ## Create alchemical state ##from openmmtools.alchemy import AlchemicalState ##alchemical_state = AlchemicalState.from_system(self.reference_thermodynamic_state.system) # Create restraint state restraint_state = RestraintState(lambda_restraints=1.0) # Create thermodynamic states to be sampled # TODO: Should we include an unbiased state? initial_time = time.time() thermodynamic_states = list() ##compound_state = states.CompoundThermodynamicState(self.reference_thermodynamic_state, composable_states=[alchemical_state, restraint_state]) compound_state = states.CompoundThermodynamicState( self.reference_thermodynamic_state, composable_states=[restraint_state]) for lambda_restraints in np.linspace(0, 1, nstates): thermodynamic_state = copy.deepcopy(compound_state) thermodynamic_state.lambda_restraints = lambda_restraints # #thermodynamic_state.lambda_sterics = 1.0 # #thermodynamic_state.lambda_electrostatics = 1.0 thermodynamic_states.append(thermodynamic_state) elapsed_time = time.time() - initial_time print('Creating thermodynamic states took %.3f s' % elapsed_time) return thermodynamic_states def _create_thermodynamic_states(self, reference_thermodynamic_state, spacing=0.25 * unit.angstroms): """ Create thermodynamic states for sampling along pore axis. Here, the porin is restrained in the (x,z) plane and the pore axis is oriented along the y-axis. Parameters ---------- reference_thermodynamic_state : openmmtools.states.ThermodynamicState Reference ThermodynamicState containing system, temperature, and pressure spacing : simtk.unit.Quantity with units compatible with angstroms Spacing between umbrellas spanning pore axis """ # Determine relevant coordinates # TODO: Allow selections to be specified as class methods axis_center_atom = self.mdtraj_topology.select( 'residue 128 and resname ALA and name CA')[0] axis_center = self._get_reference_coordinates(axis_center_atom) print('axis center (x,z) : {} : {}'.format(axis_center_atom, axis_center)) pore_top_atom = self.mdtraj_topology.select( 'residue 226 and resname SER and name CA')[0] pore_top = self._get_reference_coordinates(pore_top_atom) print('pore top (y-max): {} : {}'.format(pore_top_atom, pore_top)) pore_bottom_atom = self.mdtraj_topology.select( 'residue 88 and resname VAL and name CA')[0] pore_bottom = self._get_reference_coordinates(pore_bottom_atom) print('pore bottom (y-min): {} : {}'.format(pore_bottom_atom, pore_bottom)) # Determine ligand atoms selection = '(residue {}) and (mass > 1.5)'.format(self.ligand_resseq) print('Determining ligand atoms using "{}"...'.format(selection)) ligand_atoms = self.mdtraj_topology.select(selection) if len(ligand_atoms) == 0: raise ValueError('Ligand residue name {} not found'.format( self.ligand_resseq)) print('ligand heavy atoms: {}'.format(ligand_atoms)) # Determine protein atoms bottom_selection = '((residue 342 and resname GLY) or (residue 97 and resname ASP) or (residue 184 and resname SER)) and (name CA)' bottom_protein_atoms = self.mdtraj_topology.select(bottom_selection) top_selection = '((residue 226 and resname SER) or (residue 421 and resname LEU) or (residue 147 and resname GLU)) and (name CA)' top_protein_atoms = self.mdtraj_topology.select(top_selection) # Compute pore bottom (lambda=0) and top (lambda=1) axis_bottom = (axis_center[0], pore_bottom[1], axis_center[2]) axis_top = (axis_center[0], pore_top[1], axis_center[2]) axis_distance = pore_top[1] - pore_bottom[1] print('axis_distance: {}'.format(axis_distance)) # Compute spacing and spring constant expansion_factor = 1.3 nstates = int(expansion_factor * axis_distance / spacing) + 1 print('nstates: {}'.format(nstates)) sigma_y = axis_distance / float( nstates) # stddev of force-free fluctuations in y-axis K_y = self.kT / (sigma_y**2) # spring constant print('vertical sigma_y = {:.3f} A'.format(sigma_y / unit.angstroms)) # Compute restraint width # TODO: Come up with a better way to define pore width? scale_factor = 0.25 sigma_xz = scale_factor * unit.sqrt( (axis_center[0] - pore_bottom[0])**2 + (axis_center[2] - pore_bottom[2])** 2) # stddev of force-free fluctuations in xz-plane K_xz = self.kT / (sigma_xz**2) # spring constant print('in-plane sigma_xz = {:.3f} A'.format(sigma_xz / unit.angstroms)) dr = axis_distance * (expansion_factor - 1.0) / 1.0 rmax = axis_distance + dr rmin = -1.0 * axis_distance # Create restraint state that encodes this axis # TODO: Rework this as CustomCVForce with in-plane and along-axis deviations as separate forces so we can store them each iteration print('Creating restraint...') from yank.restraints import RestraintState energy_expression = '(K_parallel/2)*(r_parallel-r0)^2 + (K_orthogonal/2)*r_orthogonal^2;' energy_expression += 'r_parallel = r*cos(theta);' energy_expression += 'r_orthogonal = r*sin(theta);' energy_expression += 'r = distance(g1,g2);' energy_expression += 'theta = angle(g1,g2,g3);' energy_expression += 'r0 = lambda_restraints * (rmax - rmin) + rmin;' force = openmm.CustomCentroidBondForce(3, energy_expression) force.addGlobalParameter('lambda_restraints', 1.0) force.addGlobalParameter('K_parallel', K_y) force.addGlobalParameter('K_orthogonal', K_xz) force.addGlobalParameter('rmax', rmax) force.addGlobalParameter('rmin', rmin) force.addGroup([int(index) for index in ligand_atoms]) force.addGroup([int(index) for index in bottom_protein_atoms]) force.addGroup([int(index) for index in top_protein_atoms]) force.addBond([0, 1, 2], []) self.system.addForce(force) # Update reference thermodynamic state print('Updating system in reference thermodynamic state...') self.reference_thermodynamic_state.set_system(self.system, fix_state=True) # Create alchemical state #from openmmtools.alchemy import AlchemicalState #alchemical_state = AlchemicalState.from_system(self.reference_thermodynamic_state.system) # Create restraint state restraint_state = RestraintState(lambda_restraints=1.0) # Create thermodynamic states to be sampled # TODO: Should we include an unbiased state? initial_time = time.time() thermodynamic_states = list() #compound_state = states.CompoundThermodynamicState(self.reference_thermodynamic_state, composable_states=[alchemical_state, restraint_state]) compound_state = states.CompoundThermodynamicState( self.reference_thermodynamic_state, composable_states=[restraint_state]) for lambda_restraints in np.linspace(0, 1, nstates): thermodynamic_state = copy.deepcopy(compound_state) thermodynamic_state.lambda_restraints = lambda_restraints #thermodynamic_state.lambda_sterics = 1.0 #thermodynamic_state.lambda_electrostatics = 1.0 thermodynamic_states.append(thermodynamic_state) elapsed_time = time.time() - initial_time print('Creating thermodynamic states took %.3f s' % elapsed_time) return thermodynamic_states def _get_reference_coordinates(self, atom): return self.mdtraj_refpdb.xyz[0, atom, :] * unit.nanometers def _alchemically_modify_ligand(self, reference_system): """ Alchemically soften ligand. """ # Determine ligand atoms ligand_atoms = self.mdtraj_topology.select('residue {}'.format( self.ligand_resseq)) if len(ligand_atoms) == 0: raise ValueError('Ligand residue name {} not found'.format( self.ligand_resseq)) print('ligand heavy atoms: {}'.format(ligand_atoms)) # Create alchemically-modified system # TODO: Use exact PME if PME is selected from openmmtools.alchemy import AbsoluteAlchemicalFactory, AlchemicalRegion, AlchemicalState #factory = AbsoluteAlchemicalFactory(consistent_exceptions=False, alchemical_pme_treatment='exact') factory = AbsoluteAlchemicalFactory( consistent_exceptions=False, disable_alchemical_dispersion_correction=True) alchemical_region = AlchemicalRegion(alchemical_atoms=ligand_atoms) alchemical_system = factory.create_alchemical_system( reference_system, alchemical_region) # Soften ligand #alchemical_state = AlchemicalState.from_system(alchemical_system) #alchemical_state.lambda_sterics = 1.0 #alchemical_state.lambda_electrostatics = 1.0 #alchemical_state.apply_to_system(alchemical_system) # Store alchemical system return alchemical_system def _anneal_ligand(self): """Anneal ligand interactions to clean up clashes. """ alchemical_system = self._alchemically_modify_ligand(self.system) from openmmtools.alchemy import AlchemicalState alchemical_state = AlchemicalState.from_system(alchemical_system) thermodynamic_state = states.ThermodynamicState( system=alchemical_system, temperature=self.temperature, pressure=self.pressure) alchemical_thermodynamic_state = states.CompoundThermodynamicState( thermodynamic_state=thermodynamic_state, composable_states=[alchemical_state]) # Initial softened minimization print('Minimizing softened ligand...') alchemical_thermodynamic_state.lambda_sterics = 0.01 alchemical_thermodynamic_state.lambda_electrostatics = 0.0 self.sampler_state = self._minimize_sampler_state( alchemical_thermodynamic_state, self.sampler_state) # Anneal n_annealing_steps = 1000 integrator = openmm.LangevinIntegrator(self.temperature, 90.0 / unit.picoseconds, 1.0 * unit.femtoseconds) context, integrator = openmmtools.cache.global_context_cache.get_context( alchemical_thermodynamic_state, integrator) self.sampler_state.apply_to_context(context) print('Annealing sterics...') for step in progressbar.progressbar(range(n_annealing_steps)): alchemical_state.lambda_sterics = float(step) / float( n_annealing_steps) alchemical_state.lambda_electrostatics = 0.0 alchemical_state.apply_to_context(context) integrator.step(1) print('Annealing electrostatics...') for step in progressbar.progressbar(range(n_annealing_steps)): alchemical_state.lambda_sterics = 1.0 alchemical_state.lambda_electrostatics = float(step) / float( n_annealing_steps) alchemical_state.apply_to_context(context) integrator.step(1) self.sampler_state.update_from_context(context) # Compute the final energy of the system for logging. final_energy = thermodynamic_state.reduced_potential(context) logger.debug('final alchemical energy {:8.3f}kT'.format(final_energy)) # Initial softened minimization print('Minimizing real ligand...') self.sampler_state = self._minimize_sampler_state( self.reference_thermodynamic_state, self.sampler_state) def _add_barostat(self): """Add a barostat to the system if one doesn't already exist. """ # TODO: For membrane protein simulations, we should add an anisotropic barostat has_barostat = False for (index, force) in enumerate(self.system.getForces()): if force.__class__.__name__ in [ 'MonteCarloBarostat', 'MonteCarloAnisotropicBarostat' ]: has_barostat = True force_index = index if self.pressure is not None: if not has_barostat: print('Adding a barostat...') barostat = openmm.MonteCarloBarostat(self.pressure, self.temperature) self.system.addForce(barostat) else: if has_barostat: # Remove barostat self.system.removeForce(force_index) def _restrain_protein(self, protein_atoms_to_restrain): """ Restrain protein atoms in system to keep protein from moving from reference position. Parameters ---------- protein_atoms_to_restrain : list of int List of atom indices to restraint """ # TODO: Allow sigma_protein to be configured print('Adding restraints to protein...') sigma_protein = 2.0 * unit.angstroms # stddev of fluctuations of protein atoms K_protein = self.kT / (sigma_protein**2) # spring constant energy_expression = '(K_protein/2)*((x-x0)^2 + (y-y0)^2 + (z-z0)^2);' force = openmm.CustomExternalForce(energy_expression) force.addGlobalParameter('K_protein', K_protein) force.addPerParticleParameter('x0') force.addPerParticleParameter('y0') force.addPerParticleParameter('z0') for particle_index in protein_atoms_to_restrain: x0, y0, z0 = self._get_reference_coordinates( particle_index) / unit.nanometers force.addParticle(int(particle_index), [x0, y0, z0]) self.system.addForce(force) @staticmethod def _minimize_sampler_state(thermodynamic_state, sampler_state, tolerance=None, max_iterations=100): """Minimize the specified sampler state at the given thermodynamic state. Parameters ---------- thermodynamic_state : openmmtools.states.ThermodynamicState The thermodynamic state for the simulation. sampler_state : openmmtools.states.SamplerState Initial sampler state tolerance : unit.Quantity compatible with kilojoules_per_mole/nanometer, optional, default=1.0*unit.kilojoules_per_mole/unit.angstroms Minimization will be terminated when RMS force reaches this tolerance max_iterations : int, optional, default=100 Maximum number of iterations for minimization. If 0, minimization continues until converged. Returns ------- minimized_sampler_state : openmmtools.states.SamplerState Minimized sampler state """ if tolerance is None: tolerance = 1.0 * unit.kilojoules_per_mole / unit.angstroms # Use the FIRE minimizer from yank.fire import FIREMinimizationIntegrator reference_integrator = FIREMinimizationIntegrator(tolerance=tolerance) # Create context #context = thermodynamic_state.create_context(integrator) context, integrator = openmmtools.cache.global_context_cache.get_context( thermodynamic_state, reference_integrator) # DEBUG: Reset FIRE minimizer integrator print('Resetting FIRE integrator...') for index in range(reference_integrator.getNumGlobalVariables()): value = reference_integrator.getGlobalVariable(index) integrator.setGlobalVariable(index, value) # Set initial positions and box vectors. sampler_state.apply_to_context(context) # Compute the initial energy of the system for logging. initial_energy = thermodynamic_state.reduced_potential(context) logger.debug('initial energy {:8.3f}kT'.format(initial_energy)) # Minimize energy. try: if max_iterations == 0: logger.debug( 'Using FIRE: tolerance {} minimizing to convergence'. format(tolerance)) while integrator.getGlobalVariableByName('converged') < 1: integrator.step(50) else: logger.debug( 'Using FIRE: tolerance {} max_iterations {}'.format( tolerance, max_iterations)) integrator.step(max_iterations) except Exception as e: if str(e) == 'Particle coordinate is nan': logger.debug( 'NaN encountered in FIRE minimizer; falling back to L-BFGS after resetting positions' ) sampler_state.apply_to_context(context) openmm.LocalEnergyMinimizer.minimize(context, tolerance, max_iterations) else: raise e # Get the minimized positions. sampler_state.update_from_context(context) # Compute the final energy of the system for logging. final_energy = thermodynamic_state.reduced_potential(context) logger.debug('final energy {:8.3f}kT'.format(final_energy)) if (final_energy >= initial_energy): logger.debug('minimizing again since no progress was made...') #sampler_state.apply_to_context(context) initial_energy = final_energy logger.debug('initial energy {:8.3f}kT'.format(initial_energy)) openmm.LocalEnergyMinimizer.minimize(context, tolerance, max_iterations) final_energy = thermodynamic_state.reduced_potential(context) logger.debug('final energy {:8.3f}kT'.format(final_energy)) sampler_state.update_from_context(context) # Clean up the integrator #del context, integrator return sampler_state