Python ete3.Tree类代码示例

def createImg(filename, thres=0, samples=1):
    count = parseLineage(filename)
    suffix, matrix, taxo = getSuffixandMatrixandNewick(count,thres,samples)
    newick = convert(taxo,suffix)
    newick += ';'

    t = Tree(newick, format=1)
    ct = ClusterTree(t.write(),  text_array=matrix)
    addColors(ct)

    # nodes are linked to the array table
    array = ct.arraytable
    # Calculates some stats on the matrix. Needed to establish the color gradients.
    matrix_dist = [i for r in xrange(len(array.matrix))for i in array.matrix[r] if np.isfinite(i)]
    matrix_max = np.max(matrix_dist)
    matrix_min = np.min(matrix_dist)
    matrix_avg = (matrix_max+matrix_min)/2
    # Creates a profile face that will represent node's profile as a heatmap
    profileFace  = ProfileFace(matrix_max, matrix_min, matrix_avg, 200, 14, "heatmap",colorscheme=3)
    # Creates my own layout function that uses previous faces
    def mylayout(node):
        # If node is a leaf
        if node.is_leaf():
            # And a line profile
            add_face_to_node(profileFace, node, 0, aligned=True)
            node.img_style["size"]=2

    # Use my layout to visualize the tree
    ts = TreeStyle()
    ts.layout_fn = mylayout
    # ct.show(tree_style=ts)
    filedir = '/'.join(filename.split('/')[:-1])
    # t.write(format=9, outfile="output/newick/"+param+".nw")
    ct.render(filedir+'/phylo.png',tree_style=ts)

def tree_distances(file):

    t = Tree(file)
    branch_len_out = open(file + ".patristic-dist.tsv", "w")
    avg_distance_leaves = 0

    # Computing patristic distance matrix
    header = ""
    all_leaves = t.get_leaves()
    for i in all_leaves:
        header = header + "\t" + i.name

    nb_of_distances = 0
    max_len = 0
    min_len = 9999999999999999
    branch_len_out.write(header+"\n")
    for leaf1 in all_leaves:
        row = ""
        row += str(leaf1.name)
        for leaf2 in all_leaves:
            distance = np.clip(leaf1.get_distance(leaf2), 0.0, 99999999999999999999999999)
            avg_distance_leaves += distance
            row += "\t%f" % distance
            nb_of_distances += 1
            if distance > max_len:
                max_len = distance
            if distance < min_len and distance > 0:
                min_len = distance

        branch_len_out.write(row+"\n")

    branch_len_out.close()

def sanitizeByType(container, sanitizeby='tsv', onlycolumns=False):
    '''for a iterable of strings, carry out sanitizeString by:
        line, 
        tsv (all or onlycolumns), 
        fasta headers, or 
        leaf in nwk'''
    
    assert sanitizeby in set(['line', 'tsv', 'newick', 'fasta'])
    if sanitizeby=='line': 
        for line in container:
            print(sanitizeString(line.strip("\r\n"), False))
    if sanitizeby=='tsv': 
        for line in container:
            if onlycolumns: 
                newline = line.strip("\r\n").split("\t")
                for i in onlycolumns: 
                    newline[i-1]=sanitizeString(newline[i-1], False)
            else:
                newline=[sanitizeString(item.strip("\r\n"), False) for item in line.split("\t")]
            print("\t".join(newline))
    if sanitizeby=='newick':
        from ete3 import Tree
        t=Tree("".join(container))
        for l in t:
            l.name=sanitizeString(l.name, False)
        print(t.write())
    if sanitizeby=='fasta': 
        from Bio import SeqIO
        from io import StringIO
        from sys import stdout
        fasta = StringIO("".join(container))
        for seq_record in SeqIO.parse(fasta, "fasta"):
            seq_record.id=sanitizeString(seq_record.description, False)
            seq_record.description=''
            SeqIO.write(seq_record, stdout, "fasta")

def migrate(db_path):
    if db_path is None:
        raise ConfigError("No database path is given.")

    # make sure someone is not being funny
    utils.is_profile_db(db_path)

    # make sure the version is accurate
    profile_db = db.DB(db_path, None, ignore_version = True)
    if str(profile_db.get_version()) != current_version:
        raise ConfigError("Version of this profile database is not %s (hence, this script cannot really do anything)." % current_version)

    # migrate item orders
    item_orders = profile_db.get_table_as_dict(item_orders_table_name)
    for order_name in item_orders:
        if item_orders[order_name]['type'] == 'newick':
            newick = Tree(item_orders[order_name]['data'], format=1)
            newick = newick.write(format=2)
            profile_db._exec("""UPDATE %s SET "data" = ? WHERE "name" LIKE ?""" % item_orders_table_name, (newick, order_name))

    # migrate layer orders
    layer_orders = profile_db.get_table_as_dict(layer_orders_table_name)
    for order_name in layer_orders:
        if layer_orders[order_name]['data_type'] == 'newick':
            newick = Tree(layer_orders[order_name]['data_value'], format=1)
            newick = newick.write(format=2)
            profile_db._exec("""UPDATE %s SET "data_value" = ? WHERE "data_key" LIKE ?""" % layer_orders_table_name, (newick, order_name))

    # set the version
    profile_db.remove_meta_key_value_pair('version')
    profile_db.set_version(next_version)

    # bye
    profile_db.disconnect()
    progress.end()

