add PRAS to preprocess file

Author: DaniBodor

deeprank2/tools/pdbprep/preprocess.py

@@ -1,8 +1,40 @@
+from pathlib import Path
+from tempfile import TemporaryFile
+
 from pdbtools import pdb_delresname, pdb_fixinsert, pdb_keepcoord, pdb_reatom, pdb_reres, pdb_rplresname, pdb_selaltloc, pdb_sort, pdb_tidy
+from Pras_Server.RunType import InitRunType as PRAS
 
 from deeprank2.domain.aminoacidlist import amino_acids_by_code, amino_acids_by_letter
 
 
+def preprocess_pdbs(
+    pdb_path: str | Path,
+    rename_residues: dict[str, str] | None = None,
+) -> str:
+    """Preprocess a pdb file for adding/fixing hydrogens.
+
+    Args:
+        pdb_path: Path of pdb file to preprocess.
+        rename_residues: Dictionary mapping of non-standard residue names (keys) to their standard names.
+            Defaults to:
+            {
+                "MSE": "MET",
+                "HIP": "HIS",
+                "HIE": "HIS",
+                "HID": "HIS",
+                "HSE": "HIS",
+                "HSD": "HIS",
+            }
+    """
+    with Path(pdb_path).open("r") as f:
+        pdb_str = f.read()
+
+    pdb_str = _run_pdb_tools(pdb_str, rename_residues)
+    pdb_str = _add_missing_heavy_atoms(pdb_str)
+
+    return pdb_str  # noqa: RET504
+
+
 def _run_pdb_tools(
     pdb_str: str,
     rename_residues: dict[str, str] | None = None,
@@ -22,23 +54,11 @@ def _run_pdb_tools(
         8. Renumber atoms from 1.
         9. Tidy up to somewhat adhere to pdb format specifications.
 
-    Args:
-        pdb_str: string representation of pdb file.
-        rename_residues: dictionary mapping non-standard residue names (keys) to their standard names. Defaults to:
-            {
-                "MSE": "MET",
-                "HIP": "HIS",
-                "HIE": "HIS",
-                "HID": "HIS",
-                "HSE": "HIS",
-                "HSD": "HIS",
-            }
-
     Raises:
         ValueError: if an invalid amino acid (3-letter or 1-letter) code is given as a value to rename_residues.
 
     Returns:
-        str: updated pdb
+        str: Updated pdb
     """
     if not rename_residues:
         rename_residues = {
@@ -70,3 +90,22 @@ def _run_pdb_tools(
     new_pdb = pdb_tidy.run(new_pdb)  # Tidy up to somewhat adhere to pdb format specifications
 
     return "".join(list(new_pdb))
+
+
+def _add_missing_heavy_atoms(pdb_str: str) -> str:
+    """Add missing heavy atoms (usually many) using PRAS.
+
+    PRAS can only use files (no strings) as input and output, which is why this function is wrapped inside
+    TemporaryFile context managers.
+
+    Returns:
+        str: Updated pdb
+    """
+    with TemporaryFile(mode="w", suffix="pdb", encoding="utf-8") as input_pdb, TemporaryFile(mode="r", encoding="utf-8") as output_pdb:
+        input_pdb.write(pdb_str)
+
+        fixing = PRAS(ofname=output_pdb)
+        fixing.fname = input_pdb
+        fixing.ProcessOther()  # write to specified filename
+
+        return output_pdb.read()