From 455d7a0c77df12bf4dbbdd0c15bc1e84910cef0c Mon Sep 17 00:00:00 2001 From: Nick Croucher Date: Sat, 30 Nov 2024 01:46:59 +0000 Subject: [PATCH] Add extend query graph flag --- PopPUNK/visualise.py | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/PopPUNK/visualise.py b/PopPUNK/visualise.py index de2aab27..d9f10031 100644 --- a/PopPUNK/visualise.py +++ b/PopPUNK/visualise.py @@ -99,7 +99,8 @@ def get_options(): default=None) iGroup.add_argument('--extend-query-graph', help='Extend the partial query graph to include all co-clustered isolates', - default=None) + default=False, + action = 'store_true') # output options oGroup = parser.add_argument_group('Output options') @@ -419,6 +420,22 @@ def generate_visualisations(query_db, return_dict = True) isolateClustering = joinClusterDicts(isolateClustering, queryIsolateClustering) + # Add extra isolates to partial query graph if requested + if use_partial_query_graph and extend_query_graph: + # First identify the query clusteres + query_clusters = set() + cluster_types = ['Cluster'] + query_cluster_isolates = [] + for cluster_type in cluster_types: + for isolate in viz_subset: + query_clusters.add(isolateClustering[cluster_type][isolate]) + # Then identify the isolates in these clusters + for cluster_type in cluster_types: + for isolate in isolateClustering[cluster_type]: + if isolateClustering[cluster_type][isolate] in query_clusters: + query_cluster_isolates.append(isolate) + viz_subset = set(query_cluster_isolates) + #******************************# #* *# #* Process dense or sparse *#