def main(treefile, to, metric):
    with open(treefile) as fh:
        for treeline in fh:
            tree = Tree(treeline)
            tree = alphbetise_names(tree)
            tree = normalise_tree(tree, to, metric)
            print(tree.write(format=5))

def show_tree(experiment_folder):
    model = MDPD.Hierachical_MDPD(1)
    model.load(os.path.join(experiment_folder, 'model.p'))

    width, depth = model.width, model.depth

    root = Tree()

    cache = [(0, root)]

    for i in range(depth + 1):
        foo = []

        for idx, node in cache:
            paren = int((idx - 1) / width)
            kid = idx - paren * width
            face = faces.ImgFace(os.path.join(experiment_folder, 'images', '{}_{}_{}.png'.format(idx, paren, kid)))
            node.add_face(face, 0)

            if i < depth:
                for k in range(width):
                    foo.append((idx * width + k + 1, node.add_child()))

        cache = foo

    ts = TreeStyle()
    ts.mode = "c"

    root.render(os.path.join(experiment_folder, 'images', 'tree_plot.png'), tree_style=ts)
    return root

def pruning(inputtree, inputfasta, tree_outfilename):
    #This function remove sequences from a FASTA from a larger tree
    

    #Full initial tree - to be pruned
    k = open(inputtree, "r").read() 

    #ete3 Tree format
    f = Tree(inputtree)
 
    #List of IDs to be picked from the full FASTA
    IDlist=[] 
    fasta = open(inputfasta, "rU")
    record_dict = SeqIO.to_dict(SeqIO.parse(fasta, "fasta"))
    for recordID in record_dict.keys():
         print recordID
         IDlist.append(recordID)
    print IDlist

    tree_outfile=open(tree_outfilename, "w")

    print "pruning...", inputfasta
    f.prune(IDlist, preserve_branch_length=True)
    f.write(format=0, outfile=tree_outfilename)
    print "pruned", inputfasta

    def ete_draw(self, fname=None):
        """ Draws the tree and saves it to a file.  If `fname` is None,
            show the tree instead of saving it.

            Args:
                fname: filename to save to (default=None)
        """
        if Cfg.USE_ETE3:
            def layout(node):
                faces.add_face_to_node(AttrFace("name"), node, column=0,
                                       position="branch-right")

            ts = TreeStyle()
            ts.show_leaf_name = False
            ts.layout_fn = layout
            ts.rotation = 90
            
            tree = EteTree(self.ete_str(), format=8)

            if fname:
                tree.render(fname, tree_style=ts)
            else:
                tree.show(tree_style=ts)
        else:
            # TODO maybe throw an error?
            pass

def get_example_tree():
    # Random tree
    t = Tree()
    t.populate(20, random_branches=True)

    # Some random features in all nodes
    for n in t.traverse():
        n.add_features(weight=random.randint(0, 50))

    # Create an empty TreeStyle
    ts = TreeStyle()

    # Set our custom layout function
    ts.layout_fn = layout

    # Draw a tree
    ts.mode = "c"

    # We will add node names manually
    ts.show_leaf_name = False
    # Show branch data
    ts.show_branch_length = True
    ts.show_branch_support = True

    return t, ts

def main():
    args = parse_args()

    # Use the extension specified by the user if present
    if args.extension:
        ext = args.extension
    else:
        ext = '.mod.tre'

    # Load the tree
    t = Tree(args.tree)

    # Iterate over the nodes, convert to desired value
    for node in t.iter_search_nodes():
        if args.decimal:
            if node.support >= 1:
                node.support = float(node.support * 0.01)
            else:
                print >> sys.stderr, 'bootstrap value in {} is < 1, \
                        ignoring.'.format(args.tree)
        else:
            if node.support <= 1:
                node.support = int(node.support * 100)
            else:
                print >> sys.stderr, 'bootstrap value in {} is > 1, \
                        ignoring.'.format(args.tree)

    # If the replace flag is set, replace the input file with the output file.
    # Otherwise create a new file with the '.mod.tre' extension
    if args.replace:
        out = args.tree
    else:
        out = args.tree + ext

    t.write(format=0, outfile=out)

def parseTree(root):
    tree = Tree()
    tree.name = root['Name']
    tree.add_face(TextFace(root['Split'], fgcolor="red"), column=0, position="branch-bottom")
    if root['Children']:
        for child in root['Children']:
            tree.children.append(parseTree(child))
    return tree

def nhx2key(nhxtree):
    """Parse a PHYLDOG nhx file or string and create key for each node."""
    t = Tree(nhxtree)
    keyD = {}
    for node in t.traverse():
        k = "|".join(sorted([n for n in node.get_leaf_names()]))
        keyD[k] = node.ND
    return(keyD)

def revBayesTree2key(file):
    """Parse a revBayes node index tree file and create a key for each node."""
    t=Tree(file, format = 1)
    keyD = {}
    for node in t.traverse():
        k = "|".join(sorted([re.sub('\[&index=\d+\]','', n) for n in node.get_leaf_names()]))
        keyD[k] = nodeIndexFromString(node.name)
    return(keyD)

def get_example_tree():
    t = Tree()
    ts = TreeStyle()
    ts.layout_fn = layout
    ts.mode = "r"
    ts.show_leaf_name = False
    t.populate(10)
    return t, ts

def run(args):
    import random
    from ete3 import Tree

    for n in range(args.number):
        t = Tree()
        t.populate(args.size, random_branches=args.random_branches)
        dump(t)