import math
from collections import defaultdict
class MarkovChainOrder(object):
''' Simple higher-order Markov chain, specifically for DNA
sequences. User provides training data (DNA strings). '''
def __init__(self, examples, order=1):
''' Initialize the model given collection of example DNA
strings. '''
self.order = order
self.mers = defaultdict(int)
self.longestMer = longestMer = order + 1
for ex in examples:
# count up all k-mers of length 'longestMer' or shorter
for i in range(len(ex) - longestMer + 1):
for j in range(longestMer, -1, -1):
self.mers[ex[i:i+j]] += 1
def condProb(self, obs, given):
''' Return conditional probability of seeing nucleotide "obs"
given we just saw nucleotide string "given". Length of
"given" can't exceed model order. Return None if "given"
was never observed. '''
assert len(given) <= self.order
ngiven = self.mers[given]
if ngiven == 0: return None
return float(self.mers[given + obs]) / self.mers[given]
def jointProb(self, s):
''' Return joint probability of observing string s '''
cum = 1.0
for i in range(self.longestMer-1, len(s)):
obs, given = s[i], s[i-self.longestMer+1:i]
p = self.condProb(obs, given)
if p is not None:
cum *= p
# include the smaller k-mers at the very beginning
for j in range(self.longestMer-2, -1, -1):
obs, given = s[j], s[:j]
p = self.condProb(obs, given)
if p is not None:
cum *= p
return cum
def jointProbL(self, s):
''' Return log2 of joint probability of observing string s '''
cum = 0.0
for i in range(self.longestMer-1, len(s)):
obs, given = s[i], s[i-self.longestMer+1:i]
p = self.condProb(obs, given)
if p is not None:
cum += math.log(p, 2)
# include the smaller k-mers at the very beginning
for j in range(self.longestMer-2, -1, -1):
obs, given = s[j], s[:j]
p = self.condProb(obs, given)
if p is not None:
cum += math.log(p, 2)
return cum