biotaphy_scripts/cnode.pyx at master · biotaphy/biotaphy_scripts · 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
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155

cdef class Node:
    cdef public dict data
    cdef public bint istip
    cdef public double length,time_length,height,droot,lroot
    cdef public str label,note
    cdef public object parent
    cdef public list children

    def __init__(self):
        self.label = ""
        self.length = 0.0
        self.time_length = 0.0
        self.parent = None
        self.children = []
        self.data = {}
        self.istip = False
        self.height = 0
        self.droot = 0
        self.lroot = 0.0
        self.note = ""

    def add_child(self,child):
        #make sure that the child is not already in there
        assert child not in self.children
        self.children.append(child)
        child.parent = self

    def remove_child(self,child):
        #make sure that the child is in there
        assert child in self.children
        self.children.remove(child)
        child.parent = None

    def leaves(self,v=None):
        if v == None:
            v = []
        if len(self.children) == 0:
            v.append(self)
        else:
            for child in self.children:
                child.leaves(v)
        return v

    def get_leaf_by_name(self,nm):
        for i in self.leaves():
            if i.label == nm:
                return i
        return None

    def leaves_fancy(self):
        return [n for n in self.iternodes() if n.istip ]

    def lvsnms(self):
        return [n.label for n in self.iternodes() if n.istip ]

    def iternodes(self,order="preorder"):
        if order.lower() == "preorder":
            yield self
        for child in self.children:
            for d in child.iternodes(order):
                yield d
        if order.lower() == "postorder":
            yield self

    def prune(self):
        p = self.parent
        if p != None:
            p.remove_child(self)
        return p

    def get_newick_repr_paint(self):
        ret = ""
        painted_children = []
        for i in self.children:
            if "paint" in i.data:
                painted_children.append(i)
        for i in range(len(painted_children)):
            if i == 0:
                ret += "("
            ret += painted_children[i].get_newick_repr_paint()
            if i == len(painted_children)-1:
                ret += ")"
            else:
                ret += ","
        if self.label != None and "paint" in self.data:
            ret += self.label
        return ret

    def get_newick_repr(self,showbl=False):
        ret = ""
        for i in range(len(self.children)):
            if i == 0:
                ret += "("
            ret += self.children[i].get_newick_repr(showbl)
            if i == len(self.children)-1:
                ret += ")"
            else:
                ret += ","
        if self.label != None:
            ret += self.label
        if showbl == True:
            ret += ":" + str(self.length)
        return ret

    def get_newick_repr_note(self,showbl=False):
        ret = ""
        for i in range(len(self.children)):
            if i == 0:
                ret += "("
            ret += self.children[i].get_newick_repr_note(showbl)
            if i == len(self.children)-1:
                ret += ")"
            else:
                ret += ","
        if self.label != None:
            ret += self.label
        if len(self.note) > 0:
            ret += "["+self.note+"]"
        if showbl == True:
            ret += ":" + str(self.length)
        return ret

    def set_height(self):
        if len(self.children) == 0:
            self.height = 0
        else:
            tnode = self
            h = 0
            while len(tnode.children) > 0:
                if len(tnode.children) == 1:
                    tnode = tnode.children[0]
                else:
                    if tnode.children[1].length < tnode.children[0].length:
                        tnode = tnode.children[1]
                    else:
                        tnode = tnode.children[0]
                h += tnode.length
            self.height = h

    def set_dist_root(self):
        self.droot = 0
        if self.parent != None:
            cn = self
            while cn.parent != None:
                self.droot += 1
                cn = cn.parent

    def set_length_root(self):
        self.lroot = 0.0
        if self.parent != None:
            cn = self
            while cn.parent != None:
                self.lroot += cn.length
                cn = cn.parent