treesak 1.53.3__py3-none-any.whl

TreeSAK/PlotMcmcNode_old.py

@@ -0,0 +1,252 @@
+import os
+import glob
+import argparse
+import pandas as pd
+from ete3 import Tree
+import plotly.express as px
+
+
+PlotMcmcNode_usage = '''
+============================ PlotMcmcNode example commands ============================
+
+TreeSAK PlotMcmcNode -i McmcTree_op_files -n n179 -o Clock2_n179.pdf
+TreeSAK PlotMcmcNode -i McmcTree_op_files -n n161,n186 -o Clock3_n161_n186.pdf
+TreeSAK PlotMcmcNode -i McmcTree_op_files -n nodes.txt -o multi_runs_multi_nodes.pdf
+
+# File name of the mcmc.txt and the corresponding mcmc out file need to follow 
+# the rule as specified below:
+[setting_1]_mcmc.txt
+[setting_1]_out.txt
+[setting_2]_mcmc.txt
+[setting_2]_out.txt
+
+# file format (-n, tab separated)
+# leave the 2nd column blank for nodes without renaming
+setting_1    node1	Bacteria
+setting_2    node2
+setting_3    node3,node9	Archaea
+
+# Y-axis label file format (-l, tab separated)
+PA_75_DeltaLL_50_clock3_mcmc.txt	DeltaLL_50
+PA_75_DeltaLL_75_clock3_mcmc.txt	DeltaLL_75
+
+=======================================================================================
+'''
+
+def sep_path_basename_ext(file_in):
+
+    # separate path and file name
+    f_path, file_name = os.path.split(file_in)
+    if f_path == '':
+        f_path = '.'
+
+    # separate file basename and extension
+    f_base, f_ext = os.path.splitext(file_name)
+
+    return f_path, f_base, f_ext
+
+
+def mcmctree_out_to_tree_str(mamctree_out):
+
+    # get tree string from mcmctree_out
+    tree_str = ''
+    tree_line = 0
+    current_line = 1
+    for each_line in open(mamctree_out):
+        if 'Species tree for FigTree.  Branch lengths = posterior mean times; 95% CIs = labels' in each_line:
+            tree_line = current_line + 1
+        if tree_line == current_line:
+            tree_str = each_line.strip()
+        current_line += 1
+
+    tree_str_no_space = tree_str.replace(' ', '')
+
+    # rename tree nodes
+    t = Tree(tree_str_no_space, format=1)
+    for each_node in t.traverse():
+        if each_node.is_leaf():
+            node_name_new = '_'.join(each_node.name.split('_')[1:])
+        else:
+            node_name_new = 't_n%s' % each_node.name
+        each_node.name = node_name_new
+
+    tree_str_renamed = t.write(format=8)
+
+    return tree_str_renamed
+
+
+def plot_distribution(df_txt, output_plot):
+
+    df = pd.read_table(df_txt, sep=',')
+    run_id_list = df['Setting'].unique()
+    node_id_list = df['Node'].unique()
+
+    # sort dataframe by run id
+    df = df.sort_values(by='Setting', ascending=False)
+
+    plot_width  = 900
+    plot_height = len(run_id_list)*100
+    if plot_height < 360:
+        plot_height = 360
+
+    fig = px.violin(df, x="Value", y="Setting", color="Node", points=False, orientation="h", width=plot_width, height=plot_height)
+    if len(node_id_list) == 1:
+        fig.update_traces(side="positive", fillcolor='lightblue', width=1.6, opacity=0.75)
+    else:
+        fig.update_traces(side="positive", fillcolor='rgba(0,0,0,0)', width=1.6)
+
+    fig.update_traces(showlegend=True)
+    fig.layout.template = "simple_white"
+    # fig.layout.width = 700
+    # fig.layout.height = 750
+    # fig.update_xaxes(range=[40, 0])
+    # fig.update_layout(margin_t=10, title_text='Demo', title_x=0.5)
+    fig.write_image(output_plot)
+
+
+def get_internal_node_to_plot(node_txt, mo_file):
+
+    tree_str = ''
+    if os.path.isfile(mo_file):
+        tree_str = mcmctree_out_to_tree_str(mo_file)
+
+    # get nodes to plot
+    node_set = set()
+    node_rename_dict = dict()
+    if os.path.isfile(node_txt) is True:
+        for each in open(node_txt):
+            each_split = each.strip().split('\t')
+            node_str = each_split[0]
+
+            # get internal_node_to_plot
+            internal_node_to_plot = ''
+            if ',' not in node_str:
+                internal_node_to_plot = each_split[0]
+            else:
+                leaf_list = node_str.split(',')
+                if tree_str == '':
+                    print('*out.txt file not found, program exited!')
+                    exit()
+                current_lca = Tree(tree_str, format=1).get_common_ancestor(leaf_list)
+                internal_node_to_plot = current_lca.name
+
+            # add internal_node_to_plot to  node_set
+            if internal_node_to_plot != '':
+                node_set.add(internal_node_to_plot)
+
+            # read in name to show in plot
+            if len(each_split) == 2:
+                if each_split[1] != '':
+                    node_rename_dict[internal_node_to_plot] = each_split[1]
+    else:
+        node_set = node_txt.split(',')
+
+    return node_set, node_rename_dict, tree_str
+
+
+def PlotMcmcNode(args):
+
+    mcmc_in             = args['i']
+    node_txt            = args['n']
+    output_plot         = args['o']
+    specified_out_file  = args['of']
+    y_label_txt         = args['l']
+    keep_tmp_file       = args['tmp']
+
+    # check MCMCTree output file/dir
+    if os.path.isfile(mcmc_in) is True:
+        mcmc_file_list = [mcmc_in]
+    else:
+        mcmc_file_re = '%s/*_mcmc.txt' % (mcmc_in)
+        mcmc_file_list = glob.glob(mcmc_file_re)
+
+    if len(mcmc_file_list) == 0:
+        print('*mcmc.txt file not found, program exited!')
+        exit()
+
+    if specified_out_file is None:
+        missed_out_file_list = []
+        for each_mcmc_file in mcmc_file_list:
+            pwd_out_file = each_mcmc_file.replace('_mcmc.txt', '_out.txt')
+            if os.path.isfile(pwd_out_file) is False:
+                missed_out_file_list.append(pwd_out_file)
+        if len(missed_out_file_list) > 0:
+            print('The following *out.txt files are missing, program exited!')
+            print('\n'.join(sorted(missed_out_file_list)))
+            exit()
+
+    # read in y-axis label file
+    y_label_dict = dict()
+    if y_label_txt is not None:
+        for each_sample in open(y_label_txt):
+            each_sample_split = each_sample.strip().split('\t')
+            if len(each_sample_split) == 2:
+                y_label_dict[each_sample_split[0]] = each_sample_split[1]
+            else:
+                print('Format error: %s' % y_label_txt)
+                exit()
+
+    _, f_base, _ = sep_path_basename_ext(output_plot)
+    found_matched_node = False
+    op_tree_tmp  = '%s_tree.txt'  % f_base
+    op_df_tmp    = '%s_data.txt'  % f_base
+    op_label_tmp = '%s_label.txt' % f_base
+
+    op_label_tmp_handle = open(op_label_tmp, 'w')
+    op_tree_tmp_handle = open(op_tree_tmp, 'w')
+    op_df_tmp_handle = open(op_df_tmp, 'w')
+    op_df_tmp_handle.write('Value,Node,Setting\n')
+    for mcmc_file in mcmc_file_list:
+
+        mcmc_file_no_path = mcmc_file
+        if '/' in mcmc_file_no_path:
+            mcmc_file_no_path = mcmc_file_no_path.split('/')[-1]
+
+        if specified_out_file is None:
+            pwd_current_run_mcmc_out = mcmc_file.replace('_mcmc.txt', '_out.txt')
+        else:
+            pwd_current_run_mcmc_out = specified_out_file
+        node_set, node_rename_dict, tree_str = get_internal_node_to_plot(node_txt, pwd_current_run_mcmc_out)
+        op_tree_tmp_handle.write('%s\t%s\n' % (mcmc_file_no_path.replace('_mcmc.txt', ''), tree_str))
+        label_to_write = y_label_dict.get(mcmc_file_no_path, mcmc_file_no_path)
+        mcmc_df = pd.read_table(mcmc_file, index_col=0)
+        for each_col in mcmc_df:
+            if each_col in node_set:
+                node_name_to_write = node_rename_dict.get(each_col, each_col)
+                found_matched_node = True
+                value_list = mcmc_df[each_col].values
+                for each_value in value_list:
+                    op_df_tmp_handle.write('%s,%s,%s\n' % (each_value, node_name_to_write, label_to_write))
+
+                op_label_tmp_handle.write('%s\t%s\t%s\n' % (label_to_write, each_col, node_name_to_write))
+    op_df_tmp_handle.close()
+    op_label_tmp_handle.close()
+    op_tree_tmp_handle.close()
+
+    if found_matched_node is False:
+        print('Provided node(s) not found, program exited!')
+        exit()
+
+    # plot distribution
+    plot_distribution(op_df_tmp, output_plot)
+
+    # remove tmp files
+    if keep_tmp_file is False:
+        os.system('rm %s' % op_tree_tmp)
+        os.system('rm %s' % op_df_tmp)
+        os.system('rm %s' % op_label_tmp)
+
+    print('Plot exported to %s, done!' % output_plot)
+
+
+if __name__ == '__main__':
+
+    PlotMcmcNode_parser = argparse.ArgumentParser()
+    PlotMcmcNode_parser.add_argument('-i',      required=True,                       help='folder holds the *mcmc.txt and *out.txt files')
+    PlotMcmcNode_parser.add_argument('-of',     required=False, default=None,        help='the *out.txt file')
+    PlotMcmcNode_parser.add_argument('-n',      required=True,                       help='Nodes to plot')
+    PlotMcmcNode_parser.add_argument('-l',      required=False, default=None,        help='labels on y axis')
+    PlotMcmcNode_parser.add_argument('-o',      required=True,                       help='Output plot')
+    PlotMcmcNode_parser.add_argument('-tmp',    required=False, action="store_true", help='keep tmp files')
+    args = vars(PlotMcmcNode_parser.parse_args())
+    PlotMcmcNode(args)

