physcripts/make_balanced_tree.py at master · chinchliff/physcripts · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
#!/usr/bin/env python

import argparse, math, newick3, phylo3, sys

def recur_balanced(p, max_d, max_t, _d=0):

    global _j

    if max_t < 2 or _d > max_d:
        return

    _d += 1
    n_max_t = [math.floor(max_t/2),]
    n_max_t.append(max_t - n_max_t[0])
    for i in range(2):
        c = phylo3.Node()
        p.add_child(c)
        if n_max_t[i] < 2:
            c.istip = True
            c.label = "T"+str(_j)
            _j += 1
        else:
            recur_balanced(c, max_d, n_max_t[i], _d)

def get_balanced_tree(n_tips_per_tree, br_length_function, mean_br_length):

    global _j
    _j = 1

    tree_depth = math.floor(math.log(n_tips_per_tree,2))

    tree = phylo3.Node()
    recur_balanced(tree, tree_depth, n_tips_per_tree)

    return assign_branch_lengths(tree, br_length_function, mean_br_length)

def get_pectinate_tree(n_tips_per_tree, br_length_function, mean_br_length):

    tree = phylo3.Node()
    n = tree
    for i in range(n_tips_per_tree - 1):
        c = phylo3.Node()
        t = phylo3.Node()
        t.istip = True
        t.label = "T"+str(i+1)
        n.add_child(c)
        n.add_child(t)
        n = c
        if i == n_tips_per_tree - 2:
            c.istip = True
            c.label = "T"+str(i+2)

    return assign_branch_lengths(tree, br_length_function, mean_br_length)

def assign_branch_lengths(tree, br_length_function, mean_br_length):

    # NOTE: assumes a bifurcating tree
    for n in tree.iternodes(order=phylo3.POSTORDER):

        if n.istip:

            # check if this is a cherry
            p = n.parent
            for s in p.children:
                if s != n:
                    if s.istip and hasattr(s, "height"):
                        # a cherry, other tip has an assigned length, copy it
                        n.length = s.length
                    else: # not a cherry or other tip has no length, assign a length
                        n.length = br_length_function(mean_br_length)
            n.height = 0 # tips are always height 0

        else: # not a tip

            # collect lengths and heights of children
            l = (n.children[0].length, n.children[1].length)
            h = (n.children[0].height, n.children[1].height)

            # normalize the lengths of this node's children so the
            # cumulative distance to the tips is equal on both sides
            if h[0] + l[0] > h[1] + l[1]:
                m = 0
                n.children[1].length += h[0] + l[0] - (h[1] + l[1])
            else:
                m = 1
                n.children[0].length += h[1] + l[1] - (h[0] + l[0])

            # assign length and height to this node
            n.length = br_length_function(mean_br_length)
            n.height = n.children[m].height + n.children[m].length

    return tree

def validate_tree_type(t):
    if t == 'pectinate':
        return get_pectinate_tree
    elif t == 'balanced':
        return get_balanced_tree
    else:
        sys.exit('invalid tree type: ' + t)

if __name__ == '__main__':

    parser = argparse.ArgumentParser('')

    parser.add_argument('-n', '--number-of-tips', type=int, required=True, \
        help='The number of tips to be in the final tree')

    parser.add_argument('-b', '--mean-branch-length', type=float, required=True, \
        help='The length to be used for internal branches')

    parser.add_argument('-t', '--tree-generator-method', type=validate_tree_type, required=True, \
        help='The type of tree to generate. Must be either "pectinate" or "balanced".')

    args = parser.parse_args()

    sys.setrecursionlimit(args.number_of_tips+100)

    # currently just supports constant rate
    t = args.tree_generator_method(args.number_of_tips, lambda x: x, args.mean_branch_length)

    print(newick3.to_string(t) + ';')