nwck_class.py

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)
                
        

               
f=open('rosalind_ctbl.txt','rU')
s=f.readline()
f.close()
T=Newick_Tree(s)
        
base=T.taxa[0]
f=open('chartable.txt','w')
row=''
for x in T.parents:
    if x<1:
        continue
    for t in sorted(T.taxa):
       row+=str(int(T.separates(x,base,t)))
    row+='\n'

f.write(row)
f.close()