<jupyter_start><jupyter_text>HMM Example Notebook[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/lindermanlab/ssm-jax-refactor/blob/main/notebooks/gaussian-hmm-example.ipynb)<jupyter_code>try:
    import ssm
except:
    !pip install git+https://github.com/lindermanlab/ssm-jax-refactor.git -qqq
    import ssm<jupyter_output><empty_output><jupyter_text>Imports and Plotting Functions<jupyter_code>import jax.random as jr
import jax.numpy as np
import matplotlib.pyplot as plt

from ssm.hmm import GaussianHMM
from ssm.plots import gradient_cmap
from ssm.utils import find_permutation

import seaborn as sns
sns.set_style("white")
sns.set_context("talk")

color_names = [
    "windows blue",
    "red",
    "amber",
    "faded green",
    "dusty purple",
    "orange"
    ]

colors = sns.xkcd_palette(color_names)
cmap = gradient_cmap(colors)

def plot_transition_matrix(transition_matrix):
    plt.imshow(transition_matrix, vmin=0, vmax=1, cmap="Greys")
    plt.xlabel("next state")
    plt.ylabel("current state")
    plt.colorbar()
    plt.show()
    
def compare_transition_matrix(true_matrix, test_matrix):
    fig, axs = plt.subplots(1, 2)
    out = axs[0].imshow(true_matrix, vmin=0, vmax=1, cmap="Greys")
    axs[1].imshow(test_matrix, vmin=0, vmax=1, cmap="Greys")
    axs[0].set_title("True Transition Matrix")
    axs[1].set_title("Test Transition Matrix")
    cax = fig.add_axes([axs[1].get_position().x1 + 0.07,
                        axs[1].get_position().y0,
                        0.02,
                        axs[1].get_position().y1-axs[1].get_position().y0])
    plt.colorbar(out, cax=cax)
    plt.show()

def plot_hmm_data(obs, states):
    lim = 1.01 * abs(obs).max()
    time_bins, obs_dim = obs.shape
    plt.figure(figsize=(8, 3))
    plt.imshow(states[None,:],
            aspect="auto",
            cmap=cmap,
            vmin=0,
            vmax=len(colors)-1,
            extent=(0, time_bins, -lim, (obs_dim)*lim))

    for d in range(obs_dim):
        plt.plot(obs[:,d] + lim * d, '-k')

    plt.xlim(0, time_bins)
    plt.xlabel("time")
    plt.yticks(lim * np.arange(obs_dim), ["$x_{}$".format(d+1) for d in range(obs_dim)])

    plt.title("Simulated data from an HMM")

    plt.tight_layout()
    
def plot_posterior_states(Ez, states, perm):
    plt.figure(figsize=(20, 2))
    plt.imshow(Ez.T[perm], aspect="auto", interpolation="none", cmap="Greys")
    plt.plot(states, label="True State")
    plt.plot(Ez.T[perm].argmax(axis=0), '--', label="Predicted State")
    plt.xlabel("time")
    plt.ylabel("latent state")
    plt.legend(bbox_to_anchor=(1,1))
    plt.title("Predicted vs. Ground Truth Latent State")
    plt.show()<jupyter_output><empty_output><jupyter_text>1: Gaussian HMM  First, let's create a ground truth HMM model<jupyter_code># define some parameters
num_states = 3
num_emission_dims = 2

transition_matrix = 0.9 * np.eye(num_states) \
                    + 0.1 * np.roll(np.eye(num_states), 1, axis=1)

emission_means = np.arange(6, dtype=float).reshape((3, 2)) - 3
emission_covariance=0.25**2 * np.eye(num_emission_dims)

# create the model
true_hmm = GaussianHMM(
    num_states=num_states, 
    num_emission_dims=num_emission_dims,
    transition_matrix=transition_matrix,
    emission_means=emission_means,
    emission_covariances=emission_covariance)

# sample latent states, observations (data) from the model
num_timesteps = 1000
states, data = true_hmm.sample(jr.PRNGKey(0), num_timesteps)

