From 9cc6c13502a271d60a5fe5f01ba68c3769a0486f Mon Sep 17 00:00:00 2001
From: Nick Croucher <n.croucher@imperial.ac.uk>
Date: Thu, 17 Oct 2024 21:21:44 +0100
Subject: [PATCH] Quicker network plotting

---
 PopPUNK/plot.py      | 7 +++++--
 PopPUNK/visualise.py | 8 +++++---
 2 files changed, 10 insertions(+), 5 deletions(-)

diff --git a/PopPUNK/plot.py b/PopPUNK/plot.py
index 6f7a995b..a49988ca 100644
--- a/PopPUNK/plot.py
+++ b/PopPUNK/plot.py
@@ -511,7 +511,7 @@ def drawMST(mst, outPrefix, isolate_clustering, clustering_name, overwrite):
                     output=graph2_file_name, output_size=(3000, 3000))
 
 def outputsForCytoscape(G, G_mst, isolate_names, clustering, outPrefix, epiCsv, queryList = None,
-                        suffix = None, writeCsv = True):
+                        suffix = None, writeCsv = True, use_partial_query_graph = False):
     """Write outputs for cytoscape. A graphml of the network, and CSV with metadata
 
     Args:
@@ -536,6 +536,8 @@ def outputsForCytoscape(G, G_mst, isolate_names, clustering, outPrefix, epiCsv,
             (default = None)
         writeCsv (bool)
             Whether to print CSV file to accompany network
+        use_partial_query_graph (bool)
+            Whether to optimise output for outputting only partial graphs
     """
 
     # Avoid circular import
@@ -553,7 +555,8 @@ def outputsForCytoscape(G, G_mst, isolate_names, clustering, outPrefix, epiCsv,
         suffix = '_cytoscape'
     else:
         suffix = suffix + '_cytoscape'
-    save_network(G, prefix = outPrefix, suffix = suffix, use_graphml = True)
+    if not use_partial_query_graph:
+      save_network(G, prefix = outPrefix, suffix = suffix, use_graphml = True)
 
     # Save each component too (useful for very large graphs)
     component_assignments, component_hist = gt.label_components(G)
diff --git a/PopPUNK/visualise.py b/PopPUNK/visualise.py
index 8ea3030b..b39d0407 100644
--- a/PopPUNK/visualise.py
+++ b/PopPUNK/visualise.py
@@ -421,7 +421,7 @@ def generate_visualisations(query_db,
     #*                            *#
     #******************************#
 
-    if (tree == "nj" or tree == "both") or rank_fit == None:
+    if (tree == "nj" or tree == "both") or (model.type == 'lineage' and rank_fit == None):
         
         # Either calculate or read distances
         if recalculate_distances:
@@ -699,7 +699,7 @@ def generate_visualisations(query_db,
             if gpu_graph:
                 genomeNetwork = cugraph_to_graph_tool(genomeNetwork, isolateNameToLabel(all_seq))
             # Hard delete from network to remove samples (mask doesn't work neatly)
-            if include_files is not None:
+            if include_files is not None and not use_partial_query_graph:
                 genomeNetwork = remove_nodes_from_graph(genomeNetwork, all_seq, viz_subset, use_gpu = gpu_graph)
         elif rank_fit is not None:
             genomeNetwork = sparse_mat_to_network(sparse_mat, combined_seq, use_gpu = gpu_graph)
@@ -710,12 +710,14 @@ def generate_visualisations(query_db,
         # for full network
         node_labels = viz_subset if (use_partial_query_graph is not None or include_files is not None) \
                                       else combined_seq
+        sys.stderr.write('Preparing outputs for cytoscape\n')
         outputsForCytoscape(genomeNetwork,
                             mst_graph,
                             node_labels,
                             isolateClustering,
                             output,
-                            info_csv)
+                            info_csv,
+                            use_partial_query_graph)
         if model.type == 'lineage':
             sys.stderr.write("Note: Only support for output of cytoscape graph at lowest rank\n")