openmmtools
openmmtools copied to clipboard
Published
20 hours ago •
choderalab
choderalab
Add gentle equilibration code
I recently implemented the gentle equilibration protocol from this publication in OpenMM:
from openmm import unit
import pickle
import argparse
import logging
# Set up logger
_logger = logging.getLogger()
_logger.setLevel(logging.INFO)
# Read args
parser = argparse.ArgumentParser(description='run perses protein mutation on capped amino acid')
parser.add_argument('infile', type=str, help='path to input pickle')
parser.add_argument('outfile', type=str, help='path to output file (.pdb or .cif)')
args = parser.parse_args()
def run_gentle_equilibration(topology, positions, system, stages, filename, platform_name='CUDA'):
"""
Run gentle equilibration.
Parameters
----------
topology : openmm.app.Topology
topology
positions : np.array in unit.nanometer
positions
system : openmm.System
system
stages : list of dicts
each dict corresponds to a stage of equilibration and contains the equilibration parameters for that stage
equilibration parameters:
EOM : str
'minimize' or 'MD' or 'MD_interpolate' (the last one will allow interpolation between 'temperature' and 'temperature_end')
n_steps : int
number of steps of MD
temperature : openmm.unit.kelvin
temperature (kelvin)
temperature_end : openmm.unit.kelvin, optional
the temperature (kelvin) at which to finish interpolation, if 'EOM' is 'MD_interpolate'
ensemble : str or None
'NPT' or 'NVT'
restraint_selection : str or None
to be used by mdtraj to select atoms for which to apply restraints
force_constant : openmm.unit.kilocalories_per_mole/openmm.unit.angstrom**2
force constant (kcal/molA^2)
collision_rate : 1/openmm.unit.picoseconds
collision rate (1/picoseconds)
timestep : openmm.unit.femtoseconds
timestep (femtoseconds)
filename : str
path to save the equilibrated structure
platform_name : str, default 'CUDA'
name of platform to be used by OpenMM. If not specified, OpenMM will select the fastest available platform
"""
import copy
import openmm
import time
import numpy as np
import mdtraj as md
from rich.progress import track
for i, parameters in enumerate(stages):
initial_time = time.time()
print(f"Executing stage {i + 1}")
# Make a copy of the system
system_copy = copy.deepcopy(system)
# Add restraint
if parameters['restraint_selection'] is not None:
traj = md.Trajectory(positions, md.Topology.from_openmm(topology))
selection_indices = traj.topology.select(parameters['restraint_selection'])
custom_cv_force = openmm.CustomCVForce('(K_RMSD/2)*(RMSD)^2')
custom_cv_force.addGlobalParameter('K_RMSD', parameters['force_constant'] * 2)
rmsd_force = openmm.RMSDForce(positions, selection_indices)
custom_cv_force.addCollectiveVariable('RMSD', rmsd_force)
system_copy.addForce(custom_cv_force)
# Set barostat update interval to 0 (for NVT)
if parameters['ensemble'] == 'NVT':
force_dict = {force.__class__.__name__: index for index, force in enumerate(system_copy.getForces())}
system_copy.removeForce(force_dict['MonteCarloBarostat']) # TODO : change this to `system_copy.getForce(force_dict['MonteCarloBarostat']).setFrequency(0) once the next release comes out (this recently merged PR allows frequency to be 0: https://github.com/openmm/openmm/pull/3411)
elif parameters['ensemble'] == 'NPT' or parameters['ensemble'] is None:
pass
else:
raise Exception("Invalid parameter supplied for 'ensemble'")
# Set up integrator
temperature = parameters['temperature']
collision_rate = parameters['collision_rate']
timestep = parameters['timestep']
if parameters['EOM'] == 'MD_interpolate':
temperature_end = parameters['temperature_end']
integrator = openmm.LangevinMiddleIntegrator(temperature, collision_rate, timestep)
# Set up context
platform = openmm.Platform.getPlatformByName(platform_name)
if platform_name in ['CUDA', 'OpenCL']:
platform.setPropertyDefaultValue('Precision', 'mixed')
if platform_name in ['CUDA']:
platform.setPropertyDefaultValue('DeterministicForces', 'true')
context = openmm.Context(system_copy, integrator, platform)
context.setPeriodicBoxVectors(*system_copy.getDefaultPeriodicBoxVectors())
context.setPositions(positions)
context.setVelocitiesToTemperature(temperature)
# Run minimization or MD
n_steps = parameters['n_steps']
n_steps_per_iteration = 100
if parameters['EOM'] == 'minimize':
openmm.LocalEnergyMinimizer.minimize(context, maxIterations=n_steps)
elif parameters['EOM'] == 'MD':
for _ in track(range(int(n_steps/n_steps_per_iteration))):
integrator.step(n_steps_per_iteration)
elif parameters['EOM'] == 'MD_interpolate':
temperature_unit = unit.kelvin
temperatures = np.linspace(temperature/temperature_unit, temperature_end/temperature_unit, int(n_steps/n_steps_per_iteration)) * temperature_unit
for temperature in track(temperatures):
integrator.setTemperature(temperature)
integrator.step(n_steps_per_iteration)
else:
raise Exception("Invalid parameter supplied for 'EOM'")
# Retrieve positions after this stage of equil
positions = context.getState(getPositions=True).getPositions(asNumpy=True)
# Update default box vectors for next iteration
box_vectors = context.getState().getPeriodicBoxVectors()
system.setDefaultPeriodicBoxVectors(*box_vectors)
# Delete context and integrator
del context, integrator, system_copy
elapsed_time = time.time() - initial_time
print(f"\tStage {i + 1} took {elapsed_time} seconds")
# Save the final equilibrated positions
if filename.endswith('pdb'):
openmm.app.PDBFile.writeFile(topology, positions, open(filename, "w"), keepIds=True)
elif filename.endswith('cif'):
openmm.app.PDBxFile.writeFile(topology, positions, open(filename, "w"), keepIds=True)
# Save the box vectors
with open(filename[:-4] + '_box_vectors.npy', 'wb') as f:
np.save(f, box_vectors)
# Load htf
with open(args.infile, "rb") as f:
htf = pickle.load(f)
# Retrieve old positions and system
# (in PointMutationExecutor, we will want to feed the gently equilibrated positions in before geometry engine)
positions = htf.old_positions(htf.hybrid_positions)
system = htf._topology_proposal.old_system
topology = htf._topology_proposal.old_topology
stages = [
{'EOM': 'minimize', 'n_steps': 10000, 'temperature': 300*unit.kelvin, 'ensemble': None, 'restraint_selection': 'protein and not type H', 'force_constant': 1*unit.kilocalories_per_mole/unit.angstrom**2, 'collision_rate': 2/unit.picoseconds, 'timestep': 1*unit.femtoseconds},
{'EOM': 'MD_interpolate', 'n_steps': 100000, 'temperature': 100*unit.kelvin, 'temperature_end': 300*unit.kelvin, 'ensemble': 'NVT', 'restraint_selection': 'protein and not type H', 'force_constant': 1*unit.kilocalories_per_mole/unit.angstrom**2, 'collision_rate': 10/unit.picoseconds, 'timestep': 1*unit.femtoseconds},
{'EOM': 'MD', 'n_steps': 100000, 'temperature': 300, 'ensemble': 'NPT', 'restraint_selection': 'protein and not type H', 'force_constant': 1*unit.kilocalories_per_mole/unit.angstrom**2, 'collision_rate': 10/unit.picoseconds, 'timestep': 1*unit.femtoseconds},
{'EOM': 'MD', 'n_steps': 250000, 'temperature': 300*unit.kelvin, 'ensemble': 'NPT', 'restraint_selection': 'protein and not type H', 'force_constant': 0.1*unit.kilocalories_per_mole/unit.angstrom**2, 'collision_rate': 2/unit.picoseconds, 'timestep': 1*unit.femtoseconds},
{'EOM': 'minimize', 'n_steps': 10000, 'temperature': 300*unit.kelvin, 'ensemble': None, 'restraint_selection': 'protein and backbone', 'force_constant': 0.1*unit.kilocalories_per_mole/unit.angstrom**2, 'collision_rate': 2/unit.picoseconds, 'timestep': 1*unit.femtoseconds},
{'EOM': 'MD', 'n_steps': 100000, 'temperature': 300*unit.kelvin, 'ensemble': 'NPT', 'restraint_selection': 'protein and backbone', 'force_constant': 0.1*unit.kilocalories_per_mole/unit.angstrom**2, 'collision_rate': 2/unit.picoseconds, 'timestep': 1*unit.femtoseconds},
{'EOM': 'MD', 'n_steps': 100000, 'temperature': 300*unit.kelvin, 'ensemble': 'NPT', 'restraint_selection': 'protein and backbone', 'force_constant': 0.01*unit.kilocalories_per_mole/unit.angstrom**2, 'collision_rate': 2/unit.picoseconds, 'timestep': 1*unit.femtoseconds},
{'EOM': 'MD', 'n_steps': 100000, 'temperature': 300*unit.kelvin, 'ensemble': 'NPT', 'restraint_selection': 'protein and backbone', 'force_constant': 0.001*unit.kilocalories_per_mole/unit.angstrom**2, 'collision_rate': 2/unit.picoseconds, 'timestep': 1*unit.femtoseconds},
{'EOM': 'MD', 'n_steps': 2500000, 'temperature': 300*unit.kelvin, 'ensemble': 'NPT', 'restraint_selection': None, 'force_constant': 0*unit.kilocalories_per_mole/unit.angstrom**2, 'collision_rate': 2/unit.picoseconds, 'timestep': 2*unit.femtoseconds},
]
run_gentle_equilibration(topology, positions, system, stages, args.outfile)
@jchodera mentioned that we may want to capture this in openmmtools as a utility method (eg into a GentleEquilibration factory class in openmmtools/equilibration.py) in case we ever need to use this in different contexts.
I just realized that the force constants defined here:
stages = [
{'EOM': 'minimize', 'n_steps': 10000, 'temperature': 300*unit.kelvin, 'ensemble': None, 'restraint_selection': 'protein and not type H', 'force_constant': 1*unit.kilocalories_per_mole/unit.angstrom**2, 'collision_rate': 2/unit.picoseconds, 'timestep': 1*unit.femtoseconds},
{'EOM': 'MD_interpolate', 'n_steps': 100000, 'temperature': 100*unit.kelvin, 'temperature_end': 300*unit.kelvin, 'ensemble': 'NVT', 'restraint_selection': 'protein and not type H', 'force_constant': 1*unit.kilocalories_per_mole/unit.angstrom**2, 'collision_rate': 10/unit.picoseconds, 'timestep': 1*unit.femtoseconds},
{'EOM': 'MD', 'n_steps': 100000, 'temperature': 300, 'ensemble': 'NPT', 'restraint_selection': 'protein and not type H', 'force_constant': 1*unit.kilocalories_per_mole/unit.angstrom**2, 'collision_rate': 10/unit.picoseconds, 'timestep': 1*unit.femtoseconds},
{'EOM': 'MD', 'n_steps': 250000, 'temperature': 300*unit.kelvin, 'ensemble': 'NPT', 'restraint_selection': 'protein and not type H', 'force_constant': 0.1*unit.kilocalories_per_mole/unit.angstrom**2, 'collision_rate': 2/unit.picoseconds, 'timestep': 1*unit.femtoseconds},
{'EOM': 'minimize', 'n_steps': 10000, 'temperature': 300*unit.kelvin, 'ensemble': None, 'restraint_selection': 'protein and backbone', 'force_constant': 0.1*unit.kilocalories_per_mole/unit.angstrom**2, 'collision_rate': 2/unit.picoseconds, 'timestep': 1*unit.femtoseconds},
{'EOM': 'MD', 'n_steps': 100000, 'temperature': 300*unit.kelvin, 'ensemble': 'NPT', 'restraint_selection': 'protein and backbone', 'force_constant': 0.1*unit.kilocalories_per_mole/unit.angstrom**2, 'collision_rate': 2/unit.picoseconds, 'timestep': 1*unit.femtoseconds},
{'EOM': 'MD', 'n_steps': 100000, 'temperature': 300*unit.kelvin, 'ensemble': 'NPT', 'restraint_selection': 'protein and backbone', 'force_constant': 0.01*unit.kilocalories_per_mole/unit.angstrom**2, 'collision_rate': 2/unit.picoseconds, 'timestep': 1*unit.femtoseconds},
{'EOM': 'MD', 'n_steps': 100000, 'temperature': 300*unit.kelvin, 'ensemble': 'NPT', 'restraint_selection': 'protein and backbone', 'force_constant': 0.001*unit.kilocalories_per_mole/unit.angstrom**2, 'collision_rate': 2/unit.picoseconds, 'timestep': 1*unit.femtoseconds},
{'EOM': 'MD', 'n_steps': 2500000, 'temperature': 300*unit.kelvin, 'ensemble': 'NPT', 'restraint_selection': None, 'force_constant': 0*unit.kilocalories_per_mole/unit.angstrom**2, 'collision_rate': 2/unit.picoseconds, 'timestep': 2*unit.femtoseconds},
]
do not match those outlined in the screenshotted equilibration protocol -- all force constants in this list should be multiplied by 100.
We should try to add this at least in a utility part of the code that we can optionally use if desired.
