use Model instead of Structure

Author: rvhonorato

src/prodigy_prot/cli.py

@@ -8,10 +8,9 @@
 from argparse import RawTextHelpFormatter
 from pathlib import Path
 
-from Bio.PDB.Structure import Structure
+from Bio.PDB.Model import Model
 
-from prodigy_prot.modules.parsers import (get_parser, parse_structure,
-                                          validate_structure)
+from prodigy_prot.modules.parsers import parse_structure
 from prodigy_prot.modules.prodigy import Prodigy
 
 # setup logging
@@ -100,27 +99,28 @@ def main():
         sys.exit(1)
 
     for input_f in input_list:
-        structure, n_chains, n_res = parse_structure(str(input_f))
+        models: list[Model]
+        n_chains: int
+        n_res: int
+        models, n_chains, n_res = parse_structure(str(input_f))
+        identifier = Path(input_f).stem
 
-        if len(input_list) > 1:
+        for model in models:
             log.info("#" * 42)
-
-        log.info(
-            "[+] Parsed structure file {0} ({1} chains, {2} residues)".format(
-                structure.id, n_chains, n_res
+            log.info(
+                f"[+] Processing structure {identifier}_model{model.id} ({n_chains} chains, {n_res} residues)"
+            )
+            prodigy = Prodigy(model, args.selection, args.temperature)
+            prodigy.predict(
+                distance_cutoff=args.distance_cutoff, acc_threshold=args.acc_threshold
             )
-        )
-        prodigy = Prodigy(structure, args.selection, args.temperature)
-        prodigy.predict(
-            distance_cutoff=args.distance_cutoff, acc_threshold=args.acc_threshold
-        )
-        prodigy.print_prediction(quiet=args.quiet)
-
-        if args.contact_list:
-            prodigy.print_contacts(outfile=str(struct_path.with_suffix(".ic")))
-
-        if args.pymol_selection:
-            prodigy.print_pymol_script(outfile=str(struct_path.with_suffix(".pml")))
+            prodigy.print_prediction(quiet=args.quiet)
+
+            if args.contact_list:
+                prodigy.print_contacts(outfile=str(struct_path.with_suffix(".ic")))
+
+            if args.pymol_selection:
+                prodigy.print_pymol_script(outfile=str(struct_path.with_suffix(".pml")))
 
 
 if __name__ == "__main__":

src/prodigy_prot/modules/freesasa_tools.py

@@ -5,6 +5,7 @@
 import os
 
 import freesasa
+from Bio.PDB.Model import Model
 from Bio.PDB.Structure import Structure
 from freesasa import Classifier, calc, structureFromBioPDB
 
@@ -14,7 +15,7 @@
 freesasa.setVerbosity(freesasa.nowarnings)
 
 
-def execute_freesasa_api(structure: Structure) -> tuple[dict, dict]:
+def execute_freesasa_api(model: Model) -> tuple[dict, dict]:
     """
     Calls freesasa using its Python API and returns
     per-residue accessibilities.
@@ -26,9 +27,14 @@ def execute_freesasa_api(structure: Structure) -> tuple[dict, dict]:
 
     classifier = Classifier(str(NACCESS_CONFIG))
 
+    # NOTE: `structureFromBioPDB` requires a Structure object
+    #  so here build one from a model
+    s = Structure(model.id)
+    s.add(model)
+
     try:
         struct = structureFromBioPDB(
-            structure,
+            s,
             classifier,
         )
         result = calc(struct)

src/prodigy_prot/modules/parsers.py

@@ -4,12 +4,15 @@
 
 import logging
 import sys
+import typing
 import warnings
 from pathlib import Path
 from typing import Optional, Union
 
+from Bio.PDB.Atom import DisorderedAtom
 from Bio.PDB.Chain import Chain
 from Bio.PDB.MMCIFParser import MMCIFParser
+from Bio.PDB.Model import Model
 from Bio.PDB.PDBExceptions import PDBConstructionWarning
 from Bio.PDB.PDBParser import PDBParser
 from Bio.PDB.Polypeptide import PPBuilder, is_aa
@@ -26,112 +29,110 @@ def get_parser(input_f: Path) -> Union[PDBParser, MMCIFParser]:
         return PDBParser()
 
 
+def ignore(r):
+    return r.id[0][0] == "W" or r.id[0][0] == "H"
+
+
 def validate_structure(
-    s: Structure, selection: Optional[list[str]] = None, clean: bool = True
-) -> Structure:
-
-    # setup logging
-    logger = logging.getLogger("Prodigy")
-
-    # Keep first model only
-    if len(s) > 1:
-        logger.warning(
-            (
-                "[!] Structure contains more than one model."
-                " Only the first one will be kept"
-            )
-        )
-        model_one = s[0].id
-        for m in s.child_list[:]:
-            if m.id != model_one:
-                s.detach_child(m.id)
-
-    # process selected chains
-    chains: list[Chain] = list(s.get_chains())
-    chain_ids = set([c.id for c in chains])
-
-    if selection:
-        sel_chains = []
-        # Match selected chain with structure
-        for sel in selection:
-            for c_str in sel.split(","):
-                sel_chains.append(c_str)
-                if c_str not in chain_ids:
+    input_strcture_obj: Structure,
+    selection: Optional[list[str]] = None,
+    clean: bool = True,
+) -> list[Model]:
+
+    result: list[Model] = []
+    for model in [m for m in input_strcture_obj.child_list]:
+
+        # process selected chains
+        chains: list[Chain] = list(model.get_chains())
+        chain_ids = set([c.id for c in chains])
+
+        if selection:
+            sel_chains = []
+            # Match selected chain with structure
+            for sel in selection:
+                for c_str in sel.split(","):
+                    sel_chains.append(c_str)
+                    if c_str not in chain_ids:
+                        raise ValueError(
+                            f"Selected chain not present in provided structure: {c_str}"
+                        )
+
+            # Remove unselected chains
+            def _ignore_helper(x) -> bool:
+                return x.id not in sel_chains
+
+            for c in chains:
+                if _ignore_helper(c):
+                    if c.parent is not None:
+                        c.parent.detach_child(c.id)
+
+        # Double occupancy check
+        for atom in list(model.get_atoms()):
+            if atom.is_disordered():
+                atom = typing.cast(DisorderedAtom, atom)
+                residue = atom.parent
+                assert residue is not None
+                sel_at = atom.selected_child
+                assert sel_at is not None
+                sel_at.altloc = " "
+                sel_at.disordered_flag = 0
+                residue.detach_child(atom.id)
+                residue.add(sel_at)
+
+        # Insertion code check
+        for c in chains:
+            for residue in c.get_residues():
+                if residue.get_id()[2] != " ":
+                    c.detach_child(residue.id)
+
+        if clean:
+            # Remove HETATMs and solvent
+            res_list = list(model.get_residues())
+
+            for res in res_list:
+                if ignore(res):
+                    chain = res.parent
+                    assert chain is not None
+                    chain.detach_child(res.id)
+                elif not is_aa(res, standard=True):
                     raise ValueError(
-                        f"Selected chain not present in provided structure: {c_str}"
+                        "Unsupported non-standard amino acid found: {0}".format(
+                            res.resname
+                        )
                     )
 
-        # Remove unselected chains
-        def _ignore_helper(x) -> bool:
-            return x.id not in sel_chains
-
-        for c in chains:
-            if _ignore_helper(c):
-                if c.parent is not None:
-                    c.parent.detach_child(c.id)
-
-    # Double occupancy check
-    for atom in list(s.get_atoms()):
-        if atom.is_disordered():
-            residue = atom.parent
-            sel_at = atom.selected_child
-            sel_at.altloc = " "
-            sel_at.disordered_flag = 0
-            residue.detach_child(atom.id)
-            residue.add(sel_at)
-
-    # Insertion code check
-    for c in chains:
-        for residue in c.get_residues():
-            if residue.get_id()[2] != " ":
-                c.detach_child(residue.id)
-
-    if clean:
-        # Remove HETATMs and solvent
-        res_list = list(s.get_residues())
-
-        def _ignore(r):
-            return r.id[0][0] == "W" or r.id[0][0] == "H"
-
-        for res in res_list:
-            if _ignore(res):
-                chain = res.parent
-                chain.detach_child(res.id)
-            elif not is_aa(res, standard=True):
-                raise ValueError(
-                    "Unsupported non-standard amino acid found: {0}".format(res.resname)
+            # Remove Hydrogens
+            atom_list = list(model.get_atoms())
+
+            def _ignore(x):
+                return x.element == "H"
+
+            for atom in atom_list:
+                if _ignore(atom):
+                    residue = atom.parent
+                    assert residue is not None
+                    residue.detach_child(atom.name)
+
+        # Detect gaps and compare with no. of chains
+        pep_builder = PPBuilder()
+        peptides = pep_builder.build_peptides(model)
+        n_peptides = len(peptides)
+
+        if n_peptides != len(chain_ids):
+            message = "[!] Structure contains gaps:\n"
+            for i_pp, pp in enumerate(peptides):
+                message += (
+                    "\t{1.parent.id} {1.resname}{1.id[1]} < Fragment {0} > "
+                    "{2.parent.id} {2.resname}{2.id[1]}\n".format(i_pp, pp[0], pp[-1])
                 )
+            log.warning(message)
 
-        # Remove Hydrogens
-        atom_list = list(s.get_atoms())
+        result.append(model)
 
-        def _ignore(x):
-            return x.element == "H"
+    return result
 
-        for atom in atom_list:
-            if _ignore(atom):
-                residue = atom.parent
-                residue.detach_child(atom.name)
 
-    # Detect gaps and compare with no. of chains
-    pep_builder = PPBuilder()
-    peptides = pep_builder.build_peptides(s)
-    n_peptides = len(peptides)
-
-    if n_peptides != len(chain_ids):
-        message = "[!] Structure contains gaps:\n"
-        for i_pp, pp in enumerate(peptides):
-            message += (
-                "\t{1.parent.id} {1.resname}{1.id[1]} < Fragment {0} > "
-                "{2.parent.id} {2.resname}{2.id[1]}\n".format(i_pp, pp[0], pp[-1])
-            )
-        logger.warning(message)
-        # raise Exception(message)
-
-    return s
-
-
-def parse_structure(path: str) -> tuple[Structure, int, int]:
+def parse_structure(path: str) -> tuple[list[Model], int, int]:
     """Return a validated `Structure`, number of chains and number of residues"""
 
     extension = Path(path).suffix
@@ -154,13 +155,33 @@ def parse_structure(path: str) -> tuple[Structure, int, int]:
 
     assert isinstance(original_structure, Structure)
 
-    structure = validate_structure(original_structure)
+    models: list[Model] = validate_structure(original_structure)
 
     # Get number of chains
-    number_of_chains = len(set([c.id for c in structure.get_chains()]))
+    chain_dict = {}
+    res_dict = {}
+    for model in models:
+        chain_dict.update({c.id: c for c in model.get_chains()})
+        res_dict.update({r.id: r for r in model.get_residues()})
+
+    ## Make sure all models have the same chains
+    # Get chain sets for all models
+    chain_sets = [set(chain.id for chain in model.get_chains()) for model in models]
+
+    # Check if all sets are identical
+    if not all(chain_set == chain_sets[0] for chain_set in chain_sets):
+        raise ValueError(
+            "Not all models have the same chains. Found chain sets: "
+            + ", ".join(str(s) for s in chain_sets)
+        )
+
+    res_sets = [set(res.id for res in model.get_residues()) for model in models]
 
-    # Get number of residues
-    number_of_residues = len(list(structure.get_residues()))
+    if not all(res_set == res_sets[0] for res_set in res_sets):
+        raise ValueError(
+            "Not all models have the same residues. Found residue sets: "
+            + ", ".join(str(s) for s in res_sets)
+        )
 
     # structure, n_chains, n_res = parse_structure(path=str(struct_path))
-    return (structure, number_of_chains, number_of_residues)
+    return (models, len(chain_sets[0]), len(res_sets[0]))