Skip to content
Permalink
Branch: master
Find file Copy path
Find file Copy path
Fetching contributors…
Cannot retrieve contributors at this time
205 lines (176 sloc) 6.06 KB
import numpy as np
class Newick_Tree:
def __init__(self,s):
(self.tree,self.labels,self.weights,self.taxa)=self.build_tree2(s)
self.parents=np.unique(self.tree.values())
self.string=s
def tokens(self,s):
cursor=0
n=''
while cursor<len(s):
if s[cursor]=='(':
cursor+=1
yield 'OPEN_PAREN','('
continue
if s[cursor]==')':
cursor+=1
yield 'CLOSE_PAREN',')'
continue
if s[cursor]==':':
cursor+=1
yield 'COLON',':'
if s[cursor]==';':
cursor=len(s)
break
if s[cursor]==',':
cursor+=1
yield 'COMMA',','
continue
n=''
while cursor<len(s) and s[cursor] not in [',',';',')','(',':']:
n+=s[cursor]
cursor+=1
if len(n.strip())>0:
yield 'LABEL',n.strip()
def build_tree2(self,s):
tree,labels,weights={},{},{}
taxa=[]
parent,curnode,savenode=-1,0,0
new_node=1
tree[0]=-1 #introduce an artificial root note
prior_kind=None
for (kind,value) in self.tokens(s):
if kind=='OPEN_PAREN':
parent=curnode
curnode=new_node
tree[curnode]=parent
new_node=new_node+1
prior_kind=kind
continue
if kind=='LABEL':
if prior_kind=='CLOSE_PAREN':
labels[value]=savenode
elif prior_kind=='COLON':
weights[savenode]=int(value)
else:
labels[value]=new_node
tree[new_node]=curnode
savenode=new_node # needed if this label is followed by colon
taxa.append(value)
new_node+=1
prior_kind=kind
continue
if kind=='COLON':
if prior_kind=='COMMA' or prior_kind=='OPEN_PAREN':
tree[new_node]=curnode
#savenode=new_node
new_node+=1
prior_kind=kind
continue
if kind=='CLOSE_PAREN':
if prior_kind=='COMMA':
tree[new_node]=curnode
new_node+=1
savenode=curnode
curnode=parent
parent=tree[curnode]
prior_kind=kind
continue
if kind=='COMMA':
if prior_kind=='COMMA':
tree[new_node]=curnode
new_node+=1
prior_kind=kind
continue
if parent!=-1:
raise Exception('ParseError')
else:
return tree,labels,weights,taxa
def path_to_root(self,x):
tree=self.tree
start_node=labels[x]
path=[start_node]
parent=tree[start_node]
while parent!=-1:
path.append(parent)
parent=tree[parent]
return path
def common_ancestor(self,x,y):
tree,labels,weights=self.tree,self.labels,self.weights
xnode=labels[x]
ynode=labels[y]
d=0
while True:
if xnode<ynode:
d+=weights.get(ynode,1)
ynode=tree[ynode]
elif ynode<xnode:
d+=weights.get(xnode,1)
xnode=tree[xnode]
elif ynode==xnode:
return xnode
def distance(self,x,y):
tree,labels,weights=self.tree,self.labels,self.weights
xnode=labels[x]
ynode=labels[y]
d=0
while True:
if xnode<ynode:
d+=weights.get(ynode,1)
ynode=tree[ynode]
elif ynode<xnode:
d+=weights.get(xnode,1)
xnode=tree[xnode]
elif ynode==xnode:
return d
def path(self,x,y):
tree,labels,weights=self.tree,self.labels,self.weights
xnode=labels[x]
ynode=labels[y]
front=[]
back=[]
while True:
if xnode<ynode:
back.insert(0,ynode)
ynode=tree[ynode]
elif xnode>ynode:
front.append(xnode)
xnode=tree[xnode]
elif ynode==xnode:
back.insert(0,ynode)
front.extend(back)
return front
def separates(self,a,x,y):
tree,labels,weights=self.tree,self.labels,self.weights
b=tree[a]
xnode=labels[x]
ynode=labels[y]
front=[]
back=[]
has_a,has_b=False,False
while True:
if xnode==a or ynode==a:
has_a=True
if ynode==b or xnode==b:
has_b=True
if xnode<ynode:
back.insert(0,ynode)
ynode=tree[ynode]
elif xnode>ynode:
front.append(xnode)
xnode=tree[xnode]
elif ynode==xnode:
back.insert(0,ynode)
front.extend(back)
return (has_a and has_b)
def char_table(self):
base=self.taxa[0]
answer=[]
for x in self.parents:
if x<1:
continue
row=''
for t in sorted(self.taxa):
row+=str(int(self.separates(x,base,t)))
answer.append(row)
return(answer[1:])
You can’t perform that action at this time.