treesak 1.53.3__py3-none-any.whl

TreeSAK/ALE5.py

@@ -0,0 +1,210 @@
+import os
+import glob
+import operator
+from ete3 import Tree
+from itertools import chain
+from itertools import combinations
+
+
+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 powerset(iterable):
+
+    " powerset([1,2,3]) --> () (1,) (2,) (3,) (1,2) (1,3) (2,3) (1,2,3) "
+
+    s = list(iterable)  # allows duplicate elements
+    chain_obj = chain.from_iterable(combinations(s, r) for r in range(len(s)+1))
+    combo_lol = []
+    for _, combo in enumerate(chain_obj, 1):
+        if len(list(combo)) > 0:
+            combo_lol.append(list(combo))
+
+    return combo_lol
+
+
+def lca_to_two_leaves(species_tree_from_ale, internal_node_id):
+
+    # read in ale species tree
+    stree_ale = Tree(species_tree_from_ale, format=1)
+
+    # get all leaves of the internal node
+    internal_node = stree_ale.search_nodes(name=internal_node_id)[0]
+    internal_node_leaf_object = internal_node.get_leaves()
+    internal_node_leaf_set = set()
+    for each_leaf in internal_node_leaf_object:
+        internal_node_leaf_set.add(each_leaf.name)
+
+    # get the two leaves needed
+    targeted_two_leaves = []
+    leaves_found = False
+    for leaf_1 in internal_node_leaf_set:
+        for leaf_2 in internal_node_leaf_set:
+            if leaf_1 != leaf_2:
+                if leaves_found is False:
+                    current_lca_id = stree_ale.get_common_ancestor(leaf_1, leaf_2).name
+                    if current_lca_id == internal_node_id:
+                        targeted_two_leaves.append(leaf_1)
+                        targeted_two_leaves.append(leaf_2)
+                        leaves_found = True
+
+    return targeted_two_leaves[0], targeted_two_leaves[1]
+
+
+def keep_highest_rrtc(rrtc_in, rrtc_out):
+
+    rrtc_highest_prob_dict = dict()
+    for each_rrtc in open(rrtc_in):
+        rrtc_r = each_rrtc.strip().split(':')[0].split('\t')[0]
+        rrtc_d = each_rrtc.strip().split(':')[0].split('\t')[1]
+        rrtc_v = float(each_rrtc.strip().split(':')[1])
+        rrtc_key = '%s___%s' % (rrtc_r, rrtc_d)
+        if rrtc_key not in rrtc_highest_prob_dict:
+            rrtc_highest_prob_dict[rrtc_key] = rrtc_v
+        else:
+            if rrtc_v > rrtc_highest_prob_dict[rrtc_key]:
+                rrtc_highest_prob_dict[rrtc_key] = rrtc_v
+
+    with open(rrtc_out, 'w') as rrtc_out_handle:
+        for each_rrtc in sorted(rrtc_highest_prob_dict.items(), key=operator.itemgetter(1))[::-1]:
+            rrtc_r = each_rrtc[0].split('___')[0]
+            rrtc_d = each_rrtc[0].split('___')[1]
+            rrtc_v = each_rrtc[1]
+            rrtc_out_handle.write('%s\t%s:%s\n' % (rrtc_r, rrtc_d, rrtc_v))
+
+
+########################################################################################################################
+
+# file in
+ip_dir                  = '/Users/songweizhi/Desktop/DateArTree/01_HGT_ALE_with_OMA'
+species_tree_from_ale   = '/Users/songweizhi/Desktop/DateArTree/05_pRTC_wd/genome_tree.newick.ufboot.ale.stree'
+round_list              = [1, 2, 3, 4, 5]
+color_list              = ['dodgerblue', 'goldenrod1', 'darkorange1', 'seagreen3', 'orchid3']
+min_detected_times      = 2
+
+# file out
+op_dir                  = '/Users/songweizhi/Desktop/DateArTree/05_pRTC_wd/HGTs_5_rds_ALE'
+
+########################################################################################################################
+
+rscript                 = '%s/rscript.R'                        % op_dir
+plot_file               = '%s/Venn.pdf'                         % op_dir
+rtc_txt                 = '%s/rrtc.txt'                         % op_dir
+rtc_txt_highest         = '%s/rrtc_uniq_by_highest_prob.txt'    % op_dir
+
+if os.path.isdir(op_dir):
+    os.system('rm -r %s' % op_dir)
+os.system('mkdir %s' % op_dir)
+
+########################################################################################################################
+
+hgt_dict = dict()
+rd_to_hgt_dict= dict()
+for each_rd in round_list:
+
+    rd_id               = each_rd
+    current_rd_op_dir   = '%s/ALE4_op_dir_%s_0.3'   % (ip_dir, each_rd)
+    pdf_file_re         = '%s/*.%s'                 % (current_rd_op_dir, 'pdf')
+    pdf_file_list       = glob.glob(pdf_file_re)
+
+    rd_to_hgt_dict[rd_id] = set()
+
+    for each_pdf in pdf_file_list:
+        f_path, f_base, f_ext = sep_path_basename_ext(each_pdf)
+        f_base_split = f_base.split('_')
+        id_by_d_to_r = '%s_to_%s' % (f_base_split[3], f_base_split[5])
+        rd_og        = '%s_%s'    % (each_rd, f_base_split[0])
+        rd_og_value  = '%s_%s_%s' % (each_rd, f_base_split[0], f_base_split[6])
+
+        rd_to_hgt_dict[rd_id].add(id_by_d_to_r)
+
+        if id_by_d_to_r not in hgt_dict:
+            hgt_dict[id_by_d_to_r] = []
+        hgt_dict[id_by_d_to_r].append(rd_og_value)
+
+################################################### get Venn diagram ###################################################
+
+combination_list = powerset(round_list)
+
+value_str_list = []
+for each_cmbo in combination_list:
+    current_str = ''
+    if len(each_cmbo) == 1:
+        current_value = rd_to_hgt_dict[each_cmbo[0]]
+        current_str = 'area%s=%s' % (each_cmbo[0], len(current_value))
+        value_str_list.append(current_str)
+    else:
+        value_lol = []
+        for each_element in each_cmbo:
+            ele_value = rd_to_hgt_dict[each_element]
+            value_lol.append(ele_value)
+        shared = set(value_lol[0]).intersection(*value_lol)
+        current_str = 'n%s=%s' % (''.join([str(i) for i in each_cmbo]), len(shared))
+        value_str_list.append(current_str)
+
+value_str     = ', '.join(value_str_list)
+label_str     = '"' + '", "'.join([str(i) for i in round_list]) + '"'
+color_str     = '"' + '", "'.join([str(i) for i in color_list]) + '"'
+font_size_str = ', '.join(['1.2']*len(combination_list))
+
+rscript_handle = open(rscript, 'w')
+rscript_handle.write('library(futile.logger)\n')
+rscript_handle.write('library(gridBase)\n')
+rscript_handle.write('library(VennDiagram)\n')
+rscript_handle.write('pdf(file="%s")\n' % plot_file)
+rscript_handle.write('venn.plot <- draw.quintuple.venn(%s, category=c(%s), fill=c(%s), cat.col=c(%s), cat.cex=1.2, cat.dist=0.3, margin=0.05, cex=c(%s), ind=TRUE)\n' % (value_str, label_str, color_str, color_str, font_size_str))
+rscript_handle.write('dev.off()\n')
+rscript_handle.close()
+
+os.system('Rscript %s' % rscript)
+
+########################################################################################################################
+
+rtc_txt_handle = open(rtc_txt, 'w')
+qualified_hgt_num = 0
+for each_hgt in hgt_dict:
+
+    occurence_list = hgt_dict[each_hgt]
+    pdf_dir = '%s/%s_%s' % (op_dir, each_hgt, len(occurence_list))
+    if len(occurence_list) >= min_detected_times:
+
+        #################### prepare rtc file ####################
+
+        donor_id           = each_hgt.split('_to_')[0][2:]
+        recipient_id       = each_hgt.split('_to_')[1][2:]
+        d_leaf_1, d_leaf_2 = lca_to_two_leaves(species_tree_from_ale, donor_id)
+        r_leaf_1, r_leaf_2 = lca_to_two_leaves(species_tree_from_ale, recipient_id)
+
+        for each_occurence in occurence_list:
+            value = each_occurence.split('_')[-1]
+            rtc_str = '%s,%s\t%s,%s:%s' % (r_leaf_1, r_leaf_2, d_leaf_1, d_leaf_2, value)
+            rtc_txt_handle.write(rtc_str + '\n')
+
+        ##########################################################
+
+        qualified_hgt_num += 1
+        os.system('mkdir %s' % pdf_dir)
+        for each_h in occurence_list:
+            rd_id             = each_h.split('_')[0]
+            og_id             = each_h.split('_')[1]
+            value             = each_h.split('_')[2]
+            pwd_input_pdf_in  = '%s/ALE4_op_dir_%s_0.3/%s_HGT_*_%s_%s.pdf'  % (ip_dir, rd_id, og_id, each_hgt,value)
+            pwd_input_pdf_out = '%s/%s_%s_%s_%s.pdf'                        % (pdf_dir, rd_id, og_id, each_hgt,value)
+            os.system('cp %s %s' % (pwd_input_pdf_in, pwd_input_pdf_out))
+rtc_txt_handle.close()
+
+# remove redundant HGTs, keep the one with the highest probability
+keep_highest_rrtc(rtc_txt, rtc_txt_highest)
+
+print('The number of HGTs detected in >= %s runs is %s.' % (min_detected_times, qualified_hgt_num))
+print('Done!')