TreeSAK/RootTree.py

@@ -0,0 +1,101 @@
+import random
+import dendropy
+import argparse
+from ete3 import Tree
+
+
+RootTree_usage = '''
+====================== RootTree example commands ======================
+
+TreeSAK RootTree -i input.tree -og outgroup_genomes.txt -o rooted.tree
+
+=======================================================================
+'''
+
+
+def root_with_outgroup(input_tree, out_group_list, add_root_branch, tree_file_rooted):
+
+    """
+    Reroot the tree using the given outgroup.
+    modified based on: https://github.com/Ecogenomics/GTDBTk/blob/master/gtdbtk/reroot_tree.py
+
+    input_tree:  File containing Newick tree to rerooted.
+    output_tree: Name of file for rerooted tree.
+    outgroup:    Labels of taxa in outgroup.
+    """
+
+    tree = dendropy.Tree.get_from_path(input_tree, schema='newick', rooting='force-rooted', preserve_underscores=True)
+
+    outgroup_in_tree = set()
+    ingroup_leaves = set()
+    for n in tree.leaf_node_iter():
+        if n.taxon.label in out_group_list:
+            outgroup_in_tree.add(n.taxon)
+        else:
+            ingroup_leaves.add(n)
+
+    # Since finding the MRCA is a rooted tree operation, the tree is first rerooted on an ingroup taxa. This
+    # ensures the MRCA of the outgroup can be identified so long as the outgroup is monophyletic. If the
+    # outgroup is polyphyletic trying to root on it is ill-defined. To try and pick a "good" root for
+    # polyphyletic outgroups, random ingroup taxa are selected until two of them give the same size
+    # lineage. This will, likely, be the smallest bipartition possible for the given outgroup though
+    # this is not guaranteed.
+
+    mrca = tree.mrca(taxa=outgroup_in_tree)
+    mrca_leaves = len(mrca.leaf_nodes())
+    while True:
+        rnd_ingroup = random.sample(list(ingroup_leaves), 1)[0]
+        tree.reroot_at_edge(rnd_ingroup.edge, length1=0.5 * rnd_ingroup.edge_length, length2=0.5 * rnd_ingroup.edge_length)
+        mrca = tree.mrca(taxa=outgroup_in_tree)
+        if len(mrca.leaf_nodes()) == mrca_leaves:
+            break
+        mrca_leaves = len(mrca.leaf_nodes())
+
+    if mrca.edge_length is not None:
+        tree.reroot_at_edge(mrca.edge, length1=0.5 * mrca.edge_length, length2=0.5 * mrca.edge_length)
+
+        # tree.write_to_path(tree_file_rooted, schema='newick', suppress_rooting=True, unquoted_underscores=True)
+        tree_out_string = tree.as_string(schema='newick', suppress_rooting=True, unquoted_underscores=True)
+        tree_out_string = tree_out_string.replace("'", "")
+
+        # add the root bar
+        if add_root_branch is True:
+            tree_out_string = '(' + tree_out_string
+            tree_out_string = tree_out_string.replace(');', '):0.02);')
+
+        # write out tree string
+        tree_file_rooted_handle = open(tree_file_rooted, 'w')
+        tree_file_rooted_handle.write(tree_out_string)
+        tree_file_rooted_handle.close()
+
+
+def RootTree(args):
+
+    tree_file           = args['i']
+    out_group_txt       = args['og']
+    tree_file_rooted    = args['o']
+    tree_fmt            = args['fmt']
+    add_root_branch     = args['add_root']
+
+    out_group_set = set()
+    for each_og in open(out_group_txt):
+        out_group_set.add(each_og.strip())
+
+    # tre = Tree(tree_file, format=tree_fmt)
+    # out_group_lca = tre.get_common_ancestor(out_group_set)
+    # tre.set_outgroup(out_group_lca)
+    # tre.write(outfile=tree_file_rooted, format=tree_fmt)
+
+    root_with_outgroup(tree_file, out_group_set, add_root_branch, tree_file_rooted)
+
+
+if __name__ == '__main__':
+
+    RootTree_parser = argparse.ArgumentParser()
+    RootTree_parser.add_argument('-i',          required=True,                          help='input tree')
+    RootTree_parser.add_argument('-og',         required=True,                          help='out group leaves')
+    RootTree_parser.add_argument('-o',          required=True,                          help='output tree')
+    RootTree_parser.add_argument('-add_root',   required=False, action='store_true',    help='add the root branch')
+    RootTree_parser.add_argument('-fmt',        required=False, default=1, type=int,    help='tree format, default: 1')
+    args = vars(RootTree_parser.parse_args())
+    RootTree(args)