Commit f6439a19 by cassandra

[feat] add ability to accept systems with multiple ligands

parent 8e3a3715
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"
MMPBSA_ENV_DIR = Path("/path/to/your/envs/mmpbsa")
\ No newline at end of file
......@@ -2,13 +2,13 @@
# Usage: ./run_tleap.sh input.pdb
if [ $# -ne 4 ]; then
echo "Usage: $0 sandbox_dir complex_pdb ligand_name prepi_file"
echo "Usage: $0 sandbox_dir complex_pdb frcmod_file prepi_file"
exit 1
fi
SANDBOX_DIR=$1 # PDBFILE already has sandbox path
PDBFILE=$2
LIGNAME=$3
FRCMOD_FILE=$3
PREPI_FILE=$4
# Create tleap input file
......@@ -22,7 +22,7 @@ addPdbAtomMap { { "CH3" "C" } { "HH31" "H1" } { "HH32" "H2" } { "HH33" "H3" } {
# Load ligand parameters
loadamberprep ${SANDBOX_DIR}/${PREPI_FILE}
loadamberparams ${SANDBOX_DIR}/${LIGNAME}.frcmod
loadamberparams ${SANDBOX_DIR}/${FRCMOD_FILE}
# PDBFILE already has sandbox path
mol = loadpdb ${PDBFILE}
......
......@@ -34,23 +34,29 @@ def run_tleap(sandbox_dir: str, input_pdb: str, pdb_id: str) -> str:
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:
def run_tleap_ligand(sandbox_dir: str, input_pdb: str, pdb_id: str, ligand_files: str | list[str], ligand_name: str) -> str:
"""
Run tleap preparation using run_tleap.sh, for a protein-ligand complex.
"""
# make sure it's a list
if isinstance(ligand_files, str):
ligand_files = [ligand_files]
else:
ligand_files = ligand_files
# 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)
with open(f"{sandbox_dir}/{ligand_files[0]}", "r") as ligand_infile:
for line in ligand_infile:
if line.startswith("HETATM"):
outfile.write(line)
logger.info(f"Only the first ligand {ligand_files[0]} was added to complex.pdb as we will only simulate one ligand.")
outfile.write("TER\n")
outfile.write("END\n")
......@@ -59,13 +65,14 @@ def run_tleap_ligand(sandbox_dir: str, input_pdb: str, pdb_id: str, ligand_file:
# 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"
ligand_stem=Path(f"{ligand_files[0]}").stem
if Path(f"{sandbox_dir}/{ligand_stem}_fixed.prepi").exists():
prepi_file = f"{ligand_stem}_fixed.prepi"
else:
prepi_file = f"{ligand_name}.prepi"
prepi_file = f"{ligand_stem}.prepi"
result = subprocess.run(
[str(script), sandbox_dir, complex_pdb, ligand_name, prepi_file],
[str(script), sandbox_dir, complex_pdb, f"{ligand_stem}.frcmod", prepi_file],
cwd=sandbox_dir,
capture_output=True,
text=True,
......
......@@ -9,12 +9,14 @@ import time
logger = get_class_logger(__name__)
def gromacs_equil(sandbox_dir: str, input_gro: str, md_temp: str, ligand_name=None, ligand_file=None) -> str:
def gromacs_equil(sandbox_dir: str, input_gro: str, md_temp: str, ligand_name=None, ligand_files=None) -> str:
# sometimes llm passes ligands as empty strings
if not ligand_name:
ligand_name = None
if not ligand_file:
if not ligand_files:
ligand_file = None
else:
ligand_file=ligand_files[0]
# ------------ Modify topol.top to include position restraints ------------
......
......@@ -6,12 +6,82 @@ from src.utils import get_class_logger
logger = get_class_logger(__name__)
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}")
def param_ligand(sandbox_dir: str, ligand_files: str | list[str], ligand_name: str, charge_ligand: str | None = None) -> str:
if isinstance(ligand_files, str):
ligand_files = [ligand_files]
else:
ligand_files = ligand_files
ligand_file=ligand_files[0]
if len(ligand_files) > 1:
logger.info(f"For now, we will only parameterize the first ligand since we will only simulate one ligand: {ligand_file}")
logger.info(f"Parameterizing ligand file: {ligand_file}")
ligand_stem=Path(f"{ligand_file}").stem
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"
tmp_mol2file = f"{sandbox_dir}/{ligand_stem}.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)
......@@ -42,105 +112,52 @@ def param_ligand(sandbox_dir: str, ligand_file: str, ligand_name: str, charge_li
# 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}")
logger.info(f"Charge of ligand {ligand_stem} 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"
f"antechamber -i {sandbox_dir}/{ligand_file} -fi pdb -o {sandbox_dir}/{ligand_stem}.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")
logger.info(f"Mol2 file for ligand {ligand_stem} created")
cmd = shlex.split(f"sed -i 's/UNL/{ligand_name}/g' {sandbox_dir}/{ligand_name}.mol2")
cmd = shlex.split(f"sed -i 's/UNL/{ligand_name}/g' {sandbox_dir}/{ligand_stem}.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"
f"antechamber -i {sandbox_dir}/{ligand_stem}.mol2 -fi mol2 -o {sandbox_dir}/{ligand_stem}.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}"
logger.info(f"Prepi file for ligand {ligand_file} created: {ligand_stem}.prepi")
# Create frcmod file using parmchk2
cmd = shlex.split(f"parmchk2 -i {sandbox_dir}/{ligand_name}.mol2 -f mol2 -o {sandbox_dir}/{ligand_name}.frcmod")
cmd = shlex.split(f"parmchk2 -i {sandbox_dir}/{ligand_stem}.mol2 -f mol2 -o {sandbox_dir}/{ligand_stem}.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}"
logger.info(f"Frcmod file for ligand {ligand_file} created: {ligand_stem}.frcmod")
# Update charge prepi file
fix_charges(f"{sandbox_dir}/{ligand_stem}.prepi", f"{sandbox_dir}/{ligand_stem}_fixed.prepi")
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")
if Path(f"{sandbox_dir}/{ligand_stem}_fixed.prepi").exists():
prepi_file = f"{ligand_stem}_fixed.prepi"
else:
prepi_file = f"{ligand_stem}.prepi"
return "Ligand parameterisation complete. File saved to {sandbox_dir}/{ligand_name}_fixed.prepi"
if len(ligand_files) == 1:
return f"Ligand parameterisation complete. Parameters saved to {sandbox_dir}/{prepi_file} and {sandbox_dir}/{ligand_stem}.frcmod"
if len(ligand_files) > 1:
return f"Ligand parameterisation complete for the the first ligand: {ligand_files[0]}. Parameters saved to {sandbox_dir}/{ligand_stem}.frcmod and {sandbox_dir}/{prepi_file}"
......@@ -30,15 +30,15 @@ TOOL_MAP = {
"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"]),
"param_ligand": lambda s, i: param_ligand(s.sandbox_dir, i["ligand_files"] if isinstance(i["ligand_files"], list) else [i["ligand_files"]], 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"]
s.sandbox_dir, i["input_pdb"], i["pdb_id"], i["ligand_files"] if isinstance(i["ligand_files"], list) else [i["ligand_files"]], i["ligand_name"]
),
# GROMACS-related
"gromacs_equil": lambda s, i: gromacs_equil(
s.sandbox_dir, i["input_gro"], i["md_temp"], ligand_name=i.get("ligand_name"), ligand_file=i.get("ligand_file")
s.sandbox_dir, i["input_gro"], i["md_temp"], ligand_name=i.get("ligand_name"), ligand_files=i.get("ligand_files")
),
"gromacs_production": lambda s, i: gromacs_production(
s.sandbox_dir, i["input_gro"], i["npt_cpt_file"], i["md_temp"], i["md_duration"], ligand_name=i.get("ligand_name")
......
......@@ -83,81 +83,48 @@ def fetch_and_save_pdb(sandbox_dir: str, pdb_id: str, output_pdb: str) -> str:
except Exception:
return f"Error fetching PDB {pdb_id}: {traceback.format_exc()}"
def prepare_pdb_file_ligand_old(sandbox_dir: str, pdb_id: str, ligand_name: str = None) -> str:
def check_pdb_ligand(sandbox_dir: str, pdb_id: str, ligand_name: str = None) -> str:
"""
Takes input PDB file, extract ligand to ligand_name.pdb.
Removes HETATM, CONECT and MASTER lines from input_pdb and saves to prepared_pdb.
Protonates ligand and pH=7 and saves to ligand_name_h.pdb.
This is a publically available API that checks if a PDB file is valid by attempting to parse it with Bio.PDB.PDBParser.
Args:
input_pdb (str): The path where the input PDB file is located.
prepared_pdb (str): The path where we save the prepared PDB file.
ligand_name (str): The name of the ligand to extract.
ligand_pdb (str): The path where we save the extracted ligand PDB file.
ligand_pdb_h (str): The path where we save the protonated ligand PDB file.
pdb_id (str): The 4-character PDB ID (e.g., '1abc').
ligand_name (str): Optional: the name of the ligand if a protein-ligand complex should be simulated.
"""
# Extract ligand
if (ligand_name is not None) and (ligand_name != "XXX") and (ligand_name != "None") and (ligand_name != "None_h"):
with (
open(f"{sandbox_dir}/{pdb_id}.pdb", "r") as infile,
open(f"{sandbox_dir}/{ligand_name}.pdb", "w") as outfile,
):
for line in infile:
if line.startswith("HETATM") and ligand_name in line:
new_line = line.replace("UNL", ligand_name)
outfile.write(new_line)
# Prepare PDB
with (
open(f"{sandbox_dir}/{pdb_id}.pdb", "r") as infile,
open(f"{sandbox_dir}/{pdb_id}_prepared.pdb", "w") as outfile,
):
for line in infile:
if not (line.startswith("HETATM") or line.startswith("CONECT") or line.startswith("MASTER")):
outfile.write(line)
# Protonate ligand
if (ligand_name is not None) and (ligand_name != "XXX") and (ligand_name != "None") and (ligand_name != "None_h"):
cmd = shlex.split(f"obabel {sandbox_dir}/{ligand_name}.pdb -O {sandbox_dir}/{ligand_name}_h.pdb -p7")
subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
with open(f"{sandbox_dir}/{ligand_name}_h.pdb", "r") as infile:
lines = infile.readlines()
filtered_lines = [line for line in lines if not (line.startswith("CONECT") or line.startswith("MASTER"))]
new_filtered_lines = [line.replace("UNL", ligand_name) for line in filtered_lines]
with open(f"{sandbox_dir}/{ligand_name}_h.pdb", "w") as outfile:
outfile.writelines(new_filtered_lines)
# Rewrite hydrogen, and other atoms if needed, names in the ligand
pdb_file = sandbox_dir / f"{pdb_id}.pdb"
with open(f"{sandbox_dir}/{ligand_name}_h.pdb", "r") as infile:
lines = infile.readlines()
counters = defaultdict(int)
new_lines = []
for line in lines:
# Only modify atom records
if line.startswith(("HETATM", "ATOM")):
atom_name = line[12:16].strip()
element = line[76:78].strip()
# Only rename if atom_name is the same as the element (e.g. "C", "N", "H")
if atom_name == element and element.isalpha():
counters[element] += 1
new_name = f"{element}{counters[element]}"
# Reinsert new name in the proper 4-character PDB field
line = f"{line[:12]}{new_name:>4}{line[16:]}"
new_lines.append(line)
with open(f"{sandbox_dir}/{ligand_name}_h.pdb", "w") as outfile:
outfile.writelines(new_lines)
logger.info(f"Generic atoms renamed successfully. Output saved to {sandbox_dir}/{ligand_name}_h.pdb")
return f"Successfully Prepared PDB structure with a ligand and saved the extracted ligand PDB file to {sandbox_dir}/{pdb_id}_prepared.pdb and the protonated ligand PDB file to {sandbox_dir}/{ligand_name}_h.pdb"
if ligand_name is not None and ligand_name not in ["XXX", "None", "None_h"]:
# Check if ligand is present in the PDB file
with open(pdb_file, "r") as f:
lines = f.readlines()
ligand_present = any(line.startswith("HETATM") and ligand_name in line for line in lines)
if not ligand_present:
logger.info(f"Ligand {ligand_name} not found in PDB file {pdb_file}.")
raise NoLigand(f"Ligand {ligand_name} not found in PDB file {pdb_file}.")
#Check if ligand is covalent
with open(pdb_file, "r") as f:
lines = f.readlines()
ligand_covalent = any(line.startswith("LINK") and ligand_name in line for line in lines)
if ligand_covalent:
logger.info(f"Ligand {ligand_name} appears to be covalently bound in PDB file {pdb_file}. DynaMate doesn't support the parameterization of covalently bound ligands. This system cannot be processed.")
raise NoLigand(f"Ligand {ligand_name} appears to be covalently bound in PDB file {pdb_file}. DynaMate doesn't support the parameterization of covalently bound ligands. This system cannot be processed.")
# Count number of ligands
with open(pdb_file, "r") as f:
lines = f.readlines()
ligand_count = sum(1 for line in lines if line.startswith("HET ") and ligand_name in line)
logger.info(f"There is(are) {ligand_count} ligand(s) called {ligand_name} in {pdb_file}.")
return f"Successfully Prepared PDB structure without a ligand and saved the extracted PDB file to {sandbox_dir}/{pdb_id}_prepared.pdb"
# Check for modified residues
with open(pdb_file, "r") as f:
lines = f.readlines()
modified_residues = [line for line in lines if line.startswith("MODRES")]
logger.info("There are ", len(modified_residues), "modified residues, which are", modified_residues, ". This should be checked and the corresponding residues modified to standard residues. If they can't be modified to standard residues, the system can't be processed.")
return "PDB file check completed successfully."
def prepare_pdb_file_ligand(sandbox_dir: str, pdb_id: str, ligand_name: str = None) -> str:
"""
......@@ -171,17 +138,6 @@ def prepare_pdb_file_ligand(sandbox_dir: str, pdb_id: str, ligand_name: str = No
ligand_pdb (str): The path where we save the extracted ligand PDB file.
ligand_pdb_h (str): The path where we save the protonated ligand PDB file.
"""
# Extract ligand
if (ligand_name is not None) and (ligand_name != "XXX") and (ligand_name != "None") and (ligand_name != "None_h"):
with (
open(f"{sandbox_dir}/{pdb_id}.pdb", "r") as infile,
open(f"{sandbox_dir}/{ligand_name}.pdb", "w") as outfile,
):
for line in infile:
if line.startswith("HETATM") and ligand_name in line:
new_line = line.replace("UNL", ligand_name)
outfile.write(new_line)
# Prepare PDB
with (
open(f"{sandbox_dir}/{pdb_id}.pdb", "r") as infile,
......@@ -190,17 +146,74 @@ def prepare_pdb_file_ligand(sandbox_dir: str, pdb_id: str, ligand_name: str = No
for line in infile:
if not (line.startswith("HETATM") or line.startswith("CONECT") or line.startswith("MASTER")):
outfile.write(line)
logger.info(f"Prepared PDB file saved to {sandbox_dir}/{pdb_id}_prepared.pdb")
# Extract ligand
if (ligand_name is not None) and (ligand_name != "XXX") and (ligand_name != "None") and (ligand_name != "None_h"):
# Count number of ligands
resnums = set()
with open(f"{sandbox_dir}/{pdb_id}.pdb") as f:
for line in f:
if line.startswith("HETATM") and ligand_name in line:
resnum = int(line[22:26])
resnums.add(resnum)
print("resnums is: ", resnums)
num_ligands = len(resnums)
logger.info(f"IMPORTANT: Number of ligands {ligand_name} found: {num_ligands}")
if num_ligands == 0:
logger.info(f"Ligand {ligand_name} not found in PDB file {sandbox_dir}/{pdb_id}.pdb. You can either proceed without a ligand, check the ligand name provided or check the PDB file.")
return f"Ligand {ligand_name} not found in PDB file {sandbox_dir}/{pdb_id}.pdb. You can either proceed without a ligand, check the ligand name provided or check the PDB file"
# CASE 1: only one ligand
if num_ligands == 1:
ligand_pdb_file = f"{sandbox_dir}/{ligand_name}.pdb"
ligand_pdb_files_list = [ligand_pdb_file]
with open(f"{sandbox_dir}/{pdb_id}.pdb", "r") as infile, open(ligand_pdb_file, "w") as outfile:
for line in infile:
if line.startswith("HETATM") and ligand_name in line:
outfile.write(line)
logger.info(f"Extracted ligand {ligand_name} to {ligand_pdb_file}")
# CASE 2: multiple ligands (split by residue number)
else:
ligand_pdb_files_list = []
ligands = defaultdict(list)
# Collect HETATM lines by residue number
with open(f"{sandbox_dir}/{pdb_id}.pdb", "r") as infile:
for line in infile:
if line.startswith("HETATM") and ligand_name in line:
resnum = int(line[22:26]) # residue number column
ligands[resnum].append(line)
# Write one file per ligand
for i, (resnum, atom_lines) in enumerate(ligands.items(), start=1):
ligand_pdb_file = f"{sandbox_dir}/{ligand_name}_{i}.pdb"
ligand_pdb_files_list.append(ligand_pdb_file)
with open(ligand_pdb_file, "w") as outfile:
outfile.writelines(atom_lines)
logger.info(f"Extracted ligand {ligand_name} residue {resnum} to {ligand_pdb_file}")
# Protonate ligand
list_protonated_files = []
if (ligand_name is not None) and (ligand_name != "XXX") and (ligand_name != "None") and (ligand_name != "None_h"):
cmd = shlex.split(f"obabel {sandbox_dir}/{ligand_name}.pdb -O {sandbox_dir}/{ligand_name}_h.pdb -p7")
subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
with open(f"{sandbox_dir}/{ligand_name}_h.pdb", "r") as infile:
lines = infile.readlines()
filtered_lines = [line for line in lines if not (line.startswith("CONECT") or line.startswith("MASTER"))]
new_filtered_lines = [line.replace("UNL", ligand_name) for line in filtered_lines]
with open(f"{sandbox_dir}/{ligand_name}_h.pdb", "w") as outfile:
outfile.writelines(new_filtered_lines)
for ligand_pdb_file in ligand_pdb_files_list: # loop over all extracted ligands
if num_ligands == 1:
protonated_file = f"{sandbox_dir}/{ligand_name}_h.pdb"
else:
index = ligand_pdb_file.split("_")[-1].split(".")[0] # get index from filename
protonated_file = f"{sandbox_dir}/{ligand_name}_{index}_h.pdb"
list_protonated_files.append(f"{ligand_name}_{index}_h.pdb")
cmd = shlex.split(f"obabel {ligand_pdb_file} -O {protonated_file} -p7")
subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
with open(protonated_file, "r") as infile:
lines = infile.readlines()
filtered_lines = [line for line in lines if not (line.startswith("CONECT") or line.startswith("MASTER"))]
new_filtered_lines = [line.replace("UNL", ligand_name) for line in filtered_lines]
with open(protonated_file, "w") as outfile:
outfile.writelines(new_filtered_lines)
# Rewrite atoms names in the ligand
ELEMENTS = {
......@@ -304,10 +317,7 @@ def prepare_pdb_file_ligand(sandbox_dir: str, pdb_id: str, ligand_name: str = No
return e[0] + e[1].lower()
return e
input_file = f"{sandbox_dir}/{ligand_name}_h.pdb"
output_file = f"{sandbox_dir}/{ligand_name}_h.pdb"
with open(input_file, "r") as infile:
with open(protonated_file, "r") as infile:
lines = infile.readlines()
counters = defaultdict(int)
......@@ -337,12 +347,16 @@ def prepare_pdb_file_ligand(sandbox_dir: str, pdb_id: str, ligand_name: str = No
new_lines.append(line)
with open(output_file, "w") as outfile:
with open(protonated_file, "w") as outfile:
outfile.writelines(new_lines)
logger.info("Atom renaming of ligand completed.")
return f"Successfully Prepared PDB structure with a ligand and saved the extracted ligand PDB file to {sandbox_dir}/{pdb_id}_prepared.pdb and the protonated ligand PDB file to {sandbox_dir}/{ligand_name}_h.pdb. Ligand was protonated at pH=7 and atom names were cleaned (renumbered)"
if num_ligands == 1:
return f"Successfully Prepared PDB structure with a ligand and saved the extracted protein PDB file to {sandbox_dir}/{pdb_id}_prepared.pdb and the protonated ligand PDB file to {sandbox_dir}/{ligand_name}_h.pdb. Ligand was protonated at pH=7 and atom names were cleaned (renumbered)"
if num_ligands > 1:
return f"Successfully Prepared PDB structure with {num_ligands} ligands and saved the extracted protein PDB file to {sandbox_dir}/{pdb_id}.pdb and the {num_ligands} protonated ligand PDB files to {sandbox_dir}/{list_protonated_files}. This list of {num_ligands} protonated files: {list_protonated_files} is IMPORTANT and should be the input parameter for future functions. The extracted pdb file was saved to {sandbox_dir}/{pdb_id}_prepared.pdb Ligands were protonated at pH=7 and atom names were cleaned (renumbered)"
return f"Successfully Prepared PDB structure without a ligand and saved the extracted PDB file to {sandbox_dir}/{pdb_id}_prepared.pdb"
......@@ -357,6 +371,12 @@ def add_caps(sandbox_dir: str, input_pdb: str, pdb_id: str) -> str:
pdb_id (str): The PDB ID.
sandbox_dir (str): the directory where we add and modify files.
"""
with open(f"{sandbox_dir}/{input_pdb}") as pdbfile:
for line in pdbfile:
if line.startswith("HETATM") or line.startswith("CONECT") or line.startswith("MASTER"):
pdbfile.close()
logger.warning("Input PDB file contains HETATM, CONECT or MASTER lines. Please prepare the PDB file first to remove these lines. If the PDB file has already been prepared with the prepare_pdb_file_ligand function, use the correct parameters when calling this tool or check that it has been prepared correctly.")
return "Error: Input PDB file contains HETATM, CONECT or MASTER lines. Please prepare the PDB file first to remove these lines. If the PDB file has already been prepared with the prepare_pdb_file_ligand function, use the correct parameters when calling this tool or check that it has been prepared correctly."
def create_universe(n_atoms, name, resname, positions, resids, segid):
u_new = mda.Universe.empty(
......
......@@ -243,12 +243,17 @@ def create_tool_schema_md(sandbox_dir, ligand_name, pdb_id):
"type": "string",
"description": f"The directory path where input files are located and output files are produced (this is the sandbox directory): {sandbox_dir}",
},
"ligand_file": {
"type": "string",
"ligand_files": {
"type": "array",
"items": {
"type": "string"
},
"description": (
f"Input PDB file of the protonated ligand to be parameterized (without full path). "
f"The ligand file is named {ligand_name}_h.pdb if previously protonated with the prepare_pdb_file_ligand tool."
"This file must exist in the provided sandbox_dir ({sandbox_dir})."
"List of input PDB file(s) of the protonated ligand(s) to be parameterized "
"(without full path). If there is only one ligand, provide a list with one element, "
f"e.g. ['{ligand_name}_h.pdb']. If there are multiple ligands, provide a list such as "
f"['{ligand_name}_h_1.pdb', '{ligand_name}_h_2.pdb'] where the ligands are indexed starting at 1. "
f"All files must exist in the sandbox directory ({sandbox_dir})."
),
},
"ligand_name": {
......@@ -258,7 +263,7 @@ def create_tool_schema_md(sandbox_dir, ligand_name, pdb_id):
),
},
},
"required": ["sandbox_dir", "ligand_file", "ligand_name"],
"required": ["sandbox_dir", "ligand_files", "ligand_name"],
},
),
Tool(
......@@ -298,12 +303,17 @@ def create_tool_schema_md(sandbox_dir, ligand_name, pdb_id):
"type": "string",
"description": (f"The PDB ID of the structure to be prepared: {pdb_id}. "),
},
"ligand_file": {
"type": "string",
"ligand_files": {
"type": "array",
"items": {
"type": "string"
},
"description": (
f"Input PDB file of the ligand to be merged with the protein (without full path). "
f"The ligand file is named {ligand_name}_h.pdb if previously protonated with the prepare_pdb_file_ligand tool."
"This file must exist in the provided sandbox_dir."
"List of input PDB file(s) of the protonated ligand(s) to be parameterized "
"(without full path). If there is only one ligand, provide a list with one element, "
f"e.g. ['{ligand_name}_h.pdb']. If there are multiple ligands, provide a list such as "
f"['{ligand_name}_h_1.pdb', '{ligand_name}_h_2.pdb'] where the ligands are indexed starting at 1. "
f"All files must exist in the sandbox directory ({sandbox_dir})."
),
},
"ligand_name": {
......@@ -313,13 +323,7 @@ def create_tool_schema_md(sandbox_dir, ligand_name, pdb_id):
),
},
},
"required": [
"sandbox_dir",
"input_pdb",
"pdb_id",
"ligand_file",
"ligand_name",
],
"required": ["sandbox_dir", "input_pdb", "pdb_id", "ligand_files", "ligand_name"],
},
),
Tool(
......@@ -369,12 +373,17 @@ def create_tool_schema_md(sandbox_dir, ligand_name, pdb_id):
"Three-character residue name of the ligand (capital letters or numbers). If no ligand was provided by user, do not input this argument."
),
},
"ligand_file": {
"type": "string",
"ligand_files": {
"type": "array",
"items": {
"type": "string"
},
"description": (
"Optional argument for the input PDB file of the ligand to be merged with the protein (without full path). "
"The ligand file is typically named '<ligand_name>_h.pdb' if previously protonated. "
"If no ligand was provided by user, do not input this argument."
"List of input PDB file(s) of the protonated ligand(s) to be parameterized "
"(without full path). If there is only one ligand, provide a list with one element, "
f"e.g. ['{ligand_name}_h.pdb']. If there are multiple ligands, provide a list such as "
f"['{ligand_name}_h_1.pdb', '{ligand_name}_h_2.pdb']. "
f"All files must exist in the sandbox directory ({sandbox_dir})."
),
},
},
......
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