hmmc_theano.py

# https://udemy.com/unsupervised-machine-learning-hidden-markov-models-in-python
# http://lazyprogrammer.me
# Continuous-observation HMM in Theano using gradient descent.
import wave
import theano
import theano.tensor as T
import numpy as np
import matplotlib.pyplot as plt

# from theano.sandbox import solve # does not have gradient functionality
from generate_c import get_signals, big_init

def random_normalized(d1, d2):
    x = np.random.random((d1, d2))
    return x / x.sum(axis=1, keepdims=True)


class HMM:
    def __init__(self, M, K):
        self.M = M # number of hidden states
        self.K = K # number of Gaussians
    
    def fit(self, X, learning_rate=10e-3, max_iter=10):
        # train the HMM model using the Baum-Welch algorithm
        # a specific instance of the expectation-maximization algorithm

        N = len(X)
        D = X[0].shape[1] # assume each x is organized (T, D)

        pi0 = np.ones(self.M) / self.M # initial state distribution
        A0 = random_normalized(self.M, self.M) # state transition matrix
        R0 = np.ones((self.M, self.K)) / self.K # mixture proportions
        mu0 = np.zeros((self.M, self.K, D))
        for i in xrange(self.M):
            for k in xrange(self.K):
                random_idx = np.random.choice(N)
                x = X[random_idx]
                random_time_idx = np.random.choice(len(x))
                mu0[i,k] = x[random_time_idx]
        sigma0 = np.zeros((self.M, self.K, D, D))
        for j in xrange(self.M):
            for k in xrange(self.K):
                sigma0[j,k] = np.eye(D)

        thx, cost = self.set(pi0, A0, R0, mu0, sigma0)

        pi_update = self.pi - learning_rate*T.grad(cost, self.pi)
        pi_update = pi_update / pi_update.sum()

        A_update = self.A - learning_rate*T.grad(cost, self.A)
        A_update = A_update / A_update.sum(axis=1).dimshuffle(0, 'x')

        R_update = self.R - learning_rate*T.grad(cost, self.R)
        R_update = R_update / R_update.sum(axis=1).dimshuffle(0, 'x')

        updates = [
            (self.pi, pi_update),
            (self.A, A_update),
            (self.R, R_update),
            (self.mu, self.mu - learning_rate*T.grad(cost, self.mu)),
            (self.sigma, self.sigma - learning_rate*T.grad(cost, self.sigma)),
        ]

        train_op = theano.function(
            inputs=[thx],
            updates=updates,
        )

        costs = []
        for it in xrange(max_iter):
            print "it:", it
            
            for n in xrange(N):
                c = self.log_likelihood_multi(X).sum()
                print "c:", c
                costs.append(c)
                train_op(X[n])

        print "A:", self.A.get_value()
        print "mu:", self.mu.get_value()
        print "sigma:", self.sigma.get_value()
        print "R:", self.R.get_value()
        print "pi:", self.pi.get_value()

        plt.plot(costs)
        plt.show()

    def set(self, pi, A, R, mu, sigma):
        self.pi = theano.shared(pi)
        self.A = theano.shared(A)
        self.R = theano.shared(R)
        self.mu = theano.shared(mu)
        self.sigma = theano.shared(sigma)
        M, K = R.shape
        self.M = M
        self.K = K

        D = self.mu.shape[2]
        twopiD = (2*np.pi)**D

        # set up theano variables and functions
        thx = T.matrix('X') # represents a TxD matrix of sequential observations
        def mvn_pdf(x, mu, sigma):
            k = 1 / T.sqrt(twopiD * T.nlinalg.det(sigma))
            e = T.exp(-0.5*(x - mu).T.dot(T.nlinalg.matrix_inverse(sigma).dot(x - mu)))
            return k*e

        def gmm_pdf(x):
            def state_pdfs(xt):
                def component_pdf(j, xt):
                    Bj_t = 0
                    for k in xrange(self.K):
                        Bj_t += self.R[j,k] * mvn_pdf(xt, self.mu[j,k], self.sigma[j,k])
                    return Bj_t

                Bt, _ = theano.scan(
                    fn=component_pdf,
                    sequences=T.arange(self.M),
                    n_steps=self.M,
                    outputs_info=None,
                    non_sequences=[xt],
                )
                return Bt

            B, _ = theano.scan(
                fn=state_pdfs,
                sequences=x,
                n_steps=x.shape[0],
                outputs_info=None,
            )
            return B.T
        
        B = gmm_pdf(thx)
        # scale = T.zeros((thx.shape[0], 1), dtype=theano.config.floatX)
        # scale[0] = (self.pi*B[:,0]).sum()

        def recurrence(t, old_a, B):
            a = old_a.dot(self.A) * B[:, t]
            s = a.sum()
            return (a / s), s

        [alpha, scale], _ = theano.scan(
            fn=recurrence,
            sequences=T.arange(1, thx.shape[0]),
            outputs_info=[self.pi*B[:,0], None],
            n_steps=thx.shape[0]-1,
            non_sequences=[B],
        )

        cost = -T.log(scale).sum()
        self.cost_op = theano.function(
            inputs=[thx],
            outputs=cost,
        )
        return thx, cost

    def log_likelihood_multi(self, X):
        return np.array([self.cost_op(x) for x in X])


def real_signal():
    spf = wave.open('helloworld.wav', 'r')

    #Extract Raw Audio from Wav File
    # If you right-click on the file and go to "Get Info", you can see:
    # sampling rate = 16000 Hz
    # bits per sample = 16
    # The first is quantization in time
    # The second is quantization in amplitude
    # We also do this for images!
    # 2^16 = 65536 is how many different sound levels we have
    signal = spf.readframes(-1)
    signal = np.fromstring(signal, 'Int16')
    T = len(signal)
    signal = (signal - signal.mean()) / signal.std()

    hmm = HMM(5, 3)
    # signal needs to be of shape N x T(n) x D
    hmm.fit(signal.reshape(1, T, 1), learning_rate=10e-6, max_iter=20)


def fake_signal():
    signals = get_signals()
    hmm = HMM(5, 3)
    hmm.fit(signals)
    L = hmm.log_likelihood_multi(signals).sum()
    print "LL for fitted params:", L

    # test in actual params
    _, _, _, pi, A, R, mu, sigma = big_init()
    hmm.set(pi, A, R, mu, sigma)
    L = hmm.log_likelihood_multi(signals).sum()
    print "LL for actual params:", L

if __name__ == '__main__':
    # real_signal()
    fake_signal()