feat(rspr): Save cluster tree in newick format (#189)

* Save cluster tree file in newick format * refactor: Change write_tree to use Phylo.write Signed-off-by: jvfe <[email protected]> --------- Signed-off-by: jvfe <[email protected]> Co-authored-by: KARTIK KAKADIYA <[email protected]>
beiko-lab · Apr 16, 2024 · cc5c53f · cc5c53f
1 parent e47e851
commit cc5c53f
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 6 deletions.
diff --git a/bin/rspr_heatmap.py b/bin/rspr_heatmap.py
@@ -6,7 +6,7 @@
 from matplotlib import pyplot as plt
 from matplotlib.colors import LogNorm
 import seaborn as sns
-import os 
+import os
 import json
 from Bio import Phylo
 import re
@@ -42,6 +42,12 @@ def parse_args(args=None):
         dest="CLUSTER_OUTPUT",
         help="Cluster probability output file name",
     )
+    parser.add_argument(
+        "-cfo",
+        "--cluster_file_output",
+        dest="CLUSTER_FILE_OUTPUT",
+        help="Cluster probability output newick tree file name",
+    )
     parser.add_argument(
         "-mnher",
         "--min_heatmap_exact_rspr",
@@ -231,7 +237,7 @@ def generate_cluster_network(lst_tree_clusters, refer_tree):
     print("Generating cluster network")
     if not refer_tree:
         return
-      
+
     lst_leaves = [leave.name for leave in refer_tree.get_terminals()]
     leaf_mapping = {leaf: i for i, leaf in enumerate(lst_leaves)}
     dict_clstr_map = defaultdict(int)
@@ -288,13 +294,15 @@ def generate_cluster_heatmap(lst_tree_clusters, cluster_heatmap_path):
     plt.ylabel("Leaves")
     plt.savefig(cluster_heatmap_path)
 
-    
+
 def read_tree(input_path):
     with open(input_path, "r") as f:
         tree_string = f.read()
         formatted = re.sub(r";[^:]+:", ":", tree_string)
         return Phylo.read(io.StringIO(formatted), "newick")
-
+
+def write_tree(output_path, data):
+    Phylo.write(data, output_path, "newick")
 
 def get_fig_size(refer_tree):
     max_fig_size = 100
@@ -306,7 +314,7 @@ def get_fig_size(refer_tree):
 
 
 #endregion
-    
+
 def main(args=None):
     args = parse_args(args)
 
@@ -315,7 +323,7 @@ def main(args=None):
     # Generate standard heatmap
     results["exact_drSPR"] = pd.to_numeric(results["exact_drSPR"])
     make_heatmap(
-        results, 
+        results,
         args.OUTPUT,
         args.MIN_HEATMAP_RSPR_DISTANCE,
         args.MAX_HEATMAP_RSPR_DISTANCE
@@ -337,10 +345,12 @@ def main(args=None):
             lst_tree_clusters.append(json.loads(str_clstr))
 
     cluster_tree_path = args.CLUSTER_OUTPUT
+    cluster_file_path = args.CLUSTER_FILE_OUTPUT
     refer_tree_path = os.path.join("rooted_reference_tree/core_gene_alignment.tre")
     refer_tree = read_tree(refer_tree_path)
     if refer_tree:
         generate_cluster_network(lst_tree_clusters, refer_tree)
+        write_tree(cluster_file_path, refer_tree)
 
         plt.rcParams['font.size'] = '12'
         fig_size = get_fig_size(refer_tree)

diff --git a/modules/local/rspr/heatmap.nf b/modules/local/rspr/heatmap.nf
@@ -17,6 +17,7 @@ process RSPR_HEATMAP {
     path "exact_output.png", emit: png
     path "exact_group_output.png", emit: exact_group_output
     path "cluster_tree_output.png", emit: cluster_tree_output
+    path "cluster_file_output.nwk", emit: cluster_file_output
 
     when:
     task.ext.when == null || task.ext.when
@@ -30,6 +31,7 @@ process RSPR_HEATMAP {
         -o exact_output.png \\
         -go exact_group_output.png \\
         -co cluster_tree_output.png \\
+        -cfo cluster_file_output.nwk \\
         --min_heatmap_exact_rspr $min_heatmap_exact_rspr \\
         --max_heatmap_exact_rspr $max_heatmap_exact_rspr \\
         $args
@@ -39,5 +41,6 @@ process RSPR_HEATMAP {
     touch exact_output.png
     touch exact_group_output.png
     touch cluster_tree_output.png
+    touch cluster_file_output.nwk
     """
 }