Repository: haddocking/arctic3d

Issue:
add min_cluster_size parameter
this (optional) parameter must allow to filter out clusters that contain less than `min_cluster_size` residues

Comments:

Patch:
diff --git a/src/arctic3d/cli.py b/src/arctic3d/cli.py
index 8eb450c..6c8da6f 100644
--- a/src/arctic3d/cli.py
+++ b/src/arctic3d/cli.py
@@ -6,6 +6,7 @@ from pathlib import Path
 
 from arctic3d import log
 from arctic3d.modules.blast import run_blast
+from arctic3d.modules.clustering import filter_clusters
 from arctic3d.modules.cluster_interfaces import cluster_interfaces
 from arctic3d.modules.input import Input
 from arctic3d.modules.interface import (
@@ -111,6 +112,14 @@ argument_parser.add_argument(
     default="average",
 )
 
+argument_parser.add_argument(
+    "--min_clust_size",
+    help="Minimum number of residues in clusters",
+    type=int,
+    required=False,
+    default=0,
+)
+
 
 def load_args(arguments):
     """
@@ -166,6 +175,7 @@ def main(
     ligand,
     linkage_strategy,
     threshold,
+    min_clust_size,
     log_level="DEBUG",
 ):
     """Main function."""
@@ -282,6 +292,13 @@ def main(
 
         log.info(f"Clustered interfaces {cl_dict}")
         log.info(f"Clustered interface residues: {cl_residues}")
+        if min_clust_size > 0:
+            log.info(
+                f"Excluding clusters with less than {min_clust_size} residues"
+            )
+            cl_dict, cl_residues, cl_residues_probs = filter_clusters(
+                cl_dict, cl_residues, cl_residues_probs, min_clust_size
+            )
 
         make_output(
             interface_residues=interface_residues,
diff --git a/src/arctic3d/modules/clustering.py b/src/arctic3d/modules/clustering.py
index 344869b..849c968 100644
--- a/src/arctic3d/modules/clustering.py
+++ b/src/arctic3d/modules/clustering.py
@@ -221,3 +221,50 @@ def interface_clustering(
     log.info(f"Clustering performed in {elap_time} seconds")
     log.info(f"Clustering produced {len(cl_dict)} clusters")
     return cl_dict, cl_residues, cl_residues_probs
+
+
+def filter_clusters(cl_dict, cl_residues, cl_residues_probs, min_clust_size):
+    """
+    Filter clusters based on size.
+
+    Parameters
+    ----------
+    cl_dict : dict
+        dictionary of clustered interfaces
+    cl_residues : dict
+        dictionary of clustered residues
+    cl_residues_probs : dict of dicts
+        dictionary of probabilities for clustered residues
+    min_clust_size : int
+        minimum cluster size
+
+    Returns
+    -------
+    flt_cl_dict : dict
+        dictionary of clustered interfaces
+    flt_cl_residues : dict
+        dictionary of clustered residues
+    flt_cl_residues_probs : dict of dicts
+        dictionary of probabilities for clustered residues
+    """
+    # gather clusters not respecting the minimum size
+    excl_clusts = []
+    for cl in cl_residues:
+        if len(cl_residues[cl]) < min_clust_size:
+            log.info(f"Cluster {cl} has less than {min_clust_size} residues.")
+            excl_clusts.append(cl)
+    # remove clusters not respecting the minimum size
+    for cl in excl_clusts:
+        log.info(f"Removing cluster {cl}")
+        del cl_dict[cl]
+        del cl_residues[cl]
+        del cl_residues_probs[cl]
+    # renumber clusters
+    flt_cl_dict = {}
+    flt_cl_residues = {}
+    flt_cl_residues_probs = {}
+    for idx, cl in enumerate(cl_dict.keys()):
+        flt_cl_dict[idx + 1] = cl_dict[cl]
+        flt_cl_residues[idx + 1] = cl_residues[cl]
+        flt_cl_residues_probs[idx + 1] = cl_residues_probs[cl]
+    return flt_cl_dict, flt_cl_residues, flt_cl_residues_probs


Test patch:
diff --git a/tests/test_cli.py b/tests/test_cli.py
index 5d0f747..6cf8c76 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -23,6 +23,7 @@ def test_cli_empty():
         ligand=None,
         linkage_strategy=None,
         threshold=None,
+        min_clust_size=None,
     )
     os.chdir(start_cwd)
     exp_dir = Path(f"arctic3d-{target_uniprot}")
diff --git a/tests/test_clustering.py b/tests/test_clustering.py
index 7f6498a..f2ccb71 100644
--- a/tests/test_clustering.py
+++ b/tests/test_clustering.py
@@ -3,6 +3,7 @@ import pytest
 
 from arctic3d.modules.clustering import (
     cluster_similarity_matrix,
+    filter_clusters,
     get_clustering_dict,
     get_residue_dict,
 )
@@ -66,3 +67,27 @@ def test_get_res_dict():
     )
     assert expected_res_dict == observed_res_dict
     assert expected_res_probs == observed_res_probs
+
+
+def test_filter_clusters():
+    """Test correct filtering of clusters."""
+    example_cl_dict = {
+        1: ["int_1", "int_2"],
+        2: ["int_3"],
+    }
+    example_res_dict = {1: [1, 2, 3, 4, 5], 2: [27, 28, 29]}
+    example_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},
+    }
+    obs_cl_dict, obs_res_dict, obs_res_probs = filter_clusters(
+        example_cl_dict, example_res_dict, example_res_probs, 4
+    )
+    exp_cl_dict = {1: ["int_1", "int_2"]}
+    exp_res_dict = {1: [1, 2, 3, 4, 5]}
+    exp_res_probs = {
+        1: {1: 0.5, 2: 0.5, 3: 1.0, 4: 0.5, 5: 0.5},
+    }
+    assert exp_cl_dict == obs_cl_dict
+    assert exp_res_dict == obs_res_dict
+    assert exp_res_probs == obs_res_probs
diff --git a/tests/test_pdb.py b/tests/test_pdb.py
index aac2e7b..b265871 100644
--- a/tests/test_pdb.py
+++ b/tests/test_pdb.py
@@ -269,7 +269,6 @@ def test_pdb_data(inp_pdb_data):
 def test_convert_cif_to_pdbs(inp_cif_3psg):
     """Test convert_cif_to_pdbs."""
     obs_out_pdb_fnames = convert_cif_to_pdbs(inp_cif_3psg, "3psg", "P00791")
-    print(f"obs_out_pdb_fnames {obs_out_pdb_fnames}")
     exp_out_pdb_fnames = [Path("3psg-A.pdb")]
     assert exp_out_pdb_fnames == obs_out_pdb_fnames
     # inspect the pdb file