sankoff.py

from nwck_class_3 import Newick_Tree
from collections import defaultdict
sequences=dict(zip(['a','b','c','d'],['CAA','TAG','ATG','GTC']))
N=min([len(sequences[x]) for x in sequences.keys()])

x=Newick_Tree('(a,b,(c,d))')

seq=dict()
costV=dict()
cost=dict()
for taxon in x.labels.keys():
    seq[x.labels[taxon]]=sequences[taxon]
    for i in range(N):
        costV[(x.labels[taxon],i,sequences[taxon][i])]=0
        cost[(x.labels[taxon],i)]=0
        
    

alphabet=['C','A','G','T']

def sankoff(base):
    print base
    if len(x.dtree[base])==0:
        return
    for v in x.dtree[base]:
        sankoff(v)
    for i in range(N):   
        for c in alphabet:
            m=0
            for v in x.dtree[base]:                     
                if (v,i,c) in costV:
                    m=m+min(costV[(v,i,c)],cost[(v,i)]+1)
                else:
                    m=m+cost[(v,i)]+1
            costV[(base,i,c)]=m
        cost[(base,i)]=min([costV[(base,i,c)] for c in alphabet])
        
    return
    
sankoff(x.root)
print sum([costV[(x.root,i,sequences[x.taxa[x.node_to_taxon[x.root]]][i])] for i in range(N)])