from __future__ import division
import numpy as np
np.seterr(divide='ignore') # these warnings are usually harmless for this code
from matplotlib import pyplot as plt
import matplotlib
matplotlib.rcParams['font.size'] = 8

import pyhsmm
from pyhsmm.util.text import progprint_xrange

print \
'''
This demo shows how HDP-HMMs can fail when the underlying data has state
persistence without some kind of temporal regularization (in the form of a
sticky bias or duration modeling): without setting the number of states to be
the correct number a priori, lots of extra states can be intsantiated.

BUT the effect is much more relevant on real data (when the data doesn't exactly
fit the model). Maybe this demo should use multinomial emissions...
'''

#####################
#  data generation  #
#####################

# Set parameters
N = 4
T = 1000
obs_dim = 2

obs_hypparams = {'mu_0':np.zeros(obs_dim),
                'sigma_0':np.eye(obs_dim),
                'kappa_0':0.25,
                'nu_0':obs_dim+5}

dur_hypparams = {'alpha_0':2*30,
                 'beta_0':2}

true_obs_distns = [pyhsmm.distributions.Gaussian(**obs_hypparams) for state in xrange(N)]
true_dur_distns = [pyhsmm.distributions.PoissonDuration(**dur_hypparams) for state in range(N)]

true_trans_matrix = np.diag(np.repeat(0.5,N-1),-1) + np.diag(np.repeat(0.5,N-1),1)
true_trans_matrix /= true_trans_matrix.sum(1)[:,None]

truemodel = pyhsmm.models.HSMM(alpha=6.,init_state_concentration=1.,
                              obs_distns=true_obs_distns,
                              dur_distns=true_dur_distns,
                              trans_matrix=true_trans_matrix)

data, labels = truemodel.generate(T)

plt.figure()
truemodel.plot()
plt.gcf().suptitle('True model')


#########################
#  posterior inference  #
#########################

# Set the weak limit truncation level
Nmax = 25

### HDP-HMM without the sticky bias

obs_distns = [pyhsmm.distributions.Gaussian(**obs_hypparams) for state in xrange(Nmax)]
posteriormodel = pyhsmm.models.WeakLimitHDPHMM(alpha=6.,gamma=6.,init_state_concentration=1.,
                                   obs_distns=obs_distns)
posteriormodel.add_data(data)

for idx in progprint_xrange(100):
    posteriormodel.resample_model()

plt.figure()
posteriormodel.plot()
plt.gcf().suptitle('HDP-HMM sampled model after 100 iterations')

### Sticky-HDP-HMM

obs_distns = [pyhsmm.distributions.Gaussian(**obs_hypparams) for state in xrange(Nmax)]
posteriormodel = pyhsmm.models.WeakLimitStickyHDPHMM(
        kappa=50.,alpha=6.,gamma=6.,init_state_concentration=1.,
        obs_distns=obs_distns)
posteriormodel.add_data(data)

for idx in progprint_xrange(100):
    posteriormodel.resample_model()

plt.figure()
posteriormodel.plot()
plt.gcf().suptitle('Sticky HDP-HMM sampled model after 100 iterations')

plt.show()


@@ -178,16 +178,26 @@
 plotlib%0A
+import os%0A
 matplotl
@@ -770,1123 +770,451 @@
 ####
-######%0A#  data generation  #%0A#####################%0A%0A# Set parameters%0AN = 4%0AT = 1000%0Aobs_dim = 2%0A%0Aobs_hypparams = %7B'mu_0':np.zeros(obs_dim),%0A                'sigma_0':np.eye(obs_dim),%0A                'kappa_0':0.25,%0A                'nu_0':obs_dim+5%7D%0A%0Adur_hypparams = %7B'alpha_0':2*30,%0A                 'beta_0':2%7D%0A%0Atrue_obs_distns = %5Bpyhsmm.distributions.Gaussian(**obs_hypparams) for state in xrange(N)%5D%0Atrue_dur_distns = %5Bpyhsmm.distributions.PoissonDuration(**dur_hypparams) for state in range(N)%5D%0A%0Atrue_trans_matrix = np.diag(np.repeat(0.5,N-1),-1) + np.diag(np.repeat(0.5,N-1),1)%0Atrue_trans_matrix /= true_trans_matrix.sum(1)%5B:,None%5D%0A%0Atruemodel = pyhsmm.models.HSMM(alpha=6.,init_state_concentration=1.,%0A                              obs_distns=true_obs_distns,%0A                              dur_distns=true_dur_distns,%0A                              trans_matrix=true_trans_matrix)%0A%0Adata, labels = truemodel.generate(T)%0A%0Aplt.figure()%0Atruemodel.plot()%0Aplt.gcf().suptitle('True model')%0A%0A%0A#########################%0A#  posterior inference  #%0A#########################%0A%0A# Set the weak limit truncation level%0ANmax = 25
+%0A#  load data  #%0A###############%0A%0Adata = np.loadtxt(os.path.join(os.path.dirname(__file__),'example-data.txt'))%0A%0A#########################%0A#  posterior inference  #%0A#########################%0A%0A# Set the weak limit truncation level%0ANmax = 25%0A%0A# and some hyperparameters%0Aobs_dim = data.shape%5B1%5D%0Aobs_hypparams = %7B'mu_0':np.zeros(obs_dim),%0A                'sigma_0':np.eye(obs_dim),%0A                'kappa_0':0.25,%0A                'nu_0':obs_dim+2%7D
 %0A%0A##

make hmm.py load example data