Skip to content

Commit

Permalink
80-adding-solvate-option-in-setupandrun-tool (#84)
Browse files Browse the repository at this point in the history 
* adding solvate option in setup and run tool: 1. Adding the parameter in the input arg schema, 2. Adding the modeller functionality that acutally solvates the box if needed

* adding solvate in the check_parameters as part of the tool validation. 2. Change the document strings to the ('') format

* Adding modeller.addSolvent to the final script output. For simplicity, it uses padding

* Saving topology just before simulations for post analysis
  • Loading branch information
@Jgmedina95
Jgmedina95 authored Feb 14, 2024
1 parent eaa4f17 commit dae6379
Showing 1 changed file with 140 additions and 46 deletions.
186 changes: 140 additions & 46 deletions mdagent/tools/base_tools/simulation_tools/setup_and_run.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -387,6 +387,15 @@ def _setup_and_run_simulation(self, query, PathRegistry):
simulation = Simulation(modeller.topology, system, integrator)
simulation.context.setPositions(modeller.positions)
simulation.minimizeEnergy()
# save initial positions to registry
file_name = "initial_positions.pdb"
with open(file_name, "w") as f:
PDBFile.writeFile(
simulation.topology,
simulation.context.getState(getPositions=True).getPositions(),
f,
)
print("Initial Positions saved to initial_positions.pdb")
simulation.reporters.append(PDBReporter(f"{name}.pdb", 1000))
# reporter_args = {"reportInterval": 1000}
reporter_args = {}
Expand Down Expand Up @@ -579,6 +588,7 @@ class SetUpandRunFunctionInput(BaseModel):
"constraints": "None",
"rigidWater": False,
"constraintTolerance": None,
"solvate": False,
},
description=(
"Parameters for the openmm system. "
Expand All @@ -593,6 +603,8 @@ class SetUpandRunFunctionInput(BaseModel):
"None, HBonds, AllBonds or OnlyWater."
"For rigidWater, you can choose from the following:\n"
"True, False.\n"
"Finally, if you want to solvate the system, before the simulation,"
"you can set solvate to True.\n"
"Example1:\n"
"{'nonbondedMethod': 'NoCutoff',\n"
"'constraints': 'None',\n"
Expand All @@ -602,7 +614,8 @@ class SetUpandRunFunctionInput(BaseModel):
"'nonbondedCutoff': 1.0,\n"
"'constraints': 'HBonds',\n"
"'rigidWater': True,\n"
"'constraintTolerance': 0.00001} "
"'constraintTolerance': 0.00001,\n"
"'solvate': True} "
),
)
integrator_params: Dict[str, Any] = Field(
Expand Down Expand Up @@ -670,6 +683,7 @@ def __init__(
"constraints": AllBonds,
"rigidWater": True,
"constraintTolerance": 0.000001,
"solvate": False,
}
self.sim_params = self.params.get("simmulation_params", None)
if self.sim_params is None:
Expand Down Expand Up @@ -735,12 +749,12 @@ def create_simulation(self):
print("Creating simulation...")
st.markdown("Creating simulation", unsafe_allow_html=True)
self.simulation = Simulation(
self.pdb.topology,
self.modeller.topology,
self.system,
self.integrator,
Platform.getPlatformByName("CPU"),
)
self.simulation.context.setPositions(self.pdb.positions)
self.simulation.context.setPositions(self.modeller.positions)

# TEMPORARY FILE MANAGEMENT OR PATH REGISTRY MAPPING
if self.save:
Expand Down Expand Up @@ -821,6 +835,7 @@ def _create_system(
constraints="None",
rigidWater=False,
constraintTolerance=None,
solvate=False,
**kwargs,
):
# Create a dictionary to hold system parameters
Expand Down Expand Up @@ -850,8 +865,26 @@ def _create_system(

# if use_constraint_tolerance:
# constraintTolerance = system_params.pop('constraintTolerance')

system = forcefield.createSystem(pdb.topology, **system_params)
self.modeller = Modeller(pdb.topology, pdb.positions)
if solvate:
try:
self.modeller.addSolvent(forcefield)
except ValueError as e:
print("Error adding solvent", type(e).__name__, "–", e)
if "No Template for" in str(e):
raise ValueError(str(e))
except AttributeError as e:
print("Error adding solvent: ", type(e).__name__, "–", e)
print("Trying to add solvent with 1 nm padding")
if "NoneType" and "value_in_unit" in str(e):
try:
self.modeller.addSolvent(forcefield, padding=1 * nanometers)
except Exception as e:
print("Error adding solvent", type(e).__name__, "–", e)
raise (e)
system = forcefield.createSystem(self.modeller.topology, **system_params)
else:
system = forcefield.createSystem(self.modeller.topology, **system_params)

return system

Expand Down Expand Up @@ -879,6 +912,7 @@ def unit_to_string(unit):
ewaldErrorTolerance = {self.sys_params.get("ewaldErrorTolerance", 0.0005)}
constraintTolerance = self.sys_params.get("constraintTolerance", None)
hydrogenMass = self.sys_params.get("hydrogenMass", None)
solvate = self.sys_params.get("solvate", False)

integrator_type = self.int_params.get("integrator_type", "LangevinMiddle")
friction = self.int_params.get("Friction", 1.0 / picoseconds)
Expand Down Expand Up @@ -956,46 +990,51 @@ def unit_to_string(unit):
# Simulate
print('Building system...')
topology = pdb.topology
positions = pdb.positions
modeller = Modeller(pdb.topology, pdb.positions)
"""
if solvate:
script_content += (
"""modeller.addSolvent(forcefield, padding=1*nanometers)"""
)

if nonbondedMethod == NoCutoff:
if hydrogenMass:
script_content += """
system = forcefield.createSystem(topology, nonbondedMethod=nonbondedMethod,
constraints=constraints, rigidWater=rigidWater, hydrogenMass=hydrogenMass)
system = forcefield.createSystem(modeller.topology,
nonbondedMethod=nonbondedMethod, constraints=constraints,
rigidWater=rigidWater, hydrogenMass=hydrogenMass)
"""
else:
script_content += """
system = forcefield.createSystem(topology, nonbondedMethod=nonbondedMethod,
constraints=constraints, rigidWater=rigidWater)
system = forcefield.createSystem(modeller.topology,
nonbondedMethod=nonbondedMethod, constraints=constraints,
rigidWater=rigidWater)
"""
if nonbondedMethod == CutoffNonPeriodic or nonbondedMethod == CutoffPeriodic:
if hydrogenMass:
script_content += """
system = forcefield.createSystem(topology,
nonbondedMethod=nonbondedMethod,
nonbondedCutoff=nonbondedCutoff, constraints=constraints,
rigidWater=rigidWater, hydrogenMass=hydrogenMass)
system = forcefield.createSystem(modeller.topology,
nonbondedMethod=nonbondedMethod, nonbondedCutoff=nonbondedCutoff,
constraints=constraints, rigidWater=rigidWater,
hydrogenMass=hydrogenMass)
"""
else:
script_content += """
system = forcefield.createSystem(topology,
nonbondedMethod=nonbondedMethod,
nonbondedCutoff=nonbondedCutoff, constraints=constraints,
rigidWater=rigidWater)
system = forcefield.createSystem(modeller.topology,
nonbondedMethod=nonbondedMethod, nonbondedCutoff=nonbondedCutoff,
constraints=constraints, rigidWater=rigidWater)
"""
if nonbondedMethod == PME:
if hydrogenMass:
script_content += """
system = forcefield.createSystem(topology,
system = forcefield.createSystem(modeller.topology,
nonbondedMethod=nonbondedMethod,
nonbondedCutoff=nonbondedCutoff, ewaldErrorTolerance=ewaldErrorTolerance,
constraints=constraints, rigidWater=rigidWater, hydrogenMass=hydrogenMass)
"""
else:
script_content += """
system = forcefield.createSystem(topology,
system = forcefield.createSystem(modeller.topology,
nonbondedMethod=nonbondedMethod,
nonbondedCutoff=nonbondedCutoff, ewaldErrorTolerance=ewaldErrorTolerance,
constraints=constraints, rigidWater=rigidWater)
Expand All @@ -1009,14 +1048,14 @@ def unit_to_string(unit):
script_content += """
integrator = LangevinMiddleIntegrator(temperature, friction, dt)
integrator.setConstraintTolerance(constraintTolerance)
simulation = Simulation(topology, system, integrator, platform)
simulation.context.setPositions(positions)
simulation = Simulation(modeller.topology, system, integrator, platform)
simulation.context.setPositions(modeller.positions)
"""
if integrator_type == "LangevinMiddle" and constraints == "None":
script_content += """
integrator = LangevinMiddleIntegrator(temperature, friction, dt)
simulation = Simulation(topology, system, integrator, platform)
simulation.context.setPositions(positions)
simulation = Simulation(modeller.topology, system, integrator, platform)
simulation.context.setPositions(modeller.positions)
"""

script_content += """
Expand Down Expand Up @@ -1065,6 +1104,16 @@ def run(self):

self.simulation.minimizeEnergy()
print("Minimization complete!")
top_name = f"files/pdb/{self.sim_id}_initial_positions.pdb"
top_description = f"Initial positions for simulation {self.sim_id}"
with open(top_name, "w") as f:
PDBFile.writeFile(
self.simulation.topology,
self.simulation.context.getState(getPositions=True).getPositions(),
f,
)
self.path_registry.map_path(f"top_{self.sim_id}", top_name, top_description)
print("Initial Positions saved to initial_positions.pdb")
st.markdown("Minimization complete! Equilibrating...", unsafe_allow_html=True)
print("Equilibrating...")
_temp = self.int_params["Temperature"]
Expand Down Expand Up @@ -1148,7 +1197,14 @@ def _run(self, **input_args):
print("simulation set!")
st.markdown("simulation set!", unsafe_allow_html=True)
except ValueError as e:
return str(e) + f"This were the inputs {input_args}"
msg = str(e) + f"This were the inputs {input_args}"
if "No template for" in msg:
msg += (
"This error is likely due to non standard residues "
"in the protein, if you havent done it yet, try "
"cleaning the pdb file using the cleaning tool"
)
return msg
except FileNotFoundError:
return f"File not found, check File id. This were the inputs {input_args}"
except OpenMMException as e:
Expand Down Expand Up @@ -1383,9 +1439,12 @@ def _process_parameters(self, user_params, param_type="system_params"):
try:
processed_params[key] = float(value)
except TypeError as e:
error_msg += f"""Invalid ewaldErrorTolerance: {e}.
If you are using null or None, just dont include
as part of the parameters.\n"""
error_msg += (
f"Invalid ewaldErrorTolerance: {e}. "
"If you are using null or None, "
"just dont include it "
"as part of the parameters.\n"
)
if key == "constraints":
try:
if type(value) == str:
Expand All @@ -1398,14 +1457,19 @@ def _process_parameters(self, user_params, param_type="system_params"):
elif value == "HAngles":
processed_params[key] = HAngles
else:
error_msg += f"""Invalid constraints. got {value}.
Try using None, HBonds, AllBonds,
HAngles"""
error_msg += (
f"Invalid constraints: Got {value}. "
"Try using None, HBonds, AllBonds or "
"HAngles\n"
)
else:
processed_params[key] = value
except TypeError as e:
error_msg += f"""Invalid constraints: {e}. If you are using
null or None, just dont include as part of the parameters.\n"""
error_msg += (
f"Invalid constraints: {e}. If you are using "
"null or None, just dont include as "
"part of the parameters.\n"
)
if key == "rigidWater" or key == "rigidwater":
if type(value) == bool:
processed_params[key] = value
Expand All @@ -1414,17 +1478,42 @@ def _process_parameters(self, user_params, param_type="system_params"):
elif value == "False":
processed_params[key] = False
else:
error_msg += f"""Invalid rigidWater. got {value}.
Try using True or False.\n"""
error_msg += (
f"Invalid rigidWater: got {value}. "
"Try using True or False.\n"
)
if key == "constraintTolerance" or key == "constrainttolerance":
try:
processed_params[key] = float(value)
except ValueError as e:
error_msg += f"Invalid constraintTolerance. {e}."
error_msg += f"Invalid constraintTolerance: {e}."
except TypeError as e:
error_msg += f"""Invalid constraintTolerance. {e}. If
constraintTolerance is null or None,
just dont include as part of the parameters.\n"""
error_msg += (
f"Invalid constraintTolerance: {e}. If "
"constraintTolerance is null or None, "
"just dont include as part of "
"the parameters.\n"
)
if key == "solvate":
try:
if type(value) == bool:
processed_params[key] = value
elif value == "True":
processed_params[key] = True
elif value == "False":
processed_params[key] = False
else:
error_msg += (
f"Invalid solvate: got {value}. "
"Use either True or False.\n"
)
except TypeError as e:
error_msg += (
f"Invalid solvate: {e}. If solvate is null or "
"None, just dont include as part of "
"the parameters.\n"
)

return processed_params, error_msg
if param_type == "integrator_params":
for key, value in user_params.items():
Expand All @@ -1438,9 +1527,11 @@ def _process_parameters(self, user_params, param_type="system_params"):
elif value == "Brownian" or value == BrownianIntegrator:
processed_params[key] = "Brownian"
else:
error_msg += f"""\nInvalid integrator_type. got {value}.
Try using LangevinMiddle, Langevin,
Verlet, or Brownian."""
error_msg += (
f"Invalid integrator_type: got {value}. "
"Try using LangevinMiddle, Langevin, "
"Verlet, or Brownian.\n"
)
if key == "Temperature" or key == "temperature":
temperature, msg = self.parse_temperature(value)
processed_params[key] = temperature
Expand Down Expand Up @@ -1469,8 +1560,10 @@ def _process_parameters(self, user_params, param_type="system_params"):
elif value == "NVE":
processed_params[key] = "NVE"
else:
error_msg += f"""Invalid Ensemble. got {value}.
Try using NPT, NVT, or NVE."""
error_msg += (
f"Invalid Ensemble. got {value}. "
"Try using NPT, NVT, or NVE.\n"
)

if key == "Number of Steps" or key == "number of steps":
processed_params[key] = int(value)
Expand Down Expand Up @@ -1501,6 +1594,7 @@ def check_system_params(cls, values):
"constraints": AllBonds,
"rigidWater": True,
"constraintTolerance": 0.00001,
"solvate": False,
}
integrator_params = values.get("integrator_params")
if integrator_params:
Expand Down Expand Up @@ -1618,7 +1712,7 @@ def check_system_params(cls, values):
if file not in FORCEFIELD_LIST:
error_msg += "The forcefield file is not present"

save = values.get("final", False)
save = values.get("save", True)
if type(save) != bool:
error_msg += "save must be a boolean value"

Expand Down

0 comments on commit dae6379

Please sign in to comment.