Hidden Markov Model implemented in edward

ed_hmm.py

#!/usr/bin/env python
"""Hidden Markov Models

Abstract base class for HMMs and an implementation of an HMM
with discrete states and gaussian emissions.
"""

import tensorflow as tf
import edward as ed
from edward.models import Bernoulli, Categorical, Normal
import numpy as np


def categorical(ps):
    return Categorical(logits=ed.logit(ps)).value()

def flip(p):
    return tf.equal(Bernoulli(p=p), tf.constant(1))

def append(lst, x):
    return tf.concat(0, [lst, [x]])


class HMM(object):
    """A Hidden Markov Model."""

    def step(self, state):
        """Returns a new state following `state`."""
        raise NotImplementedError()

    def emit(self, state):
        """Returns an observable emission from `state`."""
        raise NotImplementedError()

    def init(self):
        """Returns an intial state."""
        raise NotImplementedError()

    def final(self, state):
        """Returns true if the model should stop in `state`."""
        raise NotImplementedError()
    
    def sample(self):
        def cond(states, emissions):
            s0 = states[-1]
            return self.final(s0) # TODO not final
        
        def body(states, emissions):
            s0 = states[-1]
            s1 = self.step(s0)
            e1 = self.emit(s1)
            return append(states, s1), append(emissions, e1)

        s0 = self.init()
        e0 = self.emit(s0)
        states = tf.convert_to_tensor([s0])
        emissions = tf.convert_to_tensor([e0])
        return tf.while_loop(
            cond, body, 
            loop_vars=[states, emissions],
            shape_invariants=[tf.TensorShape(None), tf.TensorShape(None)]
        )


class DiscreteGaussianHMM(HMM):
    """HMM with discrete transitions and gaussian emissions."""
    def __init__(self, P, mu, sigma, p_init, p_final):
        super().__init__()
        self.P = P
        self.mu = mu
        self.sigma = sigma
        self.p_init = p_init
        self.p_final = p_final

    def step(self, state):
        return categorical(self.P[state])

    def emit(self, state):
        return Normal(mu=self.mu[state], sigma=self.sigma[state]).value()

    def init(self):
        return categorical(self.p_init)

    def final(self, state):
        return flip(1 - tf.gather(self.p_final, state))
      

def demo():
    import matplotlib.pyplot as plt

    P = tf.constant(np.array([
        [.6, .4],
        [.2, .8],
    ], dtype='float32'))
    mu = tf.constant([5., -5.])
    sigma = tf.constant([1., 1.])
    model = DiscreteGaussianHMM(P, mu, sigma, [0.5, 0.5], [0.02, 0.02])

    sess = ed.get_session()
    emissions = [sess.run(model.sample()[1]) for _ in range(3)]
    for e in emissions:
        plt.plot(e)
    plt.show()

if __name__ == '__main__':
    demo()