# Autoregressive (AR) HMM Demo

[](https://colab.research.google.com/github/lindermanlab/ssm-jax-refactor/blob/main/notebooks/arhmm-example.ipynb)

This notebook illustrates the use of the _auto_regression_ observation model.  
Let $x_t$ denote the observation at time $t$. Let $z_t$ denote the corresponding discrete latent state.

The autoregressive hidden Markov model has the following likelihood,
$$
\begin{align}
x_t \mid x_{t-1}, z_t &\sim
\mathcal{N}\left(A_{z_t} x_{t-1} + b_{z_t}, Q_{z_t} \right).
\end{align}
$$
(Technically, higher-order autoregressive processes with extra linear terms from inputs are also implemented.) 


```python
try:
    import ssm
except:
    !pip install git+https://github.com/lindermanlab/ssm-jax-refactor.git -qqq
    import ssm
```


```python
import jax.numpy as np

import jax.random as jr

from tensorflow_probability.substrates import jax as tfp

from ssm.distributions.linreg import GaussianLinearRegression
from ssm.arhmm import GaussianARHMM
from ssm.utils import random_rotation
from ssm.plots import gradient_cmap #, white_to_color_cmap

import matplotlib.pyplot as plt
import seaborn as sns
```

    /Users/collinschlager/miniforge3/envs/ssmjax/lib/python3.9/site-packages/jax/_src/lib/__init__.py:33: UserWarning: JAX on Mac ARM machines is experimental and minimally tested. Please see https://github.com/google/jax/issues/5501 in the event of problems.
      warnings.warn("JAX on Mac ARM machines is experimental and minimally tested. "



```python
sns.set_style("white")
sns.set_context("talk")

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


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


```python
# Make a transition matrix
num_states = 5
transition_probs = (np.arange(num_states)**10).astype(float)
transition_probs /= transition_probs.sum()
transition_matrix = np.zeros((num_states, num_states))
for k, p in enumerate(transition_probs[::-1]):
    transition_matrix += np.roll(p * np.eye(num_states), k, axis=1)
    
plt.imshow(transition_matrix, vmin=0, vmax=1, cmap="Greys")
plt.xlabel("next state")
plt.ylabel("current state")
plt.title("transition matrix")
plt.colorbar()
```


```python
# Make observation distributions
data_dim = 2
num_lags = 1

keys = jr.split(jr.PRNGKey(0), num_states)
angles = np.linspace(0, 2 * np.pi, num_states, endpoint=False)
theta = np.pi / 25 # rotational frequency
weights = np.array([0.8 * random_rotation(key, data_dim, theta=theta) for key in keys])
biases = np.column_stack([np.cos(angles), np.sin(angles), np.zeros((num_states, data_dim - 2))])
covariances = np.tile(0.001 * np.eye(data_dim), (num_states, 1, 1))

# Compute the stationary points
stationary_points = np.linalg.solve(np.eye(data_dim) - weights, biases)
```

# Plot dynamics functions


```python
if data_dim == 2:
    lim = 5
    x = np.linspace(-lim, lim, 10)
    y = np.linspace(-lim, lim, 10)
    X, Y = np.meshgrid(x, y)
    xy = np.column_stack((X.ravel(), Y.ravel()))

    fig, axs = plt.subplots(1, num_states, figsize=(3 * num_states, 6))
    for k in range(num_states):
        A, b = weights[k], biases[k]
        dxydt_m = xy.dot(A.T) + b - xy
        axs[k].quiver(xy[:, 0], xy[:, 1],
                    dxydt_m[:, 0], dxydt_m[:, 1],
                    color=colors[k % len(colors)])

        
        axs[k].set_xlabel('$x_1$')
        axs[k].set_xticks([])
        if k == 0:
            axs[k].set_ylabel("$x_2$")
        axs[k].set_yticks([])
        axs[k].set_aspect("equal")
                        

    plt.tight_layout()
```

# Sample data from the ARHMM


```python
import warnings
warnings.filterwarnings("error")
```


```python
# Make an Autoregressive (AR) HMM
true_initial_distribution = tfp.distributions.Categorical(logits=np.zeros(num_states))
true_transition_distribution = tfp.distributions.Categorical(probs=transition_matrix)

import logging
logging.captureWarnings(True)

true_arhmm = GaussianARHMM(num_states, 
                            transition_matrix=transition_matrix,
                            emission_weights=weights,
                            emission_biases=biases,
                            emission_covariances=covariances)
time_bins = 10000
true_states, data = true_arhmm.sample(jr.PRNGKey(0), time_bins)
```

    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.



```python
fig = plt.figure(figsize=(8, 8))
for k in range(num_states):
    plt.plot(*data[true_states==k].T, 'o', color=colors[k],
         alpha=0.75, markersize=3)
    
plt.plot(*data[:1000].T, '-k', lw=0.5, alpha=0.2)
plt.xlabel("$x_1$")
plt.ylabel("$x_2$")
# plt.gca().set_aspect("equal")
```

Below, we visualize each component of of the observation variable as a time series. The colors correspond to the latent state. The dotted lines represent the stationary point of the the corresponding AR state while the solid lines are the actual observations sampled from the HMM.


```python
# Plot the data and the smoothed data
plot_slice = (0, 200)
lim = 1.05 * abs(data).max()
plt.figure(figsize=(8, 6))
plt.imshow(true_states[None, :],
           aspect="auto",
           cmap=cmap,
           vmin=0,
           vmax=len(colors)-1,
           extent=(0, time_bins, -lim, (data_dim)*lim))


Ey = np.array(stationary_points)[true_states]
for d in range(data_dim):
    plt.plot(data[:,d] + lim * d, '-k')
    plt.plot(Ey[:,d] + lim * d, ':k')

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

plt.tight_layout()

```

# Fit an ARHMM


```python
# Now fit an HMM to the data
key1, key2 = jr.split(jr.PRNGKey(0), 2)
test_num_states = num_states
initial_distribution = tfp.distributions.Categorical(logits=np.zeros(test_num_states))
transition_distribution = tfp.distributions.Categorical(logits=np.zeros((test_num_states, test_num_states)))
emission_distribution = GaussianLinearRegression(
    weights=np.tile(0.99 * np.eye(data_dim), (test_num_states, 1, 1)),
    bias=0.01 * jr.normal(key2, (test_num_states, data_dim)),
    scale_tril=np.tile(np.eye(data_dim), (test_num_states, 1, 1))
)

arhmm = GaussianARHMM(test_num_states, 
                      data_dim,
                      num_lags,
                      seed=jr.PRNGKey(0))

lps, arhmm, posterior = arhmm.fit(data, method="em")
```

    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.


    Initializing...
    Done.


    [jit compiling...]:   0%|          | 0/100 [00:00<?, ?it/s]WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    [converged] LP: 38590.152:   7%|▋         | 7/100 [00:00<00:11,  8.22it/s]



```python
# Plot the log likelihoods against the true likelihood, for comparison
true_lp = true_arhmm.marginal_likelihood(data)
plt.plot(lps, label="EM")
plt.plot(true_lp * np.ones(len(lps)), ':k', label="True")
plt.xlabel("EM Iteration")
plt.ylabel("Log Probability")
plt.legend(loc="lower right")
plt.show()
```


```python
# # Find a permutation of the states that best matches the true and inferred states
# most_likely_states = posterior.most_likely_states()
# arhmm.permute(find_permutation(true_states[num_lags:], most_likely_states))
# posterior.update()
# most_likely_states = posterior.most_likely_states()
```


```python
if data_dim == 2:
    lim = abs(data).max()
    x = np.linspace(-lim, lim, 10)
    y = np.linspace(-lim, lim, 10)
    X, Y = np.meshgrid(x, y)
    xy = np.column_stack((X.ravel(), Y.ravel()))

    fig, axs = plt.subplots(2, max(num_states, test_num_states), figsize=(3 * num_states, 6))
    for i, model in enumerate([true_arhmm, arhmm]):
        for j in range(model.num_states):
            dist = model._emissions._distribution[j]
            A, b = dist.weights, dist.bias
            dxydt_m = xy.dot(A.T) + b - xy
            axs[i,j].quiver(xy[:, 0], xy[:, 1],
                        dxydt_m[:, 0], dxydt_m[:, 1],
                        color=colors[j % len(colors)])


            axs[i,j].set_xlabel('$x_1$')
            axs[i,j].set_xticks([])
            if j == 0:
                axs[i,j].set_ylabel("$x_2$")
            axs[i,j].set_yticks([])
            axs[i,j].set_aspect("equal")
                        

    plt.tight_layout()
```


```python
# Plot the true and inferred discrete states
plot_slice = (0, 1000)
plt.figure(figsize=(8, 4))
plt.subplot(211)
plt.imshow(true_states[None,num_lags:], aspect="auto", interpolation="none", cmap=cmap, vmin=0, vmax=len(colors)-1)
plt.xlim(plot_slice)
plt.ylabel("$z_{\\mathrm{true}}$")
plt.yticks([])

plt.subplot(212)
# plt.imshow(most_likely_states[None,: :], aspect="auto", cmap=cmap, vmin=0, vmax=len(colors)-1)
plt.imshow(posterior.expected_states[0].T, aspect="auto", interpolation="none", cmap="Greys", vmin=0, vmax=1)
plt.xlim(plot_slice)
plt.ylabel("$z_{\\mathrm{inferred}}$")
plt.yticks([])
plt.xlabel("time")

plt.tight_layout()
```


```python
# Sample the fitted model
sampled_states, sampled_data = arhmm.sample(jr.PRNGKey(0), time_bins)
```

    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.
    WARNING:root:The use of `check_types` is deprecated and does not have any effect.



```python
fig = plt.figure(figsize=(8, 8))
for k in range(num_states):
    plt.plot(*sampled_data[sampled_states==k].T, 'o', color=colors[k],
         alpha=0.75, markersize=3)
    
plt.plot(*sampled_data.T, '-k', lw=0.5, alpha=0.2)
plt.xlabel("$x_1$")
plt.ylabel("$x_2$")
# plt.gca().set_aspect("equal")
```

4f942ffd8e6ed9425e756aa5573f8e1de0450624
ipynb
Jupyter Notebook
notebooks/arhmm-example.ipynb
lindermanlab/ssm-jax-refactor
879243a5b649daeacea3467fab09b5a405cb4ff9
MIT
2021-12-02T08:40:57.000Z
2022-03-08T15:23:37.000Z
notebooks/arhmm-example.ipynb
lindermanlab/ssm-jax-refactor
879243a5b649daeacea3467fab09b5a405cb4ff9
MIT
2021-11-20T00:21:58.000Z
2022-02-25T11:05:53.000Z
notebooks/arhmm-example.ipynb
lindermanlab/ssm-jax-refactor
879243a5b649daeacea3467fab09b5a405cb4ff9
MIT
2022-03-02T05:55:16.000Z
2022-03-02T05:55:16.000Z
Qwen/Qwen-72B
1. YES
2. YES