def test_viterbi(T=1000, K=20, D=2):
# Create a true HMM
A = npr.rand(K, K)
A /= A.sum(axis=1, keepdims=True)
A = 0.75 * np.eye(K) + 0.25 * A
C = npr.randn(K, D)
sigma = 0.01
# Sample from the true HMM
z = np.zeros(T, dtype=int)
y = np.zeros((T, D))
for t in range(T):
if t > 0:
z[t] = np.random.choice(K, p=A[z[t - 1]])
y[t] = C[z[t]] + np.sqrt(sigma) * npr.randn(D)
# Compare to pyhsmm answer
from pyhsmm.models import HMM as OldHMM
from pyhsmm.basic.distributions import Gaussian
oldhmm = OldHMM(
[Gaussian(mu=C[k], sigma=sigma * np.eye(D)) for k in range(K)],
trans_matrix=A,
init_state_distn="uniform")
oldhmm.add_data(y)
states = oldhmm.states_list.pop()
states.Viterbi()
z_star = states.stateseq
# Make an HMM with these parameters
hmm = HMM(K, D, observations="diagonal_gaussian")
hmm.transitions.log_Ps = np.log(A)
hmm.observations.mus = C
hmm.observations.sigmasq = sigma * np.ones((K, D))
z_star2 = hmm.most_likely_states(y)
assert np.allclose(z_star, z_star2)
]
# Fit with both SGD and EM
methods = ["sgd", "em"]
results = {}
for obs in observations:
for method in methods:
print("Fitting {} HMM with {}".format(obs, method))
model = HMM(K, D, observations=obs)
train_lls = model.fit(y, method=method)
test_ll = model.log_likelihood(y_test)
smoothed_y = model.smooth(y)
# Permute to match the true states
model.permute(find_permutation(z, model.most_likely_states(y)))
smoothed_z = model.most_likely_states(y)
results[(obs, method)] = (model, train_lls, test_ll, smoothed_z,
smoothed_y)
# Plot the inferred states
fig, axs = plt.subplots(len(observations) + 1, 1, figsize=(12, 8))
# Plot the true states
plt.sca(axs[0])
plt.imshow(z[None, :], aspect="auto", cmap="jet")
plt.title("true")
plt.xticks()
# Plot the inferred states
for i, obs in enumerate(observations):
# In[7]:
N_iters = 50
hmm = HMM(K, D, observations="gaussian")
hmm_lls = hmm.fit(y, method="em", num_em_iters=N_iters, verbose=True)
plt.plot(hmm_lls, label="EM")
plt.plot([0, N_iters], true_ll * np.ones(2), ':k', label="True")
plt.xlabel("EM Iteration")
plt.ylabel("Log Probability")
plt.legend(loc="lower right")
# In[8]:
# Find a permutation of the states that best matches the true and inferred states
hmm.permute(find_permutation(z, hmm.most_likely_states(y)))
# In[11]:
# Plot the true and inferred discrete states
hmm_z = hmm.most_likely_states(y)
plt.figure(figsize=(8, 4))
plt.subplot(211)
plt.imshow(z[None, :], aspect="auto", cmap=cmap, vmin=0, vmax=len(colors) - 1)
plt.xlim(0, T)
plt.ylabel("$z_{\\mathrm{true}}$")
plt.yticks([])
plt.subplot(212)
plt.imshow(hmm_z[None, :],
color=colors[1],
label="Struct" if n == 0 else None)
plt.legend()
plt.xlabel("time")
# # Fit an HMM to the LDS states
# In[13]:
from ssm.models import HMM
N_iters = 50
K = 15
hmm = HMM(K, D, observations="gaussian")
hmm_lls = hmm.fit(x, method="em", num_em_iters=N_iters, verbose=True)
z = hmm.most_likely_states(x)
# In[14]:
plt.plot(hmm_lls, label="EM")
plt.xlabel("EM Iteration")
plt.ylabel("Log Probability")
plt.legend(loc="lower right")
# In[15]:
# Plot the observation distributions
from hips.plotting.colormaps import white_to_color_cmap
xmins = x.min(axis=0)
xmaxs = x.max(axis=0)
npts = 100
if regularize:
arhmm_ses1 = HMM(K,N, observations="ar",
transitions="sticky",
transition_kwargs=dict(kappa=10),
observation_kwargs=dict(regularization_params=dict(type='l2',lambda_A=lambda_reg)))
regularizestring='_regl2'
else:
arhmm_ses1 = HMM(K,N, observations="ar",
transitions="sticky",
transition_kwargs=dict(kappa=10))
regularizestring=''
arhmm_em_lls_ses1 = arhmm_ses1.fit(trialdata_ses1, method="em", num_em_iters=numiters)
# Get the inferred states for session 1
trialdata_ses1_z=[arhmm_ses1.most_likely_states(trial) for trial in trialdata_ses1]
trialdata_ses1_z=np.asarray(trialdata_ses1_z)
As_ses1=[None]*K;
for k in np.arange(K):
As_ses1[k]=arhmm_ses1.params[2][0][k,:,:];
fig = plt.figure(figsize=(8, 10))
plt.imshow(trialdata_ses1_z, aspect="auto", vmin=0, vmax=K)
plt.plot([54,54],[0,trialdata_ses1_z.shape[0]],'-r')
plt.ylim([0,trialdata_ses1_z.shape[0]])
fig.savefig(os.path.join(resultsfolder,'ses1'+cca_str+constring+'K'+str(K)+whitenstring+regularizestring+'.png'))
fig = plt.figure(figsize=(8*K,8))
for k in np.arange(K):
ax = plt.subplot(1,K,k+1)
plt.imshow(As_ses1[k],cmap='jet');
sequence=np.array(sequence)
# Divide into training and testing (I didn't really end up using the testing - but can check log likelihoods to decide K)
traintrials=[trialdata[j] for j in sequence[:int(np.ceil(0.8*len(trialdata)))]]
testtrials=[trialdata[j] for j in sequence[int(np.ceil(0.8*len(trialdata))):]]
print(len(traintrials)); print(len(testtrials))
# Run the ARHMM
arhmm = HMM(K,N, observations="ar",
transitions="sticky",
transition_kwargs=dict(kappa=kappa))
arhmm_em_lls = arhmm.fit(traintrials, method="em", num_em_iters=numiters)
# Get the inferred states for train and test trials
traintrials_z=[arhmm.most_likely_states(traintrial) for traintrial in traintrials]
traintrials_z=np.asarray(traintrials_z)
testtrials_z=[arhmm.most_likely_states(testtrial) for testtrial in testtrials]
testtrials_z=np.asarray(testtrials_z)
As=[None]*K; maxvals=[None]*K
for k in np.arange(K):
As[k]=arhmm.params[2][0][k,:,:];
maxvals[k]=np.var(As[k]) # Tried to permute the states so that it would be 'no movement' --> 'movement', based on the variance of the values in the A matrix (didn't really work)
# permute the states
sortorder=np.argsort(maxvals)
sortedmaxvals=np.sort(maxvals)
print(sortorder); print(sortedmaxvals)
As=[As[i] for i in sortorder]
# Now create a new HMM and fit it to the data with EM
N_iters = 50
hmm = HMM(K, D, M,
observations="categorical", observation_kwargs=dict(C=C),
transitions="inputdriven")
# Fit
hmm_lps = hmm.fit(y, inputs=inpt, method="em", num_em_iters=N_iters)
# In[5]:
# Find a permutation of the states that best matches the true and inferred states
hmm.permute(find_permutation(z, hmm.most_likely_states(y, input=inpt)))
z_inf = hmm.most_likely_states(y, input=inpt)
# In[6]:
# Plot the log probabilities of the true and fit models
plt.plot(hmm_lps, label="EM")
plt.plot([0, N_iters], true_lp * np.ones(2), ':k', label="True")
plt.legend(loc="lower right")
plt.xlabel("EM Iteration")
plt.xlim(0, N_iters)
plt.ylabel("Log Probability")
