Merge pull request #208 from beiko-lab/fix-approx-rspr

Updated rspr_approx.py to remove double counting

Author: rbeiko

bin/rspr_approx.py

@@ -8,13 +8,15 @@
 from pathlib import Path
 import argparse
 import subprocess
+import shutil
 from ete3 import Tree
 import pandas as pd
 from collections import defaultdict
 from matplotlib import pyplot as plt
 from matplotlib.colors import LogNorm
 import seaborn as sns
 import tempfile
+import logging
 
 
 #####################################################################
@@ -98,8 +100,9 @@ def read_tree(input_path):
         tree_string = f.read()
         formatted = re.sub(r";[^:]+:", ":", tree_string)
         is_duplicated = check_formatted_tree(formatted)
+        is_small = formatted.count(",") < 3
 
-        return Tree(formatted), is_duplicated
+        return Tree(formatted), is_duplicated, is_small
 
 
 #####################################################################
@@ -111,33 +114,38 @@ def read_tree(input_path):
 #####################################################################
 
 
-def root_tree(input_path, basename, output_path):
-    tre,is_duplicated = read_tree(input_path)
+def root_one_tree(input_path, basename, output_path):
+    tre,is_duplicated,is_small = read_tree(input_path)
     midpoint = tre.get_midpoint_outgroup()
     tre.set_outgroup(midpoint)
     if is_duplicated:
         outdir = Path(output_path) / "multiple"
         Path(outdir).mkdir(exist_ok=True, parents=True)
         output_path = outdir / basename
         output_path = str(output_path).replace(".tre", ".tre.multiple")
+    elif is_small:
+        outdir = Path(output_path) / "small"
+        Path(outdir).mkdir(exist_ok=True, parents=True)
+        output_path = outdir / basename
+        output_path = str(output_path).replace(".tre", ".tre.small")
     else:
         outdir = Path(output_path) / "unique"
         Path(outdir).mkdir(exist_ok=True, parents=True)
         output_path = outdir / basename
 
     tre.write(outfile=output_path)
-    return tre.write(), len(tre.get_leaves()), output_path, is_duplicated
+    return tre.write(), len(tre.get_leaves()), output_path, is_duplicated, is_small
 
 def root_reference_tree(input_path, output_path):
-    tre, _ = read_tree(input_path)
+    tre, _, _ = read_tree(input_path)
     midpoint = tre.get_midpoint_outgroup()
     tre.set_outgroup(midpoint)
     tre.write(outfile=output_path)
     return tre.write(), len(tre.get_leaves())
 
 
 #####################################################################
-### FUNCTION ROOT_TREE
+### FUNCTION ROOT_ALL_TREES
 ### Root all the unrooted input trees in directory
 ### core_tree: path of the core tree
 ### gene_trees: path of the csv file containing all the gene tree paths
@@ -148,8 +156,7 @@ def root_reference_tree(input_path, output_path):
 #####################################################################
 
 
-def root_trees(core_tree, gene_trees_path, output_dir, results, merge_pair=False):
-    print("Rooting trees")
+def root_all_trees(core_tree, gene_trees_path, output_dir, results, merge_pair=False):
     #'''
     reference_tree = core_tree
 
