Repository: haddocking/arctic3d

Issue:
adjust error message when no PDB is found
the current error message when the program fails to retrieve the pdb is
`Could not retrieve a valid PDB for the target, please provide one using the --pdb option`
such option does not exist, as the pdb can only be the main argument

Comments:

Patch:
diff --git a/src/arctic3d/cli.py b/src/arctic3d/cli.py
index b0e4333..ea1b0a4 100644
--- a/src/arctic3d/cli.py
+++ b/src/arctic3d/cli.py
@@ -264,7 +264,7 @@ def main(
         if pdb_f is None:
             log.error(
                 "Could not retrieve a valid PDB for the target, please provide"
-                " one using the --pdb option"
+                " one as the main input argument."
             )
             sys.exit()
 
diff --git a/src/arctic3d/modules/interface_matrix.py b/src/arctic3d/modules/interface_matrix.py
index 79204fe..fd64777 100644
--- a/src/arctic3d/modules/interface_matrix.py
+++ b/src/arctic3d/modules/interface_matrix.py
@@ -91,7 +91,7 @@ def get_coupling_matrix(mdu, int_resids):
     if u.positions.shape[0] != len(int_resids):
         raise Exception(
             "shape mismatch: positions do not match input residues"
-            " {int_resids}"
+            f" {int_resids}"
         )
     distmap = cdist(u.positions, u.positions)
     exp_factor = 4 * SIGMA * SIGMA
diff --git a/src/arctic3d/modules/pdb.py b/src/arctic3d/modules/pdb.py
index 9bc154b..82b1dc7 100644
--- a/src/arctic3d/modules/pdb.py
+++ b/src/arctic3d/modules/pdb.py
@@ -9,6 +9,8 @@ from pdbecif.mmcif_io import MMCIF2Dict
 from pdbtools.pdb_selaltloc import select_by_occupancy
 from pdbtools.pdb_selchain import select_chain
 from pdbtools.pdb_tidy import tidy_pdbfile
+from pdbtools.pdb_selmodel import select_model
+
 
 from arctic3d.functions import make_request
 from arctic3d.modules.interface_matrix import filter_interfaces
@@ -177,6 +179,36 @@ def renumber_pdb_from_cif(pdb_id, uniprot_id, chain_id, pdb_fname):
     return pdb_renum_fname, cif_fname
 
 
+def fetch_pdb_files(pdb_to_fetch):
+    """
+    Fetches the pdb files from PDBe database.
+
+    Parameters
+    ----------
+    pdb_to_fetch : list
+        list of pdb hits to fetch
+
+    Returns
+    -------
+    validated_pdbs : list
+        list of tuples (pdb_file, hit)
+    """
+    validated_pdbs = []
+    valid_pdb_set = set()  # set of valid pdb IDs
+    for hit in pdb_to_fetch:
+        pdb_id = hit["pdb_id"]
+        pdb_fname = f"{pdb_id}.pdb"
+        if pdb_fname not in os.listdir():
+            pdb_f = fetch_pdb(pdb_id)
+        else:
+            pdb_f = Path(pdb_fname)
+        if pdb_f is not None:
+            validated_pdbs.append((pdb_f, hit))
+            if pdb_id not in valid_pdb_set:
+                valid_pdb_set.add(pdb_id)
+    return validated_pdbs
+
+
 def fetch_pdb(pdb_id):
     """
     Fetches the pdb from PDBe database.
@@ -227,7 +259,32 @@ def selchain_pdb(inp_pdb_f, chain):
         with open(out_pdb_fname, "w") as f:
             for line in select_chain(pdb_fh, chain):
                 f.write(line)
-            f.write(line)
+    return out_pdb_fname
+
+
+def selmodel_pdb(inp_pdb_f, model_id=1):
+    """
+    Select model from PDB file.
+
+    Parameters
+    ----------
+    inp_pdb_f : Path
+        Path to PDB file.
+    model_id : int, optional
+        Model ID, by default 1
+
+    Returns
+    -------
+    out_pdb_fname : Path
+        Path to PDB file.
+    """
+    # log.debug(f"Selecting model {model_id} from PDB file")
+    out_pdb_fname = Path(f"{inp_pdb_f.stem}-model{model_id}.pdb")
+    with open(inp_pdb_f, "r") as pdb_fh:
+        with open(out_pdb_fname, "w") as f:
+            line = ""
+            for line in select_model(pdb_fh, [model_id]):
+                f.write(line)
     return out_pdb_fname
 
 
@@ -328,48 +385,38 @@ def validate_api_hit(
     validated_pdbs : list
         List of (pdb_f, hit) tuples
     """
-    validated_pdbs = []  # list of good pdbs
-    valid_pdb_set = set()  # set of valid pdb IDs
-
-    for hit in fetch_list[:max_pdb_num]:
-        check_list = {}
+    pdbs_to_fetch = []
+    for hit in fetch_list:
+        check_list = []
         pdb_id = hit["pdb_id"]
         coverage = hit["coverage"]
         resolution = hit["resolution"]
+        exp_method = hit["experimental_method"]
 
-        pdb_fname = f"{pdb_id}.pdb"
-        if pdb_fname not in os.listdir():
-            pdb_f = fetch_pdb(pdb_id)
-        else:
-            pdb_f = Path(pdb_fname)
-
-        if pdb_f is not None:
-            check_list["pdb_f"] = True
-        else:
-            check_list["pdb_f"] = False
-
+        # check coverage value
         if coverage > coverage_cutoff:
-            check_list["cov"] = True
+            check_list.append(True)
         else:
-            check_list["cov"] = False
-
+            check_list.append(False)
+        # check resolution value
         if resolution is None:
-            check_list["res"] = False
+            if "NMR" in exp_method:
+                check_list.append(True)
+            else:
+                check_list.append(False)
         elif resolution < resolution_cutoff:
-            check_list["res"] = True
+            check_list.append(True)
         else:
-            check_list["res"] = False
+            check_list.append(False)
 
-        if all(check_list.values()):
-            validated_pdbs.append((pdb_f, hit))
-            if pdb_id not in valid_pdb_set:
-                valid_pdb_set.add(pdb_id)
+        if all(check_list):
+            pdbs_to_fetch.append(hit)
         else:
-            log.debug(f"{pdb_id} failed validation ({check_list})")
-            # pdb_f could be None or the pdb (another chain)
-            #   could be valid and the file should not be removed
-            if pdb_f is not None and pdb_id not in valid_pdb_set:
-                os.unlink(pdb_f)
+            log.debug(f"{pdb_id} failed validation")
+    log.info(f"Found {len(pdbs_to_fetch)} valid PDBs to fetch")
+    # downloading a list of good pdbs
+    validated_pdbs = fetch_pdb_files(pdbs_to_fetch[:max_pdb_num])
+    log.info(f"Found {len(pdbs_to_fetch)} valid PDBs")
     return validated_pdbs
 
 
@@ -389,7 +436,8 @@ def preprocess_pdb(pdb_fname, chain_id):
     tidy_pdb_f : Path
         preprocessed pdb file
     """
-    atoms_pdb_f = keep_atoms(pdb_fname)
+    model_pdb_f = selmodel_pdb(pdb_fname)
+    atoms_pdb_f = keep_atoms(model_pdb_f)
     chained_pdb_f = selchain_pdb(atoms_pdb_f, chain_id)
     occ_pdb_f = occ_pdb(chained_pdb_f)
     tidy_pdb_f = tidy_pdb(occ_pdb_f)
@@ -565,7 +613,6 @@ def get_best_pdb(
             return
 
     # if pdb_to_use is not None, already filter the list
-
     if pdb_to_use:
         pdb_to_use = pdb_to_use.lower()
     if chain_to_use:
@@ -591,8 +638,8 @@ def get_best_pdb(
 
     log.info(
         f"BestPDB hit for {uniprot_id}:"
-        f" {pdb_id}_{chain_id} {coverage:.2f} coverage"
-        f" {resolution:.2f} Angstrom / start {start} end {end}"
+        f" {pdb_id}_{chain_id} {coverage} coverage"
+        f" {resolution} Angstrom / start {start} end {end}"
     )
 
     processed_pdb = pdb_f.rename(f"{uniprot_id}-{pdb_id}-{chain_id}.pdb")


Test patch:
diff --git a/tests/test_pdb.py b/tests/test_pdb.py
index 1f7e252..7851263 100644
--- a/tests/test_pdb.py
+++ b/tests/test_pdb.py
@@ -9,6 +9,7 @@ from arctic3d.modules.pdb import (
     keep_atoms,
     occ_pdb,
     selchain_pdb,
+    selmodel_pdb,
     tidy_pdb,
     validate_api_hit,
 )
@@ -98,6 +99,15 @@ def good_hits():
     return hits_list
 
 
+@pytest.fixture
+def example_interfaces():
+    interfaces = {
+        "P01024": [103, 104, 105],
+        "P-dummy": [103, 104, 105, 1049, 1050],
+    }
+    return interfaces
+
+
 def test_selchain_pdb(inp_pdb):
     pdb = selchain_pdb(inp_pdb, "B")
     assert pdb.exists()
@@ -122,10 +132,19 @@ def test_keep_atoms(inp_pdb):
     pdb.unlink()
 
 
+def test_selmodel_pdb(inp_pdb):
+    pdb = selmodel_pdb(inp_pdb, "1")
+    assert pdb.exists()
+    pdb.unlink()
+
+
 def test_validate_api_hit(pdb_hit_no_resolution):
+    """Test validate_api_hit."""
     validated_pdbs = validate_api_hit([pdb_hit_no_resolution])
-    assert validated_pdbs == []
-
+    assert (
+        validated_pdbs == []
+    )  # this is empty because resolution is None and exp != NMR
+    # change resolution to 1.0
     pdb_hit_no_resolution["resolution"] = 1.0
     validated_pdbs = validate_api_hit([pdb_hit_no_resolution])
     pdb, dict = validated_pdbs[0]
@@ -133,12 +152,19 @@ def test_validate_api_hit(pdb_hit_no_resolution):
     assert dict == pdb_hit_no_resolution
 
 
-def test_get_best_pdb():
-    orig_interfaces = {
-        "P01024": [103, 104, 105],
-        "P-dummy": [103, 104, 105, 1049, 1050],
-    }
-    pdb, filtered_interfaces = get_best_pdb("P20023", orig_interfaces)
+def test_validate_api_hit_nmr(pdb_hit_no_resolution):
+    """Test validate_api_hit with NMR data."""
+    pdb_hit_no_resolution["experimental_method"] = "Solution NMR"
+    # NMR structures have no resolution but should be accepted
+    validated_pdbs = validate_api_hit([pdb_hit_no_resolution])
+    pdb, dict = validated_pdbs[0]
+    assert pdb.name == "2gsx.pdb"
+    assert dict == pdb_hit_no_resolution
+
+
+def test_get_best_pdb(example_interfaces):
+    """Test get_best_pdb."""
+    pdb, filtered_interfaces = get_best_pdb("P20023", example_interfaces)
     exp_pdb = Path("P20023-1ghq-B.pdb")
     exp_interfaces = {"P01024": [103, 104, 105]}
     assert pdb == exp_pdb
@@ -146,8 +172,8 @@ def test_get_best_pdb():
     exp_pdb.unlink()
 
 
-def test_get_maxint_pdb():
-    """Test get_maxint_pdb."""
+def test_get_maxint_pdb_empty():
+    """Test get_maxint_pdb with empty output."""
     empty_validated_pdbs = []
     pdb_f, top_hit, filtered_interfaces = get_maxint_pdb(
         empty_validated_pdbs, {}, uniprot_id=None
@@ -156,7 +182,17 @@ def test_get_maxint_pdb():
     assert top_hit is None
     assert filtered_interfaces is None
 
-    # TODO: test the non-empty case as well
+
+def test_get_maxint_pdb(good_hits, example_interfaces):
+    """Test get_maxint_pdb."""
+    validated_pdbs = validate_api_hit(good_hits)
+    pdb_f, top_hit, filtered_interfaces = get_maxint_pdb(
+        validated_pdbs, example_interfaces, "P00760"
+    )
+    assert pdb_f.name == "4xoj-model1-atoms-A-occ-tidy_renum.pdb"
+    assert top_hit["pdb_id"] == "4xoj"
+    assert top_hit["chain_id"] == "A"
+    assert filtered_interfaces == {"P01024": [103, 104, 105]}
 
 
 def test_filter_pdb_list(good_hits):