Commit 8e3a3715 by cassandra

[feat] improved parameter selection

parent 86bbaa3a
......@@ -28,75 +28,6 @@ def _ensure_api_key(env_var: str, prompt_name: str) -> str | None:
return key
def update_mdp_file(src_path, dest_path, md_temp, md_duration):
def extract_comment(part):
comment = ""
if ";" in part:
_, comment = part.split(";", 1)
comment = ";" + comment
return comment
with open(src_path, "r") as f:
lines = f.readlines()
dt = None
for line in lines:
if line.strip().startswith("dt"):
dt = float(line.split("=")[1].split(";")[0].strip())
break
if dt is None:
raise ValueError(f"dt not found in {src_path}")
# Calculate nsteps
# md_duration is in ns, dt is in ps, so convert: 1 ns = 1000 ps
nsteps = int((md_duration * 1000) / dt)
# Update the lines
new_lines = []
for line in lines:
if line.strip().startswith("ref_t"):
# replace the value with md_temp
parts = line.split("=")
comment = extract_comment(parts[1])
new_line = f"{parts[0]}= {md_temp} {md_temp} {comment}\n"
elif line.strip().startswith("nsteps"):
parts = line.split("=")
comment = extract_comment(parts[1])
new_line = f"{parts[0]}= {nsteps} {comment}\n"
else:
new_line = line
new_lines.append(new_line)
with open(dest_path, "w") as f:
f.writelines(new_lines)
def process_mdp_files(mdp_dir, sandbox_dir, md_temp, md_duration):
if not mdp_dir.exists():
raise FileNotFoundError(f"MDP files directory not found: {mdp_dir}")
mdp_files = os.listdir(mdp_dir)
for file in mdp_files:
full_file_name = os.path.join(mdp_dir, file)
dest_file_name = os.path.join(sandbox_dir, file)
if not os.path.isfile(full_file_name):
continue
if file in ("md.mdp", "npt.mdp", "nvt.mdp"):
update_mdp_file(full_file_name, dest_file_name, md_temp, md_duration)
else:
shutil.copy(full_file_name, sandbox_dir)
@dataclass
class CommandLineArgs:
"""
......@@ -112,6 +43,12 @@ class CommandLineArgs:
ligand: str | None = None
"Ligand ID (optional; defaults to no ligand)."
temp: float | None = None
"Simulation temperature in Kelvin."
duration: float | None = None
"Simulation length in nanoseconds."
model_supports_system_messages: bool = True
......@@ -123,11 +60,7 @@ def main(config: CommandLineArgs):
sandbox_dir = constants.DATA_DIR / run_name
sandbox_dir.mkdir(parents=True, exist_ok=True)
root_logger.info("autoMD - your assistant for running molecular dynamics")
root_logger.info("================")
root_logger.info(
"I can read, fetch, and prepare PDB files to run MD simulations. Is there a particular system I can help you with today?"
)
root_logger.info("DynaMate - your assistant for running molecular dynamics")
try:
_ensure_api_key("OPENROUTER_API_KEY", "OPENROUTER_API_KEY")
......@@ -135,31 +68,27 @@ def main(config: CommandLineArgs):
root_logger.error(str(e))
return
root_logger.info("\n=== 1. Starting PrepAgent (Planning & Parameter Determination) ===")
root_logger.info("\n=== Starting PrepAgent (Planning & Parameter Determination) ===")
prep_agent = PrepAgent(
model_name=config.model,
temperature=constants.TEMPERATURE,
sandbox_dir=sandbox_dir,
pdb_id=config.pdb_id,
ligand_name=config.ligand,
md_temp=config.temp,
md_duration=config.duration,
model_supports_system_messages=config.model_supports_system_messages,
)
# Start the prep agent to understand request from user and generate a plan
prep_agent.setup_tools()
pdb_file_path, ligand_name, plan, llm_cost = prep_agent.run()
root_logger.info("PrepAgent completed. Plan generated.")
# Copy required mdp_files into sandbox for the agent and update ref_t and nsteps
md_temp, md_duration = plan["parameters"]["temperature_k"], plan["parameters"]["duration_ns"]
md_duration=0.01
root_logger.info(f"Applying parameters: Temp={md_temp}K, Duration={md_duration}ns")
process_mdp_files(constants.MDP_FILES, sandbox_dir, md_temp, md_duration)
root_logger.info("\n=== 2. Starting MDAgent (Execution & Tool Loop) ===")
root_logger.info("\n=== Starting MDAgent (Execution & Tool Loop) ===")
# Copy PDB into run directory
md_agent = MDAgent(
model_name=config.model,
temperature=constants.TEMPERATURE,
......@@ -167,6 +96,8 @@ def main(config: CommandLineArgs):
structure_path=pdb_file_path,
pdb_id=Path(pdb_file_path).stem,
ligand_name=ligand_name,
md_temp=md_temp,
md_duration=md_duration,
model_supports_system_messages=config.model_supports_system_messages,
plan=plan,
)
......@@ -175,14 +106,12 @@ def main(config: CommandLineArgs):
# Run the MD pipeline (handles user input and retries internally)
result = md_agent.run()
# Print result and summary
if not result:
root_logger.error("=== MD Pipeline failed or incomplete ===")
else:
root_logger.info("=== MD Pipeline completed successfully ===")
if __name__ == "__main__":
config = tyro.cli(CommandLineArgs)
main(config)
#!/bin/bash
# load your environment
source ~/miniforge3/bin/activate dynagent
source ~/miniforge3/bin/activate dynamate
# Activate GROMACS environment
if [ -f /usr/local/gromacs/bin/GMXRC ]; then
......
......@@ -28,6 +28,8 @@ class BaseAgent(ABC):
sandbox_dir: str,
pdb_id: str | None = None,
ligand_name: str | None = None,
md_temp: float | None = None,
md_duration: float | None = None,
model_supports_system_messages: bool = True,
):
self.model_name = model_name
......@@ -35,6 +37,8 @@ class BaseAgent(ABC):
self.sandbox_dir = Path(sandbox_dir)
self.pdb_id = pdb_id
self.ligand_name = ligand_name
self.md_temp = md_temp
self.md_duration = md_duration
self.model_supports_system_messages = model_supports_system_messages
......@@ -66,7 +70,6 @@ class BaseAgent(ABC):
if not messages:
return []
# Start from the end
block = []
i = len(messages) - 1
......@@ -77,11 +80,7 @@ class BaseAgent(ABC):
role = m.get("role")
# Stop conditions for different last-message types:
# 1. If the last message is a 'tool' message, we must also include the
# assistant tool call that triggered it.
if role == "tool":
# include until we find the matching assistant tool call
i -= 1
while i >= 0 and messages[i].get("role") != "assistant":
block.insert(0, messages[i])
......@@ -89,22 +88,13 @@ class BaseAgent(ABC):
if i >= 0:
block.insert(0, messages[i])
return block
# 2. If last message is assistant tool call (function call),
# we keep only that one so the model knows what it's continuing.
if role == "assistant" and m.get("tool_calls"):
return block
# 3. If the last message is a normal assistant message,
# keep it and the preceding user message.
if role == "assistant":
# include the previous user message
if i - 1 >= 0 and messages[i-1].get("role") == "user":
block.insert(0, messages[i-1])
return block
# 4. If the last message is a user message,
# include previous assistant message.
if role == "user":
if i - 1 >= 0 and messages[i-1].get("role") == "assistant":
block.insert(0, messages[i-1])
......@@ -168,7 +158,6 @@ class BaseAgent(ABC):
tool_output = None
# try:
self._validate_tool_path(tool_input)
func = TOOL_MAP.get(tool_name)
......@@ -179,15 +168,6 @@ class BaseAgent(ABC):
passed = self._additional_check_for_errors_tool_output(tool_name, tool_output)
return {"ok": passed, "output": tool_output}
# except Exception as e:
# tb = traceback.format_exc()
# error_msg = (
# f"Tool '{tool_name}' failed with exception:\n{e}\n\n"
# f"Traceback:\n{tb}\n"
# "Decide whether to retry, fix the issue, or skip this step."
# )
# self.logger.error(error_msg)
# return {"ok": False, "output": error_msg}
def _format_tool_usage_ouput(self, id_, tool_name, arguments, output):
......@@ -203,7 +183,7 @@ class BaseAgent(ABC):
def to_dict_safe(msg):
if isinstance(msg, dict):
return msg
elif hasattr(msg, "model_dump"): # Pydantic/LiteLLM object
elif hasattr(msg, "model_dump"):
return msg.model_dump()
else:
return str(msg)
......
......@@ -7,7 +7,7 @@ from src.agents.agent import BaseAgent, ToolOutputError
from src.prompts import MD_SYSTEM_PROMPT
from src.tools import tool_schema
litellm.drop_params = True # avoid problems with setting temp on GPT-5
litellm.drop_params = True
class MDAgent(BaseAgent):
......@@ -31,11 +31,20 @@ class MDAgent(BaseAgent):
structure_path,
pdb_id=None,
ligand_name=None,
md_temp=None,
md_duration=None,
model_supports_system_messages=True,
plan: Dict[str, Any] = None,
):
super().__init__(
model_name, temperature, sandbox_dir, pdb_id, ligand_name, model_supports_system_messages
model_name,
temperature,
sandbox_dir,
pdb_id,
ligand_name,
md_temp,
md_duration,
model_supports_system_messages
)
self.structure_path = Path(structure_path)
......@@ -249,7 +258,7 @@ class MDAgent(BaseAgent):
success = self._pipeline_successful()
if self.ligand_name and success:
user_prompt = "Would you like me to calculate the free energy of binding for your protein-ligand system using the MMPBSA tool? (yes/no) \n"
user_prompt = "\n==========\nWould you like me to calculate the free energy of binding for your protein-ligand system using the MMPBSA tool? (yes/no) \n\n"
user_answer = input(user_prompt).strip().lower()
......
......@@ -9,7 +9,7 @@ from src.agents.agent import BaseAgent
from src.prompts import PREP_SYSTEM_PROMPT
litellm.drop_params = True # avoid problems with setting temp on GPT-5
litellm.drop_params = True
class Tool(BaseModel):
......@@ -26,10 +26,19 @@ class PrepAgent(BaseAgent):
sandbox_dir,
pdb_id=None,
ligand_name=None,
md_temp=None,
md_duration=None,
model_supports_system_messages=True,
):
super().__init__(
model_name, temperature, sandbox_dir, pdb_id, ligand_name, model_supports_system_messages
model_name,
temperature,
sandbox_dir,
pdb_id,
ligand_name,
md_temp,
md_duration,
model_supports_system_messages
)
self.messages: List[Dict[str, Any]] = []
......@@ -122,7 +131,7 @@ class PrepAgent(BaseAgent):
sys.exit(1)
def _find_simulation_temperature(self):
temperature = None
temperature = self.md_temp
while temperature is None:
user_input = "Please pick a suitable temperature (in Kelvins) for running the molecular dynamics simulation. Also provide a rational for why you selected this temperature."
......@@ -146,7 +155,7 @@ class PrepAgent(BaseAgent):
return temperature
def _calculate_duration(self):
duration = None
duration = self.md_duration
while duration is None:
user_input = "Please pick a suitable simulation duration (in nanoseconds) for running a short molecular dynamics of this system. Also provide a rational for why you selected this duration. Note that dynamics simulations of 10 ns take 1 hour to complete, so keep the experiment brief (less than 1 ns duration)."
......@@ -218,18 +227,20 @@ class PrepAgent(BaseAgent):
prompt = f"I would like to run molecular dynamics for the system {user_input}. If a PDB has not been uploaded, use the tools available to fetch and prepare the PDB for {user_input}."
self.logger.info(f"User input: {prompt}")
self.pdb_file_path = self._get_pdb_file_path(prompt)
self.logger.info(f"Thank you, I now have access to the structure information for protein {self.pdb_file_path}")
self.logger.info(f"I now have access to the structure information for protein {self.pdb_file_path}")
self._find_ligand()
temperature = self._find_simulation_temperature()
self.logger.info(f"Using simulation temperature: {temperature} K")
if self.md_temp is None:
self.md_temp = self._find_simulation_temperature()
self.logger.info(f"Using simulation temperature: {self.md_temp} K")
duration = self._calculate_duration()
self.logger.info(f"Using simulation duration: {duration} ns")
if self.md_duration is None:
self.md_duration = self._calculate_duration()
self.logger.info(f"Using simulation duration: {self.md_duration} ns")
# Build plan steps depending on ligand
plan = self._generate_plan(temperature, duration)
plan = self._generate_plan(self.md_temp, self.md_duration)
self.agent_plan = json.dumps(plan, indent=2)
self.logger.info(f"Generated plan: {self.agent_plan}")
......
......@@ -15,3 +15,5 @@ DATA_DIR = Path(__file__).resolve().parent.parent / "sandbox"
AGENT_LOGS = Path(__file__).resolve().parent.parent / "agent_logs"
JSON_LOG_FILE = AGENT_LOGS / "agent_runs.jsonl"
MMPBSA_ENV_DIR = Path("/path/to/your/envs/mmpbsa")
\ No newline at end of file
title = Protein-ligand complex MD simulation
; Run parameters
integrator = md ; leap-frog integrator
nsteps = 5000 ; 2 * 50,000 = 100 ps (0.01 ns)
nsteps = 50000 ; 2 * 500,000 = 1000 ps (0.1 ns)
dt = 0.002 ; 2 fs
; Output control
nstenergy = 5000 ; save energies every 10.0 ps
......
#!/bin/bash
if [ "$#" -lt 2 ]; then
echo "Usage: $0 input_gro npt_cpt_file"
if [ "$#" -lt 3 ]; then
echo "Usage: $0 input_gro npt_cpt_file log_file [ligand_name]"
exit 1
fi
......@@ -11,6 +11,65 @@ LOG_FILE="$3"
> $LOG_FILE
# Optional fourth argument
if [ "$#" -ge 4 ]; then
LIGNAME="$4"
else
LIGNAME=""
fi
#------- EDIT MD.MDP FILE ------------
# if no ions are present, update md.mdp such that Water_and_ions group is changed to Water only
# if ligand, Protein Water becomes Protein_ligand Water
if ! grep -q "Cl-" index.ndx && ! grep -q "Na+" index.ndx; then
md_file="md.mdp"
original="Protein Water_and_ions"
if [ -n "$LIGNAME" ]; then
replacement="Protein_$LIGNAME Water"
if grep "Protein_$LIGNAME" index.ndx; then
echo "Protein_$LIGNAME already in index.ndx"
else
echo -e "1 | 13\nq" | $GMX make_ndx -f em.gro -n index.ndx -o index.ndx
fi
else
replacement="Protein Water"
fi
if grep "$original" "$md_file"; then
sed -i "s|$original|$replacement|" "$md_file"
echo "$replacement added successfully to tc-grps group in $md_file."
else
echo "tc-grps line was not found in $md_file."
fi
else
echo "Ions present. Keeping Water_and_ions group."
fi
# if ions are present, but ligand exists, update md.mdp accordingly
if grep -q "Cl-" index.ndx || grep -q "Na+" index.ndx; then
if [ -n "$LIGNAME" ]; then
md_file="md.mdp"
original="Protein Water_and_ions"
replacement="Protein_$LIGNAME Water_and_ions"
if grep "Protein_$LIGNAME" index.ndx; then
echo "Protein_$LIGNAME already in index.ndx"
else
echo -e "1 | 13\nq" | $GMX make_ndx -f em.gro -n index.ndx -o index.ndx
fi
if grep "$original" "$md_file"; then
sed -i "s|$original|$replacement|" "$md_file"
echo "Protein_$LIGNAME added successfully to tc-grps group in $md_file."
else
echo "tc-grps line "Protein Water_and_ions" was not found in $md_file."
fi
fi
fi
#------- PRODUCTION MD ------------
if ! ls md.gro 1> /dev/null 2>&1; then
$GMX grompp -f md.mdp -c $INPUT_GRO -t $NPT_CPT_FILE -p topol.top -n index.ndx -o md.tpr >> $LOG_FILE 2>&1
......
......@@ -6,18 +6,19 @@ import parmed as pmd # type: ignore
import sys
import re
import os
from src import constants
def run_gmxMMPBSA(sandbox_dir: str, pdb_id: str, nsteps:str, nstxout_compressed:str, temp=str) -> str:
def run_gmxMMPBSA(sandbox_dir: str, pdb_id: str, nsteps:str, nstxout_compressed:str, md_temp:str) -> str:
nframes=int(nsteps)/int(nstxout_compressed)
os.makedirs(f"{sandbox_dir}/gmx_MMPBSA", exist_ok=True)
os.chdir(f"{sandbox_dir}/gmx_MMPBSA")
mmpbsa_infile = open('mmpbsa.in', 'w' )
MMPBSA_dir=f"{sandbox_dir}/gmx_MMPBSA"
mmpbsa_infile = open(f"{MMPBSA_dir}/mmpbsa.in", 'w' )
mmpbsa_infile.write(f'''&general
sys_name={pdb_id}
startframe=1
endframe={int(float(nframes))}
interval=5
temperature={int(float(temp))}
temperature={int(float(md_temp))}
verbose=2
/
&pb
......@@ -68,17 +69,13 @@ verbose=2
''')
mmpbsa_infile.close()
#run_gmxMMPBSA("6JJ3","10000000","5000","300")
tpr_file=f"{sandbox_dir}/md.tpr"
xtc_file=f"{sandbox_dir}/md_noPBC.xtc"
index_file=f"{sandbox_dir}/index.ndx"
topol_file=f"{sandbox_dir}/topol.top"
GMXMMPBSA_PATH = "/home/hackathon/miniforge3/envs/gmxMMPBSA/bin/gmx_MMPBSA"
cmd = [
GMXMMPBSA_PATH,
constants.MMPBSA_ENV_DIR,
"-O",
"-i", "mmpbsa.in",
"-cs", tpr_file,
......@@ -91,5 +88,25 @@ verbose=2
"-nogui"
]
subprocess.run(cmd, check=True)
return f"MMPBSA complete! Files created: {sandbox_dir}/gmx_MMPBSA/FINAL_RESULTS_MMPBSA.dat and {sandbox_dir}/gmx_MMPBSA/FINAL_RESULTS_MMPBSA.csv"
result = subprocess.run(cmd, cwd=MMPBSA_dir, stdout=sys.stdout, stderr=sys.stderr, text=True)
MMPBSA_output = ""
log_file_path = Path(MMPBSA_dir) / "gmx_MMPBSA.log"
if log_file_path.exists():
try:
MMPBSA_output = log_file_path.read_text(encoding="utf-8")
except Exception as e:
MMPBSA_output = f"Could not read gmx_MMPBSA log file: {e}"
if result.returncode != 0:
return (f"Equilibration script failed with return code {result.returncode}.\n"
f"--- Full gmx_MMPBSA Log ---\n"
f"{MMPBSA_output}\n"
f"--- Shell Script Stderr ---\n"
f"{result.stderr or 'None captured directly'}")
else:
return (f"MMPBSA complete! Files created: {MMPBSA_dir}/FINAL_RESULTS_MMPBSA.dat and {MMPBSA_dir}/FINAL_RESULTS_MMPBSA.csv."
f"Full gmx_MMPBSA output:\n"
f"{MMPBSA_output}")
\ No newline at end of file
......@@ -38,15 +38,15 @@ TOOL_MAP = {
),
# GROMACS-related
"gromacs_equil": lambda s, i: gromacs_equil(
s.sandbox_dir, i["input_gro"], ligand_name=i.get("ligand_name"), ligand_file=i.get("ligand_file")
s.sandbox_dir, i["input_gro"], i["md_temp"], ligand_name=i.get("ligand_name"), ligand_file=i.get("ligand_file")
),
"gromacs_production": lambda s, i: gromacs_production(
s.sandbox_dir, i["input_gro"], i["npt_cpt_file"], ligand_name=i.get("ligand_name")
s.sandbox_dir, i["input_gro"], i["npt_cpt_file"], i["md_temp"], i["md_duration"], ligand_name=i.get("ligand_name")
),
"gromacs_analysis": lambda s, i: gromacs_analysis(s.sandbox_dir, i["input_xtc"], ligand_name=i.get("ligand_name")),
# MMPBSA-related
"run_gmxMMPBSA": lambda s, i: run_gmxMMPBSA(
s.sandbox_dir, i["pdb_id"], i["nsteps"], i["nstxout_compressed"], i["temp"],
s.sandbox_dir, i["pdb_id"], i["nsteps"], i["nstxout_compressed"], i["md_temp"],
),
# # RAG tools
"search_papers": lambda _, i: search_papers(i["query"]),
......
......@@ -342,7 +342,7 @@ def prepare_pdb_file_ligand(sandbox_dir: str, pdb_id: str, ligand_name: str = No
logger.info("Atom renaming of ligand completed.")
return f"Successfully Prepared PDB structure with a ligand and saved the extracted ligand PDB file to {sandbox_dir}/{pdb_id}_prepared.pdb and the protonated ligand PDB file to {sandbox_dir}/{ligand_name}_h.pdb. Ligand was protonated at pH=7 and atom names were cleaned (renumbered)."
return f"Successfully Prepared PDB structure with a ligand and saved the extracted ligand PDB file to {sandbox_dir}/{pdb_id}_prepared.pdb and the protonated ligand PDB file to {sandbox_dir}/{ligand_name}_h.pdb. Ligand was protonated at pH=7 and atom names were cleaned (renumbered)"
return f"Successfully Prepared PDB structure without a ligand and saved the extracted PDB file to {sandbox_dir}/{pdb_id}_prepared.pdb"
......
......@@ -251,12 +251,6 @@ def create_tool_schema_md(sandbox_dir, ligand_name, pdb_id):
"This file must exist in the provided sandbox_dir ({sandbox_dir})."
),
},
# "charge_ligand": {
# "type": ["integer", "null"],
# "description": (
# "The net charge of the ligand. If not provided, the charge will be calculated from the structure."
# ),
# },
"ligand_name": {
"type": "string",
"description": (
......@@ -363,6 +357,12 @@ def create_tool_schema_md(sandbox_dir, ligand_name, pdb_id):
"This file must be located within sandbox_dir."
),
},
"md_temp": {
"type": "string",
"description": (
"Temperature in Kelvin for the MD simulation."
),
},
"ligand_name": {
"type": "string",
"description": (
......@@ -378,7 +378,7 @@ def create_tool_schema_md(sandbox_dir, ligand_name, pdb_id):
),
},
},
"required": ["sandbox_dir", "input_gro"],
"required": ["sandbox_dir", "input_gro", "md_temp"],
},
),
Tool(
......@@ -416,6 +416,18 @@ def create_tool_schema_md(sandbox_dir, ligand_name, pdb_id):
"Must be located within sandbox_dir."
),
},
"md_temp": {
"type": "string",
"description": (
"Temperature in Kelvin for the MD simulation."
),
},
"md_duration": {
"type": "string",
"description": (
"Duration of the MD simulation in nanoseconds."
),
},
"ligand_name": {
"type": "string",
"description": (
......@@ -423,7 +435,7 @@ def create_tool_schema_md(sandbox_dir, ligand_name, pdb_id):
),
}
},
"required": ["sandbox_dir", "input_gro", "npt_cpt_file"],
"required": ["sandbox_dir", "input_gro", "npt_cpt_file", "md_temp", "md_duration"],
},
),
Tool(
......@@ -502,12 +514,12 @@ def create_tool_schema_md(sandbox_dir, ligand_name, pdb_id):
"type": ["string"],
"description": ("Number of MD steps that elapse between writing position coordinates using lossy compression, found in the md.mdp file located in sandbox_dir."),
},
"temp": {
"md_temp": {
"type": ["string"],
"description": ("Temperature used during the MD simulation, found in the md.mdp file located in sandbox_dir. This value is an integer with base 10."),
},
},
"required": ["sandbox_dir", "pdb_id", "nsteps", "nstxout_compressed", "temp"],
"required": ["sandbox_dir", "pdb_id", "nsteps", "nstxout_compressed", "md_temp"],
},
),
Tool(
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment