"""

Hidden Markov Model Class

Parameters:

-----------

- S: Number of states.

- T: Transition matrix of size S by S

stores probability from state i to state j.

- E: Emission matrix of size S by N (number of observations)

stores the probability of observing O_j from state S_i.

- T0: Initial state probabilities of size S.

"""

def __init__(self, T, E, T0, epsilon = 0.001, maxStep = 10):

with tf.name_scope('Inital_Parameters'):

with tf.name_scope('Scalar_constants'):

# Max number of iteration

self.maxStep = maxStep

# convergence criteria

self.epsilon = epsilon

# Number of possible states

self.S = T.shape[0]

# Number of possible observations

self.O = E.shape[0]

self.prob_state_1 = []

with tf.name_scope('Model_Parameters'):

# Emission probability

self.E = tf.Variable(E, dtype=tf.float64, name='emission_matrix')

# Transition matrix

self.T = tf.Variable(T, dtype=tf.float64, name='transition_matrix')

# Initial state vector

self.T0 = tf.Variable(tf.constant(T0, dtype=tf.float64, name='inital_state_vector'))

def initialize_viterbi_variables(self, shape):

pathStates = tf.Variable(tf.zeros(shape, dtype=tf.int64), name='States_matrix')

pathScores = tf.Variable(tf.zeros(shape, dtype=tf.float64), name='Score_matrix')

states_seq = tf.Variable(tf.zeros([shape[0]], dtype=tf.int64), name='States_sequence')

return pathStates, pathScores, states_seq

def belief_propagation(self, scores):

scores_reshape = tf.reshape(scores, (-1,1))

return tf.add(scores_reshape, tf.log(self.T))

def viterbi_inference(self, obs_seq):

# length of observed sequence

self.N = len(obs_seq)

# shape path Variables

shape = [self.N, self.S]

# observed sequence

x = tf.constant(obs_seq, dtype=tf.int32, name='observation_sequence')

with tf.name_scope('Init_viterbi_variables'):

# Initialize variables

pathStates, pathScores, states_seq = self.initialize_viterbi_variables(shape)

with tf.name_scope('Emission_seq_'):

# log probability of emission sequence

obs_prob_seq = tf.log(tf.gather(self.E, x))

obs_prob_list = tf.split(0, self.N, obs_prob_seq)

with tf.name_scope('Starting_log-priors'):

# initialize with state starting log-priors

pathScores = tf.scatter_update(pathScores, 0, tf.log(self.T0) + tf.squeeze(obs_prob_list[0]))

with tf.name_scope('Belief_Propagation'):

for step, obs_prob in enumerate(obs_prob_list[1:]):

with tf.name_scope('Belief_Propagation_step_%s' %step):

# propagate state belief

belief = self.belief_propagation(pathScores[step, :])

# the inferred state by maximizing global function

# and update state and score matrices

pathStates = tf.scatter_update(pathStates, step + 1, tf.argmax(belief, 0))

pathScores = tf.scatter_update(pathScores, step + 1, tf.reduce_max(belief, 0) + tf.squeeze(obs_prob))

with tf.name_scope('Max_Likelyhood_update'):

# infer most likely last state

states_seq = tf.scatter_update(states_seq, self.N-1, tf.argmax(pathScores[self.N-1, :], 0))

with tf.name_scope('Backtrack'):

for step in range(self.N - 1, 0, -1):

with tf.name_scope('Back_track_step_%s' %step):

# for every timestep retrieve inferred state

state = states_seq[step]

idx = tf.reshape(tf.pack([step, state]), [1, -1])

state_prob = tf.gather_nd(pathStates, idx)

states_seq = tf.scatter_update(states_seq, step - 1, state_prob[0])

return states_seq, tf.exp(pathScores) # turn scores back to probabilities

def run_viterbi(self, obs_seq, summary=False):

state_graph, state_prob_graph = self.viterbi_inference(obs_seq)

with tf.Session() as sess:

sess.run(tf.initialize_all_variables())

states_seq, state_prob = sess.run([state_graph, state_prob_graph])

if summary:

# Instantiate a SummaryWriter to output summaries and the Graph.

summary_writer = tf.train.SummaryWriter('logs/', graph=sess.graph)

return states_seq, state_prob

def initialize_forw_back_variables(self, shape):

self.forward = tf.Variable(tf.zeros(shape, dtype=tf.float64), name='forward')

self.backward = tf.Variable(tf.zeros(shape, dtype=tf.float64), name='backward')

self.posterior = tf.Variable(tf.zeros(shape, dtype=tf.float64), name='posteriror')