@@ -165,11 +172,11 @@ def root_trees(core_tree, gene_trees_path, output_dir, results, merge_pair=False
     rooted_gene_trees_path = os.path.join(output_dir, "rooted_gene_trees")
     for filename in df_gene_trees["path"]:
         basename = Path(filename).name
-        gene_content, gene_tree_size, gene_tree_path, is_duplicated = root_tree(
+        gene_content, gene_tree_size, gene_tree_path, is_duplicated, is_small = root_one_tree(
                                                                 filename,
                                                                 basename,
                                                                 rooted_gene_trees_path)
-        if not is_duplicated:
+        if not (is_duplicated or is_small):
             results.loc[basename, "tree_size"] = gene_tree_size
         if merge_pair:
             with open(gene_tree_path, "w") as f2:
@@ -205,6 +212,9 @@ def extract_approx_distance(text):
 
 def run_approx_rspr(results, input_file, lst_filename, rspr_path):
     input_file.seek(0)
+
+    command_exists = shutil.which(rspr_path[0])
+
     result = subprocess.run(
         rspr_path, stdin=input_file, capture_output=True, text=True
     )
@@ -231,7 +241,6 @@ def run_approx_rspr(results, input_file, lst_filename, rspr_path):
 def approx_rspr(
     rooted_gene_trees_path, results, min_branch_len=0, max_support_threshold=0.7
 ):
-    print("Calculating approx distance")
     rspr_path = [
         "rspr",
         "-approx",
@@ -245,20 +254,73 @@ def approx_rspr(
     lst_filename = []
     with tempfile.TemporaryFile(mode='w+') as temp_file:
         for filename in os.listdir(rooted_gene_trees_path):
-            if cur_count == group_size:
-                run_approx_rspr(results, temp_file, lst_filename, rspr_path)
-                temp_file.seek(0)
-                temp_file.truncate()
-                lst_filename.clear()
-                cur_count = 0
-
-            gene_tree_path = os.path.join(rooted_gene_trees_path, filename)
-            with open(gene_tree_path, "r") as infile:
-                temp_file.write(infile.read() + "\n")
-                lst_filename.append(filename)
-                cur_count += 1
-        if cur_count > 0:
-            run_approx_rspr(results, temp_file, lst_filename, rspr_path)
+            if str(filename) in results.index:
+                print("Found " + str(filename))
+                if cur_count == group_size:
+                    run_approx_rspr(results, temp_file, lst_filename, rspr_path)
+                    temp_file.seek(0)
+                    temp_file.truncate()
+                    lst_filename.clear()
+                    cur_count = 0
+
+                gene_tree_path = os.path.join(rooted_gene_trees_path, filename)
+                with open(gene_tree_path, "r") as infile:
+                    lines = infile.readlines()
+                    if len(lines) < 2:
+                        print(f"File {filename} does not have enough lines.")
+                        continue
+                    tree = Tree(lines[1].strip())
+                    # Calculate N: number of nodes at or above the support threshold
+                    # num_resolved = sum(1 for node in tree.traverse() if node.support >= max_support_threshold and not node.is_leaf())
+                    num_resolved = -1
+                    for node in tree.traverse():
+                        if node.support is not None and node.support >= max_support_threshold and not node.is_leaf():
+                            num_resolved += 1
+
+                    tree_size = len(tree.get_leaves())
+                    results.loc[filename, "Num resolved"] = num_resolved
+                    results.loc[filename, "N/tree_size"] = num_resolved / tree_size if tree_size > 0 else 0
+                    lst_filename.append(filename)
+                    temp_file.write(lines[0].strip() + "\n" + lines[1].strip() + "\n")
+                    cur_count += 1
+                if cur_count > 0:
+                    run_approx_rspr(results, temp_file, lst_filename, rspr_path)
+
+    # Add the approx_drSPR/N column
+    results["approx_drSPR/N"] = results.apply(lambda row: float(row["approx_drSPR"]) / row["Num resolved"] if row["Num resolved"] > 0 else 0, axis=1)
+    print("CBA " + str(results))
+
+#def approx_rspr_old(
+#    rooted_gene_trees_path, results, min_branch_len=0, max_support_threshold=0.7
+#):
+#    print("Calculating approx distance")
+#    rspr_path = [
+#        "rspr",
+#        "-approx",
+#        "-multifurcating",
+#        "-length " + str(min_branch_len),
+#        "-support " + str(max_support_threshold),
+#    ]
+#
+#    group_size = 10000
+#    cur_count = 0
+#    lst_filename = []
+#    with tempfile.TemporaryFile(mode='w+') as temp_file:
+#        for filename in os.listdir(rooted_gene_trees_path):
+#            if cur_count == group_size:
+#                run_approx_rspr(results, temp_file, lst_filename, rspr_path)
+#                temp_file.seek(0)
+#                temp_file.truncate()
+#                lst_filename.clear()
+#                cur_count = 0
+#
+#            gene_tree_path = os.path.join(rooted_gene_trees_path, filename)
+#            with open(gene_tree_path, "r") as infile:
+#                temp_file.write(infile.read() + "\n")
+#                lst_filename.append(filename)
+#                cur_count += 1
+#        if cur_count > 0:
+#            run_approx_rspr(results, temp_file, lst_filename, rspr_path)
 
 
 #####################################################################
@@ -289,7 +351,6 @@ def generate_heatmap(freq_table, output_path, log_scale=False):
 #####################################################################
 
 def make_heatmap(results, output_path, min_distance, max_distance):
-    print("Generating heatmap")
 
     # create sub dataframe
     sub_results = results[(results["approx_drSPR"] >= min_distance)]
@@ -306,7 +367,6 @@ def make_heatmap(results, output_path, min_distance, max_distance):
 
 
 def make_heatmap_from_tsv(input_path, output_path, min_distance, max_distance):
-    print("Generating heatmap from CSV")
     results = pd.read_table(input_path)
     make_heatmap(results, output_path, min_distance, max_distance)
 
@@ -339,7 +399,6 @@ def get_heatmap_group_size(all_values, max_groups=15):
 #####################################################################
 
 def make_group_heatmap(results, output_path, min_distance, max_distance):
-    print("Generating group heatmap")
 
     # create sub dataframe
     sub_results = results[(results["approx_drSPR"] >= min_distance)]
@@ -383,7 +442,7 @@ def make_group_heatmap(results, output_path, min_distance, max_distance):
 ### RETURN groups of trees
 #####################################################################
 
-def generate_group_sizes(target_sum, max_groups=500):
+def generate_group_sizes(target_sum, max_groups=1000):
     degree = 1
     current_sum = 0
     group_sizes = []
@@ -410,7 +469,6 @@ def generate_group_sizes(target_sum, max_groups=500):
 #####################################################################
 
 def make_groups_v1(results, min_limit=10):
-    print("Generating groups")
     min_group = results[results["approx_drSPR"] <= min_limit]["file_name"].tolist()
     groups = defaultdict()
     first_group = "group_0"
@@ -438,7 +496,6 @@ def make_groups_v1(results, min_limit=10):
 #####################################################################
 
 def make_groups(results, min_limit=10):
-    print("Generating groups")
     min_group = results[results["approx_drSPR"] <= min_limit]["file_name"].tolist()
     groups = defaultdict()
     first_group = "group_0"
@@ -463,7 +520,6 @@ def make_groups(results, min_limit=10):
 
 
 def make_groups_from_csv(input_df, min_limit):
-    print("Generating groups from CSV")
     groups = make_groups_v1(input_df, min_limit)
     tidy_data = [
         (key, val)
@@ -476,6 +532,24 @@ def make_groups_from_csv(input_df, min_limit):
     return merged
 
 
+# def join_annotation_data(df, annotation_data):
+#    ann_df = pd.read_table(annotation_data, dtype={"genome_id": "str"})
+#    ann_df.columns = map(str.lower, ann_df.columns)
+#    ann_df.columns = ann_df.columns.str.replace(" ", "_")
+#    ann_subset = ann_df[["gene", "product"]]
+#
+#    df["tree_name"] = [f.split(".")[0] for f in df["file_name"]]
+#
+#    merged = df.merge(ann_subset, how="left", left_on="tree_name", right_on="gene")
+#
+#    if merged["gene"].isnull().all():
+#        ann_subset = ann_df[["locus_tag", "gene", "product"]]
+#        merged = df.merge(
+#            ann_subset, how="left", left_on="tree_name", right_on="locus_tag"
+#        )
+#
+#    return merged.fillna(value="NULL").drop("tree_name", axis=1).drop_duplicates()
+
 def join_annotation_data(df, annotation_data):
     ann_df = pd.read_table(annotation_data, dtype={"genome_id": "str"})
     ann_df.columns = map(str.lower, ann_df.columns)
@@ -492,8 +566,23 @@ def join_annotation_data(df, annotation_data):
             ann_subset, how="left", left_on="tree_name", right_on="locus_tag"
         )
 
-    return merged.fillna(value="NULL").drop("tree_name", axis=1).drop_duplicates()
+    merged = merged.fillna("NULL").drop("tree_name", axis=1)
 
+    # Group by all columns except 'product' and aggregate 'product'
+    grouped = (
+        merged.groupby(list(merged.columns.difference(['product'])))
+        .agg({'product': lambda x: '||'.join(sorted(set(x)))})
+        .reset_index()
+    )
+
+        # Reorder columns
+    desired_order = [
+        "file_name", "gene", "tree_size", "product", "N/tree_size",
+        "Num resolved", "approx_drSPR", "approx_drSPR/N"
+    ]
+    grouped = grouped[desired_order]
+
+    return grouped.drop_duplicates()
 
 def main(args=None):
     args = parse_args(args)
@@ -502,7 +591,7 @@ def main(args=None):
     #'''
     results = pd.DataFrame(columns=["file_name", "tree_size", "approx_drSPR"])
     results.set_index("file_name", inplace=True)
-    rooted_paths = root_trees(
+    rooted_paths = root_all_trees(
         args.CORE_TREE, args.GENE_TREES, args.OUTPUT_DIR, results, True
     )
     approx_rspr(
@@ -512,7 +601,10 @@ def main(args=None):
         args.MAX_SUPPORT_THRESHOLD,
     )
 
+    #exit(11)
+
     # Generate standard heatmap
+    # results["approx_drSPR"] = pd.to_numeric(results["approx_drSPR"]).fillna(1000000)
     results["approx_drSPR"] = pd.to_numeric(results["approx_drSPR"])
     fig_path = os.path.join(args.OUTPUT_DIR, "output.png")
     make_heatmap(