-
Notifications
You must be signed in to change notification settings - Fork 1
/
node.py
85 lines (75 loc) · 2.33 KB
/
node.py
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
class Node:
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.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 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(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 set_height(self):
if len(self.children) == 0:
self.height = 0
else:
tnode = self
h = 0
while len(tnode.children) > 0:
if tnode.children[1].length < tnode.children[0].length:
tnode = tnode.children[1]
else:
tnode = tnode.children[0]
h += tnode.length
self.height = h