Updates to MST

PopPUNK/plot.py

@@ -414,7 +414,7 @@ def distHistogram(dists, rank, outPrefix):
                 "_rank_" + str(rank) + "_histogram.png")
     plt.close()
 
-def drawMST(mst, outPrefix, isolate_clustering, overwrite):
+def drawMST(mst, outPrefix, isolate_clustering, clustering_name, overwrite):
     """Plot a layout of the minimum spanning tree
 
     Args:
@@ -424,6 +424,8 @@ def drawMST(mst, outPrefix, isolate_clustering, overwrite):
             Output prefix for save files
         isolate_clustering (dict)
             Dictionary of ID: cluster, used for colouring vertices
+        clustering_name (str)
+            Name of clustering scheme to be used for colouring
         overwrite (bool)
             Overwrite existing output files
     """
@@ -446,12 +448,12 @@ def drawMST(mst, outPrefix, isolate_clustering, overwrite):
                             output=graph1_file_name, output_size=(3000, 3000))
         if overwrite or not os.path.isfile(graph2_file_name):
             cluster_fill = {}
-            for cluster in set(isolate_clustering['Cluster'].values()):
+            for cluster in set(isolate_clustering[clustering_name].values()):
                 cluster_fill[cluster] = list(np.random.rand(3)) + [0.9]
             plot_color = mst.new_vertex_property('vector<double>')
             mst.vertex_properties['plot_color'] = plot_color
             for v in mst.vertices():
-                plot_color[v] = cluster_fill[isolate_clustering['Cluster'][mst.vp.id[v]]]
+                plot_color[v] = cluster_fill[isolate_clustering[clustering_name][mst.vp.id[v]]]
 
             gt.graph_draw(mst, pos=pos, vertex_fill_color=mst.vertex_properties['plot_color'],
                     output=graph2_file_name, output_size=(3000, 3000))

PopPUNK/sparse_mst.py

@@ -32,6 +32,7 @@ def get_options():
     iGroup.add_argument('--distances', required=True, help='Prefix of input pickle of pre-calculated distances (required)')
     iGroup.add_argument('--rank-fit', required=True, help='Location of rank fit, a sparse matrix (*_rank*_fit.npz)')
     iGroup.add_argument('--previous-clustering', help='CSV file with cluster definitions')
+    iGroup.add_argument('--display-cluster', default=None, help='Column of clustering CSV to use for plotting')
 
     # output options
     oGroup = parser.add_argument_group('Output options')
@@ -116,7 +117,7 @@ def main():
     if not args.no_plot:
         if args.previous_clustering != None:
             mode = "clusters"
-            if re.match(r"_lineages\.csv$", args.previous_clustering):
+            if args.previous_clustering.endswith('_lineages.csv'):
                 mode = "lineages"
             isolateClustering = readIsolateTypeFromCsv(args.previous_clustering,
                                                        mode = mode,
@@ -127,7 +128,20 @@ def main():
             for v in mst.vertices:
                 isolateClustering['Cluster'][mst.vp.id[v]] = '0'
 
-        drawMST(mst, args.output, isolateClustering, True)
+        # Check selecting clustering type is in CSV
+        clustering_name = 'Cluster'
+        if args.display_cluster != None and args.previous_clustering != None:
+            if args.display_cluster not in isolateClustering.keys():
+                sys.stderr.write('Unable to find clustering column ' + args.display_cluster + ' in file ' +
+                                 args.previous_clustering + '\n')
+                sys.exit()
+            else:
+                clustering_name = args.display_cluster
+        else:
+            clustering_name = list(isolateClustering.keys())[0]
+
+        # Draw MST
+        drawMST(mst, args.output, isolateClustering, clustering_name, True)
 
     sys.exit(0)
 

PopPUNK/visualise.py

@@ -53,15 +53,18 @@ def get_options():
                              'to clusters [default = reference database directory]',
                         type = str)
     iGroup.add_argument('--previous-clustering',
-                        help='Directory containing previous cluster definitions '
+                        help='File containing previous cluster definitions '
                              'and network [default = use that in the directory '
                              'containing the model]',
                         type = str)
     iGroup.add_argument('--previous-query-clustering',
-                        help='Directory containing previous cluster definitions '
+                        help='File containing previous cluster definitions '
                              'from poppunk_assign [default = use that in the directory '
-                             'containing the model]',
+                             'of the query database]',
                         type = str)
+    iGroup.add_argument('--display-cluster',
+                        help='Column of clustering CSV to use for plotting',
+                        default=None)
 
     # output options
     oGroup = parser.add_argument_group('Output options')
@@ -150,6 +153,7 @@ def generate_visualisations(query_db,
                             overwrite,
                             core_only,
                             accessory_only,
+                            display_cluster,
                             web):
 
     from .models import loadClusterFit
