<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 <a href="https
Step1:  Imports and Plotting Functions
Step2:  Bernoulli HMM
Let's create a true model
Step3:  From the true model, we can sample synthetic data
Step4:  Let's view the synthetic data
Step5:  Fit HMM using exact EM update
Step6:  Fit Bernoulli Over Multiple Trials
<ASSISTANT_TASK:>
Python Code:
!pip install  git+git://github.com/lindermanlab/ssm-jax-refactor.git
import ssm
Explanation: <a href="https://colab.research.google.com/github/probml/probml-notebooks/blob/main/notebooks/bernoulli_hmm_example.ipynb" target="_parent"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/></a>
Bernoulli HMM Example Notebook
Modified from https://github.com/lindermanlab/ssm-jax-refactor/blob/main/notebooks/bernoulli-hmm-example.ipynb
End of explanation
import jax.random as jr
import jax.numpy as np
import matplotlib.pyplot as plt
from tensorflow_probability.substrates import jax as tfp
from ssm.hmm import BernoulliHMM
from ssm.plots import gradient_cmap
from ssm.utils import find_permutation
import warnings
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, figsize=(10, 5))
    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=(25, 5))
    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()
Explanation: Imports and Plotting Functions
End of explanation
num_states = 5
num_channels = 10
transition_matrix = 0.90 * np.eye(num_states) + 0.10 * np.ones((num_states, num_states)) / num_states
true_hmm = BernoulliHMM(
    num_states, num_emission_dims=num_channels, transition_matrix=transition_matrix, seed=jr.PRNGKey(0)
)
plot_transition_matrix(true_hmm.transition_matrix)
Explanation: Bernoulli HMM
Let's create a true model
End of explanation
rng = jr.PRNGKey(0)
num_timesteps = 500
states, data = true_hmm.sample(rng, num_timesteps)
Explanation: From the true model, we can sample synthetic data
End of explanation
fig, axs = plt.subplots(2, 1, sharex=True, figsize=(20, 8))
axs[0].imshow(data.T, aspect="auto", interpolation="none")
# axs[0].set_ylabel("neuron")
axs[0].set_title("Observations")
axs[1].plot(states)
axs[1].set_title("Latent State")
axs[1].set_xlabel("time")
axs[1].set_ylabel("state")
plt.savefig("bernoulli-hmm-data.pdf")
plt.savefig("bernoulli-hmm-data.png")
plt.show()
Explanation: Let's view the synthetic data
End of explanation
test_hmm = BernoulliHMM(num_states, num_channels, seed=jr.PRNGKey(32))
lps, test_hmm, posterior = test_hmm.fit(data, method="em", tol=-1)
# Plot the log probabilities
plt.plot(lps)
plt.xlabel("iteration")
plt.ylabel("log likelihood")
test_hmm.transition_matrix
# Compare the transition matrices
compare_transition_matrix(true_hmm.transition_matrix, test_hmm.transition_matrix)
plt.savefig("bernoulli-hmm-transmat-comparison.pdf")
# Posterior distribution
Ez = posterior.expected_states.reshape(-1, num_states)
perm = find_permutation(states, np.argmax(Ez, axis=-1))
plot_posterior_states(Ez, states, perm)
plt.savefig("bernoulli-hmm-state-est-comparison.pdf")
plt.savefig("bernoulli-hmm-state-est-comparison.png")
plt.show()
Explanation: Fit HMM using exact EM update
End of explanation
rng = jr.PRNGKey(0)
num_timesteps = 500
num_trials = 5
all_states, all_data = true_hmm.sample(rng, num_timesteps, num_samples=num_trials)
# Now we have a batch dimension of size `num_trials`
print(all_states.shape)
print(all_data.shape)
lps, test_hmm, posterior = test_hmm.fit(all_data, method="em", tol=-1)
# plot marginal log probabilities
plt.title("Marginal Log Probability")
plt.ylabel("lp")
plt.xlabel("idx")
plt.plot(lps / data.size)
compare_transition_matrix(true_hmm.transition_matrix, test_hmm.transition_matrix)
# For the first few trials, let's see how good our predicted states are
for trial_idx in range(3):
    print("=" * 5, f"Trial: {trial_idx}", "=" * 5)
    Ez = posterior.expected_states[trial_idx]
    states = all_states[trial_idx]
    perm = find_permutation(states, np.argmax(Ez, axis=-1))
    plot_posterior_states(Ez, states, perm)
Explanation: Fit Bernoulli Over Multiple Trials
End of explanation
<END_TASK>