Skip to content

Commit

Permalink
Change GMX settings to reproduce OpenMM (#12)
Browse files Browse the repository at this point in the history 
* Separate solvation and parameter assignment

* Add espaloma model

* Parametrize solute with espaloma

* Linters

* Change GMX settings to reproduce OpenMM

* Fix minor variable name error

* Fix minor settings issues

---------

Co-authored-by: Anika Friedman <[email protected]>
  • Loading branch information
@chapincavender @ajfriedman22
chapincavender and ajfriedman22 authored Oct 19, 2023
1 parent db9e984 commit 53b78f9
Show file tree
Hide file tree
Showing 5 changed files with 291 additions and 244 deletions.
274 changes: 152 additions & 122 deletions proteinbenchmark/gmx_simulation.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,8 +5,6 @@
import numpy
from openff.units import unit

from proteinbenchmark.utilities import exists_and_not_empty, read_xml


class GMXSimulation:
"""A class representing a simulation in GROMACS."""
Expand All @@ -16,17 +14,16 @@ def __init__(
gmx_executable: str,
parametrized_system: str,
initial_coords_file: str,
save_state_prefix: str,
setup_prefix: str,
output_prefix: str,
nonbonded_cutoff: unit.Quantity,
vdw_switch_width: unit.Quantity,
temperature: unit.Quantity,
pressure: unit.Quantity,
langevin_friction: unit.Quantity,
barostat_constant: unit.Quantity,
thermostat_constant: unit.Quantity,
timestep: unit.Quantity,
traj_length: unit.Quantity,
frame_length: unit.Quantity,
restraints_present: bool,
load_state_prefix: str = None,
):
"""
Initializes the simulation parameters and checks units.
Expand All @@ -40,16 +37,22 @@ def __init__(
The path to the parametrized system as a top file.
initial_coords_file
Path to GRO file used to set initial coordinates.
save_state_prefix
Path prefix for files to write serialized simulation states.
output_prefix
Path prefix for GMX output files.
nonbonded_cutoff
The cutoff for the Lennard-Jones potential and PME direct space
summation.
vdw_switch_width
The distance from the nonbonded cutoff at which to apply the
switching function.
temperature
The target temperature of the Langevin thermostat.
pressure
The target pressure of the Monte Carlo barostat.
langevin_friction
The collision frequency of the Langevin integrator.
barostat_constant
The time constant for pressure coupling
thermostat_constant
The time constant for temperature coupling
timestep
The timestep for the Langevin integrator.
traj_length
Expand All @@ -64,17 +67,19 @@ def __init__(
self.gmx_executable = gmx_executable
self.parametrized_system = parametrized_system
self.initial_coords_file = initial_coords_file
self.save_state_prefix = save_state_prefix
self.setup_prefix = setup_prefix
self.load_prefix = load_state_prefix
self.output_prefix = output_prefix

# Check units of arguments
if not nonbonded_cutoff.is_compatible_with(unit.nanometer):
raise ValueError("nonbonded_cutoff does not have units of Length")
if not vdw_switch_width.is_compatible_with(unit.nanometer):
raise ValueError("vdw_switch_width does not have units of Length")
if not temperature.is_compatible_with(unit.kelvin):
raise ValueError("temperature does not have units of Temperature")

if not pressure.is_compatible_with(unit.atmosphere):
raise ValueError("pressure does not have units of Mass Length^-1 Time^-2")

if not langevin_friction.is_compatible_with(unit.picosecond**-1):
raise ValueError("langevin_friction does not have units of Time^-1")
if not timestep.is_compatible_with(unit.picosecond):
raise ValueError("timestep does not have units of Time")
if not traj_length.is_compatible_with(unit.picosecond):
Expand All @@ -83,162 +88,187 @@ def __init__(
raise ValueError("frame_length does not have units of Time")

# Ensure proper units for output
self.nonbonded_cutoff = nonbonded_cutoff.m_as(unit.nanometer)
self.switch_distance = (nonbonded_cutoff - vdw_switch_width).m_as(
unit.nanometer
)
self.temperature = temperature.m_as(unit.kelvin) # Kelvin
self.pressure = pressure.m_as(unit.atmosphere) # atm
self.thermostat_constant = thermostat_constant.m_as(unit.picosecond) # ps
self.pressure = pressure.m_as(unit.bar) # atm
self.thermostat_constant = (1.0 / langevin_friction).m_as(unit.picosecond) # ps
self.barostat_constant = barostat_constant.m_as(unit.picosecond) # ps
self.timestep = timestep.m_as(unit.picosecond) # fs
self.n_steps = int(numpy.round(traj_length / timestep)) # steps
self.output_frequency = int(numpy.round(frame_length / timestep)) # steps
self.restraints_present = restraints_present # str

def setup_simulation(self):
def setup_simulation(self, continuation: bool = False):
"""
Set up a GROMACS simulation.
Set up a GROMACS simulation by writing the MDP file.
Parameters
----------
continuation
Continue from an existing simulation.
"""

# create mdp file
mdp_file = self.save_state_prefix + ".mdp"
mdp_file = self.output_prefix + ".mdp"

# Write settings to MDP File
with open(mdp_file, "w") as mdp_file_w:
# Apply position restraints for equilibration
if self.restraints_present == "NPT":
mdp_file_w.write("define=-DPOSRES\n" + "continuation=no\n")
if continuation:
mdp_file_w.write("continuation=yes\ngen_vel = no\n")

else:
mdp_file_w.write("continuation=yes\n")
mdp_file_w.write(
"continuation=no\n"
"gen_vel = yes\n"
f"gen_temp = {self.temperature}\n"
"gen_seed = -1\n"
)

mdp_file_w.write(
f"integrator = md\nnsteps={self.n_steps}\n"
f"dt = {self.timestep}\nnstenergy = 5000\nnstlog = 5000\n"
f"integrator = sd\n"
f"nsteps={self.n_steps}\n"
f"dt = {self.timestep}\n"
f"nstenergy = {self.output_frequency}\n"
f"nstlog = {self.output_frequency}\n"
f"nstxout-compressed = {self.output_frequency}\n"
"constraint_algorithm = lincs\nlincs-warnangle = 45\n"
"constraints = h-bonds\nlincs_iter = 2\nlincs_order = 4\n"
"cutoff-scheme = Verlet\nns_type = grid\nnstlist = 40\n"
"rlist = 0.9\nvdwtype = cutoff\nvdw-modifier = force-switch\n"
"rvdw-switch = 0.88\nrvdw = 0.9\ncoulombtype = PME\n"
"rcoulomb = 0.9\npme_order = 4\nfourierspacing = 0.16\n"
"tcoupl = V-rescale\ntc-grps = System\n"
"constraint_algorithm = lincs\n"
"lincs-warnangle = 45\n"
"constraints = h-bonds\n"
"lincs_iter = 2\n"
"lincs_order = 4\n"
"cutoff-scheme = Verlet\n"
"nstlist = 40\n"
"vdwtype = cut-off\n"
"vdw-modifier = potential-switch\n"
f"rvdw = {self.nonbonded_cutoff}\n"
f"rvdw-switch = {self.switch_distance}\n"
"coulombtype = PME\n"
f"rcoulomb = {self.nonbonded_cutoff}\n"
"pme_order = 4\n"
"fourierspacing = 0.16\n"
"tc-grps = System\n"
f"tau_t = {self.thermostat_constant}\n"
f"ref_t = {self.temperature}\npbc = xyz\nDispCorr = EnerPres\n"
f"pcoupltype = isotropic\ntau_p = {self.barostat_constant}\n"
f"ref_p = {self.pressure}\ncompressibility = 4.5e-5\n"
f"ref_t = {self.temperature}\n"
"pbc = xyz\n"
"DispCorr = EnerPres\n"
"pcoupl = C-rescale\n"
"pcoupltype = isotropic\n"
f"tau_p = {self.barostat_constant}\n"
f"ref_p = {self.pressure}\n"
"compressibility = 4.5e-5\n"
"refcoord_scaling = com\n"
)
if self.restraints_present == "NPT":
mdp_file_w.write(
f"gen_vel = yes\ngen_temp = {self.temperature}\n"
"gen_seed = -1\npcoupl = C-rescale\n"
)
else:
mdp_file_w.write("gen_vel = no\npcoupl = Parrinello-Rahman\n")

return mdp_file

def run(self):
def start_from_pdb(self):
"""
Start a new simulation initializing positions from a PDB and velocities
to random samples from a Boltzmann distribution at the simulation
temperature.
"""

# Generate MDP
mdp_file = self.setup_simulation()
mdp_file = self.setup_simulation(continuation=False)

# Generate TPR
if self.restraints_present == "NPT":
grompp = subprocess.run(
[
self.gmx_executable,
"grompp",
"-f",
mdp_file,
"-p",
self.parametrized_system,
"-c",
self.initial_coords_file,
"-r",
self.initial_coords_file,
"-maxwarn",
"2",
"-o",
str(self.save_state_prefix),
]
)

else:
print(self.gmx_executable)
grompp = subprocess.run(
[
self.gmx_executable,
"grompp",
"-f",
mdp_file,
"-p",
self.parametrized_system,
"-c",
self.initial_coords_file,
"-t",
f"{self.load_prefix}.cpt",
"-maxwarn",
"2",
"-o",
str(self.save_state_prefix),
]
)
grompp = subprocess.run(
[
self.gmx_executable,
"grompp",
"-f",
mdp_file,
"-p",
self.parametrized_system,
"-c",
self.initial_coords_file,
"-maxwarn",
"2",
"-o",
str(self.output_prefix),
]
)

# Run Simulation
if self.restraints_present == "NPT":
simulation = subprocess.run(
[
self.gmx_executable,
"mdrun",
"-deffnm",
str(self.save_state_prefix),
"-ntmpi",
"1",
]
)
# simulation = gmx.mdrun(
# input=grompp.output.file['-o'], runtime_args={'-ntmpi': '1'}
# )
else:
simulation = subprocess.run(
[
self.gmx_executable,
"mdrun",
"-deffnm",
str(self.save_state_prefix),
"-ntomp",
"1",
]
)
# simulation = gmx.mdrun(
# input=grompp.output.file['-o'], runtime_args={'-ntomp': '1'}
# )
# Run simulation
simulation = subprocess.run(
[
self.gmx_executable,
"mdrun",
"-deffnm",
str(self.output_prefix),
"-ntmpi",
"1",
]
)

def start_from_save_state(
self,
save_state_file: str,
checkpoint_file: str,
):
"""
Start a new simulation initializing positions and velocities from a
serialized OpenMM state.
GMX checkpoint.
Parameters
----------
save_state_file
Path to the serialized simulation state.
checkpoint_file
Path to the checkpoint file.
"""

# Generate MDP
mdp_file = self.setup_simulation(continuation=True)

# Generate TPR
grompp = subprocess.run(
[
self.gmx_executable,
"grompp",
"-f",
mdp_file,
"-p",
self.parametrized_system,
"-c",
self.initial_coords_file,
"-t",
checkpoint_file,
"-maxwarn",
"2",
"-o",
str(self.output_prefix),
]
)

# Run Simulation
simulation = subprocess.run(
[
self.gmx_executable,
"mdrun",
"-deffnm",
str(self.output_prefix),
"-ntomp",
"1",
]
)
# simulation = gmx.mdrun(
# input=grompp.output.file['-o'], runtime_args={'-ntomp': '1'}
# )

def resume_from_checkpoint(self):
"""
Resume an existing simulation from a GMX checkpoint.
"""

# Run Simulation
simulation = subprocess.run(
[
self.gmx_executable,
"mdrun",
"-deffnm",
str(self.save_state_prefix),
str(self.output_prefix),
"-cpi",
save_state_file,
f"{self.output_prefix}.cpt",
"-ntomp",
"1",
]
Expand Down
2 changes: 0 additions & 2 deletions proteinbenchmark/openmm_simulation.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -78,10 +78,8 @@ def __init__(
# Check units of arguments
if not temperature.unit.is_compatible(unit.kelvin):
raise ValueError("temperature does not have units of Temperature")

if not pressure.unit.is_compatible(unit.atmosphere):
raise ValueError("pressure does not have units of Mass Length^-1 Time^-2")

if not langevin_friction.unit.is_compatible(unit.picosecond**-1):
raise ValueError("langevin_friction does not have units of Time^-1")
if not timestep.unit.is_compatible(unit.picosecond):
Expand Down
Loading

0 comments on commit 53b78f9

Please sign in to comment.