def _forward(self, obs_prob_list):

with tf.name_scope('init_scaling_factor'):

self.scale = tf.Variable(tf.zeros([self.N], tf.float64)) #scale factors

with tf.name_scope('forward_first_step'):

# initialize with state starting priors

init_prob = tf.mul(self.T0, tf.squeeze(obs_prob_list[0]))

# scaling factor at t=0

self.scale = tf.scatter_update(self.scale, 0, 1.0 / tf.reduce_sum(init_prob))

# scaled belief at t=0

self.forward = tf.scatter_update(self.forward, 0, self.scale[0] * init_prob)

# propagate belief

for step, obs_prob in enumerate(obs_prob_list[1:]):

with tf.name_scope('time_step-%s' %step):

# previous state probability

prev_prob = tf.expand_dims(self.forward[step, :], 0)

# transition prior

prior_prob = tf.matmul(prev_prob, self.T)

# forward belief propagation

forward_score = tf.mul(prior_prob, tf.squeeze(obs_prob))

forward_prob = tf.squeeze(forward_score)

# scaling factor

self.scale = tf.scatter_update(self.scale, step+1, 1.0 / tf.reduce_sum(forward_prob))

# Update forward matrix

self.forward = tf.scatter_update(self.forward, step+1, self.scale[step+1] * forward_prob)

def _backward(self, obs_prob_list):

with tf.name_scope('backward_last_step'):

# initialize with state ending priors

self.backward = tf.scatter_update(self.backward, 0, self.scale[self.N-1] * tf.ones([self.S], dtype=tf.float64))

# propagate belief

for step, obs_prob in enumerate(obs_prob_list[:-1]):

with tf.name_scope('time_step-%s' %step):

# next state probability

next_prob = tf.expand_dims(self.backward[step, :], 1)

# observation emission probabilities

obs_prob_d = tf.diag(tf.squeeze(obs_prob))

# transition prior

prior_prob = tf.matmul(self.T, obs_prob_d)

# backward belief propagation

backward_score = tf.matmul(prior_prob, next_prob)

backward_prob = tf.squeeze(backward_score)

# Update backward matrix

self.backward = tf.scatter_update(self.backward, step+1, self.scale[self.N-2-step] * backward_prob)

self.backward = tf.assign(self.backward, tf.reverse(self.backward, [True, False]))

def _posterior(self):

# posterior score

self.posterior = tf.mul(self.forward, self.backward)

marginal = tf.reduce_sum(self.posterior, 1)

self.posterior = self.posterior / tf.expand_dims(marginal, 1)

def re_estimate_emission(self, x):

states_marginal = tf.reduce_sum(self.gamma, 0)

seq_one_hot = tf.one_hot(tf.cast(x, tf.int64), self.O, 1, 0)

emission_score = tf.matmul(tf.cast(seq_one_hot, tf.float64), self.gamma, transpose_a=True)

return emission_score / states_marginal

def re_estimate_transition(self, x):

with tf.name_scope('Init_3D_tensor'):

self.M = tf.Variable(tf.zeros((self.N-1, self.S, self.S), tf.float64))

with tf.name_scope('3D_tensor_transition'):

for t in range(self.N - 1):

with tf.name_scope('time_step-%s' %t):

tmp_0 = tf.matmul(tf.expand_dims(self.forward[t, :], 0), self.T)

tmp_1 = tf.mul(tmp_0, tf.expand_dims(tf.gather(self.E, x[t+1]), 0))

denom = tf.squeeze(tf.matmul(tmp_1, tf.expand_dims(self.backward[t+1, :], 1)))

with tf.name_scope('Init_new_transition'):

trans_re_estimate = tf.Variable(tf.zeros((self.S, self.S), tf.float64))

for i in range(self.S):

with tf.name_scope('State-%s' %i):

numer = self.forward[t, i] * self.T[i, :] * tf.gather(self.E, x[t+1]) * self.backward[t+1, :]

trans_re_estimate = tf.scatter_update(trans_re_estimate, i, numer / denom)

self.M = tf.scatter_update(self.M, t, trans_re_estimate)

with tf.name_scope('Smooth_gamma'):

self.gamma = tf.squeeze(tf.reduce_sum(self.M, 2))

T_new = tf.reduce_sum(self.M, 0) / tf.expand_dims(tf.reduce_sum(self.gamma, 0), 1)

with tf.name_scope('New_init_states_prob'):

T0_new = self.gamma[0,:]

with tf.name_scope('Append_gamma_final_time_step'):

prod = tf.expand_dims(tf.mul(self.forward[self.N-1, :], self.backward[self.N-1, :]), 0)

