Repository: haddocking/arctic3d

Issue:
include probability in the output
when several interfaces are clustered together, the list of clustered residues can be quite long.

An example is can be found by running
`arctic3d example/1ppe_E.fasta --db db/swissprot`

48 interfaces can be found. 46 of them are clustered together. They are indeed similar, although they span a big set of 41 residues:
`1: [43, 45, 46, 47, 63, 65, 66, 67, 68, 89, 91, 98, 99, 100, 101, 102, 103, 104, 114, 149, 150, 152, 153, 154, 156, 157, 159, 178, 194, 195, 197, 198, 200, 215, 217, 218, 219, 220, 222, 223, 225]`
Of course not all of them are present in every of the 46 interfaces. Does it make sense to include their probability to be in an interface in the output to facilitate the understanding?

Discussion:

Patch:
diff --git a/src/arctic3d/cli.py b/src/arctic3d/cli.py
index da52d4c..9612f52 100644
--- a/src/arctic3d/cli.py
+++ b/src/arctic3d/cli.py
@@ -12,7 +12,7 @@ from arctic3d.modules.input import Input
 from arctic3d.modules.interface import get_interface_residues, read_interface_residues
 
 # from arctic3d.modules.output import make_output
-from arctic3d.modules.pdb import get_best_pdb
+from arctic3d.modules.pdb import get_best_pdb, output_pdb
 from arctic3d.modules.sequence import to_fasta
 
 # from arctic3d.modules.sequence import load_seq
@@ -151,18 +151,20 @@ def main(input_arg, db, interface_file, out_uniprot, out_pdb, pdb_to_use):
 
         # cluster interfaces
         if filtered_interfaces:
