Skip to content

D
DynaMate
Commit 8e3a3715 by cassandra
Options

[feat] improved parameter selection

parent 86bbaa3a
Showing with 460 additions and 298 deletions
main.py
......@@ -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)
setup.sh
#!/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
......
src/agents/agent.py
......@@ -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)
......
src/agents/md_agent.py
......@@ -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()
......
src/agents/prep_agent.py
......@@ -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}")
......
src/constants.py
......@@ -14,4 +14,6 @@ ENV_FILE = Path(__file__).resolve().parent.parent / ".env"
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"
\ No newline at end of file
JSON_LOG_FILE = AGENT_LOGS / "agent_runs.jsonl"
MMPBSA_ENV_DIR = Path("/path/to/your/envs/mmpbsa")
\ No newline at end of file
src/scripts/equil_Gromacs.sh
#!/bin/bash
if [ "$#" -lt 2 ]; then
echo "Usage: $0 sandbox_dir input_gro [ligand_name] [ligand_file] [ligand_gro]"
if [ "$#" -lt 3 ]; then
echo "Usage: $0 sandbox_dir input_gro log_file [ligand_name] [ligand_file] [ligand_gro]"
exit 1
fi
......@@ -30,13 +30,13 @@ if ! ls em.gro 1> /dev/null 2>&1; then
if [ -f em.gro ]; then
echo "'em.gro' created"
echo "11 0" | $GMX energy -f em.edr -o potential.xvg
echo "11 0" | $GMX energy -f em.edr -o potential.xvg >> $LOG_FILE 2>&1
else
echo "Error: Failed to create 'em.gro'"
echo "Error: Failed to create 'em.gro'" >> $LOG_FILE 2>&1
exit 1
fi
else
echo "'em.gro' already exists. Skipping energy minimisation."
echo "'em.gro' already exists. Skipping energy minimisation." >> $LOG_FILE 2>&1
fi
#----------Create posres files-----------
......@@ -44,16 +44,16 @@ fi
# Step 1: count NME residues (6 per chain)
n_nme=$(grep -c "NME" em.gro)
chains=$((n_nme / 6))
echo "Detected $chains chains based on NME residues."
echo "Detected $chains chains based on NME residues." >> $LOG_FILE 2>&1
if ! ls index.ndx 1> /dev/null 2>&1; then
echo "q" | $GMX make_ndx -f em.gro -o index.ndx
echo "q" | $GMX make_ndx -f em.gro -o index.ndx >> $LOG_FILE 2>&1
fi
# Step 2: if 1 chain, create posre.itp file
if [ "$chains" -eq "1" ]; then
if ! ls posre.itp 1> /dev/null 2>&1;then
echo ""Protein-H"" | $GMX genrestr -f em.gro -n index.ndx -o posre.itp -fc 1000 1000 1000
echo ""Protein-H"" | $GMX genrestr -f em.gro -n index.ndx -o posre.itp -fc 1000 1000 1000 >> $LOG_FILE 2>&1
fi
else
# Step 2: if more than 1 chain, create posre_chain(i).itp files
......@@ -64,7 +64,7 @@ else
for ((i=5; i<${#nme_residues[@]}; i+=6)); do
chain_end_residues+=("${nme_residues[i]}")
done
echo "Chain end residues: ${chain_end_residues[@]}"
echo "Chain end residues: ${chain_end_residues[@]}" >> $LOG_FILE 2>&1
# Step 3: compute chain residue ranges
start=1
......@@ -76,7 +76,7 @@ else
start=$((end + 1))
done
echo "Residue ranges per chain: ${ranges[@]}"
echo "Residue ranges per chain: ${ranges[@]}" >> $LOG_FILE 2>&1
# Step 4: create index groups for each chain
# Start from existing index.ndx or create new
......@@ -87,20 +87,30 @@ EOF
fi
# Step 5: add groups per chain
i=1
if grep -E "system1 +2" topol.top; then #special case for two identical chains named system1
echo "You have two identical chains named system1, therefore only one position restraint file for the first chain will be created." >> $LOG_FILE 2>&1
echo "Creating group for residues ${ranges[0]}" >> $LOG_FILE 2>&1
echo -e "ri ${ranges[0]}\n2 & \"r_${ranges[0]}\"\nq" | $GMX make_ndx -f em.gro -n index.ndx -o index.ndx >> $LOG_FILE 2>&1
else
i=1
for range in "${ranges[@]}"; do
echo "Creating group for residues $range..."
echo -e "ri $range\n2 & \"r_${range}\"\nq" | $GMX make_ndx -f em.gro -n index.ndx -o index.ndx
echo "Creating group for residues $range..." >> $LOG_FILE 2>&1
echo -e "ri $range\n2 & \"r_${range}\"\nq" | $GMX make_ndx -f em.gro -n index.ndx -o index.ndx >> $LOG_FILE 2>&1
((i++))
done
fi
# Step 6: generate posre.itp for each chain
if grep -E "system1 +2" topol.top; then #special case for two identical chains named system1
group_name="Protein-H_&_r_${ranges[0]}"
echo "$group_name" | $GMX genrestr -f em.gro -n index.ndx -o "posre.itp" -fc 1000 1000 1000 >> $LOG_FILE 2>&1
else
i=1
for range in "${ranges[@]}"; do
if [ ! -f posre_chain${i}.itp ]; then
echo "Generating position restraints for chain${i}"
echo "Generating position restraints for chain${i}" >> $LOG_FILE 2>&1
group_name="Protein-H_&_r_${range}"
echo "$group_name" | $GMX genrestr -f em.gro -n index.ndx -o "posre_chain${i}.itp" -fc 1000 1000 1000
echo "$group_name" | $GMX genrestr -f em.gro -n index.ndx -o "posre_chain${i}.itp" -fc 1000 1000 1000 >> $LOG_FILE 2>&1
if [ "$i" -gt "1" ]; then
#Adjust atom indices so first = 1
......@@ -123,66 +133,31 @@ EOF
((i++))
fi
done
fi
fi
if [ -n "$LIGNAME" ]; then
#obabel $LIGFILE -O $LIGGRO
if ! ls "posre_$LIGNAME.itp" 1> /dev/null 2>&1;then
echo -e "0 & ! a H*\nq" | $GMX make_ndx -f $LIGGRO -o "index_$LIGNAME.ndx"
echo "3" | $GMX genrestr -f $LIGGRO -n "index_$LIGNAME.ndx" -o "posre_$LIGNAME.itp" -fc 1000 1000 1000
echo "Generating position restraints for ligand $LIGNAME" >> $LOG_FILE 2>&1
echo -e "0 & ! a H* \n q" | $GMX make_ndx -f $LIGGRO -o "index_$LIGNAME.ndx" >> $LOG_FILE 2>&1
echo "3" | $GMX genrestr -f $LIGGRO -n "index_$LIGNAME.ndx" -o "posre_$LIGNAME.itp" -fc 1000 1000 1000 >> $LOG_FILE 2>&1
fi
fi
# #-------- UPDATE TOPOLOGY FILE -----------
# TMP_FILE="topol.tmp"
# # Check if block already exists
# if [ -n "$LIGNAME" ]; then
# if grep -q "^; posre_$LIGNAME.itp" topol.top; then
# echo "Position restraint block already present in topol.top. No changes made."
# else
# awk -v ligname="$LIGNAME" '
# BEGIN { in_dihedrals=0; inserted=0 }
# {
# if ($0 ~ /^\[ dihedrals \]/ && in_dihedrals==0) {
# in_dihedrals=1
# }
# else if (in_dihedrals==1 && $0 ~ /^$/) {
# # Insert blank line + restraint block
# print ""
# print "; Include Position restraint file"
# print "#ifdef POSRES"
# print "#include \"posre_" ligname ".itp\""
# print "#endif"
# inserted=1
# in_dihedrals=0
# }
# print
# }
# ' topol.top > "$TMP_FILE"
# # Replace the original file with the modified one
# cp "topol.top" topol_old.top
# mv "$TMP_FILE" "topol.top"
# echo "Updated topol.top with position restraint block."
# fi
# fi
# Create group Water_and_ions if not exists
if grep -q "Water_and_ions" index.ndx; then
echo "Group Water_and_ions already exists in index.ndx"
else
if grep -q "Cl-" index.ndx; then
echo -e '"WAT" | "Cl-" \n q' | $GMX make_ndx -f em.gro -n index.ndx -o index.ndx
echo -e '"WAT" | "Cl-" \n q' | $GMX make_ndx -f em.gro -n index.ndx -o index.ndx >> $LOG_FILE 2>&1
sed -i 's/Water_Cl-/Water_and_ions/g' index.ndx
echo "Group Water_and_ions created in index.ndx"
echo "Group Water_and_ions created in index.ndx" >> $LOG_FILE 2>&1
fi
if grep -q "Na+" index.ndx; then
echo -e '"WAT" | "Na+" \n q' | $GMX make_ndx -f em.gro -n index.ndx -o index.ndx
echo -e '"WAT" | "Na+" \n q' | $GMX make_ndx -f em.gro -n index.ndx -o index.ndx >> $LOG_FILE 2>&1
sed -i 's/Water_Na+/Water_and_ions/g' index.ndx
echo "Group Water_and_ions created in index.ndx"
echo "Group Water_and_ions created in index.ndx" >> $LOG_FILE 2>&1
fi
fi
......@@ -190,16 +165,14 @@ fi
if ! grep -q "Cl-" index.ndx && ! grep -q "Na+" index.ndx; then
nvt_file="nvt.mdp"
npt_file="npt.mdp"
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"
echo "Protein_$LIGNAME already in index.ndx" >> $LOG_FILE 2>&1
else
echo -e "1 | 13\nq" | $GMX make_ndx -f em.gro -n index.ndx -o index.ndx
echo -e "1 | 13\nq" | $GMX make_ndx -f em.gro -n index.ndx -o index.ndx >> $LOG_FILE 2>&1
fi
else
replacement="Protein Water"
......@@ -207,101 +180,83 @@ if ! grep -q "Cl-" index.ndx && ! grep -q "Na+" index.ndx; then
if grep "$original" "$nvt_file"; then
sed -i "s|$original|$replacement|" "$nvt_file"
echo "$replacement added successfully to tc-grps group in $nvt_file."
echo "$replacement added successfully to tc-grps group in $nvt_file." >> $LOG_FILE 2>&1
else
echo "tc-grps line was not found in $nvt_file."
echo "tc-grps line was not found in $nvt_file." >> $LOG_FILE 2>&1
fi
if grep "$original" "$npt_file"; then
sed -i "s|$original|$replacement|" "$npt_file"
echo "$replacement added successfully to tc-grps group in $npt_file."
echo "$replacement added successfully to tc-grps group in $npt_file." >> $LOG_FILE 2>&1
else
echo "tc-grps line was not found in $npt_file."
echo "tc-grps line was not found in $npt_file." >> $LOG_FILE 2>&1
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."
echo "Ions present. Keeping Water_and_ions group." >> $LOG_FILE 2>&1
fi
#-------- UPDATE TEMP GROUPS NPT, NVT, MD.MDP FILES -----
#cp "${MDP_FILES}/nvt.mdp" .
#cp "${MDP_FILES}/npt.mdp" .
#cp "${MDP_FILES}/md.mdp" .
#echo "mdp files copied"
#-------- UPDATE TEMP GROUPS NPT, NVT, MD.MDP FILES if there are ions -----
if [ -n "$LIGNAME" ]; then
nvt_file="nvt.mdp"
npt_file="npt.mdp"
md_file="md.mdp"
if grep -q "Cl-" index.ndx || grep -q "Na+" index.ndx; then
if [ -n "$LIGNAME" ]; then
nvt_file="nvt.mdp"
npt_file="npt.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" "$nvt_file"; then
sed -i "s|$original|$replacement|" "$nvt_file"
echo "Protein_$LIGNAME added successfully to tc-grps group in $nvt_file."
else
echo "tc-grps line was not found in $nvt_file."
fi
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" >> $LOG_FILE 2>&1
else
echo -e "1 | 13\nq" | $GMX make_ndx -f em.gro -n index.ndx -o index.ndx >> $LOG_FILE 2>&1
fi
if grep "$original" "$npt_file"; then
sed -i "s|$original|$replacement|" "$npt_file"
echo "Protein_$LIGNAME added successfully to tc-grps group in $npt_file."
else
echo "tc-grps line was not found in $npt_file."
fi
if grep "$original" "$nvt_file"; then
sed -i "s|$original|$replacement|" "$nvt_file"
echo "Protein_$LIGNAME added successfully to tc-grps group in $nvt_file." >> $LOG_FILE 2>&1
else
echo "tc-grps line "Protein Water_and_ions" was not found in $nvt_file." >> $LOG_FILE 2>&1
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 was not found in $md_file."
if grep "$original" "$npt_file"; then
sed -i "s|$original|$replacement|" "$npt_file"
echo "Protein_$LIGNAME added successfully to tc-grps group in $npt_file." >> $LOG_FILE 2>&1
else
echo "tc-grps line "Protein Water_and_ions" was not found in $npt_file." >> $LOG_FILE 2>&1
fi
fi
fi
#--------------- NVT --------------------
if ! ls nvt.gro 1> /dev/null 2>&1; then
$GMX grompp -f nvt.mdp -c em.gro -r em.gro -p topol.top -o nvt.tpr -n index.ndx -maxwarn 2
$GMX mdrun -v -deffnm nvt
$GMX grompp -f nvt.mdp -c em.gro -r em.gro -p topol.top -o nvt.tpr -n index.ndx -maxwarn 2 >> $LOG_FILE 2>&1
$GMX mdrun -v -deffnm nvt >> $LOG_FILE 2>&1
if [ -f nvt.gro ]; then
echo "'nvt.gro' created"
echo -e "Temperature \n 0" | $GMX energy -f nvt.edr -o temperature.xvg
echo "'nvt.gro' created" >> $LOG_FILE 2>&1
echo -e "Temperature \n 0" | $GMX energy -f nvt.edr -o temperature.xvg >> $LOG_FILE 2>&1
else
echo "Error: Failed to create 'nvt.gro'"
echo "Error: Failed to create 'nvt.gro'" >> $LOG_FILE 2>&1
exit 1
fi
else
echo "'nvt.gro' already exists. Skipping NVT."
echo "'nvt.gro' already exists. Skipping NVT." >> $LOG_FILE 2>&1
fi
#--------------- NPT --------------------
if ! ls npt.gro 1> /dev/null 2>&1; then
$GMX grompp -f npt.mdp -c nvt.gro -t nvt.cpt -r nvt.gro -p topol.top -o npt.tpr -n index.ndx -maxwarn 2
$GMX mdrun -v -deffnm npt
$GMX grompp -f npt.mdp -c nvt.gro -t nvt.cpt -r nvt.gro -p topol.top -o npt.tpr -n index.ndx -maxwarn 2 >> $LOG_FILE 2>&1
$GMX mdrun -v -deffnm npt >> $LOG_FILE 2>&1
if [ -f npt.gro ]; then
echo "'npt.gro' created"
echo -e "Pressure \n 0" | $GMX energy -f npt.edr -o pressure.xvg
echo -e "Density \n 0" | $GMX energy -f npt.edr -o density.xvg
echo "'npt.gro' created" >> $LOG_FILE 2>&1
echo -e "Pressure \n 0" | $GMX energy -f npt.edr -o pressure.xvg >> $LOG_FILE 2>&1
echo -e "Density \n 0" | $GMX energy -f npt.edr -o density.xvg >> $LOG_FILE 2>&1
else
echo "Error: Failed to create 'npt.gro'"
echo "Error: Failed to create 'npt.gro'" >> $LOG_FILE 2>&1
exit 1
fi
else
echo "'npt.gro' already exists. Skipping NPT."
echo "'npt.gro' already exists. Skipping NPT." >> $LOG_FILE 2>&1
fi
\ No newline at end of file
src/scripts/mdp_files/md.mdp
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
......@@ -43,4 +43,4 @@ pbc = xyz ; 3-D PBC
; Dispersion correction is not used for proteins with the C36 additive FF
DispCorr = no
; Velocity generation
gen_vel = no ; continuing from NPT equilibration
gen_vel = no ; continuing from NPT equilibration
\ No newline at end of file
src/scripts/prod_Gromacs.sh
#!/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
......
src/tools/MMPBSA.py
......@@ -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
src/tools/gromacs_tools.py
......@@ -9,13 +9,15 @@ import time
logger = get_class_logger(__name__)
def gromacs_equil(sandbox_dir: str, input_gro: str, ligand_name=None, ligand_file=None) -> str:
def gromacs_equil(sandbox_dir: str, input_gro: str, md_temp: str, ligand_name=None, ligand_file=None) -> str:
# sometimes llm passes ligands as empty strings
if not ligand_name:
ligand_name = None
if not ligand_file:
ligand_file = None
# ------------ Modify topol.top to include position restraints ------------
input_path = Path(f"{sandbox_dir}/topol.top")
backup_path = Path(f"{sandbox_dir}/topol_without_posre.top")
......@@ -36,8 +38,12 @@ def gromacs_equil(sandbox_dir: str, input_gro: str, ligand_name=None, ligand_fil
if num_systems == 0:
logger.warning("No system entries found. No changes made.")
sys.exit(0)
logger.info(f"Detected {num_systems} chain(s): {', '.join(systems)}")
with open(f"{sandbox_dir}/topol_without_posre.top", "r") as f:
if re.search(r"system1\s+2", f.read()):
logger.info(f"Detected 2 chain(s) because of 2 identical chains named system1 in topol.top")
else:
logger.info(f"Detected {num_systems} chain(s): {', '.join(systems)}")
# --- Create posre include block ---
def make_posre_block(posre_file):
......@@ -68,7 +74,7 @@ def gromacs_equil(sandbox_dir: str, input_gro: str, ligand_name=None, ligand_fil
inserted_blocks = []
inserted_ligand = False
# --- Loop over each [ moleculetype ] block ---
# Loop over each [ moleculetype ] block
for i in range(len(positions) - 1):
seg = text[positions[i] : positions[i + 1]]
......@@ -105,27 +111,158 @@ def gromacs_equil(sandbox_dir: str, input_gro: str, ligand_name=None, ligand_fil
input_path.write_text(modified_text, encoding="utf-8")
logger.info(f"Added position restraints for: {', '.join(inserted_blocks) or 'none'}")
## Position restraints added in topol.top
# -------------- Create em.mdp, nvt.mdp, npt.mdp files --------------
em_mdp_infile = open(f'{sandbox_dir}/em.mdp', 'w' )
em_mdp_infile.write(f'''; LINES STARTING WITH ';' ARE COMMENTS
title = Minimization ; Title of run
; Parameters describing what to do, when to stop and what to save
integrator = steep ; Algorithm (steep = steepest descent minimization)
emtol = 1000.0 ; Stop minimization when the maximum force < 10.0 kJ/mol
emstep = 0.01 ; Energy step size
nsteps = 50000 ; Maximum number of (minimization) steps to perform
; Parameters describing how to find the neighbors of each atom and how to calculate the interactions
nstlist = 1 ; Frequency to update the neighbor list and long range forces
cutoff-scheme = Verlet
ns_type = grid ; Method to determine neighbor list (simple, grid)
rlist = 1.2 ; Cut-off for making neighbor list (short range forces)
coulombtype = PME ; Treatment of long range electrostatic interactions
rcoulomb = 1.2 ; long range electrostatic cut-off
vdwtype = cutoff
vdw-modifier = force-switch
rvdw-switch = 1.0
rvdw = 1.2 ; long range Van der Waals cut-off
pbc = xyz ; Periodic Boundary Conditions
DispCorr = no
''')
em_mdp_infile.close()
nvt_mdp_infile = open(f'{sandbox_dir}/nvt.mdp', 'w' )
nvt_mdp_infile.write(f'''title = Protein-ligand complex NPT equilibration
define = -DPOSRES ; position restrain the protein and ligand
; Run parameters
integrator = md ; leap-frog integrator
nsteps = 5000 ; 2 * 5000 = 10 ps
dt = 0.002 ; 2 fs
; Output control
nstenergy = 500 ; save energies every 1.0 ps
nstlog = 500 ; update log file every 1.0 ps
nstxout-compressed = 500 ; save coordinates every 1.0 ps
; Bond parameters
continuation = yes ; continuing from NVT
constraint_algorithm = lincs ; holonomic constraints
constraints = h-bonds ; bonds to H are constrained
lincs_iter = 1 ; accuracy of LINCS
lincs_order = 4 ; also related to accuracy
; Neighbor searching and vdW
cutoff-scheme = Verlet
ns_type = grid ; search neighboring grid cells
nstlist = 20 ; largely irrelevant with Verlet
rlist = 1.2
vdwtype = cutoff
vdw-modifier = force-switch
rvdw-switch = 1.0
rvdw = 1.2 ; short-range van der Waals cutoff (in nm)
; Electrostatics
coulombtype = PME ; Particle Mesh Ewald for long-range electrostatics
rcoulomb = 1.2
pme_order = 4 ; cubic interpolation
fourierspacing = 0.16 ; grid spacing for FFT
; Temperature coupling
tcoupl = V-rescale ; modified Berendsen thermostat
tc-grps = Protein Water_and_ions ; two coupling groups - more accurate
tau_t = 0.1 0.1 ; time constant, in ps
ref_t = {float(md_temp)} {float(md_temp)} ; reference temperature, one for each group, in K
; Pressure coupling
pcoupl = Berendsen ; pressure coupling is on for NPT
pcoupltype = isotropic ; uniform scaling of box vectors
tau_p = 2.0 ; time constant, in ps
ref_p = 1.0 ; reference pressure, in bar
compressibility = 4.5e-5 ; isothermal compressibility of water, bar^-1
refcoord_scaling = com
; Periodic boundary conditions
pbc = xyz ; 3-D PBC
; Dispersion correction is not used for proteins with the C36 additive FF
DispCorr = no
; Velocity generation
gen_vel = no ; velocity generation off after NVT
''')
nvt_mdp_infile.close()
npt_mdp_infile = open(f'{sandbox_dir}/npt.mdp', 'w' )
npt_mdp_infile.write(f'''title = Protein-ligand complex NPT equilibration
define = -DPOSRES ; position restrain the protein and ligand
; Run parameters
integrator = md ; leap-frog integrator
nsteps = 5000 ; 2 * 5000 = 10 ps
dt = 0.002 ; 2 fs
; Output control
nstenergy = 500 ; save energies every 1.0 ps
nstlog = 500 ; update log file every 1.0 ps
nstxout-compressed = 500 ; save coordinates every 1.0 ps
; Bond parameters
continuation = yes ; continuing from NVT
constraint_algorithm = lincs ; holonomic constraints
constraints = h-bonds ; bonds to H are constrained
lincs_iter = 1 ; accuracy of LINCS
lincs_order = 4 ; also related to accuracy
; Neighbor searching and vdW
cutoff-scheme = Verlet
ns_type = grid ; search neighboring grid cells
nstlist = 20 ; largely irrelevant with Verlet
rlist = 1.2
vdwtype = cutoff
vdw-modifier = force-switch
rvdw-switch = 1.0
rvdw = 1.2 ; short-range van der Waals cutoff (in nm)
; Electrostatics
coulombtype = PME ; Particle Mesh Ewald for long-range electrostatics
rcoulomb = 1.2
pme_order = 4 ; cubic interpolation
fourierspacing = 0.16 ; grid spacing for FFT
; Temperature coupling
tcoupl = V-rescale ; modified Berendsen thermostat
tc-grps = Protein Water_and_ions ; two coupling groups - more accurate
tau_t = 0.1 0.1 ; time constant, in ps
ref_t = {float(md_temp)} {float(md_temp)} ; reference temperature, one for each group, in K
; Pressure coupling
pcoupl = Berendsen ; pressure coupling is on for NPT
pcoupltype = isotropic ; uniform scaling of box vectors
tau_p = 2.0 ; time constant, in ps
ref_p = 1.0 ; reference pressure, in bar
compressibility = 4.5e-5 ; isothermal compressibility of water, bar^-1
refcoord_scaling = com
; Periodic boundary conditions
pbc = xyz ; 3-D PBC
; Dispersion correction is not used for proteins with the C36 additive FF
DispCorr = no
; Velocity generation
gen_vel = no ; velocity generation off after NVT
''')
npt_mdp_infile.close()
# -------------- Run equil_Gromacs.sh script --------------
script = constants.SCRIPTS_DIR / "equil_Gromacs.sh"
log_file_path = Path(f"{sandbox_dir}/gromacs_equil.log")
cmd = [str(script), sandbox_dir, input_gro, log_file_path]
print("I will convert 89G_h.pdb to 89.gro using obabel") # Debug print
#time.sleep(10)
if ligand_file is not None:
obabel_cmd = f"obabel {ligand_file} -O {sandbox_dir}/{ligand_name}.gro"
obabel_result = subprocess.run(obabel_cmd, cwd=sandbox_dir, capture_output=True, text=True, shell=True)
if obabel_result.returncode != 0:
error_text = "\n".join(filter(None, [obabel_result.stderr, obabel_result.stdout]))
return f"Equilibration failed with error: {obabel_result.stderr}"
return f"Equilibration failed with error: {error_text}"
ligand_gro = f"{ligand_name}.gro"
#time.sleep(10)
cmd.append(ligand_name)
cmd.append(ligand_file)
cmd.append(ligand_gro)
print(cmd) # Debug print
print(cmd)
result = subprocess.run(cmd, cwd=sandbox_dir, stdout=sys.stdout, stderr=sys.stderr, text=True)
......@@ -138,22 +275,75 @@ def gromacs_equil(sandbox_dir: str, input_gro: str, ligand_name=None, ligand_fil
gromacs_output = f"Could not read GROMACS log file: {e}"
if result.returncode != 0:
# Report failure and include the captured GROMACS output for debugging
return (f"Equilibration script failed with return code {result.returncode}.\n"
f"--- Full GROMACS Log ---\n"
f"{gromacs_output}\n"
f"--- Shell Script Stderr ---\n"
f"{result.stderr or 'None captured directly.'}") # Note: result.stderr will be empty if we set stderr=sys.stderr, but we keep it here for safety.
f"{result.stderr or 'None captured directly'}")
else:
# Report success and return the captured GROMACS output
return (f"Equilibration ran successfully. Full GROMACS output:\n"
f"{gromacs_output}")
def gromacs_production(sandbox_dir: str, input_gro: str, npt_cpt_file: str, ligand_name=None) -> str:
def gromacs_production(sandbox_dir: str, input_gro: str, npt_cpt_file: str, md_temp: str, md_duration: str, ligand_name=None) -> str:
"""
Run production MD with GROMACS using prod_Gromacs.sh.
"""
# ---------- Create md.mdp file --------------
nsteps = int(((float(md_duration)) * 1000000) / 2) # Convert ns to number of steps (2 fs per step)
md_mdp_infile = open(f'{sandbox_dir}/md.mdp', 'w' )
md_mdp_infile.write(f'''title = Protein-ligand complex MD simulation
; Run parameters
integrator = md ; leap-frog integrator
nsteps = {nsteps} ; 2 * 500,000 = 1000 ps (0.1 ns)
dt = 0.002 ; 2 fs
; Output control
nstenergy = 5000 ; save energies every 10.0 ps
nstlog = 5000 ; update log file every 10.0 ps
nstxout-compressed = 5000 ; save coordinates every 10.0 ps
; Bond parameters
continuation = yes ; continuing from NPT
constraint_algorithm = lincs ; holonomic constraints
constraints = h-bonds ; bonds to H are constrained
lincs_iter = 1 ; accuracy of LINCS
lincs_order = 4 ; also related to accuracy
; Neighbor searching and vdW
cutoff-scheme = Verlet
ns_type = grid ; search neighboring grid cells
nstlist = 20 ; largely irrelevant with Verlet
rlist = 1.2
vdwtype = cutoff
vdw-modifier = force-switch
rvdw-switch = 1.0
rvdw = 1.2 ; short-range van der Waals cutoff (in nm)
; Electrostatics
coulombtype = PME ; Particle Mesh Ewald for long-range electrostatics
rcoulomb = 1.2
pme_order = 4 ; cubic interpolation
fourierspacing = 0.16 ; grid spacing for FFT
; Temperature coupling
tcoupl = V-rescale ; modified Berendsen thermostat
tc-grps = Protein Water_and_ions ; two coupling groups - more accurate
tau_t = 0.1 0.1 ; time constant, in ps
ref_t = {float(md_temp)} {float(md_temp)} ; reference temperature, one for each group, in K
; Pressure coupling
pcoupl = Parrinello-Rahman ; pressure coupling is on for NPT
pcoupltype = isotropic ; uniform scaling of box vectors
tau_p = 2.0 ; time constant, in ps
ref_p = 1.0 ; reference pressure, in bar
compressibility = 4.5e-5 ; isothermal compressibility of water, bar^-1
; Periodic boundary conditions
pbc = xyz ; 3-D PBC
; Dispersion correction is not used for proteins with the C36 additive FF
DispCorr = no
; Velocity generation
gen_vel = no ; continuing from NPT equilibration
''')
md_mdp_infile.close()
# ---------- Run prod_Gromacs.sh script --------------
script = constants.SCRIPTS_DIR / "prod_Gromacs.sh"
log_file_path = Path(f"{sandbox_dir}/gromacs_production.log")
......@@ -174,14 +364,12 @@ def gromacs_production(sandbox_dir: str, input_gro: str, npt_cpt_file: str, liga
gromacs_output = f"Could not read GROMACS log file: {e}"
if result.returncode != 0:
# Report failure and include the captured GROMACS output for debugging
return (f"Equilibration script failed with return code {result.returncode}.\n"
f"--- Full GROMACS Log ---\n"
f"{gromacs_output}\n"
f"--- Shell Script Stderr ---\n"
f"{result.stderr or 'None captured directly.'}") # Note: result.stderr will be empty if we set stderr=sys.stderr, but we keep it here for safety.
f"{result.stderr or 'None captured directly'}")
else:
# Report success and return the captured GROMACS output
return (f"Equilibration ran successfully. Full GROMACS output:\n"
f"{gromacs_output}")
......@@ -210,13 +398,11 @@ def gromacs_analysis(sandbox_dir: str, input_xtc: str, ligand_name=None) -> str:
gromacs_output = f"Could not read GROMACS log file: {e}"
if result.returncode != 0:
# Report failure and include the captured GROMACS output for debugging
return (f"Equilibration script failed with return code {result.returncode}.\n"
f"--- Full GROMACS Log ---\n"
f"{gromacs_output}\n"
f"--- Shell Script Stderr ---\n"
f"{result.stderr or 'None captured directly.'}") # Note: result.stderr will be empty if we set stderr=sys.stderr, but we keep it here for safety.
f"{result.stderr or 'None captured directly'}")
else:
# Report success and return the captured GROMACS output
return (f"Equilibration ran successfully. Full GROMACS output:\n"
f"{gromacs_output}")
\ No newline at end of file
src/tools/map.py
......@@ -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"]),
......
src/tools/pdb_tools.py
......@@ -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"
......
src/tools/tool_schema.py
......@@ -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