s= prod/ tf.reduce_sum(prod)

self.gamma = tf.concat(0, [self.gamma, s])

self.prob_state_1.append(self.gamma[:, 0])

return T0_new, T_new

def check_convergence(self, new_T0, new_transition, new_emission):

delta_T0 = tf.reduce_max(tf.abs(self.T0 - new_T0)) < self.epsilon

delta_T = tf.reduce_max(tf.abs(self.T - new_transition)) < self.epsilon

delta_E = tf.reduce_max(tf.abs(self.E - new_emission)) < self.epsilon

return tf.logical_and(tf.logical_and(delta_T0, delta_T), delta_E)

def forward_backward(self, obs_prob_seq):

"""

runs forward backward algorithm on observation sequence

Arguments

---------

- obs_seq : matrix of size N by S, where N is number of timesteps and

S is the number of states

Returns

-------

- forward : matrix of size N by S representing

the forward probability of each state at each time step

- backward : matrix of size N by S representing

the backward probability of each state at each time step

- posterior : matrix of size N by S representing

the posterior probability of each state at each time step

"""

obs_prob_list_for = tf.split(0, self.N, obs_prob_seq)

with tf.name_scope('forward_belief_propagation'):

# forward belief propagation

self._forward(obs_prob_list_for)

obs_prob_seq_rev = tf.reverse(obs_prob_seq, [True, False])

obs_prob_list_back = tf.split(0, self.N, obs_prob_seq_rev)

with tf.name_scope('backward_belief_propagation'):

# backward belief propagation

self._backward(obs_prob_list_back)

def expectation_maximization_step(self, x):

# probability of emission sequence

obs_prob_seq = tf.gather(self.E, x)

with tf.name_scope('Forward_Backward'):

self.forward_backward(obs_prob_seq)

with tf.name_scope('Re_estimate_transition'):

new_T0, new_transition = self.re_estimate_transition(x)

with tf.name_scope('Re_estimate_emission'):

new_emission = self.re_estimate_emission(x)

with tf.name_scope('Check_Convergence'):

converged = self.check_convergence(new_T0, new_transition, new_emission)

with tf.name_scope('Update_parameters'):

self.T0 = tf.assign(self.T0, new_T0)

self.E = tf.assign(self.E, new_emission)

self.T = tf.assign(self.T, new_transition)

#self.count = tf.assign_add(self.count, 1)

with tf.name_scope('histogram_summary'):

_ = tf.histogram_summary(self.T0.name, self.T0)

_ = tf.histogram_summary(self.T.name, self.T)

_ = tf.histogram_summary(self.E.name, self.E)

return converged

def Baum_Welch_EM(self, obs_seq):

with tf.name_scope('Input_Observed_Sequence'):

# length of observed sequence

self.N = len(obs_seq)

# shape of Variables

shape = [self.N, self.S]

# observed sequence

x = tf.constant(obs_seq, dtype=tf.int32, name='observation_sequence')

with tf.name_scope('Initialize_variables'):

# initialize variables

self.initialize_forw_back_variables(shape)

converged = tf.cast(False, tf.bool)

#self.count = tf.Variable(tf.constant(0))

with tf.name_scope('Train_Baum_Welch'):

for i in range(self.maxStep):

with tf.name_scope('EM_step-%s' %i):

converged = self.expectation_maximization_step(x)

# TF while_loop op is buggy, should be fixed in future release

# def loop_conditions(converged, obs_seq):

# cond_1 = tf.logical_not(converged)

# cond_2 = tf.less(self.count, self.maxStep)

# return tf.logical_or(cond_1, cond_2)

# def body(converged, obs_seq):

# return self.expectation_maximization_step(obs_seq)

# while_params = [converged, obs_seq]

# c = tf.while_loop(loop_conditions, body, while_params)

return converged

def run_Baum_Welch_EM(self, obs_seq, summary=False, monitor_state_1=False):

converged = self.Baum_Welch_EM(obs_seq)

# Build the summary operation based on the TF collection of Summaries.

summary_op = tf.merge_all_summaries()

with tf.Session() as sess:

sess.run(tf.initialize_all_variables())

trans0, transition, emission, c = sess.run([self.T0, self.T, self.E, converged])

if monitor_state_1:

self.state_summary = np.array([sess.run(g) for g in self.prob_state_1])

if summary:

# Instantiate a SummaryWriter to output summaries and the Graph.

summary_writer = tf.train.SummaryWriter('logs/', graph=sess.graph)

summary_str = sess.run(summary_op)

summary_writer.add_summary(summary_str)

return trans0, transition, emission, c