@@ -295,34 +299,37 @@ def generate_visualisations(query_db,
         sys.exit(1)
 
     # Load previous clusters
-    mode = "clusters"
-    suffix = "_clusters.csv"
-    if model.type == "lineage":
-        mode = "lineages"
-        suffix = "_lineages.csv"
-    if model.indiv_fitted:
-        sys.stderr.write("Note: Individual (core/accessory) fits found, but "
-                         "visualisation only supports combined boundary fit\n")
-
-    # Set directories of previous fit
     if previous_clustering is not None:
         prev_clustering = previous_clustering
+        mode = "clusters"
+        suffix = "_clusters.csv"
+        if prev_clustering.endswith('_lineages.csv'):
+            mode = "lineages"
+            suffix = "_lineages.csv"
     else:
-        prev_clustering = os.path.dirname(model_file)
-    cluster_file = prev_clustering + '/' + os.path.basename(prev_clustering) + suffix
-    isolateClustering = readIsolateTypeFromCsv(cluster_file,
+        # Identify type of clustering based on model
+        mode = "clusters"
+        suffix = "_clusters.csv"
+        if model.type == "lineage":
+            mode = "lineages"
+            suffix = "_lineages.csv"
+        if model.indiv_fitted:
+            sys.stderr.write("Note: Individual (core/accessory) fits found, but "
+                             "visualisation only supports combined boundary fit\n")
+        prev_clustering = os.path.dirname(model_file) + '/' + os.path.basename(model_file) + suffix
+    isolateClustering = readIsolateTypeFromCsv(prev_clustering,
                                                mode = mode,
                                                return_dict = True)
 
     # Join clusters with query clusters if required
     if not self:
         if previous_query_clustering is not None:
-            prev_query_clustering = previous_query_clustering + '/' + os.path.basename(previous_query_clustering)
+            prev_query_clustering = previous_query_clustering
         else:
-            prev_query_clustering = query_db
+            prev_query_clustering = os.path.dirname(query_db) + '/' + os.path.basename(query_db) + suffix
 
         queryIsolateClustering = readIsolateTypeFromCsv(
-                prev_query_clustering + suffix,
+                prev_query_clustering,
                 mode = mode,
                 return_dict = True)
         isolateClustering = joinClusterDicts(isolateClustering, queryIsolateClustering)
@@ -348,7 +355,19 @@ def generate_visualisations(query_db,
                                  weights_type=mst_distances,
                                  summarise=False)
             mst_graph = generate_minimum_spanning_tree(G)
-            drawMST(mst_graph, output, isolateClustering, overwrite)
+            # Check selecting clustering type is in CSV
+            clustering_name = 'Cluster'
+            if display_cluster != None:
+                if display_cluster not in isolateClustering.keys():
+                    sys.stderr.write('Unable to find clustering column ' + display_cluster + ' in file ' +
+                                     prev_clustering + '\n')
+                    sys.exit()
+                else:
+                    clustering_name = display_cluster
+            else:
+                clustering_name = list(isolateClustering.keys())[0]
+            # Draw MST
+            drawMST(mst_graph, output, isolateClustering, clustering_name, overwrite)
             mst_tree = mst_to_phylogeny(mst_graph, isolateNameToLabel(combined_seq))
         else:
             mst_tree = existing_tree
@@ -443,6 +462,7 @@ def main():
                             args.overwrite,
                             args.core_only,
                             args.accessory_only,
+                            args.display_cluster,
                             web = False)
 
 if __name__ == '__main__':

test/run_test.py

@@ -62,7 +62,7 @@
 subprocess.run("python ../poppunk_visualise-runner.py --ref-db example_db --output example_viz --grapetree", shell=True, check=True)
 subprocess.run("python ../poppunk_visualise-runner.py --ref-db example_db --output example_viz_subset --microreact --include-files subset.txt", shell=True, check=True)
 subprocess.run("python ../poppunk_visualise-runner.py --ref-db example_db --query-db example_query --output example_viz_query --microreact", shell=True, check=True)
-subprocess.run("python ../poppunk_visualise-runner.py --ref-db example_db --previous-clustering example_lineages --model-dir example_lineages --output example_lineage_viz --microreact", shell=True, check=True)
+subprocess.run("python ../poppunk_visualise-runner.py --ref-db example_db --previous-clustering example_lineages/example_dbscan_clusters.csv --model-dir example_lineages --output example_lineage_viz --microreact", shell=True, check=True)
 subprocess.run("python ../poppunk_visualise-runner.py --distances example_query/example_query.dists --ref-db example_db --model-dir example_lineages --query-db example_lineage_query --output example_viz_query_lineages --microreact", shell=True, check=True)
 
 # MST