Skip to content

D
DynaMate
Commit 523fdad9 by cassandra
Options

[release] DynaMate!

parents
Showing with 3318 additions and 0 deletions
.dockerignore 0 → 100644
# Python-generated files
__pycache__/
*.py[oc]
build/
dist/
wheels/
*.egg-info
# Virtual environments
.venv/
# API keys and secrets
.env
# Logger
logs/
# The sandbox
sandbox/
# vscode
.vscode
# agent logs
agent_logs/
# software
software/*
# !software/ambertools25_src
\ No newline at end of file
.gitignore 0 → 100644
# Python-generated files
__pycache__/
*.py[oc]
build/
dist/
wheels/
*.egg-info
.ipynb_checkpoints/
# Virtual environments
.venv
# API keys and secrets
.env
# The sandbox
sandbox/
# vscode
.vscode
# agent logs
agent_logs/
# literature papers
my_papers/
my_docs.pkl
software/
Launch/
Plotting/
extra_logs/
\ No newline at end of file
README.md 0 → 100644
<p align="center">
<img src="assets/logo_pink.png" alt="drawing" width="250"/>
</p>
DynaMate is your reliable mate that can run molecular dynamics simulations of protein-ligand and protein-only systems. It is built using LiteLLM and equipt with a collection of tools. Quality checks throughout the pipeline trigger re-tries when something goes wrong, allowing the agent to correct course and save you time on debugging.
### Software setup
The tools used by the agent require that you have a local installation of the following software. We provide a Docker image with all dependencies pre-installed (recommended), or you can install everything manually if you prefer
#### Docker Setup (Recommended)
1. Build a docker image:
```
docker build -t dynamate -f ./docker/Dockerfile .
```
2. Create an `.env` file to store sensitive data like API keys:
```
OPENROUTER_API_KEY=your_key_here
# Add other keys as needed
```
3. Run the agent:
```
docker run --env-file .env dynamate --pdb-id <pdb-id> --model <model_name>
```
4. Interactive mode (for debugging or exploration):
```
docker run -it --rm --env-file .env dynamate /bin/bash
python main.py --pdb-id <pdb-id> --model <model_name>
```
Happy molecular dynamics simulations! 🧬
#### Manual Setup
We recommend that you install in a separate `~/softwares` directory, **not inside the project**.:
```bash
mkdir ~/softwares
cd ~/softwares
```
### CMake
You will need `cmake` locally if you don't have the module available to load directly.
1. Download the pre-compiled binary from the official site
```bash
wget https://github.com/Kitware/CMake/releases/download/v3.27.8/cmake-3.27.8-linux-x86_64.sh
chmod +x cmake-3.27.8-linux-x86_64.sh
./cmake-3.27.8-linux-x86_64.sh --prefix=$HOME/cmake --skip-license
```
2. Add it to your PATH
```bash
echo 'export PATH=$HOME/cmake/bin:$PATH' >> ~/.bashrc
source ~/.bashrc
```
### GROMACS
1. Download the source code. You can use `wget` or `curl`:
```bash
wget https://ftp.gromacs.org/pub/gromacs/gromacs-2023.tar.gz
```
You can use a newer version if you want, but IMPORTANT to note:
* To run MM-PB(GB)SA calculations, you will need a GROMACS version inferior than 2023.4.
* The analysis script (/src/scripts/analysis_Gromacs.sh) has been written for GROMACS 2023. You will need to update the echo synthax if you use a different version of GROMACS.
2. Unpack the archive
```bash
tar -xvzf gromacs-2023.tar.gz
cd gromacs-2023
```
3. Create a build directory
```bash
mkdir build
cd build
```
4. Make sure `cmake` is in your PATH. If you installed it with `pip`, add it to your PATH
```bash
export PATH=$HOME/.local/bin:$PATH
```
5. Configure the build with GPU support (make sure you have the appropriate CUDA Toolkit for your system)
```bash
cmake .. -DGMX_GPU=on -DGMX_GPU=CUDA -DCMAKE_INSTALL_PREFIX=$HOME/softwares/gromacs-2023
```
6. Build and install
```bash
make -j 4
make install
```
7. Source GROMACS in the current session
```bash
source /usr/local/gromacs/bin/GMXRC
```
8. Verify
```bash
gmx --version
```
### PDBFixer
1. Download the source file
```bash
wget https://github.com/openmm/pdbfixer/archive/refs/tags/v1.11.tar.gz
```
2. Unpack and enter it
```bash
tar -xvzf v1.11.tar.gz
cd pdbfixer-1.11
```
4. Install in editable mode (optional) or normally
```bash
pip install -e .
```
5. Verify
```bash
python -c "import pdbfixer; print(pdbfixer.__version__)"
```
### AmberTools25
Navigate [here](https://ambermd.org/GetAmber.php#ambertools) to obtain the source code in tar format. Copy this into the `~/softwares` directory.
1. Unpack the archive
```bash
tar -xvzf ambertools25.tar.bz2
cd ambertools25_src/build
```
2. Run the `cmake` script with MPI and CUDA enabled
```bash
./run_cmake -DMPI=TRUE -DCUDA=TRUE
```
3. If the cmake build report looks OK, you should now do the following:
```bash
make -j 4
make install
source /home/softwares/ambertools25/amber.sh
```
### Conda
If you don't have conda, install it
You can use `wget` or `curl`:
```bash
wget "https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-$(uname)-$(uname -m).sh"
```
## Clone the repository
```bash
git clone https://github.com/schwallergroup/DynaMate.git
cd DynaMate
```
## Environment setup
Setup the conda env
```bash
conda env create -f environment.yml
```
## Activate your environment
```bash
conda activate dynamate
```
## Export python path so you can load the modules
At the root of the project run:
```bash
export PYTHONPATH=.
```
## Run the setup script
After setting up your project environment, make sure to run the setup script if you don't want to load gromacs each time. This will load both the environment and the softwares
```bash
source setup.sh
```
## Usage
To launch the script specify the PDB, possible ligand name, and model name in the command line arguments. For example, to launch an MD run with the protein 5UEZ, ligand 89G, and model GPT-5 mini (don't forget to specify the online status to allow the models to perform web searches):
```bash
python main.py --pdb_id 5UEZ --ligand 89G --model openrouter/anthropic/claude-4.5-sonnet:online
```
And again, happy molecular dynamics simulations! 🧬
<p align="center">
<img src="assets/MDAgent-Tools-workflow.png" alt="drawing" width="900"/>
</p>
### License
This work is licensed under the [MIT License](https://opensource.org/license/mit)
docker/Dockerfile 0 → 100644
FROM md_software:latest
RUN apt-get update && apt-get install -y \
openbabel \
build-essential python3-dev gfortran git wget curl \
&& rm -rf /var/lib/apt/lists/*
# Force Python to use unbuffered output
ENV PYTHONUNBUFFERED=1
# Download Conda
ARG CONDA_DIR=/opt/conda
RUN wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O /tmp/miniconda.sh && \
bash /tmp/miniconda.sh -b -p $CONDA_DIR && \
rm /tmp/miniconda.sh
ENV PATH="$CONDA_DIR/bin:$PATH"
# Accept Conda's Terms and Conditions
RUN conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/main && \
conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/r
COPY environment.yml /tmp/environment.yml
RUN conda env create -f /tmp/environment.yml && \
conda clean -afy
ARG CONDA_ENV_NAME=dynagent
ENV CONDA_DEFAULT_ENV=$CONDA_ENV_NAME
ENV PATH="$CONDA_DIR/envs/$CONDA_ENV_NAME/bin:$PATH"
# Create a new user for the agent
ARG AGENT_USER=beautifulagent
RUN useradd -m $AGENT_USER
WORKDIR /app
ENV SANDBOX=/app/sandbox
ENV AGENT_LOGS=/app/agent_logs
COPY . /app
RUN chmod -R a+rX,u-w,go-w /app
# Create a sandbox for the agent
RUN mkdir -p $SANDBOX && chown -R $AGENT_USER:$AGENT_USER $SANDBOX
RUN mkdir -p $AGENT_LOGS && chown -R $AGENT_USER:$AGENT_USER $AGENT_LOGS
RUN mkdir -p .venv && chown -R $AGENT_USER:$AGENT_USER .venv
COPY docker/entrypoint.sh /entrypoint.sh
RUN chmod +x /entrypoint.sh
ENTRYPOINT ["/entrypoint.sh"]
CMD ["python", "/app/main.py"]
\ No newline at end of file
docker/Dockerfile.base 0 → 100644
# Stage 1, Builder
FROM nvidia/cuda:12.4.0-devel-ubuntu22.04 AS builder
# Define build-time variables
ARG SOFTWARE_FOLDER=/opt/software
ARG GROMACS_VERSION=2023
ARG UV_DIR=/opt/uv
ARG GROMACS_INSTALL_DIR=$SOFTWARE_FOLDER/gromacs
ENV DEBIAN_FRONTEND=noninteractive
# Install core build dependencies
RUN apt-get update && apt-get install -y --no-install-recommends \
cmake gfortran flex bison wget build-essential curl ca-certificates \
zlib1g-dev \
libbz2-dev \
&& rm -rf /var/lib/apt/lists/*
# Install UV
RUN curl -LsSf https://astral.sh/uv/install.sh | env UV_INSTALL_DIR=$UV_DIR sh
ENV PATH="$UV_DIR:$PATH"
SHELL ["/bin/bash", "-c"]
# Copy local software
ARG LOCAL_SOFTWARE_FOLDER=software
COPY $LOCAL_SOFTWARE_FOLDER $SOFTWARE_FOLDER
# Install AmberTools25
WORKDIR $SOFTWARE_FOLDER/ambertools25_src/build
RUN ./run_cmake && make -j 4 && make install
# Install GROMACS dependencies (updated CMake)
RUN wget https://github.com/Kitware/CMake/releases/download/v3.29.3/cmake-3.29.3-linux-x86_64.sh -O /tmp/cmake.sh && \
chmod +x /tmp/cmake.sh && \
/tmp/cmake.sh --skip-license --prefix=/usr/local && \
rm /tmp/cmake.sh
# Download and Install GROMACS
WORKDIR /build
RUN wget https://ftp.gromacs.org/gromacs/gromacs-${GROMACS_VERSION}.tar.gz -O /tmp/gromacs-${GROMACS_VERSION}.tar.gz && \
tar xfz /tmp/gromacs-${GROMACS_VERSION}.tar.gz -C $SOFTWARE_FOLDER && \
rm /tmp/gromacs-${GROMACS_VERSION}.tar.gz
WORKDIR $SOFTWARE_FOLDER/gromacs-${GROMACS_VERSION}/build
RUN cmake .. \
-DGMX_GPU=CUDA \
-DCMAKE_INSTALL_PREFIX=${GROMACS_INSTALL_DIR} \
-DGMX_BUILD_OWN_FFTW=ON \
-DREGRESSIONTEST_DOWNLOAD=ON && \
make -j$(nproc) && make check && make install
# Stage 2, runtime, set ups the environments, copies software from the builder
FROM nvidia/cuda:12.4.0-runtime-ubuntu22.04 AS runtime
# Define common arguments/variables
ARG SOFTWARE_FOLDER=/opt/software
ARG GROMACS_INSTALL_DIR=$SOFTWARE_FOLDER/gromacs
ARG UV_DIR=/opt/uv
ARG AGENT_USER=beautifulagent
ENV DEBIAN_FRONTEND=noninteractive
RUN apt-get update && apt-get install -y \
openbabel \
build-essential python3-dev gfortran git \
&& rm -rf /var/lib/apt/lists/*
COPY --from=builder $UV_DIR $UV_DIR
ARG AMBERTOOLS_DIR=$SOFTWARE_FOLDER/ambertools25
ARG GROMACS_INSTALL_DIR=$SOFTWARE_FOLDER/gromacs
COPY --from=builder $AMBERTOOLS_DIR $AMBERTOOLS_DIR
COPY --from=builder $GROMACS_INSTALL_DIR $GROMACS_INSTALL_DIR
# Set up environment variables for the final image
ENV PATH="$UV_DIR:${GROMACS_INSTALL_DIR}/bin:/usr/local/bin:$PATH"
ENV LD_LIBRARY_PATH="${GROMACS_INSTALL_DIR}/lib:/usr/local/lib:$LD_LIBRARY_PATH"
# Metadata
LABEL maintainer="liac"
LABEL description="Molecular dynamics software: AmberTools25 + GROMACS + openbabel + UV (Final Agent Image)"
LABEL version="0.0.1"
\ No newline at end of file
docker/entrypoint.sh 0 → 100644
#!/bin/bash
source /opt/conda/etc/profile.d/conda.sh
conda activate dynagent
source /opt/software/ambertools25/amber.sh
exec "$@"
\ No newline at end of file
environment.yml 0 → 100644
name: dynagent
channels:
- conda-forge
- defaults
dependencies:
- python=3.12
- pip
- ambertools
- mdanalysis
- numpy=1.26.4
- parmed
- pip:
- openai>=2.0.0
- pydantic>=2.11.9
- litellm==1.79.3
- bio>=1.8.0
- dotenv>=0.9.9
- anthropic>=0.69.0
- mdanalysis>=2.9.0
- ruff>=0.14.0
- parmed>=4.3.0
- pdbfixer @ git+https://github.com/openmm/pdbfixer
- paper-qa>=5.29.1
- tyro>=0.9.35
- python-dotenv>=1.1.1
# dev deps:
- ipykernel>=6.30.1
- tyro>=0.9.35
\ No newline at end of file
gmx_MMPBSA/env.yml 0 → 100644
name: gmxMMPBSA
channels:
- defaults
- conda-forge
dependencies:
- python=3.9
- pip
- ambertools<=23.3
- mpi4py<=3.1.5
- gromacs<=2023.4
- git
- pip:
- pyqt5<=6.6.1
- gmx-mmpbsa
- pandas==1.2.2
- seaborn<0.12
- scipy>=1.6.1
- matplotlib==3.5.2
- tqdm
gmx_MMPBSA/gmx_MMPBSA_test.log 0 → 100644
[INFO ] Cloning gmx_MMPBSA repository in /home/hackathon/autoMD_salome/gmx_MMPBSA/gmx_MMPBSA_test
[INFO ] Cloning gmx_MMPBSA repository...Done.
[INFO ] Example STATE
[INFO ] Protein-Ligand (Single trajectory approximation) RUNNING
[INFO ] Protein-Ligand (Single trajectory approximation) [ 1/ 9] DONE
[INFO ] Protein-Protein RUNNING
[INFO ] Protein-Protein [ 2/ 9] DONE
[INFO ] Protein-DNA RUNNING
[INFO ] Protein-DNA [ 3/ 9] DONE
[INFO ] Protein-Glycan RUNNING
[INFO ] Protein-Glycan [ 4/ 9] DONE
[INFO ] Comp_receptor RUNNING
[INFO ] Comp_receptor [ 5/ 9] DONE
[INFO ] Alanine Scanning RUNNING
[INFO ] Alanine Scanning [ 6/ 9] DONE
[INFO ] Stability calculation RUNNING
[INFO ] Stability calculation [ 7/ 9] DONE
[INFO ] Decomposition Analysis RUNNING
[INFO ] Decomposition Analysis [ 8/ 9] DONE
[INFO ] Interaction Entropy approximation RUNNING
[INFO ] Interaction Entropy approximation [ 9/ 9] DONE
[INFO ] Opening gmx_MMPBSA_ana...
log_parser.py 0 → 100644
import re
import csv
from datetime import datetime
from pathlib import Path
LOG_DIR = Path("agent_logs")
OUT_FILE = "run_summary.csv"
MODEL_NAME = "claude-3-5-haiku-20241022" # update this for each model run
def parse_timestamp(line):
try:
return datetime.strptime(line.split(" - ")[0], "%Y-%m-%d %H:%M:%S,%f")
except Exception:
return None
def parse_logs():
prep_log = (LOG_DIR / "PrepAgent.log").read_text()
md_log = (LOG_DIR / "MDAgent.log").read_text()
# --- Step 1: Identify runs and their metadata ---
runs = []
current_run = {}
for line in prep_log.splitlines():
if "PrepAgent initialized." in line:
if current_run:
runs.append(current_run)
current_run = {"tools": set(), "start_time": parse_timestamp(line)}
elif "User input:" in line:
current_run["protein"] = line.split("User input:")[1].strip()
elif "User requested ligand:" in line:
current_run["ligand"] = line.split("User requested ligand:")[1].strip()
elif "Executing tool:" in line:
tool_match = re.search(r"Executing tool:\s*(\w+)", line)
if tool_match:
current_run["tools"].add(tool_match.group(1))
if current_run:
runs.append(current_run)
# --- Step 2: Parse MDAgent log ---
md_lines = md_log.splitlines()
md_runs = []
current = None
for line in md_lines:
if "MDAgent initialized." in line:
if current:
md_runs.append(current)
current = {
"start_time": parse_timestamp(line),
"iterations": 0,
"tools_called": set(),
"attempted": 0,
"success": 0,
}
elif "Logging agent iteration" in line and current:
current["iterations"] += 1
elif "Executing tool:" in line and current:
current["attempted"] += 1
match = re.search(r"Executing tool:\s*(\w+)", line)
if match:
current["tools_called"].add(match.group(1))
elif "Tool result:" in line and current:
if "Error" not in line:
current["success"] += 1
elif "MD Pipeline completed successfully" in line and current:
current["end_time"] = parse_timestamp(line)
if current:
md_runs.append(current)
# --- Step 3: Merge PrepAgent + MDAgent runs by order ---
rows = []
for i, run in enumerate(runs):
md = md_runs[i] if i < len(md_runs) else {}
start = md.get("start_time")
end = md.get("end_time")
total_time = (end - start).total_seconds() if start and end else None
rows.append(
{
"Model": MODEL_NAME,
"Protein": run.get("protein", "Unknown"),
"Ligand": run.get("ligand", "None"),
"Iterations": md.get("iterations", 0),
"Total Time (s)": total_time,
"Subtasks Attempted": md.get("attempted", 0),
"Subtasks Successful": md.get("success", 0),
"Tools Called": ", ".join(sorted(md.get("tools_called", run["tools"]))),
}
)
# --- Step 4: Write to CSV ---
fieldnames = [
"Model",
"Protein",
"Ligand",
"Iterations",
"Total Time (s)",
"Subtasks Attempted",
"Subtasks Successful",
"Tools Called",
]
with open(OUT_FILE, "w", newline="") as f:
writer = csv.DictWriter(f, fieldnames=fieldnames)
writer.writeheader()
writer.writerows(rows)
print(f"Summary written to {OUT_FILE}")
if __name__ == "__main__":
parse_logs()
main.py 0 → 100644
import os
import glob
from pathlib import Path
from dotenv import load_dotenv, set_key
import shutil
import tyro
from dataclasses import dataclass
from src.agents import MDAgent, PrepAgent
from src import utils
from src import constants
if constants.ENV_FILE.exists():
load_dotenv(dotenv_path=constants.ENV_FILE)
def _ensure_api_key(env_var: str, prompt_name: str) -> str | None:
"""Ensures a required API key is set, prompting the user if necessary."""
key = os.environ.get(env_var)
if not key:
print(f"--- Missing API Key: {prompt_name} ---")
key = input(f"Please enter your {prompt_name} API key: ").strip()
if key:
os.environ[env_var] = key
set_key(str(constants.ENV_FILE), env_var, key)
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:
"""
Tyro automatically generates a command line interface from this class.
"""
pdb_id: str
"PDB ID."
model: str
"Model name to use for the MD pipeline."
ligand: str | None = None
"Ligand ID (optional; defaults to no ligand)."
model_supports_system_messages: bool = True
def main(config: CommandLineArgs):
root_logger = utils.get_class_logger("Main")
# create a run directory inside of sandbox
run_name = f"run_{utils.time_now()}"
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?"
)
try:
_ensure_api_key("OPENROUTER_API_KEY", "OPENROUTER_API_KEY")
except ValueError as e:
root_logger.error(str(e))
return
root_logger.info("\n=== 1. 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,
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) ===")
# Copy PDB into run directory
md_agent = MDAgent(
model_name=config.model,
temperature=constants.TEMPERATURE,
sandbox_dir=sandbox_dir,
structure_path=pdb_file_path,
pdb_id=Path(pdb_file_path).stem,
ligand_name=ligand_name,
model_supports_system_messages=config.model_supports_system_messages,
plan=plan,
)
md_agent.setup_tools()
# 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 0 → 100755
#!/bin/bash
# load your environment
source /home/miniforge3/bin/activate dynagent
# Activate GROMACS environment
if [ -f /usr/local/gromacs/bin/GMXRC ]; then
source /usr/local/gromacs/bin/GMXRC
else
echo "Error: /usr/local/gromacs/bin/GMXRC not found!"
return 1 2>/dev/null || exit 1
fi
# Export Python path to current project
export PYTHONPATH=$(pwd)
echo "PYTHONPATH set to: $PYTHONPATH"
echo "Environment setup complete!"
src/agents/__init__.py 0 → 100644
from .md_agent import MDAgent
from .prep_agent import PrepAgent
from .agent import BaseAgent, ToolOutputError
__all__ = ["MDAgent", "PrepAgent", "BaseAgent", "ToolOutputError"]
\ No newline at end of file
src/agents/agent.py 0 → 100644
import os
from pathlib import Path
from typing import Dict, Any, List
import json
from litellm import completion
from abc import ABC, abstractmethod
import traceback
from src import constants
from src.tools.map import TOOL_MAP
from src import utils, constants
import tiktoken
ENC = tiktoken.get_encoding("cl100k_base")
class ToolOutputError(Exception):
"""Custom exception raised when a tool returns a known failure string."""
pass
class BaseAgent(ABC):
def __init__(
self,
model_name: str,
temperature: float,
sandbox_dir: str,
pdb_id: str | None = None,
ligand_name: str | None = None,
model_supports_system_messages: bool = True,
):
self.model_name = model_name
self.temperature = temperature
self.sandbox_dir = Path(sandbox_dir)
self.pdb_id = pdb_id
self.ligand_name = ligand_name
self.model_supports_system_messages = model_supports_system_messages
self.tool_schemas = None
self.messages: List[Dict[str, Any]] = []
self.llm_cost = 0
self.logger = utils.get_class_logger(self.__class__.__name__)
@abstractmethod
def setup_tools(self):
pass
@abstractmethod
def _additional_check_for_errors_tool_output(self, tool_name, tool_call) -> bool:
pass
import json
def _count_tokens(self, text):
if not isinstance(text, str):
text = json.dumps(text)
return len(ENC.encode(text))
def _find_recent_block(self, messages):
"""
Return the smallest suffix of messages that forms a logically valid block.
"""
if not messages:
return []
# Start from the end
block = []
i = len(messages) - 1
while i >= 0:
m = messages[i]
block.insert(0, m)
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])
i -= 1
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])
return block
i -= 1
return block
def _summarize_history(self, messages):
recent_messages = self._find_recent_block(messages)
history_to_summarize = messages[:-len(recent_messages)]
# Turn history into text
history_text = "\n".join(
f"{m['role']}: {m['content']}"
for m in history_to_summarize
)
# Summarize
summary_response = completion(
model=self.model_name,
temperature=0.1,
messages=[
{"role": "system", "content": "Summarize the conversation concisely but fully."},
{"role": "user", "content": history_text}
],
max_tokens=constants.SUMMARY_OUTPUT_TOKENS,
)
summary_text = summary_response.choices[0].message["content"]
# NEW conversation = summary + recent logical block
return [
{"role": "assistant", "content": f"[Conversation Summary]\n{summary_text}"},
*recent_messages
]
def _validate_tool_path(self, tool_input) -> None:
# if "path" in tool_input and not utils.is_path_child_dir(tool_input["path"], self.sandbox_dir):
# raise PermissionError(f"Access outside sandbox not allowed: {tool_input['path']}, {self.sandbox_dir}")
# if "sandbox_dir" in tool_input and not utils.is_path_child_dir(tool_input["sandbox_dir"], self.sandbox_dir):
# raise PermissionError(f"Access outside sandbox not allowed: {tool_input['sandbox_dir']}, {self.sandbox_dir}")
pass
def _safe_execute_tool(self, tool_name: str, tool_input: Dict[str, Any]) -> dict:
"""Executes a tool and catches any errors to pass back to the LLM."""
self.logger.info(f"Executing tool: {tool_name} with input: {tool_input}")
if not isinstance(tool_input, dict):
try:
tool_input = json.loads(tool_input)
except Exception:
error_msg = f"Invalid tool input format (not a dict): {tool_input}"
self.logger.error(error_msg)
return {"ok": False, "output": error_msg}
tool_output = None
# try:
self._validate_tool_path(tool_input)
func = TOOL_MAP.get(tool_name)
if not func:
raise ValueError(f"Unknown tool: {tool_name}")
tool_output = func(self, tool_input)
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):
return {
"tool_call_id": id_,
"role": "tool",
"name": tool_name,
"arguments": arguments,
"content": output,
}
def _create_logs(self):
def to_dict_safe(msg):
if isinstance(msg, dict):
return msg
elif hasattr(msg, "model_dump"): # Pydantic/LiteLLM object
return msg.model_dump()
else:
return str(msg)
safe_messages = [to_dict_safe(m) for m in self.messages]
created_files = [f for f in os.listdir(self.sandbox_dir) if os.path.isfile(os.path.join(self.sandbox_dir, f))]
# Build run log
run_data = {
"timestamp": utils.time_now(),
"model": self.model_name,
"temperature": self.temperature,
"tools": self.tool_schemas,
"messages": safe_messages,
"files_created": created_files,
}
utils.append_jsonl(run_data, constants.JSON_LOG_FILE)
def _call_llm(self, messages):
# periodically summarize
total_tokens = 0
for m in messages:
total_tokens += self._count_tokens(m.get("content", ""))
if total_tokens > constants.MAX_CONTEXT_TOKENS and len(messages) > 3:
messages = self._summarize_history(messages)
self.messages = messages
response = completion(
model=self.model_name,
temperature=self.temperature,
supports_system_message=self.model_supports_system_messages,
messages=messages,
web_search_options={"search_context_size": "medium"},
tools=self.tool_schemas,
tool_choice="auto",
)
self.llm_cost += response._hidden_params["response_cost"]
message = response.choices[0].message
return message
def _prompt_llm(self, prompt):
self.messages.append({"role": "user", "content": prompt})
return self._call_llm(self.messages)
def _final_log(self, llm_cost):
final_logs = {
"timestamp_final": utils.time_now(),
"total_completion_cost": llm_cost,
}
utils.append_jsonl(final_logs, constants.JSON_LOG_FILE)
def _process_tool_call(self, tool_call):
function_name = tool_call.function.name
raw_args = tool_call.function.arguments
if not raw_args:
self.logger.warning(f"Tool '{function_name}' returned empty arguments. Using empty dict.")
function_args = {}
else:
try:
function_args = json.loads(raw_args)
except json.JSONDecodeError:
self.logger.warning(f"Tool '{function_name}' returned invalid JSON: {raw_args}. Using empty dict.")
function_args = {}
exec_result = self._safe_execute_tool(function_name, function_args)
tool_output = utils.truncate_string(exec_result["output"])
exec_result["output"] = tool_output
self.logger.info(f"Tool result: {tool_output}.")
tool_message = self._format_tool_usage_ouput(tool_call.id, function_name, function_args, tool_output)
self.messages.append(tool_message)
return exec_result
@abstractmethod
def _reset_pipeline(self) -> None:
pass
@abstractmethod
def run(self):
pass
\ No newline at end of file
src/agents/prep_agent.py 0 → 100644
import re
import json
import sys
from typing import Dict, Any, List
from pydantic import BaseModel
import litellm
from src.tools import tool_schema
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
class Tool(BaseModel):
name: str
description: str
parameters: Dict[str, Any]
class PrepAgent(BaseAgent):
def __init__(
self,
model_name,
temperature,
sandbox_dir,
pdb_id=None,
ligand_name=None,
model_supports_system_messages=True,
):
super().__init__(
model_name, temperature, sandbox_dir, pdb_id, ligand_name, model_supports_system_messages
)
self.messages: List[Dict[str, Any]] = []
self.pdb_file_path = None
self.user_temp = None
self.agent_plan = ""
self.logger.info(f"PrepAgent initialized.")
def setup_tools(self):
self.tool_schemas = tool_schema.create_tool_schema_prep(self.sandbox_dir)
def _additional_check_for_errors_tool_output(self, tool_name, tool_call) -> bool:
# No additional domain-specific checks needed for PrepAgent's tools based on output string
return True
def _reset_pipeline(self):
self.agent_plan = ""
def _setup_system_prompt(self) -> None:
system_prompt_text = PREP_SYSTEM_PROMPT.format(sandbox_dir=self.sandbox_dir)
system_prompt = {
"role": "system",
"content": system_prompt_text,
}
self.messages.append(system_prompt)
def _get_pdb_file_path(self, prompt):
pdb_file_path = None
while pdb_file_path is None:
response = self._prompt_llm(prompt)
self.logger.info(f"Response: {response}")
tool_calls = response.tool_calls
if tool_calls:
self.logger.info(f"Length of tool calls: {len(tool_calls)}")
self.messages.append(response)
for tool_call in tool_calls:
self._process_tool_call(tool_call)
else:
assistant_message = {"role": "assistant", "content": response.content}
self.messages.append(assistant_message)
# Check if a PDB file exists in sandbox
pdb_files = list(self.sandbox_dir.glob("*.pdb"))
if pdb_files:
pdb_file_path = str(pdb_files[0])
return pdb_file_path
def _find_ligand(self):
lig_response = None
while lig_response is None:
user_input = self.ligand_name
if not user_input:
lig_response = True
self.logger.info("Defining system without a ligand")
else:
lig_name = re.search(r"^[A-Z0-9]{3}$", user_input)
if lig_name:
self.logger.info(f"User requested ligand: {lig_name.group()}")
self.ligand_name = lig_name.group()
lig_response = True
lig_found = False
with open(self.pdb_file_path, "r") as infile:
for line in infile:
if line.startswith("HETATM") and self.ligand_name in line:
lig_found = True
break
if not lig_found:
self.logger.error(
"The ligand name could not be identified. Carefully enter the three character identifier for the ligand."
)
sys.exit(1)
else:
self.logger.error(
"The ligand name could not be identified. Carefully enter the three character identifier for the ligand."
)
sys.exit(1)
def _find_simulation_temperature(self):
temperature = None
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."
response = self._prompt_llm(user_input)
tool_calls = response.tool_calls
if tool_calls:
self.logger.info(f"Length of tool calls: {len(tool_calls)}")
self.messages.append(response)
for tool_call in tool_calls:
self._process_tool_call(tool_call)
else:
assistant_message = {"role": "assistant", "content": response.content}
self.messages.append(assistant_message)
match = re.search(r"(\d+\.?\d*)\s+K", response.content)
temperature = float(match.group(1)) if match else 310.0
self.logger.info(f"Rational for picking simulation temperature {temperature}: {response.content}")
return temperature
def _calculate_duration(self):
duration = None
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)."
response = self._prompt_llm(user_input)
tool_calls = response.tool_calls
if tool_calls:
self.logger.info(f"Length of tool calls: {len(tool_calls)}")
self.messages.append(response)
for tool_call in tool_calls:
self._process_tool_call(tool_call)
else:
assistant_message = {"role": "assistant", "content": response.content}
self.messages.append(assistant_message)
matches = re.findall(r"(\d+\.?\d*)\s+ns", response.content)
duration = float(matches[-1]) if matches else 0.01
self.logger.info(f"Rational for picking simulation duration {duration}: {response.content}")
return duration
def _generate_plan(self, temperature, duration):
if self.ligand_name:
steps = [
{
"step": "prepare_pdb_file_ligand",
"description": "Clean and preprocess PDB file for protein-ligand system.",
},
{"step": "add_caps", "description": "Add N- and C-terminal capping groups."},
{"step": "rename_histidines", "description": "Rename HIS to HIE, HIP or HID."},
{"step": "param_ligand", "description": "Generate ligand parameters using antechamber or acpype."},
{"step": "run_tleap_ligand", "description": "Build system topology and solvate complex using tleap."},
{"step": "gromacs_equil", "description": "Perform energy minimization and equilibration."},
{"step": "gromacs_production", "description": "Run production MD simulation."},
{"step": "gromacs_analysis", "description": "Analyse production MD simulation."},
]
else:
steps = [
{
"step": "prepare_pdb_file_ligand",
"description": "Clean and preprocess PDB file for protein-only system.",
},
{"step": "add_caps", "description": "Add N- and C-terminal capping groups."},
{"step": "rename_histidines", "description": "Rename HIS to HIE, HIP or HID."},
{"step": "run_tleap", "description": "Build system topology and solvate complex using tleap."},
{"step": "gromacs_equil", "description": "Perform energy minimization and equilibration."},
{"step": "gromacs_production", "description": "Run production MD simulation."},
{"step": "gromacs_analysis", "description": "Analyse production MD simulation."},
]
plan = {
"sandbox_dir": str(self.sandbox_dir),
"pdb_file_path": self.pdb_file_path,
"ligand": self.ligand_name,
"plan": steps,
"parameters": {"temperature_k": float(temperature), "duration_ns": float(duration)},
}
return plan
def run(self):
self._reset_pipeline()
if not self.tool_schemas:
raise NameError("Tool schema is not defined, call setup_tools() first.")
self._setup_system_prompt()
user_input = self.pdb_id
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._find_ligand()
temperature = self._find_simulation_temperature()
self.logger.info(f"Using simulation temperature: {temperature} K")
duration = self._calculate_duration()
self.logger.info(f"Using simulation duration: {duration} ns")
# Build plan steps depending on ligand
plan = self._generate_plan(temperature, duration)
self.agent_plan = json.dumps(plan, indent=2)
self.logger.info(f"Generated plan: {self.agent_plan}")
self._create_logs()
return self.pdb_file_path, self.ligand_name, plan, self.llm_cost
src/constants.py 0 → 100644
from pathlib import Path
MAX_CHARACTERS_TO_LOG = 5000
SUMMARY_OUTPUT_TOKENS = 6000
MAX_CONTEXT_TOKENS = 32000
PAPER_DIR = Path(__file__).resolve().parent.parent / "my_papers"
MODEL_NAME = "openrouter/openai/gpt-4.1-2025-04-14"
TEMPERATURE = 0.1
SCRIPTS_DIR = Path(__file__).resolve().parent / "scripts"
MDP_FILES = SCRIPTS_DIR / "mdp_files"
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
src/prompts/__init__.py 0 → 100644
from .default import PREP_SYSTEM_PROMPT, MD_SYSTEM_PROMPT
__all__ = ["PREP_SYSTEM_PROMPT", "MD_SYSTEM_PROMPT"]
src/prompts/default.py 0 → 100644
PREP_SYSTEM_PROMPT = """"You are a helpful science assistant designed to fetch information about protein
structures and ligands, and make helpful suggestions regarding molecular systems.
Classify the user request and prepare the input files for an appropriate molecular dynamics pipeline.
The user will either specify a PDB ID or upload the file into {sandbox_dir}.
Depending on the user inputs you should define what a sucessful MD pipeline would involve.
Call the relevant tools when needed to prepare the system for molecular dynamics. Do not ask the user anything."""
MD_SYSTEM_PROMPT = """You are an MD execution assistant. You have access to tools that prepare and
run molecular dynamics (MD) simulations using GROMACS.
The PDB structure file has been provided at {pdb_path} and the
necessary MDP files to run GROMACS. They are available in the sandbox directory
located at {sandbox_dir}, and you should use them.
You should use the tools to solvate, equilibrate, and run MD in sequence,
starting with a simulation using GROMACS and the Amber force field ff14sb.
First, if the system is a protein-ligand complex, separate the protein and ligands into two separate pdb files. If a ligand is present, protonate it at the appropriate pH. Check that the protonation is accurate based on a literature search.
Then, if a ligand is present, parameterise the ligand using antechamber with Amber.
If the system is a protein-ligand complex, merge the two into complex.pdb, otherwise if the system is a protein alone, keep only the protein atoms.
Next, parameterise the system using tleap, which allows to create a box, solvate, add ions to neutralise and prepare the 'topol.top' file. In the tleap step, the protein will be protonated. Check that the protonation is accurate by doing a literature search, and that the protein was protonated at the optimum pH.
Then, perform a short energy minimization using GROMACS.
Next, equilibrate with short NVT and NPT runs.
Finally, perform the production run. Use a default of 0.1 ns unless the user specifies otherwise.
After production run is complete, perform a basic analysis of the trajectory including RMSD, RMSF calculations, radius of gyration, and hydrogen bond analysis.
The analysis of these plots should be saved as a text file named "analysis.txt" in the sandbox directory.
If any step fails, retry after analyzing the provided error message and make
corrections to the inputs for the current step.
"""
src/scripts/analysis_Gromacs.sh 0 → 100755
#!/bin/bash
if [ "$#" -lt 1 ]; then
echo "Usage: $0 input_xtc [ligand_name]"
exit 1
fi
GMX='gmx'
INPUT_XTC="$1"
LOG_FILE="$2"
> $LOG_FILE
FILENAME="${INPUT_XTC%.*}"
# Optional fourth argument
if [ "$#" -ge 2 ]; then
LIGNAME="$2"
else
LIGNAME=""
fi
#------ ANALYSIS ------------------
# Remove PBC
echo -e "Protein \n System" | $GMX trjconv -s $FILENAME.tpr -f $FILENAME.xtc -o "$FILENAME"_noPBC.xtc -pbc mol -center >> $LOG_FILE 2>&1
# RMSD to initial structure
echo -e "Backbone \n Backbone" | $GMX rms -s $FILENAME.tpr -f "$FILENAME"_noPBC.xtc -o rmsd.xvg -tu ns >> $LOG_FILE 2>&1
# RMSD to crystal structure
echo -e "Backbone \n Backbone" | $GMX rms -s em.tpr -f "$FILENAME"_noPBC.xtc -o rmsd_xtal.xvg -tu ns >> $LOG_FILE 2>&1
# RMSF
echo -e "C-alpha" | $GMX rmsf -s $FILENAME.tpr -f "$FILENAME"_noPBC.xtc -o rmsf.xvg -res >> $LOG_FILE 2>&1
# Radius of gyration
echo -e "Protein" | $GMX gyrate -s $FILENAME.tpr -f "$FILENAME"_noPBC.xtc -o gyrate.xvg >> $LOG_FILE 2>&1
# Hydrogen bonds
echo -e "MainChain+H \n MainChain+H" | $GMX hbond -s md.tpr -f md_noPBC.xtc -tu ns -num hbnum_mainchain.xvg >> $LOG_FILE 2>&1
echo -e "SideChain \n SideChain" | $GMX hbond -s $FILENAME.tpr -f "$FILENAME"_noPBC.xtc -tu ns -num hbnum_sidechain.xvg >> $LOG_FILE 2>&1
echo -e "Protein \n Water" | $GMX hbond -s $FILENAME.tpr -f "$FILENAME"_noPBC.xtc -tu ns -num hbnum_prot_wat.xvg >> $LOG_FILE 2>&1
if [ -n "$LIGNAME" ]; then
# Ligand-Protein hydrogen bonds
echo -e "1 \n 13" | $GMX hbond -s $FILENAME.tpr -f "$FILENAME"_noPBC.xtc -tu ns -num hbnum_prot_lig.xvg >> $LOG_FILE 2>&1
fi
# echo -e "Protein" \n "System" | gmx trjconv -s md.tpr -f md.xtc -o md_noPBC.xtc -pbc mol -center >> log 2>&1
# # RMSD to initial structure
# echo -e "Backbone" \n "Backbone" | gmx rms -s md.tpr -f md_noPBC.xtc -o rmsd.xvg -tu ns >> log.log 2>&1
# # RMSD to crystal structure
# echo -e "Backbone" \n "Backbone" | $GMX rms -s em.tpr -f "$FILENAME"_noPBC.xtc -o rmsd_xtal.xvg -tu ns >> $LOG_FILE 2>&1
# echo -e "Backbone" \n "Backbone" | gmx rms -s em.tpr -f md_noPBC.xtc -o rmsd_xtal.xvg -tu ns
# # RMSF
# echo -e "C-alpha" | $GMX rmsf -s $FILENAME.tpr -f "$FILENAME"_noPBC.xtc -o rmsf.xvg -res >> $LOG_FILE 2>&1
# # Radius of gyration
# echo -e "Protein" | $GMX gyrate -s $FILENAME.tpr -f "$FILENAME"_noPBC.xtc -o gyrate.xvg >> $LOG_FILE 2>&1
# # Hydrogen bonds
# #echo -e "MainChain+H \n MainChain+H" | $GMX hbond -s md.tpr -f md_noPBC.xtc -tu ns -num hbnum_mainchain.xvg
# echo -e "SideChain \n SideChain" | gmx hbond -s md.tpr -f md_noPBC.xtc -tu ns -num hbnum_sidechain.xvg
# echo -e "Protein \n Water" | $GMX hbond -s $FILENAME.tpr -f "$FILENAME"_noPBC.xtc -tu ns -num hbnum_prot_wat.xvg >> $LOG_FILE 2>&1
# echo -e "Protein \n Water" | $GMX hbond -s md.tpr -f md_noPBC.xtc -tu ns -num hbnum_prot_wat.xvg
# echo -e "MainChain+H \n MainChain+H" | gmx hbond -s md.tpr -f md_noPBC.xtc -tu ns -num hbnum_mainchain.xvg
\ No newline at end of file
src/scripts/equil_Gromacs.sh 0 → 100755
#!/bin/bash
if [ "$#" -lt 2 ]; then
echo "Usage: $0 sandbox_dir input_gro [ligand_name] [ligand_file] [ligand_gro]"
exit 1
fi
GMX='gmx'
SANDBOX_DIR="$1"
INPUT_GRO="$2"
LOG_FILE="$3"
> $LOG_FILE
echo "Starting GROMACS Equilibration Log" >> $LOG_FILE 2>&1
# Optional fourth argument
if [ "$#" -ge 5 ]; then
LIGNAME="$4"
LIGFILE="$5"
LIGGRO="$6"
else
LIGNAME=""
LIGFILE=""
LIGGRO=""
fi
#------- ENERGY MINIMISATION ------------
if ! ls em.gro 1> /dev/null 2>&1; then
$GMX grompp -f em.mdp -c $INPUT_GRO -p topol.top -o em.tpr >> $LOG_FILE 2>&1
$GMX mdrun -v -deffnm em >> $LOG_FILE 2>&1
if [ -f em.gro ]; then
echo "'em.gro' created"
echo "11 0" | $GMX energy -f em.edr -o potential.xvg
else
echo "Error: Failed to create 'em.gro'"
exit 1
fi
else
echo "'em.gro' already exists. Skipping energy minimisation."
fi
#----------Create posres files-----------
# 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."
if ! ls index.ndx 1> /dev/null 2>&1; then
echo "q" | $GMX make_ndx -f em.gro -o index.ndx
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
fi
else
# Step 2: if more than 1 chain, create posre_chain(i).itp files
# Get last residue number of each chain, by extracting residue numbers of all NME lines and selecting every 6th occurrence (last line of each NME block)
nme_residues=($(awk '/NME/ {resnum = substr($0,1,5); gsub(/ /,"",resnum); print resnum}' em.gro))
chain_end_residues=()
for ((i=5; i<${#nme_residues[@]}; i+=6)); do
chain_end_residues+=("${nme_residues[i]}")
done
echo "Chain end residues: ${chain_end_residues[@]}"
# Step 3: compute chain residue ranges
start=1
end=0
ranges=()
for end_residue in "${chain_end_residues[@]}"; do
end=$((end + end_residue))
ranges+=("${start}-${end}")
start=$((end + 1))
done
echo "Residue ranges per chain: ${ranges[@]}"
# Step 4: create index groups for each chain
# Start from existing index.ndx or create new
if [ ! -f index.ndx ]; then
$GMX make_ndx -f em.gro -o index.ndx << EOF
q
EOF
fi
# Step 5: add groups per chain
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
((i++))
done
# Step 6: generate posre.itp for each chain
i=1
for range in "${ranges[@]}"; do
if [ ! -f posre_chain${i}.itp ]; then
echo "Generating position restraints for chain${i}"
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
if [ "$i" -gt "1" ]; then
#Adjust atom indices so first = 1
awk '
/^\[ position_restraints \]/ { in_section=1; first_index=0; shift=0; print; next }
/^\[/ && !/\[ position_restraints \]/ { in_section=0 }
{
if (in_section && /^[0-9]/) {
if (first_index == 0) {
first_index = $1
if (first_index != 1) shift = first_index - 2
}
$1 = $1 - shift
}
print
}
' "posre_chain${i}.itp" > "posre_chain${i}_renum.itp" && mv "posre_chain${i}_renum.itp" "posre_chain${i}.itp"
fi
((i++))
fi
done
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
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
sed -i 's/Water_Cl-/Water_and_ions/g' index.ndx
echo "Group Water_and_ions created in index.ndx"
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
sed -i 's/Water_Na+/Water_and_ions/g' index.ndx
echo "Group Water_and_ions created in index.ndx"
fi
fi
#------ Update Water_and_ions is no ions present -----
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"
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" "$nvt_file"; then
sed -i "s|$original|$replacement|" "$nvt_file"
echo "$replacement added successfully to tc-grps group in $nvt_file."
else
echo "tc-grps line was not found in $nvt_file."
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."
else
echo "tc-grps line was not found in $npt_file."
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
#-------- 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"
if [ -n "$LIGNAME" ]; then
nvt_file="nvt.mdp"
npt_file="npt.mdp"
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" "$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
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" "$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."
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
if [ -f nvt.gro ]; then
echo "'nvt.gro' created"
echo -e "Temperature \n 0" | $GMX energy -f nvt.edr -o temperature.xvg
else
echo "Error: Failed to create 'nvt.gro'"
exit 1
fi
else
echo "'nvt.gro' already exists. Skipping NVT."
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
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
else
echo "Error: Failed to create 'npt.gro'"
exit 1
fi
else
echo "'npt.gro' already exists. Skipping NPT."
fi
\ No newline at end of file
src/scripts/mdp_files/em.mdp 0 → 100755
; 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
src/scripts/mdp_files/ions.mdp 0 → 100755
; 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.0 ; Cut-off for making neighbor list (short range forces)
coulombtype = cutoff ; Treatment of long range electrostatic interactions
rcoulomb = 1.0 ; long range electrostatic cut-off
rvdw = 1.0 ; long range Van der Waals cut-off
pbc = xyz ; Periodic Boundary Conditions
src/scripts/mdp_files/md.mdp 0 → 100755
title = Protein-ligand complex MD simulation
; Run parameters
integrator = md ; leap-frog integrator
nsteps = 5000 ; 2 * 50,000 = 100 ps (0.01 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 = 300 300 ; 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
src/scripts/mdp_files/npt.mdp 0 → 100755
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 = 300 300 ; 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
src/scripts/mdp_files/nvt.mdp 0 → 100755
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 = 300 300 ; 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
src/scripts/prod_Gromacs.sh 0 → 100755
#!/bin/bash
if [ "$#" -lt 2 ]; then
echo "Usage: $0 input_gro npt_cpt_file"
exit 1
fi
GMX='gmx'
INPUT_GRO="$1"
NPT_CPT_FILE="$2"
LOG_FILE="$3"
> $LOG_FILE
#------- 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
echo "y" | $GMX mdrun -v -deffnm md >> $LOG_FILE 2>&1
else
echo "'md.gro' already exists. Skipping production MD."
fi
\ No newline at end of file
src/scripts/run_tleap.sh 0 → 100755
#!/bin/bash
# Usage: ./run_tleap.sh input.pdb
if [ $# -ne 3 ]; then
echo "Usage: $0 sandbox_dir input.pdb pdb_id"
exit 1
fi
SANDBOX_DIR=$1
PDBFILE=$2
PDB_ID=$3
# Create tleap input file
cat > leap.in << EOF
source leaprc.protein.ff14SB
source leaprc.water.tip3p
# Map PDB atom names to template atom names
addPdbAtomMap { { "CH3" "C" } { "HH31" "H1" } { "HH32" "H2" } { "HH33" "H3" } }
mol = loadpdb ${SANDBOX_DIR}/${PDBFILE}
solvatebox mol TIP3PBOX 16
addions mol Cl- 0 # Neutralize system
addions mol Na+ 0 # Neutralize system
saveamberparm mol ${SANDBOX_DIR}/${PDB_ID}.prmtop ${SANDBOX_DIR}/${PDB_ID}.inpcrd
savepdb mol ${SANDBOX_DIR}/${PDB_ID}_tleap.pdb
quit
EOF
# Run tleap
tleap -f leap.in
# if output files not in directory, echo that tleap failed, otyherwise confirm success
ls ${SANDBOX_DIR}/${PDB_ID}.prmtop ${SANDBOX_DIR}/${PDB_ID}.inpcrd ${SANDBOX_DIR}/${PDB_ID}_tleap.pdb > /dev/null 2>&1
if [ $? -ne 0 ]; then
echo "tleap failed to generate output files."
else
echo "${PDBFILE} processed. Generated ${PDB_ID}.prmtop, ${PDB_ID}.inpcrd, and ${PDB_ID}_tleap.pdb."
fi
\ No newline at end of file
src/scripts/run_tleap_ligand.sh 0 → 100755
#!/bin/bash
# Usage: ./run_tleap.sh input.pdb
if [ $# -ne 4 ]; then
echo "Usage: $0 sandbox_dir complex_pdb ligand_name prepi_file"
exit 1
fi
SANDBOX_DIR=$1 # PDBFILE already has sandbox path
PDBFILE=$2
LIGNAME=$3
PREPI_FILE=$4
# Create tleap input file
cat > leap.in << EOF
source leaprc.protein.ff14SB
source leaprc.gaff
source leaprc.water.tip3p
# Map PDB atom names to template atom names
addPdbAtomMap { { "CH3" "C" } { "HH31" "H1" } { "HH32" "H2" } { "HH33" "H3" } { "CL1" "Cl1" } { "CL2" "Cl2" } }
# Load ligand parameters
loadamberprep ${SANDBOX_DIR}/${PREPI_FILE}
loadamberparams ${SANDBOX_DIR}/${LIGNAME}.frcmod
# PDBFILE already has sandbox path
mol = loadpdb ${PDBFILE}
solvatebox mol TIP3PBOX 16
addions mol Cl- 0 # Neutralize system
addions mol Na+ 0 # Neutralize system
saveamberparm mol ${SANDBOX_DIR}/complex.prmtop ${SANDBOX_DIR}/complex.inpcrd
savepdb mol ${SANDBOX_DIR}/complex_tleap.pdb
quit
EOF
# Run tleap
tleap -f leap.in
\ No newline at end of file
src/scripts/solvate_Gromacs_ff14sb.sh 0 → 100755
#!/bin/bash
# location input mdp files
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
MDP_FILES="$SCRIPT_DIR/../../sandbox/experiments/inputs_unconstrained_MD"
# get input file
if [ $# -ne 1 ]; then
echo "Usage: $0 input.pdb"
exit 1
fi
PDBFILE=$(basename "$1" .pdb)
gmx editconf -f "$PDBFILE" -o box.gro -c -d 1.2 -bt cubic
gmx solvate -cp box.gro -cs spc216.gro -o solv.gro -p topol.top
gmx grompp -f "$MDP_FILES"/ions.mdp -c solv.gro -p topol.top -o ions.tpr
echo "13" | gmx genion -s ions.tpr -o solv_ions.gro -p topol.top -pname NA -nname CL -neutral
\ No newline at end of file
src/scripts/solvate_Gromacs_ff99sb.sh 0 → 100755
#!/bin/bash
# location input mdp files
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
MDP_FILES="$SCRIPT_DIR/../../experiments/inputs_unconstrained_MD"
# get input file
if [ $# -ne 1 ]; then
echo "Usage: $0 input.pdb"
exit 1
fi
PDBFILE=$(basename "$1" .pdb)
echo "5 0 1 1" | gmx pdb2gmx -f "$PDBFILE" -o processed.gro -ter
gmx editconf -f processed.gro -o box.gro -c -d 1.2 -bt cubic
gmx solvate -cp box.gro -cs spc216.gro -o solv.gro -p topol.top
gmx grompp -f "$MDP_FILES"/ions.mdp -c solv.gro -p topol.top -o ions.tpr
echo "13" | gmx genion -s ions.tpr -o solv_ions.gro -p topol.top -pname NA -nname CL -neutral
\ No newline at end of file
src/tools/MMPBSA.py 0 → 100644
import subprocess
from pathlib import Path
import subprocess, shlex
from tkinter import constants
import parmed as pmd # type: ignore
import sys
import re
import os
def run_gmxMMPBSA(sandbox_dir: str, pdb_id: str, nsteps:str, nstxout_compressed:str, 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_infile.write(f'''&general
sys_name={pdb_id}
startframe=1
endframe={int(float(nframes))}
interval=5
temperature={int(float(temp))}
verbose=2
/
&pb
ipb = 2 # Dielectric model for PB
inp = 1 # Nonpolar solvation method
sander_apbs = 0 # Use sander.APBS?
indi = 1.0 # Internal dielectric constant
exdi = 80.0 # External dielectric constant
emem = 4.0 # Membrane dielectric constant
smoothopt = 1 # Set up dielectric values for finite-difference grid edges that are located across the solute/solvent dielectric boundary
istrng = 0.0 # Ionic strength (M)
radiopt = 1 # Use optimized radii?
prbrad = 1.4 # Probe radius
iprob = 2.0 # Mobile ion probe radius (Angstroms) for ion accessible surface used to define the Stern layer
sasopt = 0 # Molecular surface in PB implict model
arcres = 0.25 # The resolution (Å) to compute solvent accessible arcs
memopt = 0 # Use PB optimization for membrane
poretype = 1 # Use exclusion region for channel proteins
npbopt = 0 # Use NonLinear PB solver?
solvopt = 1 # Select iterative solver
accept = 0.001 # Sets the iteration convergence criterion (relative to the initial residue)
linit = 1000 # Number of SCF iterations
fillratio = 4.0 # Ratio between the longest dimension of the rectangular finite-difference grid and that of the solute
scale = 2.0 # 1/scale = grid spacing for the finite difference solver (default = 1/2 Å)
nbuffer = 0.0 # Sets how far away (in grid units) the boundary of the finite difference grid is away from the solute surface
nfocus = 2 # Electrostatic focusing calculation
fscale = 8 # Set the ratio between the coarse and fine grid spacings in an electrostatic focussing calculation
npbgrid = 1 # Sets how often the finite-difference grid is regenerated
bcopt = 5 # Boundary condition option
eneopt = 2 # Compute electrostatic energy and forces
frcopt = 0 # Output for computing electrostatic forces
scalec = 0 # Option to compute reaction field energy and forces
cutfd = 5.0 # Cutoff for finite-difference interactions
cutnb = 0.0 # Cutoff for nonbonded interations
nsnba = 1 # Sets how often atom-based pairlist is generated
decompopt = 2 # Option to select different decomposition schemes when INP = 2
use_rmin = 1 # The option to set up van der Waals radii
sprob = 0.557 # Solvent probe radius for SASA used to compute the dispersion term
vprob = 1.3 # Solvent probe radius for molecular volume (the volume enclosed by SASA)
rhow_effect = 1.129 # Effective water density used in the non-polar dispersion term calculation
use_sav = 1 # Use molecular volume (the volume enclosed by SASA) for cavity term calculation
cavity_surften = 0.0378 # Surface tension
cavity_offset = -0.5692 # Offset for nonpolar solvation calc
maxsph = 400 # Approximate number of dots to represent the maximum atomic solvent accessible surface
maxarcdot = 1500 # Number of dots used to store arc dots per atom
npbverb = 0 # Option to turn on verbose mode
/
''')
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,
"-O",
"-i", "mmpbsa.in",
"-cs", tpr_file,
"-ct", xtc_file,
"-ci", index_file,
"-cg", "1", "13",
"-cp", topol_file,
"-o", "FINAL_RESULTS_MMPBSA.dat",
"-eo", "FINAL_RESULTS_MMPBSA.csv",
"-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"
src/tools/RAG_tools.py 0 → 100644
from paperqa import Docs, Settings
from tqdm import tqdm
import os
import os.path
import contextlib
import pickle
from src import constants
documents: Docs | None = None
def _load_documents() -> Docs:
docs = Docs()
pdf_files = list(constants.PAPER_DIR.glob("*.pdf"))
total_files = len(pdf_files)
pickled_docs = "my_docs.pkl"
if not os.path.exists(pickled_docs):
with tqdm(
total=total_files,
desc="Loading PDFs",
bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt} ({percentage:.1f}%) [ time left: {remaining}, time spent: {elapsed}]",
) as pbar:
for file_path in pdf_files:
# suppress output from docs.add()
with open(os.devnull, "w") as fnull:
with contextlib.redirect_stdout(fnull), contextlib.redirect_stderr(fnull):
docs.add(file_path)
pbar.update(1)
# save embeddings of the documents
with open(pickled_docs, "wb") as f:
pickle.dump(docs, f)
else:
# load embeddings of the documents
with open(pickled_docs, "rb") as f:
docs = pickle.load(f)
return docs
def search_papers(query: dict):
global documents
if not documents:
documents = _load_documents()
if isinstance(query, dict):
query = query.get("query")
if not isinstance(query, str):
raise ValueError(f"search_papers expected a string query, got: {type(query)}")
paper_directory = constants.PAPER_DIR
if paper_directory is None:
raise ValueError(
"'paper_dir' is None. To use this tool, the user must provide a directory with PDFs at the start."
)
settings = Settings(
# Retrieval size — more is NOT always better
evidence_k=8, # retrieve top 8 chunks per query
max_chunk_size=800, # avoid huge chunks; MD details are often local
rerank_k=20, # lightly expand initial search before reranking
# LLM settings
temperature=0.1, # scientific, deterministic tone
answer_temperature=0.0, # final answers must be strict, non-creative
# Trust/scientific correctness
require_citations=True, # every claim must have a supporting doc
max_tokens=4096, # MD methods can be verbose
summary_length=5, # keep evidence chunks tight
# Error-handling / agent behavior
cohere_reranker=False, # use built-in reranker (fast, good enough)
retries=2, # avoid failures during batch queries
timeout=120, # MD queries can be long
)
result = documents.query(query, settings=settings)
answer = result.formatted_answer
if "I cannot answer." in answer:
answer += f" Check to ensure there's papers in {paper_directory}"
return answer
src/tools/amber_tools.py 0 → 100644
import subprocess
from pathlib import Path
import parmed as pmd # type: ignore
from src import constants
from src.utils import get_class_logger
logger = get_class_logger(__name__)
def run_tleap(sandbox_dir: str, input_pdb: str, pdb_id: str) -> str:
"""
Run tleap preparation using run_tleap.sh.
"""
script = constants.SCRIPTS_DIR / "run_tleap.sh"
result = subprocess.run(
[str(script), sandbox_dir, input_pdb, pdb_id], cwd=sandbox_dir, capture_output=True, text=True
)
if result.returncode != 0:
# tleap often puts errors in stdout
error_text = "\n".join(filter(None, [result.stderr, result.stdout]))
return f"tleap run failed with error:\n{error_text}"
else:
# Parmed
try:
prmtop_path = f"{sandbox_dir}/{pdb_id}.prmtop"
inpcrd_path = f"{sandbox_dir}/{pdb_id}.inpcrd"
parmed_cm = pmd.load_file(prmtop_path, inpcrd_path)
parmed_cm.save(f"{sandbox_dir}/topol.top")
parmed_cm.save(f"{sandbox_dir}/{pdb_id}.gro")
except Exception as e:
return f"ParmEd failed: {type(e).__name__}: {e}, {result}"
return f"tleap ran successfully with output: {result.stdout}. \n New files added: {sandbox_dir}/topol.top, {sandbox_dir}/{pdb_id}.gro"
def run_tleap_ligand(sandbox_dir: str, input_pdb: str, pdb_id: str, ligand_file: str, ligand_name: str) -> str:
"""
Run tleap preparation using run_tleap.sh, for a protein-ligand complex.
"""
# complex.pdb
with (
open(f"{sandbox_dir}/{input_pdb}", "r") as pdb_infile,
open(f"{sandbox_dir}/{ligand_file}", "r") as ligand_infile,
open(f"{sandbox_dir}/complex.pdb", "w") as outfile,
):
for line in pdb_infile:
if not line.startswith("END"):
outfile.write(line)
for line in ligand_infile:
if line.startswith("HETATM"):
outfile.write(line)
outfile.write("TER\n")
outfile.write("END\n")
complex_pdb = f"{sandbox_dir}/complex.pdb"
# tleap with ligand
script = constants.SCRIPTS_DIR / "run_tleap_ligand.sh"
if Path(f"{sandbox_dir}/{ligand_name}_fixed.prepi").exists():
prepi_file = f"{ligand_name}_fixed.prepi"
else:
prepi_file = f"{ligand_name}.prepi"
result = subprocess.run(
[str(script), sandbox_dir, complex_pdb, ligand_name, prepi_file],
cwd=sandbox_dir,
capture_output=True,
text=True,
)
if result.returncode != 0:
# tleap often puts errors in stdout
error_text = "\n".join(filter(None, [result.stderr, result.stdout]))
return f"tleap run failed with error:\n{error_text}"
else:
# Parmed
try:
prmtop_path = f"{sandbox_dir}/complex.prmtop"
inpcrd_path = f"{sandbox_dir}/complex.inpcrd"
parmed_cm = pmd.load_file(prmtop_path, inpcrd_path)
parmed_cm.save(f"{sandbox_dir}/topol.top")
parmed_cm.save(f"{sandbox_dir}/complex.gro")
except Exception as e:
return f"ParmEd failed: {type(e).__name__}: {e}"
return f"tleap ran successfully with output: {result.stdout}. \n New files added: {sandbox_dir}/topol.top, {sandbox_dir}/complex.gro"
src/tools/coding_tools.py 0 → 100644
import os
from pathlib import Path
def find_input(directory: Path) -> str:
try:
items = []
for file in sorted(os.listdir(directory)):
item_path = os.path.join(directory, file)
if os.path.isdir(item_path):
pass
else:
items.append(file)
if not items:
return f"Empty directory: {directory}. You forgot to upload your PDB file"
pdb_files = list(directory.glob("*.pdb"))
pdb_files_str = "\n".join(str(f) for f in pdb_files)
return f"User uploads to {directory}:\n{pdb_files_str}"
except Exception as e:
return f"Error listing files: {str(e)}"
def read_file(path: str) -> str:
try:
with open(path, "r", encoding="utf-8") as f:
content = f.read()
return f"File contents of {path}:\n{content}"
except FileNotFoundError:
return f"File not found: {path}"
except Exception as e:
return f"Error reading file: {str(e)}"
def list_files(path: str) -> str:
try:
if not os.path.exists(path):
return f"Path not found: {path}"
items = []
for item in sorted(os.listdir(path)):
item_path = os.path.join(path, item)
if os.path.isdir(item_path):
items.append(f"[DIR] {item}/")
else:
items.append(f"[FILE] {item}")
if not items:
return f"Empty directory: {path}"
return f"Contents of {path}:\n" + "\n".join(items)
except Exception as e:
return f"Error listing files: {str(e)}"
def edit_file(path: str, old_text: str, new_text: str) -> str:
try:
if os.path.exists(path) and old_text:
with open(path, "r", encoding="utf-8") as f:
content = f.read()
if old_text not in content:
return f"Text not found in file: {old_text}"
content = content.replace(old_text, new_text)
with open(path, "w", encoding="utf-8") as f:
f.write(content)
return f"Successfully edited {path}"
else:
# Only create directory if path contains subdirectories
dir_name = os.path.dirname(path)
if dir_name:
os.makedirs(dir_name, exist_ok=True)
with open(path, "w", encoding="utf-8") as f:
f.write(new_text)
return f"Successfully created {path}"
except Exception as e:
return f"Error editing file: {str(e)}"
src/tools/gromacs_tools.py 0 → 100644
import subprocess
from pathlib import Path
import re
import shutil
import sys
from src import constants
from src.utils import get_class_logger
import time
logger = get_class_logger(__name__)
def gromacs_equil(sandbox_dir: str, input_gro: 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
input_path = Path(f"{sandbox_dir}/topol.top")
backup_path = Path(f"{sandbox_dir}/topol_without_posre.top")
if not input_path.exists():
raise FileNotFoundError(f"{input_path} not found.")
# Make a backup
shutil.copyfile(input_path, backup_path)
logger.info(f"Backup created: {backup_path}")
text = input_path.read_text(encoding="utf-8", errors="replace")
# --- Detect all system names (system, system1, system2, etc.) ---
system_matches = re.findall(r"^\s*(system\d*?)\b", text, flags=re.M)
systems = sorted(set(system_matches), key=lambda x: int(re.search(r"\d*$", x).group() or 0))
num_systems = len(systems)
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)}")
# --- Create posre include block ---
def make_posre_block(posre_file):
return f'; Include Position restraint file\n#ifdef POSRES\n#include "{posre_file}"\n#endif\n\n'
if ligand_file is not None:
ligand_posre_block = (
f'; Include Position restraint file\n#ifdef POSRES\n#include "posre_{ligand_name}.itp"\n#endif\n\n'
)
# --- Split by [ moleculetype ] sections ---
header_re = re.compile(r"^\[\s*moleculetype\s*\]", flags=re.I | re.M)
headers = list(header_re.finditer(text))
# If there are no [ moleculetype ] sections, exit
if not headers:
logger.warning("No [ moleculetype ] sections found. No changes made.")
sys.exit(0)
# Get the position of the first non-moleculetype section (e.g., [ system ])
next_non_moleculetype = re.search(r"^\[\s*(system|molecules)\s*\]", text, flags=re.I | re.M)
end_of_moleculetype = next_non_moleculetype.start() if next_non_moleculetype else len(text)
# Define the boundaries for moleculetype sections
positions = [m.start() for m in headers if m.start() < end_of_moleculetype] + [end_of_moleculetype]
preamble = text[: positions[0]]
segments = []
inserted_blocks = []
inserted_ligand = False
# --- Loop over each [ moleculetype ] block ---
for i in range(len(positions) - 1):
seg = text[positions[i] : positions[i + 1]]
# Only modify real moleculetype blocks
match = re.search(r"^\s*(system\d*|system)\b\s+\d+", seg, flags=re.M)
if match:
system_name = match.group(1)
chain_index = re.search(r"\d+$", system_name)
chain_num = int(chain_index.group()) if chain_index else 1
posre_file = f"posre_chain{chain_num}.itp" if num_systems > 1 else "posre.itp"
posre_block = make_posre_block(posre_file)
if posre_block not in seg:
seg = seg.rstrip() + "\n\n" + posre_block
inserted_blocks.append(system_name)
# Check for ligand_name
if ligand_file is not None:
if not inserted_ligand:
if re.search(r"^\s*" + re.escape(ligand_name) + r"\b\s+\d+", seg, flags=re.M):
include_lig = f'#include "posre_{ligand_name}.itp"'
if include_lig not in seg:
seg = seg.rstrip() + "\n\n" + ligand_posre_block
inserted_ligand = True
segments.append(seg)
# Reassemble modified part + untouched rest
modified_text = preamble + "".join(segments) + text[end_of_moleculetype:]
# Write result
input_path.write_text(modified_text, encoding="utf-8")
logger.info(f"Added position restraints for: {', '.join(inserted_blocks) or 'none'}")
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}"
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
result = subprocess.run(cmd, cwd=sandbox_dir, stdout=sys.stdout, stderr=sys.stderr, text=True)
gromacs_output = ""
if log_file_path.exists():
try:
gromacs_output = log_file_path.read_text(encoding="utf-8")
except Exception as e:
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.
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:
"""
Run production MD with GROMACS using prod_Gromacs.sh.
"""
script = constants.SCRIPTS_DIR / "prod_Gromacs.sh"
log_file_path = Path(f"{sandbox_dir}/gromacs_production.log")
cmd = [str(script), input_gro, npt_cpt_file, log_file_path]
if ligand_name is not None:
cmd.append(ligand_name)
cmd.append(f"{sandbox_dir}/{ligand_name}.gro")
result = subprocess.run(cmd, cwd=sandbox_dir, stdout=sys.stdout, stderr=sys.stderr, text=True)
gromacs_output = ""
if log_file_path.exists():
try:
gromacs_output = log_file_path.read_text(encoding="utf-8")
except Exception as e:
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.
else:
# Report success and return the captured GROMACS output
return (f"Equilibration ran successfully. Full GROMACS output:\n"
f"{gromacs_output}")
def gromacs_analysis(sandbox_dir: str, input_xtc: str, ligand_name=None) -> str:
"""
Run production MD with GROMACS using prod_Gromacs.sh.
"""
script = constants.SCRIPTS_DIR / "analysis_Gromacs.sh"
log_file_path = Path(f"{sandbox_dir}/gromacs_analysis.log")
cmd = [str(script), input_xtc, log_file_path]
if ligand_name is not None:
cmd.append(ligand_name)
cmd.append(f"{sandbox_dir}/{ligand_name}.gro")
result = subprocess.run(cmd, cwd=sandbox_dir, stdout=sys.stdout, stderr=sys.stderr, text=True)
gromacs_output = ""
if log_file_path.exists():
try:
gromacs_output = log_file_path.read_text(encoding="utf-8")
except Exception as e:
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.
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/ligand_tools.py 0 → 100644
import subprocess
from pathlib import Path
import subprocess, shlex
import re
from src.utils import get_class_logger
logger = get_class_logger(__name__)
def param_ligand(sandbox_dir: str, ligand_file: str, ligand_name: str, charge_ligand: str | None = None) -> str:
# Find charge of ligand if not provided
charges = []
tmp_mol2file = f"{sandbox_dir}/ligand_tmp.mol2"
# Create temporary mol2 file using obabel
cmd = shlex.split(f"obabel {sandbox_dir}/{ligand_file} -O {tmp_mol2file}")
run_1 = subprocess.run(cmd, cwd=sandbox_dir, capture_output=True, text=True)
if run_1.returncode != 0:
error_text = "\n".join(filter(None, [run_1.stderr, run_2.stdout]))
return f"Ligand parameterization failed with error: {error_text}"
# Read the mol2 file to find the charge
in_atom_section = False
with open(tmp_mol2file, "r") as f:
for line in f:
line = line.strip()
if line.startswith("@<TRIPOS>ATOM"):
in_atom_section = True
continue
elif line.startswith("@<TRIPOS>") and in_atom_section:
break
elif in_atom_section and line:
try:
charge = float(line.split()[-1])
charges.append(charge)
except ValueError:
pass # skip malformed lines
total_charge = sum(charges)
charge_ligand = round(total_charge)
# Clean up temporary mol2 file
Path(tmp_mol2file).unlink(missing_ok=True)
logger.info(f"Charge of ligand {ligand_file} determined to be {charge_ligand}")
# Create mol2 file using antechamber
cmd = shlex.split(
f"antechamber -i {sandbox_dir}/{ligand_file} -fi pdb -o {sandbox_dir}/{ligand_name}.mol2 -fo mol2 -c bcc -nc {charge_ligand} -s 2"
)
run_2 = subprocess.run(cmd, cwd=sandbox_dir, capture_output=True, text=True)
if run_2.returncode != 0:
error_text = "\n".join(filter(None, [run_2.stderr, run_2.stdout]))
return f"Ligand parameterization failed with error: {error_text}"
logger.info(f"Mol2 file for ligand {ligand_name} created")
cmd = shlex.split(f"sed -i 's/UNL/{ligand_name}/g' {sandbox_dir}/{ligand_name}.mol2")
run_3 = subprocess.run(cmd, cwd=sandbox_dir, capture_output=True, text=True)
if run_3.returncode != 0:
error_text = "\n".join(filter(None, [run_3.stderr, run_3.stdout]))
return f"Ligand parameterization failed with error: {error_text}"
# Parmed to make total charge an integer
# pmd.load_file(f"{sandbox_dir}/{ligand_name}.mol2").fix_charges(precision=4).save(f"{sandbox_dir}/{ligand_name}_fixed.mol2")
# Create prepi file using antechamber
cmd = shlex.split(
f"antechamber -i {sandbox_dir}/{ligand_name}.mol2 -fi mol2 -o {sandbox_dir}/{ligand_name}.prepi -fo prepi -c bcc"
)
run_4 = subprocess.run(cmd, cwd=sandbox_dir, capture_output=True, text=True)
if run_4.returncode != 0:
error_text = "\n".join(filter(None, [run_4.stderr, run_4.stdout]))
return f"Ligand parameterization failed with error: {error_text}"
# Create frcmod file using parmchk2
cmd = shlex.split(f"parmchk2 -i {sandbox_dir}/{ligand_name}.mol2 -f mol2 -o {sandbox_dir}/{ligand_name}.frcmod")
run_5 = subprocess.run(cmd, cwd=sandbox_dir, capture_output=True, text=True, check=True)
if run_5.returncode != 0:
error_text = "\n".join(filter(None, [run_5.stderr, run_5.stdout]))
return f"Ligand parameterization failed with error: {error_text}"
# Update charge prepi file
def fix_charges(input_file, output_file=None):
"""
Adjust charges so the total is an integer.
Only modifies one atom's charge (last atom in coordinate/charge section).
"""
with open(input_file) as f:
lines = f.readlines()
# Identify section that looks like atom definitions (before LOOP/IMPROPER)
atom_section_end = len(lines)
for i, line in enumerate(lines):
if re.match(r"^\s*(LOOP|IMPROPER|DONE|STOP)\b", line):
atom_section_end = i
break
atom_lines = []
charges = []
float_pattern = re.compile(r"([-+]?\d*\.\d+|\d+)(?!.*\S)")
for i, line in enumerate(lines[:atom_section_end]):
tokens = line.split()
if len(tokens) >= 8: # typical atom line length
last = tokens[-1]
try:
charge = float(last)
atom_lines.append(i)
charges.append(charge)
except ValueError:
pass
if not charges:
logger.warning("No atom charges found.")
return
total = sum(charges)
target = round(total)
delta = target - total
if abs(delta) < 1e-6:
logger.warning(f"Already integer total ({total:.6f}). No change.")
return
# Adjust last atom charge in section
last_atom_idx = atom_lines[-1]
old_charge = charges[-1]
new_charge = old_charge + delta
lines[last_atom_idx] = float_pattern.sub(f"{new_charge:.6f}", lines[last_atom_idx])
if output_file is None:
output_file = Path(input_file).with_name(Path(input_file).stem + "_fixed.res")
with open(output_file, "w") as f:
f.writelines(lines)
logger.info(f"Adjusted total charge: {total:.6f} → {target}")
logger.info(f"Atom line {last_atom_idx + 1}: {old_charge:.6f} → {new_charge:.6f}")
logger.info(f"Saved to: {output_file}")
fix_charges(f"{sandbox_dir}/{ligand_name}.prepi", f"{sandbox_dir}/{ligand_name}_fixed.prepi")
return "Ligand parameterisation complete. File saved to {sandbox_dir}/{ligand_name}_fixed.prepi"
src/tools/map.py 0 → 100644
import os
from src.tools.amber_tools import run_tleap, run_tleap_ligand
from src.tools.gromacs_tools import gromacs_equil, gromacs_production, gromacs_analysis
from src.tools.pdb_tools import fix_pdb_file
from src.tools.ligand_tools import param_ligand
from src.tools.pdb_tools import prepare_pdb_file_ligand, add_caps, rename_histidines, fetch_and_save_pdb
from src.tools.coding_tools import read_file, edit_file, list_files, find_input
from src.tools.RAG_tools import search_papers
from src.tools.MMPBSA import run_gmxMMPBSA
from src import constants
def truncate_file_output(full_content: str) -> str:
"""Truncate long file outputs to prevent LLM token overload."""
if len(full_content) <= 2 * constants.MAX_CHARACTERS_TO_LOG:
return full_content
return f"{full_content[:constants.MAX_CHARACTERS_TO_LOG]}... truncated ... {full_content[-constants.MAX_CHARACTERS_TO_LOG:]}" if full_content else ""
TOOL_MAP = {
# Basic File Operations
"read_file": lambda _, i: truncate_file_output(read_file(i['path'])),
"list_files": lambda s, _: list_files(s.sandbox_dir),
"find_input": lambda s, _: find_input(s.sandbox_dir),
"edit_file": lambda _, i: edit_file(i["path"], i["old_text"], i["new_text"]),
# Protein prep
"fetch_and_save_pdb": lambda s, i: fetch_and_save_pdb(s.sandbox_dir, i["pdb_id"], i["output_pdb"]),
"fix_pdb_file": lambda _, i: fix_pdb_file(i["input_pdb"], f"{os.path.splitext(i['input_pdb'])[0]}_fixed.pdb"),
"prepare_pdb_file_ligand": lambda s, i: prepare_pdb_file_ligand(s.sandbox_dir, i["pdb_id"], i["ligand_name"]),
"add_caps": lambda s, i: add_caps(s.sandbox_dir, i["input_pdb"], i["pdb_id"]),
"rename_histidines": lambda s, i: rename_histidines(s.sandbox_dir, i["input_pdb"], i["pdb_id"]),
# Ligand handling
"param_ligand": lambda s, i: param_ligand(s.sandbox_dir, i["ligand_file"], i["ligand_name"]),
# AMBER-related
"run_tleap": lambda s, i: run_tleap(s.sandbox_dir, i["input_pdb"], i["pdb_id"]),
"run_tleap_ligand": lambda s, i: run_tleap_ligand(
s.sandbox_dir, i["input_pdb"], i["pdb_id"], i["ligand_file"], i["ligand_name"]
),
# 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")
),
"gromacs_production": lambda s, i: gromacs_production(
s.sandbox_dir, i["input_gro"], i["npt_cpt_file"], 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"],
),
# # RAG tools
"search_papers": lambda _, i: search_papers(i["query"]),
}
src/tools/topology_tools.py 0 → 100644
import subprocess
from pathlib import Path
def fix_topology_negative(topfile: str, sandbox_dir: str) -> str:
"""
This script should be used if ff14sb AMBER force field is used in tleap.
This script should be used if topfile has a net negative charge.
This script adds the missing water and ions parameters to the topology file.
Args:
topfile (str): Path to the topology file to be fixed.
Returns:
topol.top (str): Fixed topology file.
"""
output_file = Path(sandbox_dir) / "topol.top"
# Read the original file
with open(topfile, "r") as f:
lines = f.readlines()
# Lines to insert
insert_lines = [
"HW 1 1.008 0.0000 A 0.00000e+00 0.00000e+00\n"
"OW 8 16.00 0.0000 A 3.15061e-01 6.36386e-01\n"
"\n",
"; Include topology for water\n",
'#include "amber99.ff/tip3p.itp"\n',
"\n",
"[ moleculetype ]\n",
"; molname nrexcl\n",
"NA 1\n",
"\n",
"[ atoms ]\n",
"; id at type res nr residu name at name cg nr charge\n",
"1 NA 1 NA NA 1 1.00000\n",
]
# Find [ atomtypes ] section
start_index = None
for i, line in enumerate(lines):
if line.strip() == "[ atomtypes ]":
start_index = i
break
if start_index is None:
raise ValueError("No [ atomtypes ] section found in the file.")
# Find last atom line (last line before blank line)
end_index = start_index + 1
while end_index < len(lines) and lines[end_index].strip() != "":
end_index += 1
# Insert lines
lines = lines[:end_index] + insert_lines + lines[end_index:]
# Write to output
with open(output_file, "w") as f:
f.writelines(lines)
return "Successfully added missing water ions and fixed the topology. Output file is saved to topol.top"
def fix_topology_positive(topfile: str, sandbox_dir: str) -> str:
"""
This script should be used if ff14sb AMBER force field is used in tleap.
This script should be used if topfile has a net positive charge.
This script adds the missing water and ions parameters to the topology file.
Args:
topfile (str): Path to the topology file to be fixed.
Returns:
topol.top (str): Fixed topology file.
"""
output_file = sandbox_dir / "topol.top"
# Read the original file
with open(topfile, "r") as f:
lines = f.readlines()
# Lines to insert
insert_lines = [
"HW 1 1.008 0.0000 A 0.00000e+00 0.00000e+00\n"
"OW 8 16.00 0.0000 A 3.15061e-01 6.36386e-01\n"
"\n",
"; Include topology for water\n",
'#include "amber99.ff/tip3p.itp"\n',
"\n",
"[ moleculetype ]\n",
"; molname nrexcl\n",
"CL 1\n",
"\n",
"[ atoms ]\n",
"; id at type res nr residu name at name cg nr charge\n",
"1 CL 1 CL CL 1 -1.00000",
]
# Find [ atomtypes ] section
start_index = None
for i, line in enumerate(lines):
if line.strip() == "[ atomtypes ]":
start_index = i
break
if start_index is None:
raise ValueError("No [ atomtypes ] section found in the file.")
# Find last atom line (last line before blank line)
end_index = start_index + 1
while end_index < len(lines) and lines[end_index].strip() != "":
end_index += 1
# Insert lines
lines = lines[:end_index] + insert_lines + lines[end_index:]
# Write to output
with open(output_file, "w") as f:
f.writelines(lines)
return "Successfully added missing water ions and fixed the topology. Output file is saved to topol.top"
def analyze_Gromacs(sandbox_dir: str) -> None:
"""
This script analyzes the trajectory from the production MD simulation using Gromacs.
It calculates RMSD, RMSF, Radius of Gyration, and secondary structure content.
It uses the previously generated md.xtc and topol.top files as inputs.
The script is ran after the production_Gromacs script and requires md.gro and md.xtc files.
Args:
None
Returns:
rmsd.xvg (str): RMSD plot data.
rmsf.xvg (str): RMSF plot data.
gyrate.xvg (str): Radius of Gyration plot data.
"""
rmsd_command = "echo 4 4 | gmx rms -s md.tpr -f md.xtc -o rmsd.xvg -tu ns"
subprocess.run([rmsd_command], cwd=sandbox_dir, check=True)
rmsf_command = "echo 4 | gmx rmsf -s md.tpr -f md.xtc -o rmsf.xvg"
subprocess.run([rmsf_command], cwd=sandbox_dir, check=True)
rg_command = "echo 4 | gmx gyrate -s md.tpr -f md.xtc -o gyrate.xvg"
subprocess.run([rg_command], cwd=sandbox_dir, check=True)
return None
src/utils.py 0 → 100644
import logging
import sys
from pathlib import Path
from src import constants
import json
import pathlib
from datetime import datetime
def get_class_logger(class_name: str, log_dir: Path = None) -> logging.Logger:
"""
Create or retrieve a logger specific to a class.
Each class writes to its own log file inside agent_logs/.
"""
if log_dir is None:
log_dir = Path(__file__).resolve().parent.parent / "agent_logs"
log_dir.mkdir(exist_ok=True)
log_file = log_dir / f"{class_name}.log"
logger = logging.getLogger(class_name)
logger.setLevel(logging.INFO)
# Avoid adding duplicate handlers
if not logger.handlers:
file_handler = logging.FileHandler(log_file)
formatter = logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")
file_handler.setFormatter(formatter)
logger.addHandler(file_handler)
# also print to stdout
stream_handler = logging.StreamHandler(sys.stdout)
stream_handler.setFormatter(formatter)
logger.addHandler(stream_handler)
return logger
def append_jsonl(data, filename):
"""Append one JSON record per line."""
with open(filename, "a", encoding="utf-8") as f:
f.write(json.dumps(data, ensure_ascii=False) + "\n")
def truncate_string(string):
if not string:
return ""
if len(string) <= 2 * constants.MAX_CHARACTERS_TO_LOG:
return string
return f"{string[:constants.MAX_CHARACTERS_TO_LOG]}... truncated ... {string[-constants.MAX_CHARACTERS_TO_LOG:]}"
def is_path_child_dir(potential_child_dir: str | Path, dir: str | Path) -> bool:
"""Ensure that the requested path stays within the sandbox."""
if isinstance(potential_child_dir, str):
potential_child_dir = pathlib.Path(potential_child_dir)
if isinstance(dir, str):
dir = pathlib.Path(dir)
abs_potential_child = potential_child_dir.resolve()
abs_dir = dir.resolve()
return abs_potential_child.is_relative_to(abs_dir)
def time_now(time_format: str = "%Y%m%d_%H%M%S"):
return datetime.now().strftime(time_format)
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