def load_de_numericized_newick_tree(tree_in,before="'",after="'",root=False):
    from cogent.core.tree import PhyloNode
    from cogent import LoadTree
    import os.path
    
    if os.path.isfile(tree_in):
        tree = LoadTree(tree_in)
    else:
        tree = LoadTree(treestring=tree_in)
    terminals = tree.getTipNames()
    rename_dict = {}
    for tip in terminals:
        rename_dict[tip] = before + str(tip) + after
    tree.reassignNames(rename_dict)
    if root:
        tree = tree.rootAtMidpoint()
    treestring = tree.getNewick(with_distances=True)
    
    return treestring
Esempio n. 2
0
def main():

    # Parse input to get parameters 
    option_parser, opts, args =\
        parse_command_line_parameters(**script_info)
    
    tree_file = opts.input_tree
    trait_table_fp = opts.input_trait_table
    verbose = opts.verbose 
    
    #Handle parameters with more complex defaults
    if opts.output_table:
        output_table_fp = opts.output_table  
    else:
        output_table_fp = add_to_filename(trait_table_fp,"reformatted")
    
    if opts.output_tree:
        output_tree_fp = opts.output_tree  
    else:
        output_tree_fp = add_to_filename(trait_table_fp,"reformatted")

     

    if verbose:
        print "Running with options:"
        print "%s:%s" %("Tree file",tree_file)
        print "%s:%s" %("Trait table",trait_table_fp)
        print "%s:%s" %("Output tree",output_tree_fp)
        print "%s:%s" %("Output trait table",output_table_fp)
        print "%s:%s" %("Add branch length to root",opts.add_branch_length_to_root)
        print "%s:%s" %("Convert to NEXUS?",opts.convert_to_nexus)
         
    # Open output files
    output_trait_table_file = open(output_table_fp,"w+")
    output_tree_file  = open(output_tree_fp,"w+")
    

    # Begin reformatting
    
    # TODO: convert parameter candidates to params
    #delimiter = "\t"
    
    
    delimiter = " "
    min_branch_length = 0.0001
    root_name = "root"
    
    if format_for_bayestraits:
        convert_to_nexus = True
        convert_to_bifurcating = True
        filter_table_by_tree_tips = True
        filter_tree_by_table_entries = True
        enforce_min_branch_length = True
        convert_trait_floats_to_ints = True

    #Load inputs
    input_tree = LoadTree(tree_file)
    
    trait_table = open(trait_table_fp,"U")
    trait_table_lines = trait_table.readlines()
    
    #if opts.verbose:
        #print "Tree prior to reformatting:"
        #print input_tree.getNewick(with_distances=True)
        #print "Tree nodes prior to reformatting:"
        #print_node_summary_table(input_tree)
    
    # Remove taxa missing in tree tips
    if filter_table_by_tree_tips: 
        trait_table_lines = filter_table_by_presence_in_tree(input_tree,trait_table_lines,delimiter=delimiter)
    
    #Convert floating point values to ints
    if convert_trait_floats_to_ints:
        trait_table_lines = convert_trait_values(\
            trait_table_lines,conversion_fn = int,delimiter=delimiter)
   
    #Write out results
    output_trait_table_file.writelines(trait_table_lines)
    trait_table.close()
    output_trait_table_file.close()
   
    # Use the reformatted trait table to filter tree
    trait_table = open(output_table_fp,"U")
    trait_table_lines = trait_table.readlines()
    

    if filter_tree_by_table_entries:
        input_tree = filter_tree_tips_by_presence_in_table(input_tree,trait_table_lines,delimiter=delimiter) 
    
    # Tree reformatting
    
    
    if convert_to_bifurcating:
        input_tree = input_tree.bifurcating() # Required by most ancSR programs
    
   
    if convert_to_bifurcating:
        input_tree = ensure_root_is_bifurcating(input_tree)
        # The below nutty-looking re-filtering step is necessary
        # When ensuring the root is bifurcating, internal nodes can get moved to the tips
        # So without additional filtering we get unannotated tip nodes
        if filter_tree_by_table_entries:
            input_tree = filter_tree_tips_by_presence_in_table(input_tree,trait_table_lines,delimiter=delimiter) 
    
    if enforce_min_branch_length:
        input_tree = set_min_branch_length(input_tree,min_length = min_branch_length)
    
    if opts.add_branch_length_to_root:
        input_tree = add_branch_length_to_root(input_tree,root_name=input_tree.Name,root_length=min_branch_length)
    
    input_tree.prune() 
    
    #if opts.verbose:
    #    print "Tree nodes after reformatting:"
    #    print_node_summary_table(input_tree)
    #    print "Convert to nexus is True?",convert_to_nexus is True
    if convert_to_nexus is True:
        lines = nexus_lines_from_tree(input_tree)
        output_tree_file.write("\n".join(map(str,lines)))
    else:
        output_tree_file.write(input_tree.getNewick(with_distances=True))
    
    output_tree_file.close()