Repository: haddocking/arctic3d

Issue:
add maxclust criterion in resclust cli
adding the `maxclust` criterion can be useful when we already know that we want to divide the structure in n clusters of residues

Discussion:

Patch:
diff --git a/src/arctic3d/cli_resclust.py b/src/arctic3d/cli_resclust.py
index d5f3b12..fa7e31c 100644
--- a/src/arctic3d/cli_resclust.py
+++ b/src/arctic3d/cli_resclust.py
@@ -14,6 +14,10 @@ Input arguments:
     `threshold` : the number to be used as threshold for hierarchical clustering.
 
     `chain` : the chain ID to be used.
+
+    `linkage` : the linkage strategy.
+
+    `criterion` : the criterion to extract the clusters.
 """
 import argparse
 import logging
@@ -56,6 +60,15 @@ argument_parser.add_argument(
     default=10.0,
 )
 
+argument_parser.add_argument(
+    "--criterion",
+    help="Criterion for clustering",
+    type=str,
+    required=False,
+    choices=["distance", "maxclust"],
+    default="distance",
+)
+
 argument_parser.add_argument(
     "--linkage",
     help="Linkage strategy for clustering",
@@ -109,7 +122,7 @@ def maincli():
     cli(argument_parser, main)
 
 
-def main(input_arg, residue_list, chain, threshold, linkage):
+def main(input_arg, residue_list, chain, threshold, linkage, criterion):
     """Main function."""
     log.setLevel("INFO")
 
@@ -152,12 +165,18 @@ def main(input_arg, residue_list, chain, threshold, linkage):
         sys.exit(1)
 
     # do the clustering
+    if criterion == "maxclust":
+        threshold = int(threshold)
     log.info(
-        f"Clustering distance matrix with linkage {linkage} and threshold {threshold}"
+        f"Clustering distance matrix with linkage {linkage}, threshold {threshold}, and criterion {criterion}"
     )
     distmap = pdist(u.positions)
     clusters = cluster_similarity_matrix(
-        distmap, unique_sorted_resids, threshold=threshold, linkage_strategy=linkage
+        distmap,
+        unique_sorted_resids,
+        threshold=threshold,
+        linkage_strategy=linkage,
+        crit=criterion,
     )
 
     cl_dict = get_clustering_dict(clusters, unique_sorted_resids)
diff --git a/src/arctic3d/modules/clustering.py b/src/arctic3d/modules/clustering.py
index 0b4dd4e..560834f 100644
--- a/src/arctic3d/modules/clustering.py
+++ b/src/arctic3d/modules/clustering.py
@@ -53,7 +53,12 @@ def plot_dendrogram(linkage_matrix, entries, filename, threshold, max_entries=50
 
 
 def cluster_similarity_matrix(
-    int_matrix, entries, linkage_strategy="average", threshold=0.866, plot=False
+    int_matrix,
+    entries,
+    linkage_strategy="average",
+    threshold=0.866,
+    crit="distance",
+    plot=False,
 ):
     """
     Does the clustering.
@@ -82,7 +87,7 @@ def cluster_similarity_matrix(
         dendrogram_figure_filename = "dendrogram_" + linkage_strategy + ".png"
         plot_dendrogram(Z, entries, dendrogram_figure_filename, threshold)
     # clustering
-    clusters = fcluster(Z, t=threshold, criterion="distance")
+    clusters = fcluster(Z, t=threshold, criterion=crit)
     log.info("Dendrogram created and clustered.")
     log.debug(f"Clusters = {clusters}")
     return clusters


Test patch:
diff --git a/tests/test_cli_resclust.py b/tests/test_cli_resclust.py
index e0e17ac..32580bc 100644
--- a/tests/test_cli_resclust.py
+++ b/tests/test_cli_resclust.py
@@ -14,11 +14,22 @@ def example_pdbpath():
 
 
 def test_resclust_cli(example_pdbpath):
-    main(example_pdbpath, "100,101,102,133,134,135", None, 7.0, "average")
+    main(example_pdbpath, "100,101,102,133,134,135", None, 7.0, "average", "distance")
 
 
 def test_wrong_residue_list(example_pdbpath):
     with pytest.raises(SystemExit) as e:
-        main(example_pdbpath, "100,101,102,133,134,135%", None, 9.0, "average")
+        main(
+            example_pdbpath,
+            "100,101,102,133,134,135%",
+            None,
+            9.0,
+            "average",
+            "distance",
+        )
     assert e.type == SystemExit
     assert e.value.code == 1
+
+
+def test_resclust_maxclust(example_pdbpath):
+    main(example_pdbpath, "100,101,102,133,134,135", None, 2, "average", "maxclust")