# view the sampled data
plot_transition_matrix(transition_matrix)
plot_hmm_data(data, states)<jupyter_output>WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)<jupyter_text>Acessing model parametersIn this case, the HMM was constructed with given means and covariance matrices. How do we access them? They are available through model properties.<jupyter_code>print("Emission means (num states x num emission_dims):")
print(true_hmm.emission_means)
print("")
print("Emission covariance matrices (num states x num emission_dims x num emission_dims):")
print(true_hmm.emission_covariance_matrices)<jupyter_output>Emission means (num states x num emission_dims):
[[-3. -2.]
 [-1.  0.]
 [ 1.  2.]]

Emission covariance matrices (num states x num emission_dims x num emission_dims):
[[[0.00390625 0.        ]
  [0.         0.00390625]]

 [[0.00390625 0.        ]
  [0.         0.00390625]]

 [[0.00390625 0.        ]
  [0.         0.00390625]]]<jupyter_text>Fit Gaussian HMM using EM on a single trial<jupyter_code># we initialize a random HMM which we will fit to our data
seed = jr.PRNGKey(1)
num_states_est = 3

test_hmm = GaussianHMM(
    num_states=num_states, 
    num_emission_dims=num_emission_dims,
    seed=seed)
# fit the model using EM (for Gaussian HMM, we can perform exact EM)
log_probs, fitted_hmm, posterior = test_hmm.fit(data, method="em", tol=1e-5)

# plot the marginal log probability of the data as a function of number of fit steps
plt.plot(log_probs)
plt.ylabel("Log probability")
plt.xlabel("Iteration")
plt.show()<jupyter_output>Initializing...
Done.<jupyter_text>Let's inspect the fitted model<jupyter_code># We can obtain the expected states from the posterior over latent states
Ez = posterior.expected_states[0]

# For discrete states, we can freely permute the state identity
# here, we find the best permutation to match up to our known states
# so that it's easier to compare our inferred states with the ground truth.
perm = find_permutation(states, np.argmax(Ez, axis=-1))
# Here, we compare the true and inferred transition matrix
compare_transition_matrix(true_hmm.transition_matrix, fitted_hmm.transition_matrix)
# Here, we plot the inferred state against the true state as a function of time. 
# The black rectangles represent the probability of the inferred state. We can 
# see that the model is very confident about the different state assignments.

plot_posterior_states(Ez, states, perm)<jupyter_output><empty_output><jupyter_text>Fit Gaussian HMM on Multiple Trials<jupyter_code>rng = jr.PRNGKey(0)
num_steps = 1000
num_trials = 5

all_states, all_data = true_hmm.sample(rng, num_steps=num_steps, num_samples=num_trials)
# here, we fit across the entire dataset using EM
# this is implemented by vmapping over the batch dim of our dataset
lps, fitted_hmm, posteriors = test_hmm.fit(all_data, method="em")
# notice that our posteriors object has a batch dim for it's properties
# (one for each of the trials in our dataset)
print(posteriors.expected_states.shape)<jupyter_output>(5, 1000, 3)<jupyter_text>Let's inspect our fitted modelWith the additional trial examples, we notice that our model has recovered more of the true transition matrix!<jupyter_code># Here, we compare the true and inferred transition matrix
Ez = posteriors.expected_states
perm = find_permutation(all_states[0], np.argmax(Ez[0], axis=-1))
compare_transition_matrix(true_hmm.transition_matrix, fitted_hmm.transition_matrix[np.ix_(perm, perm)])
# For the first few trials, let's see how good our predicted states are
for i in range(3):
    print("="*5, f"trial {i}", "="*5)
    states = all_states[i]
    perm = find_permutation(states, np.argmax(Ez[i], axis=-1))
    plot_posterior_states(Ez[i], states, perm)<jupyter_output>===== trial 0 =====
{'fraction_non_alphanumeric': 0.11663714402507154, 'fraction_numerical': 0.021392560294318026, 'mean_word_length': 4.412979351032448, 'pattern_counts': {'":': 0, '<': 28, '<?xml version=': 0, '>': 28, 'https://': 3, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '8', 'dir': 'jupyter-structured-clean-dedup', 'id': '161633', 'license': 'MIT', 'n_tokens_mistral': 2673, 'n_tokens_neox': 2509, 'n_words': 646, 'path': 'notebooks/gaussian-hmm-example.ipynb', 'repo_name': 'lindermanlab/ssm-jax-refactor', 'size': '7339'}