diff get_clusters.py @ 1:d9f8cc3258f9 draft

"planemo upload for repository https://github.com/galaxycomputationalchemistry/galaxy-tools-compchem/ commit 1b23e024af45cc0999d9142d07de6897d4189ec2"
author chemteam
date Mon, 24 Aug 2020 06:09:11 -0400
parents
children 9ca30ad95444
line wrap: on
line diff
--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/get_clusters.py	Mon Aug 24 06:09:11 2020 -0400
@@ -0,0 +1,39 @@
+import argparse
+import collections
+import json
+
+import numpy as np
+
+from scipy.cluster.hierarchy import fcluster
+
+
+def separate_clusters(Z_fpath, threshold, min_members, output):
+    Z = np.loadtxt(Z_fpath)
+    branch_assignments = fcluster(Z, threshold, criterion='distance')
+    cluster_dict = collections.defaultdict(list)
+    for n, val in enumerate(branch_assignments):
+        cluster_dict[branch_assignments[n]].append(n)
+    cluster_dict = {int(k): v for k, v in cluster_dict.items()
+                    if len(v) >= min_members}
+    with open(output, 'w') as f:
+        json.dump(cluster_dict, f, indent=4, sort_keys=True)
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--Z', required=True,
+                        help='File for cluster linkage array.')
+    parser.add_argument('--threshold', type=int, required=True,
+                        help='Distance cutoff.')
+    parser.add_argument('--min-members', type=int, required=True,
+                        help='Minimum number of members of the cluster.')
+    parser.add_argument('--output', required=True,
+                        help='Output file.')
+    args = parser.parse_args()
+
+    separate_clusters(args.Z, args.threshold,
+                      args.min_members, args.output)
+
+
+if __name__ == "__main__":
+    main()