def setup_fah_run(destination_path,
protein_pdb_filename,
oemol=None,
cache=None,
restrain_rmsd=False):
"""
Prepare simulation
Parameters
----------
destination_path : str
The path to the RUN to be created
protein_pdb_filename : str
Path to protein PDB file
oemol : openeye.oechem.OEMol, optional, default=None
The molecule to parameterize, with SDData attached
If None, don't include the small molecule
restrain_rmsd : bool, optional, default=False
If True, restrain RMSD during first equilibration phase
"""
# Parameters
from simtk import unit, openmm
protein_forcefield = 'amber14/protein.ff14SB.xml'
solvent_forcefield = 'amber14/tip3p.xml'
small_molecule_forcefield = 'openff-1.2.0'
water_model = 'tip3p'
solvent_padding = 10.0 * unit.angstrom
ionic_strength = 70 * unit.millimolar # assay buffer: 20 mM HEPES pH 7.3, 1 mM TCEP, 50 mM NaCl, 0.01% Tween-20, 10% glycerol
pressure = 1.0 * unit.atmospheres
collision_rate = 1.0 / unit.picoseconds
temperature = 300.0 * unit.kelvin
timestep = 4.0 * unit.femtoseconds
iterations = 1000 # 1 ns equilibration
nsteps_per_iteration = 250
# Prepare phases
import os
system_xml_filename = os.path.join(destination_path, 'system.xml.bz2')
integrator_xml_filename = os.path.join(destination_path,
'integrator.xml.bz2')
state_xml_filename = os.path.join(destination_path, 'state.xml.bz2')
# Check if we can skip setup
openmm_files_exist = os.path.exists(
system_xml_filename) and os.path.exists(
state_xml_filename) and os.path.exists(integrator_xml_filename)
if openmm_files_exist:
return
# Create barostat
barostat = openmm.MonteCarloBarostat(pressure, temperature)
# Create RUN directory if it does not yet exist
os.makedirs(destination_path, exist_ok=True)
# Load any molecule(s)
molecule = None
if oemol is not None:
from openforcefield.topology import Molecule
molecule = Molecule.from_openeye(oemol, allow_undefined_stereo=True)
molecule.name = 'MOL' # Ensure residue is MOL
print([res for res in molecule.to_topology().to_openmm().residues()])
# Create SystemGenerator
import os
from simtk.openmm import app
forcefield_kwargs = {
'removeCMMotion': False,
'hydrogenMass': 3.0 * unit.amu,
'constraints': app.HBonds,
'rigidWater': True
}
periodic_kwargs = {
'nonbondedMethod': app.PME,
'ewaldErrorTolerance': 2.5e-04
}
forcefields = [protein_forcefield, solvent_forcefield]
from openmmforcefields.generators import SystemGenerator
openmm_system_generator = SystemGenerator(
forcefields=forcefields,
molecules=molecule,
small_molecule_forcefield=small_molecule_forcefield,
cache=cache,
barostat=barostat,
forcefield_kwargs=forcefield_kwargs,
periodic_forcefield_kwargs=periodic_kwargs)
# Read protein
print(f'Reading protein from {protein_pdb_filename}...')
pdbfile = app.PDBFile(protein_pdb_filename)
modeller = app.Modeller(pdbfile.topology, pdbfile.positions)
if oemol is not None:
# Add small molecule to the system
modeller.add(molecule.to_topology().to_openmm(),
molecule.conformers[0])
# DEBUG : Check residue name
with open(os.path.join(destination_path, 'initial-complex.pdb'),
'wt') as outfile:
app.PDBFile.writeFile(modeller.topology, modeller.positions,
outfile)
# Add solvent
print('Adding solvent...')
kwargs = {'padding': solvent_padding}
modeller.addSolvent(openmm_system_generator.forcefield,
model='tip3p',
ionicStrength=ionic_strength,
**kwargs)
# Create an OpenMM system
print('Creating OpenMM system...')
system = openmm_system_generator.create_system(modeller.topology)
# Add a virtual bond between protein and ligand to make sure they are not imaged separately
if oemol is not None:
import mdtraj as md
mdtop = md.Topology.from_openmm(
modeller.topology) # excludes solvent and ions
for res in mdtop.residues:
print(res)
protein_atom_indices = mdtop.select(
'(protein and name CA)') # protein CA atoms
ligand_atom_indices = mdtop.select(
'((resname MOL) and (mass > 1))') # ligand heavy atoms
protein_atom_index = int(protein_atom_indices[0])
ligand_atom_index = int(ligand_atom_indices[0])
force = openmm.CustomBondForce('0')
force.addBond(protein_atom_index, ligand_atom_index, [])
system.addForce(force)
# Add RMSD restraints if requested
if restrain_rmsd:
print('Adding RMSD restraint...')
kB = unit.AVOGADRO_CONSTANT_NA * unit.BOLTZMANN_CONSTANT_kB
kT = kB * temperature
import mdtraj as md
mdtop = md.Topology.from_openmm(
pdbfile.topology) # excludes solvent and ions
#heavy_atom_indices = mdtop.select('mass > 1') # heavy solute atoms
rmsd_atom_indices = mdtop.select(
'(protein and (name CA)) or ((resname MOL) and (mass > 1))'
) # CA atoms and ligand heavy atoms
rmsd_atom_indices = [int(index) for index in rmsd_atom_indices]
custom_cv_force = openmm.CustomCVForce('(K_RMSD/2)*RMSD^2')
custom_cv_force.addGlobalParameter('K_RMSD', kT / unit.angstrom**2)
rmsd_force = openmm.RMSDForce(modeller.positions, rmsd_atom_indices)
custom_cv_force.addCollectiveVariable('RMSD', rmsd_force)
force_index = system.addForce(custom_cv_force)
# Create OpenM Context
platform = openmm.Platform.getPlatformByName('OpenCL')
platform.setPropertyDefaultValue('Precision', 'mixed')
from openmmtools import integrators
integrator = integrators.LangevinIntegrator(temperature, collision_rate,
timestep)
context = openmm.Context(system, integrator, platform)
context.setPositions(modeller.positions)
# Report initial potential energy
state = context.getState(getEnergy=True)
print(
f'Initial potential energy is {state.getPotentialEnergy()/unit.kilocalories_per_mole:.3f} kcal/mol'
)
# Store snapshots in MDTraj trajectory to examine RMSD
import mdtraj as md
import numpy as np
mdtop = md.Topology.from_openmm(pdbfile.topology)
atom_indices = mdtop.select('all') # all solute atoms
protein_atom_indices = mdtop.select(
'protein and (mass > 1)') # heavy solute atoms
if oemol is not None:
ligand_atom_indices = mdtop.select(
'(resname MOL) and (mass > 1)') # ligand heavy atoms
trajectory = md.Trajectory(
np.zeros([iterations + 1, len(atom_indices), 3], np.float32), mdtop)
trajectory.xyz[0, :, :] = context.getState(getPositions=True).getPositions(
asNumpy=True)[atom_indices] / unit.nanometers
# Minimize
print('Minimizing...')
openmm.LocalEnergyMinimizer.minimize(context)
# Equilibrate (with RMSD restraint if needed)
import numpy as np
from rich.progress import track
import time
initial_time = time.time()
for iteration in track(range(iterations), 'Equilibrating...'):
integrator.step(nsteps_per_iteration)
trajectory.xyz[iteration + 1, :, :] = context.getState(
getPositions=True).getPositions(
asNumpy=True)[atom_indices] / unit.nanometers
elapsed_time = (time.time() - initial_time) * unit.seconds
ns_per_day = (context.getState().getTime() /
elapsed_time) / (unit.nanoseconds / unit.day)
print(f'Performance: {ns_per_day:8.3f} ns/day')
if restrain_rmsd:
# Disable RMSD restraint
context.setParameter('K_RMSD', 0.0)
print('Minimizing...')
openmm.LocalEnergyMinimizer.minimize(context)
for iteration in track(range(iterations),
'Equilibrating without RMSD restraint...'):
integrator.step(nsteps_per_iteration)
# Retrieve state
state = context.getState(getPositions=True,
getVelocities=True,
getEnergy=True,
getForces=True)
system.setDefaultPeriodicBoxVectors(*state.getPeriodicBoxVectors())
modeller.topology.setPeriodicBoxVectors(state.getPeriodicBoxVectors())
print(
f'Final potential energy is {state.getPotentialEnergy()/unit.kilocalories_per_mole:.3f} kcal/mol'
)
# Equilibrate again if we restrained the RMSD
if restrain_rmsd:
print('Removing RMSD restraint from system...')
system.removeForce(force_index)
#if oemol is not None:
# # Check final RMSD
# print('checking RMSD...')
# trajectory.superpose(trajectory, atom_indices=protein_atom_indices)
# protein_rmsd = md.rmsd(trajectory, trajectory[-1], atom_indices=protein_atom_indices)[-1] * 10 # Angstroms
# oechem.OESetSDData(oemol, 'equil_protein_rmsd', f'{protein_rmsd:.2f} A')
# ligand_rmsd = md.rmsd(trajectory, trajectory[-1], atom_indices=ligand_atom_indices)[-1] * 10 # Angstroms
# oechem.OESetSDData(oemol, 'equil_ligand_rmsd', f'{ligand_rmsd:.2f} A')
# print('RMSD after equilibration: protein {protein_rmsd:8.2f} A | ligand {ligand_rmsd:8.3f} A')
# Save as OpenMM
print('Exporting for OpenMM FAH simulation...')
import bz2
with bz2.open(integrator_xml_filename, 'wt') as f:
f.write(openmm.XmlSerializer.serialize(integrator))
with bz2.open(state_xml_filename, 'wt') as f:
f.write(openmm.XmlSerializer.serialize(state))
with bz2.open(system_xml_filename, 'wt') as f:
f.write(openmm.XmlSerializer.serialize(system))
with bz2.open(os.path.join(destination_path, 'equilibrated-all.pdb.gz'),
'wt') as f:
app.PDBFile.writeFile(modeller.topology, state.getPositions(), f)
with open(os.path.join(destination_path, 'equilibrated-solute.pdb'),
'wt') as f:
import mdtraj
mdtraj_topology = mdtraj.Topology.from_openmm(modeller.topology)
mdtraj_trajectory = mdtraj.Trajectory(
[state.getPositions(asNumpy=True) / unit.nanometers],
mdtraj_topology)
selection = mdtraj_topology.select('not water')
mdtraj_trajectory = mdtraj_trajectory.atom_slice(selection)
app.PDBFile.writeFile(mdtraj_trajectory.topology.to_openmm(),
mdtraj_trajectory.openmm_positions(0), f)
if oemol is not None:
# Write molecule as SDF, SMILES, and mol2
for extension in ['sdf', 'mol2', 'smi', 'csv']:
filename = os.path.join(destination_path, f'molecule.{extension}')
with oechem.oemolostream(filename) as ofs:
oechem.OEWriteMolecule(ofs, oemol)
# Clean up
del context, integrator
def _anneal_ligand(self, structure, index):
"""
Anneal ligand interactions to clean up clashes.
Returns
-------
positions : unit.Quantity
Positions of all atoms after annealing the ligand
"""
reference_system = copy.deepcopy(self.phases[index].system)
protein = self.phases[index].mdtraj_top.select(
f'protein and backbone and type CA').tolist()
rmsd = openmm.RMSDForce(
self.phases[index].structure.positions,
protein + self.phases[index].lg1_idx + self.phases[index].lg2_idx)
energy_expression = 'step(dRMSD) * (K_RMSD/2)*dRMSD^2; dRMSD = (RMSD-RMSD0);'
restraint_force = openmm.CustomCVForce(energy_expression)
restraint_force.addCollectiveVariable('RMSD', rmsd)
restraint_force.addGlobalParameter(
'K_RMSD', 2 * unit.kilocalories_per_mole / (unit.angstroms**2))
restraint_force.addGlobalParameter('RMSD0', 2 * unit.angstroms)
reference_system.addForce(restraint_force)
alchemical_system = self._alchemically_modify_ligand(
reference_system, index)
from openmmtools.alchemy import AlchemicalState
alchemical_state_zero = mmtools.alchemy.AlchemicalState.from_system(
alchemical_system, parameters_name_suffix='zero')
alchemical_state_one = mmtools.alchemy.AlchemicalState.from_system(
alchemical_system, parameters_name_suffix='one')
thermodynamic_state = states.ThermodynamicState(
system=alchemical_system, temperature=300 * unit.kelvin)
composable_states = [alchemical_state_zero, alchemical_state_one]
compound_states = states.CompoundThermodynamicState(
thermodynamic_state, composable_states=composable_states)
sampler_state = states.SamplerState(
positions=structure.positions,
box_vectors=structure.topology.getPeriodicBoxVectors())
# Anneal
n_annealing_steps = 1000
integrator = openmm.LangevinIntegrator(300 * unit.kelvin,
90.0 / unit.picoseconds,
1.0 * unit.femtoseconds)
context, integrator = mmtools.cache.global_context_cache.get_context(
compound_states, integrator)
sampler_state.apply_to_context(context)
compound_states.lambda_sterics_one = 0.0
compound_states.lambda_electrostatics_one = 0.0
compound_states.apply_to_context(context)
print('Annealing sterics of ligand 1...')
for step in progressbar.progressbar(range(n_annealing_steps)):
compound_states.lambda_sterics_zero = float(step) / float(
n_annealing_steps)
compound_states.lambda_electrostatics_zero = 0.0
compound_states.apply_to_context(context)
integrator.step(1)
print('Annealing electrostatics of ligand 1...')
for step in progressbar.progressbar(range(n_annealing_steps)):
compound_states.lambda_sterics_zero = 1.0
compound_states.lambda_electrostatics_zero = float(step) / float(
n_annealing_steps)
compound_states.apply_to_context(context)
integrator.step(1)
compound_states.lambda_sterics_zero = 1.0
compound_states.lambda_electrostatics_zero = 1.0
compound_states.apply_to_context(context)
print('Annealing sterics of ligand 2...')
for step in progressbar.progressbar(range(n_annealing_steps)):
compound_states.lambda_sterics_one = float(step) / float(
n_annealing_steps)
compound_states.lambda_electrostatics_one = 0.0
compound_states.apply_to_context(context)
integrator.step(1)
print('Annealing electrostatics of ligand 2...')
for step in progressbar.progressbar(range(n_annealing_steps)):
compound_states.lambda_sterics_one = 1.0
compound_states.lambda_electrostatics_one = float(step) / float(
n_annealing_steps)
compound_states.apply_to_context(context)
integrator.step(1)
compound_states.apply_to_context(context)
sampler_state.update_from_context(context)
# Compute the final energy of the system.
final_energy = thermodynamic_state.reduced_potential(context)
print('final alchemical energy {:8.3f}kT'.format(final_energy))
return sampler_state.positions
print('Loading %s' % pdb_filename)
pdb = app.PDBFile(pdb_filename)
print("Loading forcefield: %s" % ffxml_filenames)
forcefield = app.ForceField(*ffxml_filenames)
# Create the system
print('Creating OpenMM System...')
system = forcefield.createSystem(pdb.topology,
nonbondedMethod=app.PME,
constraints=app.HBonds,
removeCMMotion=True)
# Add forces
print("Adding CV force...")
rmsd = openmm.RMSDForce(pdb.positions, list(range(702)))
cv_force = openmm.CustomCVForce("rmsd")
cv_force.addCollectiveVariable('rmsd', rmsd)
bv = mtd.BiasVariable(cv_force, 0, 2.5, 0.1, False)
# Add a barostat
# print('Adding barostat...')
# barostat = openmm.MonteCarloBarostat(pressure, temperature)
# system.addForce(barostat)
# Create simulation
print("Creating Simulation...")
integrator = openmm.LangevinIntegrator(temperature, collision_rate, timestep)
meta = mtd.Metadynamics(system, [bv], temperature, 3,
1.2 * unit.kilojoules_per_mole, 100)
simulation = app.Simulation(pdb.topology, system, integrator)