-            clustered_interface_residues = cluster_interfaces(
+            clustered_interface_residues, cl_residues_probs = cluster_interfaces(
                 filtered_interfaces, pdb_f
             )
         else:
-            clustered_interface_residues = cluster_interfaces(interface_residues, pdb_f)
+            clustered_interface_residues, cl_residues_probs = cluster_interfaces(
+                interface_residues, pdb_f
+            )
 
         log.info(f"Clustered interface residues: {clustered_interface_residues}")
+
+        output_pdb(pdb_f, cl_residues_probs)
     else:
         log.info("No interfaces found.")
 
-    # _ = make_output(best_pdb, clustered_interface_residues)
-
 
 if __name__ == "__main__":
     sys.exit(maincli())
diff --git a/src/arctic3d/modules/clustering.py b/src/arctic3d/modules/clustering.py
index c76465f..8086075 100644
--- a/src/arctic3d/modules/clustering.py
+++ b/src/arctic3d/modules/clustering.py
@@ -5,6 +5,7 @@ import os
 import time
 
 import matplotlib.pyplot as plt
+import numpy as np
 from scipy.cluster.hierarchy import dendrogram, fcluster, linkage
 
 from arctic3d.modules.interface_matrix import read_int_matrix
@@ -33,7 +34,7 @@ def cluster_similarity_matrix(int_matrix, entries, threshold=THRESHOLD, plot=Fal
     clusters : list
         list of clusters ID, each one associated to an entry
     """
-    log.info(f"clustering with threshold {threshold}")
+    log.info(f"Clustering with threshold {threshold}")
     Z = linkage(int_matrix, LINKAGE)
     if plot:
         dendrogram_figure_filename = "dendrogram_" + LINKAGE + ".png"
@@ -112,22 +113,33 @@ def get_residue_dict(cl_dict, interface_dict):
 
     Returns
     -------
-    res_dict : dict
+    clustered_residues : dict
         dictionary of clustered residues
         example { 1 : [1,2,3,5,6,8] ,
                   2 : [29,30,31],
                   ...
                 }
+    cl_residues_probs : dict of dicts
+        dictionary of probabilities for clustered residues
+        example { 1 : {1:0.7, 2:0.2, 3:0.4 ...}
+                  ...
+                }
     """
     clustered_residues = {}
+    cl_residues_probs = {}
     for key in cl_dict.keys():
+        denom = len(cl_dict[key])
         residues = []
         for int_id in cl_dict[key]:
             residues.extend(interface_dict[int_id])
-        unique_cl_residues = list(set(residues))
-        unique_cl_residues.sort()
-        clustered_residues[key] = unique_cl_residues
-    return clustered_residues
+        unique_res = np.unique(residues, return_counts=True)
+        cl_residues_probs[key] = {}
+        # assign probabilities
+        for res_idx, res in enumerate(unique_res[0]):
+            res_prob = unique_res[1][res_idx] / denom
+            cl_residues_probs[key][res] = res_prob
+        clustered_residues[key] = list(unique_res[0])
+    return clustered_residues, cl_residues_probs
 
 
 def write_residues(res_dict, res_filename):
@@ -140,11 +152,6 @@ def write_residues(res_dict, res_filename):
         dictionary of clustered residues
     res_filename : str or Path
         output filename
-
-    Returns
-    -------
-    cl_residues : dict
-        dictionary of clustered residues
     """
     # write to file
     with open(res_filename, "w") as wfile:
@@ -153,6 +160,31 @@ def write_residues(res_dict, res_filename):
             wfile.write(f"Cluster {key} -> " + cl_string + os.linesep)
 
 
+def write_residues_probs(cl_residues_probs, res_probs_filename):
+    """
+    Writes clustered residues to file with their probability.
+
+    Parameters
+    ----------
+    res_dict : dict
+        dictionary of clustered residues
+    res_probs_filename : str or Path
+        output filename
+    """
+    # write to file
+    with open(res_probs_filename, "w") as wfile:
+        for key in cl_residues_probs.keys():
+            cl_string = f"Cluster {key} : {len(cl_residues_probs[key].keys())} residues{os.linesep}"
+            cl_string += f"rank\tresid\tprobability{os.linesep}"
+            sorted_probs = sorted(
+                cl_residues_probs[key].items(), key=lambda x: x[1], reverse=True
+            )
+            for pair_idx, pair in enumerate(sorted_probs, start=1):
+                cl_string += f"{pair_idx}\t{pair[0]}\t{pair[1]:.3f}{os.linesep}"
+            cl_string += os.linesep
+            wfile.write(cl_string)
+
+
 def interface_clustering(interface_dict, matrix_filename):
     """
     Clusters the interface matrix.
@@ -184,9 +216,11 @@ def interface_clustering(interface_dict, matrix_filename):
     write_clusters(cl_dict, cl_filename)
     # write clustered residues
     res_filename = "clustered_residues.out"
-    clustered_residues = get_residue_dict(cl_dict, interface_dict)
-    write_residues(clustered_residues, res_filename)
+    cl_residues, cl_residues_probs = get_residue_dict(cl_dict, interface_dict)
+    write_residues(cl_residues, res_filename)
+    res_probs_filename = "clustered_residues_probs.out"
+    write_residues_probs(cl_residues_probs, res_probs_filename)
     # write time
     elap_time = round((time.time() - start_time), 3)
     log.info(f"Clustering performed in {elap_time} seconds")
-    return clustered_residues
+    return cl_residues, cl_residues_probs
diff --git a/src/arctic3d/modules/pdb.py b/src/arctic3d/modules/pdb.py
index 74d0f54..068fe4f 100644
--- a/src/arctic3d/modules/pdb.py
+++ b/src/arctic3d/modules/pdb.py
@@ -353,3 +353,53 @@ def get_best_pdb(uniprot_id, interface_residues, pdb_to_use=None):
     processed_pdb = tidy_pdb_f.rename(f"{uniprot_id}-{pdb_id}-{chain_id}.pdb")
 
     return processed_pdb, filtered_interfaces
+
+
+def output_pdb(pdb_f, cl_residues_probs):
+    """Outputs pdb containing probabilities
+
+    Parameters
+    ----------
+    pdb_f : str or Path
+        Path to PDB file.
+
+    cl_residues_probs : dict of dicts
+        dictionary of probabilities for clustered residues
+        example { 1 : {1:0.7, 2:0.2, 3:0.4 ...}
+                  ...
+                }
+
+    Returns
+    -------
+    output_files : list of Paths
+        List of Paths to output PDB files
+    """
+    output_files = []  # list of paths
+    log.info("Creating output pdb files...")
+    st_beta = 50.0  # baseline value for b-factors of the identified residues
+    # read original file
+    original_content = open(pdb_f, "r").read().split(os.linesep)
+    # iterate over clusters
+    for cl_id in cl_residues_probs.keys():
+        # creating new file
+        new_filename = f"{pdb_f.stem}_cl{cl_id}{pdb_f.suffix}"
+        if new_filename in os.listdir():
+            raise Exception(f"Existing pdb file {new_filename}")
+        output_files.append(Path(new_filename))
+        log.info(f"Creating output file {new_filename}")
+        # writing new file
+        with open(new_filename, "w") as wfile:
+            for ln in original_content:
+                if ln.startswith("ATOM"):
+                    resid = int(ln[23:27].strip())
+                    new_beta = 0.00
+                    if resid in cl_residues_probs[cl_id].keys():
+                        new_beta = (
+                            st_beta + (100 - st_beta) * cl_residues_probs[cl_id][resid]
+                        )
+                    n_blank = 3 - len(str(new_beta).split(".")[0])
+                    new_line = f"{ln[:60]}{' '*n_blank}{new_beta:.2f}{ln[66:]}"
+                    wfile.write(f"{new_line}{os.linesep}")
+                else:
+                    wfile.write(f"{ln}{os.linesep}")
+    return output_files


Test patch:
diff --git a/tests/test_clustering.py b/tests/test_clustering.py
index dfd75a6..69a8815 100644
--- a/tests/test_clustering.py
+++ b/tests/test_clustering.py
@@ -9,6 +9,7 @@ from arctic3d.modules.clustering import (
     get_residue_dict,
     write_clusters,
     write_residues,
+    write_residues_probs,
 )
 
 
@@ -54,8 +55,13 @@ def test_get_res_dict():
     interface_dict = {"int_1": [1, 2, 3], "int_2": [3, 4, 5], "int_3": [27, 28, 29]}
     cl_dict = {1: ["int_1", "int_2"], 2: ["int_3"]}
     expected_res_dict = {1: [1, 2, 3, 4, 5], 2: [27, 28, 29]}
-    observed_res_dict = get_residue_dict(cl_dict, interface_dict)
+    expected_res_probs = {
+        1: {1: 0.5, 2: 0.5, 3: 1.0, 4: 0.5, 5: 0.5},
+        2: {27: 1.0, 28: 1.0, 29: 1.0},
+    }
+    observed_res_dict, observed_res_probs = get_residue_dict(cl_dict, interface_dict)
     assert expected_res_dict == observed_res_dict
+    assert expected_res_probs == observed_res_probs
 
 
 def test_write_clusters(reference_cl_dict):
@@ -80,3 +86,25 @@ def test_write_residues(reference_res_dict):
     observed_content = open(res_filename, "r").read()
     assert expected_content == observed_content
     os.unlink(res_filename)
+
+
+def test_write_res_probs():
+    """Test write_residue_probs."""
+    example_res_probs = {1: {3: 0.2, 4: 0.75}, 2: {27: 1.0, 28: 1.0}}
+    expected_content = (
+        f"Cluster 1 : 2 residues{os.linesep}"
+        f"rank\tresid\tprobability{os.linesep}"
+        f"1\t4\t0.750{os.linesep}"
+        f"2\t3\t0.200{os.linesep}"
+        f"{os.linesep}"
+        f"Cluster 2 : 2 residues{os.linesep}"
+        f"rank\tresid\tprobability{os.linesep}"
+        f"1\t27\t1.000{os.linesep}"
+        f"2\t28\t1.000{os.linesep}"
+        f"{os.linesep}"
+    )
+    res_probs_filename = "residues_probs_test.out"
+    write_residues_probs(example_res_probs, res_probs_filename)
+    observed_content = open(res_probs_filename, "r").read()
+    assert expected_content == observed_content
+    os.unlink(res_probs_filename)
diff --git a/tests/test_pdb.py b/tests/test_pdb.py
index 6e21f38..3989d83 100644
--- a/tests/test_pdb.py
+++ b/tests/test_pdb.py
@@ -1,3 +1,4 @@
+import os
 from pathlib import Path
 
 import pytest
@@ -9,6 +10,7 @@ from arctic3d.modules.pdb import (
     get_maxint_pdb,
     keep_atoms,
     occ_pdb,
+    output_pdb,
     selchain_pdb,
     tidy_pdb,
     validate_api_hit,
@@ -136,3 +138,21 @@ def test_filter_pdb_list(good_hits):
     observed_red_list = filter_pdb_list(good_hits, "6sy3")
     expected_red_list = [good_hits[1]]
     assert observed_red_list == expected_red_list
+
+
+def test_output_pdb(inp_pdb):
+    """Test output_pdb."""
+    example_res_probs = {1: {3: 0.2, 4: 0.75}}
+    output_files = output_pdb(inp_pdb, example_res_probs)
+    original_content = open(inp_pdb, "r").read().split(os.linesep)
+    original_content = [el for el in original_content if el.startswith("ATOM")]
+    # check file existence
+    assert Path.exists(output_files[0])
+    observed_content = open(output_files[0], "r").read().split(os.linesep)
+    observed_content = [el for el in observed_content if el.startswith("ATOM")]
+    # assert equal length
+    assert len(original_content) == len(observed_content)
+    # assert equal content (except for b factors)
+    for ln_id in range(len(observed_content)):
+        assert original_content[ln_id][:60] == observed_content[ln_id][:60]
+    os.unlink(output_files[0])