TreeSAK/ALE6.py

@@ -0,0 +1,401 @@
+import os
+import glob
+import argparse
+from Bio import SeqIO
+from ete3 import Tree
+
+
+ALE6_usage = '''
+====================================== ALE6 example commands ======================================
+
+# This module is developed to faa ancestral genomes based on ALE outputs
+TreeSAK ALE6 -1 ALE1_op_dir -3 ALE3_op_dir_30 -s species_tree.rooted.treefile -o ALE6_op_dir_30 -n 380 -cog BioSAK_arCOG_wd -kegg BioSAK_KEGG_wd
+TreeSAK ALE6 -1 ALE1_op_dir -3 ALE3_op_dir_30 -s species_tree.rooted.treefile -o ALE6_op_dir_30 -n 294,309,380,404
+TreeSAK ALE6 -1 ALE1_op_dir -3 ALE3_op_dir_30 -s species_tree.rooted.treefile -o ALE6_op_dir_30 -n interested_nodes.txt
+
+# Needed input files:
+-1: faa files
+-3: GeneContent.txt, SpeciesTreeRef.newick and Transfer_propensity.txt
+-s: the tree used as input for ALE2, to rename leafs back
+
+# To be added:
+1. A dereplication step to the produced faa file.
+
+===================================================================================================
+'''
+
+
+def sep_path_basename_ext(file_in):
+
+    f_path, f_name = os.path.split(file_in)
+    if f_path == '':
+        f_path = '.'
+    f_base, f_ext = os.path.splitext(f_name)
+
+    return f_name, f_path, f_base, f_ext[1:]
+
+
+def get_internal_node_leaves(ale_species_tree_file, internal_node_id):
+
+    ale_species_tree = Tree(ale_species_tree_file, format=1)
+    internal_node = ale_species_tree.search_nodes(name=internal_node_id)[0]
+    internal_node_leaf_object = internal_node.get_leaves()
+    internal_node_leaf_set = set()
+    for each_leaf in internal_node_leaf_object:
+        internal_node_leaf_set.add(each_leaf.name)
+
+    return internal_node_leaf_set
+
+
+def rename_leaves(tree_file_in, rename_dict, tree_format, tree_file_out):
+
+    t = Tree(tree_file_in, format=tree_format)
+    for leaf in t:
+        leaf_name_new = rename_dict.get(leaf.name, leaf.name)
+        leaf.name = leaf_name_new
+    t.write(format=tree_format, outfile=tree_file_out)
+
+
+def ALE6(args):
+
+    ale1_op_dir                 = args['1']
+    ale3_op_dir                 = args['3']
+    op_dir                      = args['o']
+    genome_tree_file_rooted     = args['s']
+    force_create_op_dir         = args['f']
+    interested_internal_nodes   = args['n']
+    cog_annotation_wd           = args['cog']
+    kegg_annotation_wd          = args['kegg']
+
+    GeneContent_txt             = '%s/GeneContent.txt'          % ale3_op_dir
+    SpeciesTreeRef              = '%s/SpeciesTreeRef.newick'    % ale3_op_dir
+    transfer_propensity_txt     = '%s/Transfer_propensity.txt'  % ale3_op_dir
+
+    ################################################## define op files #################################################
+
+    _, _, tree_base, tree_ext = sep_path_basename_ext(genome_tree_file_rooted)
+
+    faa_dir                                         = '%s/faa_files'                                    % op_dir
+    cog_dir                                         = '%s/annotation_COG'                               % op_dir
+    kegg_dir                                        = '%s/annotation_KEGG'                              % op_dir
+    cog_df_txt                                      = '%s/annotation_COG.txt'                           % op_dir
+    cog_df_desc_txt                                 = '%s/annotation_COG_desc.txt'                      % op_dir
+    kegg_df_txt                                     = '%s/annotation_KEGG.txt'                          % op_dir
+    kegg_df_desc_txt                                = '%s/annotation_KEGG_desc.txt'                     % op_dir
+    fun_transfer_propensity_txt                     = '%s/function_transfer_propensity_weighted.txt'    % op_dir
+    genome_tree_file_rooted_with_ale_internal_node  = '%s/%s_with_ALE_internal_nodes.%s'                % (op_dir, tree_base, tree_ext)
+
+    ########################################## get the id of nodes to process ##########################################
+
+    gnm_name_dict_ale_fmt_to_original_fmt = dict()
+    for leaf in Tree(genome_tree_file_rooted, format=1):
+        leaf_name = leaf.name
+        leaf_name_new = leaf_name.replace('_', '')
+        leaf.name = leaf_name_new
+        gnm_name_dict_ale_fmt_to_original_fmt[leaf_name_new] = leaf_name
+
+    overall_internal_node_set = set()
+    line_num_index = 0
+    for each_line in open(GeneContent_txt):
+        line_num_index += 1
+        line_split = each_line.strip().split('\t')
+        if line_num_index > 1:
+            node_id = line_split[0]
+            if '(' not in node_id:
+                overall_internal_node_set.add(node_id)
+
+    internal_nodes_to_process = set()
+    if interested_internal_nodes is None:
+        internal_nodes_to_process = overall_internal_node_set
+    else:
+        if os.path.isfile(interested_internal_nodes) is False:
+            if ',' in interested_internal_nodes:
+                internal_nodes_to_process = interested_internal_nodes.split(',')
+            else:
+                internal_nodes_to_process.add(interested_internal_nodes)
+        else:
+            for each_node in open(interested_internal_nodes):
+                internal_nodes_to_process.add(each_node.strip())
+
+    #################### an addiitonal step (add ALE added internal node names to the rooted tree) #####################
+
+    # create output directory
+    if force_create_op_dir is True:
+        if os.path.isdir(op_dir) is True:
+            os.system('rm -r %s' % op_dir)
+    os.system('mkdir %s' % op_dir)
+
+    rename_leaves(SpeciesTreeRef, gnm_name_dict_ale_fmt_to_original_fmt, 1, genome_tree_file_rooted_with_ale_internal_node)
+
+    ####################################################################################################################
+
+    os.system('mkdir %s' % faa_dir)
+
+    if cog_annotation_wd is not None:
+        os.system('mkdir %s' % cog_dir)
+
+    if kegg_annotation_wd is not None:
+        os.system('mkdir %s' % kegg_dir)
+
+    branch_to_leaf_dict = dict()
+    branch_to_content_dict = dict()
+    col_header_list = []
+    line_num_index = 0
+    for each_line in open(GeneContent_txt):
+        line_num_index += 1
+        line_split = each_line.strip().split('\t')
+        if line_num_index == 1:
+            col_header_list = line_split
+        else:
+            branch_id = line_split[0]
+            if branch_id in internal_nodes_to_process:
+                branch_to_content_dict[branch_id] = []
+                branch_child_leaf_set = get_internal_node_leaves(SpeciesTreeRef, branch_id)
+                branch_to_leaf_dict[branch_id] = branch_child_leaf_set
+                for (id, pa) in zip(col_header_list[1:], line_split[1:]):
+                    if pa != '0':
+                        branch_to_content_dict[branch_id].append(id)
+
+    branch_to_gene_dict = dict()
+    for each_branch in branch_to_content_dict:
+        branch_faa = '%s/%s.faa' % (faa_dir, each_branch)
+        branch_content = branch_to_content_dict[each_branch]
+        branch_child_set = branch_to_leaf_dict[each_branch]
+        branch_child_set_original_name = {gnm_name_dict_ale_fmt_to_original_fmt[i] for i in branch_child_set}
+        branch_faa_handle = open(branch_faa, 'w')
+        branch_to_gene_dict[each_branch] = set()
+        for each_prot_family in branch_content:
+            each_prot_family_faa = '%s/%s.faa' % (ale1_op_dir, each_prot_family)
+            for each_seq in SeqIO.parse(each_prot_family_faa, 'fasta'):
+                seq_id = each_seq.id
+                seq_gnm = '_'.join(seq_id.split('_')[:-1])
+                if seq_gnm in branch_child_set_original_name:
+                    branch_faa_handle.write('>%s %s\n' % (each_seq.id, each_prot_family))
+                    branch_faa_handle.write('%s\n' % each_seq.seq)
+                    branch_to_gene_dict[each_branch].add(each_seq.id)
+        branch_faa_handle.close()
+
+    ####################################################################################################################
+
+    # Read in COG annotation results
+    fun_to_gene_dict = dict()
+    annotation_dict_cog = dict()
+    fun_id_to_desc_dict = dict()
+    if cog_annotation_wd is not None:
+
+        print('Reading in COG annotation results')
+        file_re = '%s/*COG_wd/*_query_to_cog.txt' % (cog_annotation_wd)
+        file_list = glob.glob(file_re)
+
+        if len(file_list) == 0:
+            print('COG annotation file not detected, program exited!')
+            exit()
+
+        for each_file in file_list:
+            gnm_id = each_file.split('/')[-1].split('_query_to_cog')[0]
+            if gnm_id not in annotation_dict_cog:
+                annotation_dict_cog[gnm_id] = dict()
+            line_index = 0
+            for each_line in open(each_file):
+                if line_index > 0:
+                    each_line_split = each_line.strip().split('\t')
+                    if len(each_line_split) == 4:
+                        gene_id = each_line_split[0]
+                        cog_id = each_line_split[1]
+                        cog_desc = each_line_split[3]
+                        annotation_dict_cog[gnm_id][gene_id] = cog_id
+                        fun_id_to_desc_dict[cog_id] = cog_desc
+                        if cog_id not in fun_to_gene_dict:
+                            fun_to_gene_dict[cog_id] = set()
+                        fun_to_gene_dict[cog_id].add(gene_id)
+                line_index += 1
+
+    # Read in KEGG annotation results
+    annotation_dict_kegg = dict()
+    if kegg_annotation_wd is not None:
+
+        print('Reading in KEGG annotation results')
+        file_re = '%s/*KEGG_wd/*_ko_assignment_ABCD.txt' % (kegg_annotation_wd)
+        file_list = glob.glob(file_re)
+
+        if len(file_list) == 0:
+            print('KEGG annotation file not detected, program exited!')
+            exit()
+
+        for each_file in file_list:
+            gnm_id = each_file.split('/')[-1].split('_ko_assignment_ABCD')[0]
+            if gnm_id not in annotation_dict_kegg:
+                annotation_dict_kegg[gnm_id] = dict()
+
+            line_index = 0
+            for each_line in open(each_file):
+                if line_index > 0:
+                    each_line_split = each_line.strip().split('\t')
+                    if len(each_line_split) == 9:
+                        gene_id = each_line_split[0]
+                        ko_d_id = each_line_split[4][2:]
+                        ko_d_desc = each_line_split[8]
+                        annotation_dict_kegg[gnm_id][gene_id] = ko_d_id
+                        fun_id_to_desc_dict[ko_d_id] = ko_d_desc
+                        if ko_d_id not in fun_to_gene_dict:
+                            fun_to_gene_dict[ko_d_id] = set()
+                        fun_to_gene_dict[ko_d_id].add(gene_id)
+                line_index += 1
+
+    cog_dod = dict()
+    kegg_dod = dict()
+    all_identified_cog_set = set()
+    all_identified_kegg_set = set()
+    if (cog_annotation_wd is not None) or (kegg_annotation_wd is not None):
+        for each_branch in branch_to_gene_dict:
+            branch_gene_content = branch_to_gene_dict[each_branch]
+            branch_cog_set = set()
+            branch_kegg_set = set()
+            for each_gene in branch_gene_content:
+                gnm_id = '_'.join(each_gene.split('_')[:-1])
+                cog_fun = annotation_dict_cog[gnm_id].get(each_gene, 'na')
+                kegg_fun = annotation_dict_kegg[gnm_id].get(each_gene, 'na')
+                if cog_fun != 'na':
+                    branch_cog_set.add(cog_fun)
+                    all_identified_cog_set.add(cog_fun)
+                if kegg_fun != 'na':
+                    branch_kegg_set.add(kegg_fun)
+                    all_identified_kegg_set.add(kegg_fun)
+
+            cog_dod[each_branch] = branch_cog_set
+            kegg_dod[each_branch] = branch_kegg_set
+
+            # write out annotation
+            if len(branch_cog_set) > 0:
+                cog_annotation_txt = '%s/%s_COG.txt' % (cog_dir, each_branch)
+                cog_annotation_txt_handle = open(cog_annotation_txt, 'w')
+                for each_cog in sorted(list(branch_cog_set)):
+                    cog_annotation_txt_handle.write('%s\t%s\n' % (each_cog, fun_id_to_desc_dict[each_cog]))
+                cog_annotation_txt_handle.close()
+
+            if len(branch_kegg_set) > 0:
+                kegg_annotation_txt = '%s/%s_KEGG.txt' % (kegg_dir, each_branch)
+                kegg_annotation_txt_handle = open(kegg_annotation_txt, 'w')
+                for each_kegg in sorted(list(branch_kegg_set)):
+                    kegg_annotation_txt_handle.write('%s\t%s\n' % (each_kegg, fun_id_to_desc_dict[each_kegg]))
+                kegg_annotation_txt_handle.close()
+
+    all_identified_cog_list_sorted       = sorted(list(all_identified_cog_set))
+    all_identified_kegg_list_sorted      = sorted(list(all_identified_kegg_set))
+    all_identified_cog_list_sorted_desc  = [('%s__%s' % (i, fun_id_to_desc_dict[i])) for i in all_identified_cog_list_sorted]
+    all_identified_kegg_list_sorted_desc = [('%s__%s' % (i, fun_id_to_desc_dict[i])) for i in all_identified_kegg_list_sorted]
+
+    # write out COG dataframe
+    if len(all_identified_cog_set) > 0:
+        cog_df_txt_handle = open(cog_df_txt, 'w')
+        cog_df_txt_handle.write('\t%s\n' % '\t'.join(all_identified_cog_list_sorted))
+        cog_df_desc_txt_handle = open(cog_df_desc_txt, 'w')
+        cog_df_desc_txt_handle.write('\t%s\n' % '\t'.join(all_identified_cog_list_sorted_desc))
+        for each_branch in sorted(list(cog_dod.keys())):
+            branch_cogs = cog_dod[each_branch]
+            cog_pa_list = [each_branch]
+            for each_cog in all_identified_cog_list_sorted:
+                if each_cog in branch_cogs:
+                    cog_pa_list.append('1')
+                else:
+                    cog_pa_list.append('0')
+            cog_df_txt_handle.write('\t'.join(cog_pa_list) + '\n')
+            cog_df_desc_txt_handle.write('\t'.join(cog_pa_list) + '\n')
+        cog_df_txt_handle.close()
+        cog_df_desc_txt_handle.close()
+        print('Annotation matrix exported to: %s' % cog_df_txt)
+
+    # write out KEGG dataframe
+    if len(all_identified_kegg_set) > 0:
+        kegg_df_txt_handle = open(kegg_df_txt, 'w')
+        kegg_df_txt_handle.write('\t%s\n' % '\t'.join(all_identified_kegg_list_sorted))
+        kegg_df_desc_txt_handle = open(kegg_df_desc_txt, 'w')
+        kegg_df_desc_txt_handle.write('\t%s\n' % '\t'.join(all_identified_kegg_list_sorted_desc))
+        for each_branch in sorted(list(kegg_dod.keys())):
+            branch_keggs = kegg_dod[each_branch]
+            kegg_pa_list = [each_branch]
+            for each_kegg in all_identified_kegg_list_sorted:
+                if each_kegg in branch_keggs:
+                    kegg_pa_list.append('1')
+                else:
+                    kegg_pa_list.append('0')
+            kegg_df_txt_handle.write('\t'.join(kegg_pa_list) + '\n')
+            kegg_df_desc_txt_handle.write('\t'.join(kegg_pa_list) + '\n')
+        kegg_df_txt_handle.close()
+        kegg_df_desc_txt_handle.close()
+        print('Annotation matrix exported to: %s' % kegg_df_txt)
+
+    ################################# get transfer propensity of individual functions ##################################
+
+    print('Getting transfer propensity of individual function')
+
+    # get gene_to_oma_dict
+    faa_file_re = '%s/*.faa' % ale1_op_dir
+    faa_file_list = glob.glob(faa_file_re)
+    gene_to_oma_dict = dict()
+    for faa_file in faa_file_list:
+        _, _, faa_base, _ = sep_path_basename_ext(faa_file)
+        for each_seq in SeqIO.parse(faa_file, 'fasta'):
+            seq_id = each_seq.id
+            gene_to_oma_dict[seq_id] = faa_base
+
+    # get oma_to_transfer_propensity_dict
+    oma_to_transfer_propensity_dict = dict()
+    line_index = 0
+    for each_oma in open(transfer_propensity_txt):
+        if line_index > 0:
+            each_oma_split = each_oma.strip().split('\t')
+            oma_id = each_oma_split[0]
+            transfer_propensity = float(each_oma_split[1])
+            oma_to_transfer_propensity_dict[oma_id] = transfer_propensity
+        line_index += 1
+
+    fun_transfer_propensity_txt_handle = open(fun_transfer_propensity_txt, 'w')
+    fun_transfer_propensity_txt_handle.write('ID\tWeighted_transfer_propensity\tDescription\n')
+    oma_weighted_transfer_propensity_dict = dict()
+    for fun_id in sorted(list(fun_to_gene_dict.keys())):
+        current_fun_gene_set = fun_to_gene_dict[fun_id]
+        current_fun_oma_stats_dict = dict()
+        for gene_id in current_fun_gene_set:
+            gene_oma = gene_to_oma_dict.get(gene_id, 'na')
+            if gene_oma not in current_fun_oma_stats_dict:
+                current_fun_oma_stats_dict[gene_oma] = 1
+            else:
+                current_fun_oma_stats_dict[gene_oma] += 1
+
+        total_transfer_propensity = 0
+        total_oma_num = 0
+        for oma_id in current_fun_oma_stats_dict:
+            oma_num = current_fun_oma_stats_dict[oma_id]
+            oma_transfer_propensity = oma_to_transfer_propensity_dict.get(oma_id, 'na')
+            if oma_transfer_propensity != 'na':
+                total_transfer_propensity += (oma_num*oma_transfer_propensity)
+                total_oma_num += oma_num
+
+        oma_transfer_propensity_weighted = 'na'
+        if total_oma_num != 0:
+            oma_transfer_propensity_weighted = total_transfer_propensity/total_oma_num
+            oma_transfer_propensity_weighted = float("{0:.3f}".format(oma_transfer_propensity_weighted))
+
+        if oma_transfer_propensity_weighted != 'na':
+            oma_weighted_transfer_propensity_dict[fun_id] = oma_transfer_propensity_weighted
+            fun_transfer_propensity_txt_handle.write('%s\t%s\t%s\n' % (fun_id, oma_transfer_propensity_weighted, fun_id_to_desc_dict[fun_id]))
+
+    fun_transfer_propensity_txt_handle.close()
+
+    print('Done!')
+
+
+if __name__ == '__main__':
+
+    ALE6_parser = argparse.ArgumentParser()
+    ALE6_parser.add_argument('-1',      required=True,                          help='ALE1 output directory')
+    ALE6_parser.add_argument('-3',      required=True,                          help='ALE3 output directory')
+    ALE6_parser.add_argument('-s',      required=True,                          help='rooted species tree')
+    ALE6_parser.add_argument('-n',      required=False, default=None,           help='interested internal node(s)')
+    ALE6_parser.add_argument('-cog',    required=False, default=None,           help='COG annotation results')
+    ALE6_parser.add_argument('-kegg',   required=False, default=None,           help='KEGG annotation results')
+    ALE6_parser.add_argument('-o',      required=True,                          help='output directory')
+    ALE6_parser.add_argument('-f',      required=False, action="store_true",    help='force overwrite')
+    args = vars(ALE6_parser.parse_args())
+    ALE6(args)