Package obitools :: Package tree
[hide private]
[frames] | no frames]

Source Code for Package obitools.tree

  1  import re 
  2   
  3       
4 -class Tree(set):
5 - def registerNode(self,node):
6 assert isinstance(node, TreeNode) 7 self.add(node)
8
9 - def childNodeIterator(self,node):
10 assert isinstance(node, TreeNode) 11 return (x for x in self if x._parent==node)
12
13 - def subTreeSize(self,node):
14 n=1 15 for subnode in self.childNodeIterator(node): 16 n+=self.subTreeSize(subnode) 17 return n
18
19 - def getRoot(self):
20 roots = [x for x in self if x._parent is None] 21 assert len(roots)==1,'Tree cannot have several root node' 22 return roots[0]
23
24 - def ancestorNodeIterator(self,node):
25 assert isinstance(node, TreeNode) 26 while node._parent is not None: 27 yield node 28 node=node._parent 29 yield node
30
31 - def terminalNodeIterator(self):
32 return (x for x in self if x._isterminal)
33
34 - def commonAncestor(self,node1,node2):
35 anc1 = set(x for x in self.ancestorNodeIterator(node1)) 36 rep = [x for x in self.ancestorNodeIterator(node2) 37 if x in anc1] 38 assert len(rep)>=1 39 return rep[0]
40
41 - def getDist(self,node1,node2):
42 ca = self.commonAncestor(node1, node2) 43 dist = 0 44 while node1 != ca: 45 dist+=node1._dist 46 node1=node1._parent 47 while node2 != ca: 48 dist+=node2._dist 49 node2=node2._parent 50 return dist
51
52 - def farestNodes(self):
53 dmax=0 54 n1=None 55 n2=None 56 for node1 in self.terminalNodeIterator(): 57 for node2 in self.terminalNodeIterator(): 58 d = self.getDist(node1, node2) 59 if d > dmax: 60 dmax = d 61 n1=node1 62 n2=node2 63 return node1,node2,dmax
64
65 - def setRoot(self,node,dist):
66 assert node in self 67 assert node._parent and node._dist > dist 68 69 newroot = TreeNode(self) 70 parent = node._parent 71 node._parent = newroot 72 compdist = node._dist - dist 73 node._dist=dist 74 node = parent 75 76 while node: 77 parent = node._parent 78 if parent: 79 dist = node._dist 80 81 node._parent = newroot 82 node._dist = compdist 83 84 newroot = node 85 node = parent 86 87 if node: 88 compdist=dist 89 90 for child in self.childNodeIterator(newroot): 91 child._parent = newroot._parent 92 child._dist += newroot._dist 93 94 self.remove(newroot)
95 96
97 -class TreeNode(object):
98 - def __init__(self,tree,name=None,dist=None,bootstrap=None,**info):
99 self._parent=None 100 self._name=name 101 self._dist=dist 102 self._bootstrap=bootstrap 103 self._info=info 104 tree.registerNode(self) 105 self._isterminal=True
106 107
108 - def linkToParent(self,parent):
109 assert isinstance(parent, TreeNode) or parent is None 110 self._parent=parent 111 if parent is not None: 112 parent._isterminal=False
113