{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(sec:hmm-ex)=\n",
    "# Hidden Markov Models\n",
    "\n",
    "In this section, we introduce Hidden Markov Models (HMMs)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Boilerplate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install necessary libraries\n",
    "\n",
    "try:\n",
    "    import jax\n",
    "except:\n",
    "    # For cuda version, see https://github.com/google/jax#installation\n",
    "    %pip install --upgrade \"jax[cpu]\" \n",
    "    import jax\n",
    "\n",
    "try:\n",
    "    import jsl\n",
    "except:\n",
    "    %pip install git+https://github.com/probml/jsl\n",
    "    import jsl\n",
    "\n",
    "try:\n",
    "    import rich\n",
    "except:\n",
    "    %pip install rich\n",
    "    import rich\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import standard libraries\n",
    "\n",
    "import abc\n",
    "from dataclasses import dataclass\n",
    "import functools\n",
    "import itertools\n",
    "\n",
    "from typing import Any, Callable, NamedTuple, Optional, Union, Tuple\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "from jax import lax, vmap, jit, grad\n",
    "from jax.scipy.special import logit\n",
    "from jax.nn import softmax\n",
    "from functools import partial\n",
    "from jax.random import PRNGKey, split\n",
    "\n",
    "import inspect\n",
    "import inspect as py_inspect\n",
    "from rich import inspect as r_inspect\n",
    "from rich import print as r_print\n",
    "\n",
    "def print_source(fname):\n",
    "    r_print(py_inspect.getsource(fname))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Utility code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "def normalize(u, axis=0, eps=1e-15):\n",
    "    '''\n",
    "    Normalizes the values within the axis in a way that they sum up to 1.\n",
    "    Parameters\n",
    "    ----------\n",
    "    u : array\n",
    "    axis : int\n",
    "    eps : float\n",
    "        Threshold for the alpha values\n",
    "    Returns\n",
    "    -------\n",
    "    * array\n",
    "        Normalized version of the given matrix\n",
    "    * array(seq_len, n_hidden) :\n",
    "        The values of the normalizer\n",
    "    '''\n",
    "    u = jnp.where(u == 0, 0, jnp.where(u < eps, eps, u))\n",
    "    c = u.sum(axis=axis)\n",
    "    c = jnp.where(c == 0, 1, c)\n",
    "    return u / c, c"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(sec:casino-ex)=\n",
    "## Example: Casino HMM\n",
    "\n",
    "We first create the \"Ocassionally dishonest casino\" model from {cite}`Durbin98`.\n",
    "\n",
    "```{figure} /figures/casino.png\n",
    ":scale: 50%\n",
    ":name: casino-fig\n",
    "\n",
    "Illustration of the casino HMM.\n",
    "```\n",
    "\n",
    "There are 2 hidden states, each of which emit 6 possible observations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)\n"
     ]
    }
   ],
   "source": [
    "# state transition matrix\n",
    "A = np.array([\n",
    "    [0.95, 0.05],\n",
    "    [0.10, 0.90]\n",
    "])\n",
    "\n",
    "# observation matrix\n",
    "B = np.array([\n",
    "    [1/6, 1/6, 1/6, 1/6, 1/6, 1/6], # fair die\n",
    "    [1/10, 1/10, 1/10, 1/10, 1/10, 5/10] # loaded die\n",
    "])\n",
    "\n",
    "pi, _ = normalize(np.array([1, 1]))\n",
    "pi = np.array(pi)\n",
    "\n",
    "\n",
    "(nstates, nobs) = np.shape(B)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's make a little data structure to store all the parameters.\n",
    "We use NamedTuple rather than dataclass, since we assume these are immutable.\n",
    "(Also, standard python dataclass does not work well with JAX, which requires parameters to be\n",
    "pytrees, as discussed in https://github.com/google/jax/issues/2371)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "HMM(trans_mat=array([[0.95, 0.05],\n",
      "       [0.1 , 0.9 ]]), obs_mat=array([[0.16666667, 0.16666667, 0.16666667, 0.16666667, 0.16666667,\n",
      "        0.16666667],\n",
      "       [0.1       , 0.1       , 0.1       , 0.1       , 0.1       ,\n",
      "        0.5       ]]), init_dist=array([0.5, 0.5], dtype=float32))\n",
      "<class 'numpy.ndarray'>\n",
      "HMM(trans_mat=DeviceArray([[0.95, 0.05],\n",
      "             [0.1 , 0.9 ]], dtype=float32), obs_mat=DeviceArray([[0.16666667, 0.16666667, 0.16666667, 0.16666667, 0.16666667,\n",
      "              0.16666667],\n",
      "             [0.1       , 0.1       , 0.1       , 0.1       , 0.1       ,\n",
      "              0.5       ]], dtype=float32), init_dist=DeviceArray([0.5, 0.5], dtype=float32))\n",
      "<class 'jaxlib.xla_extension.DeviceArray'>\n"
     ]
    }
   ],
   "source": [
    "Array = Union[np.array, jnp.array]\n",
    "\n",
    "class HMM(NamedTuple):\n",
    "    trans_mat: Array  # A : (n_states, n_states)\n",
    "    obs_mat: Array  # B : (n_states, n_obs)\n",
    "    init_dist: Array  # pi : (n_states)\n",
    "\n",
    "params_np = HMM(A, B, pi)\n",
    "print(params_np)\n",
    "print(type(params_np.trans_mat))\n",
    "\n",
    "\n",
    "params = jax.tree_map(lambda x: jnp.array(x), params_np)\n",
    "print(params)\n",
    "print(type(params.trans_mat))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sampling from the joint\n",
    "\n",
    "Let's write code to sample from this model. \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Numpy version\n",
    "\n",
    "First we code it in numpy using a for loop."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def hmm_sample_np(params, seq_len, random_state=0):\n",
    "    np.random.seed(random_state)\n",
    "    trans_mat, obs_mat, init_dist = params.trans_mat, params.obs_mat, params.init_dist\n",
    "    n_states, n_obs = obs_mat.shape\n",
    "    state_seq = np.zeros(seq_len, dtype=int)\n",
    "    obs_seq = np.zeros(seq_len, dtype=int)\n",
    "    for t in range(seq_len):\n",
    "        if t==0:\n",
    "            zt = np.random.choice(n_states, p=init_dist)\n",
    "        else:\n",
    "            zt = np.random.choice(n_states, p=trans_mat[zt])\n",
    "        yt = np.random.choice(n_obs, p=obs_mat[zt])\n",
    "        state_seq[t] = zt\n",
    "        obs_seq[t] = yt\n",
    "\n",
    "    return state_seq, obs_seq"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 1 1 1 0 1 1 1 1 1 0 0 0 0 0 0 0 0 0 0 0 0\n",
      " 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      " 0 0 0 0 0 1 1 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "[4 1 0 2 3 4 5 4 3 1 5 4 5 0 5 2 5 3 5 4 5 5 4 2 1 4 1 0 0 4 2 2 3 3 3 0 4\n",
      " 0 2 4 3 2 5 5 3 5 3 1 3 3 3 2 3 5 5 0 4 4 5 0 0 1 3 5 1 5 0 1 2 4 0 0 0 4\n",
      " 0 5 1 4 3 5 4 5 0 2 3 5 2 4 1 2 1 0 4 3 5 0 4 5 1 5]\n"
     ]
    }
   ],
   "source": [
    "seq_len = 100\n",
    "state_seq, obs_seq = hmm_sample_np(params_np, seq_len, random_state=1)\n",
    "print(state_seq)\n",
    "print(obs_seq)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### JAX version\n",
    "\n",
    "Now let's write a JAX version using jax.lax.scan (for the inter-dependent states) and vmap (for the observations).\n",
    "This is harder to read than the numpy version, but faster."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "#@partial(jit, static_argnums=(1,))\n",
    "def markov_chain_sample(rng_key, init_dist, trans_mat, seq_len):\n",
    "    n_states = len(init_dist)\n",
    "\n",
    "    def draw_state(prev_state, key):\n",
    "        state = jax.random.choice(key, n_states, p=trans_mat[prev_state])\n",
    "        return state, state\n",
    "\n",
    "    rng_key, rng_state = jax.random.split(rng_key, 2)\n",
    "    keys = jax.random.split(rng_state, seq_len - 1)\n",
    "    initial_state = jax.random.choice(rng_key, n_states, p=init_dist)\n",
    "    final_state, states = jax.lax.scan(draw_state, initial_state, keys)\n",
    "    state_seq = jnp.append(jnp.array([initial_state]), states)\n",
    "\n",
    "    return state_seq"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "#@partial(jit, static_argnums=(1,))\n",
    "def hmm_sample(rng_key, params, seq_len):\n",
    "\n",
    "    trans_mat, obs_mat, init_dist = params.trans_mat, params.obs_mat, params.init_dist\n",
    "    n_states, n_obs = obs_mat.shape\n",
    "    rng_key, rng_obs = jax.random.split(rng_key, 2)\n",
    "    state_seq = markov_chain_sample(rng_key, init_dist, trans_mat, seq_len)\n",
    "\n",
    "    def draw_obs(z, key):\n",
    "        obs = jax.random.choice(key, n_obs, p=obs_mat[z])\n",
    "        return obs\n",
    "\n",
    "    keys = jax.random.split(rng_obs, seq_len)\n",
    "    obs_seq = jax.vmap(draw_obs, in_axes=(0, 0))(state_seq, keys)\n",
    "    \n",
    "    return state_seq, obs_seq"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "#@partial(jit, static_argnums=(1,))\n",
    "def hmm_sample2(rng_key, params, seq_len):\n",
    "\n",
    "    trans_mat, obs_mat, init_dist = params.trans_mat, params.obs_mat, params.init_dist\n",
    "    n_states, n_obs = obs_mat.shape\n",
    "\n",
    "    def draw_state(prev_state, key):\n",
    "        state = jax.random.choice(key, n_states, p=trans_mat[prev_state])\n",
    "        return state, state\n",
    "\n",
    "    rng_key, rng_state, rng_obs = jax.random.split(rng_key, 3)\n",
    "    keys = jax.random.split(rng_state, seq_len - 1)\n",
    "    initial_state = jax.random.choice(rng_key, n_states, p=init_dist)\n",
    "    final_state, states = jax.lax.scan(draw_state, initial_state, keys)\n",
    "    state_seq = jnp.append(jnp.array([initial_state]), states)\n",
    "\n",
    "    def draw_obs(z, key):\n",
    "        obs = jax.random.choice(key, n_obs, p=obs_mat[z])\n",
    "        return obs\n",
    "\n",
    "    keys = jax.random.split(rng_obs, seq_len)\n",
    "    obs_seq = jax.vmap(draw_obs, in_axes=(0, 0))(state_seq, keys)\n",
    "\n",
    "    return state_seq, obs_seq"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 1 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      " 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 1\n",
      " 1 1 1 1 1 1 1 1 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "[5 5 2 2 0 0 0 1 3 3 2 2 5 1 5 1 0 2 2 4 2 5 1 5 5 0 0 4 2 4 3 2 3 4 1 0 5\n",
      " 2 2 2 1 4 3 2 2 2 4 1 0 3 5 2 5 1 4 2 5 2 5 0 5 4 4 4 2 2 0 4 5 2 2 0 1 5\n",
      " 1 3 4 5 1 5 0 5 1 5 1 2 4 5 3 4 5 4 0 4 0 2 4 5 3 3]\n"
     ]
    }
   ],
   "source": [
    "\n",
    "key = PRNGKey(2)\n",
    "seq_len = 100\n",
    "\n",
    "state_seq, obs_seq = hmm_sample(key, params, seq_len)\n",
    "print(state_seq)\n",
    "print(obs_seq)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Check correctness by computing empirical pairwise statistics\n",
    "\n",
    "We will compute the number of i->j transitions, and check that it is close to the true \n",
    "A[i,j] transition probabilites."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0 0 1 1 1 1 0 0 1 1 1 0 1 0 0 1 1 1 1 0 1 0 0 0 0 1 0 0 0 0 0 1 0 0 0 0 1\n",
      " 1 0 0 0 0 1 0 1 0 0 0 0 1 0 0 1 1 0 1 1 0 1 1 0 1 1 1 0 0 1 1 0 1 0 0 1 0\n",
      " 1 0 0 0 0 0 0 0 0 0 0 1 1 1 1 0 0 0 1 1 1 0 1 0 0 0 0 1 0 0 0 0 1 1 0 0 0\n",
      " 0 0 1 1 1 1 1 1 0 0 0 1 1 0 0 0 0 0 1 0 0 0 1 0 1 1 0 1 1 0 0 0 0 0 0 1 0\n",
      " 0 0 0 1 0 0 0 0 0 0 0 1 0 0 0 0 0 0 0 0 1 1 1 1 0 1 1 0 0 0 0 0 1 0 0 0 0\n",
      " 0 0 0 0 0 1 1 0 0 0 0 0 0 0 0 0 0 0 0 1 0 1 1 0 0 0 1 1 1 0 0 0 1 1 0 0 0\n",
      " 0 0 0 1 1 1 0 0 0 0 1 0 0 1 1 1 0 1 1 1 1 1 0 1 1 0 0 0 1 1 0 1 0 0 1 0 0\n",
      " 0 0 0 1 0 0 0 1 0 1 0 0 0 0 1 0 0 1 0 0 0 1 1 0 0 0 0 0 0 0 1 0 0 1 1 1 1\n",
      " 1 1 0 0 0 0 0 1 1 0 0 0 0 0 0 1 0 0 1 0 0 0 0 0 0 1 0 0 0 0 1 0 1 0 0 0 1\n",
      " 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 1 0 1 0 0 1 0 0 1 1 0 1 0 0 0\n",
      " 0 0 0 0 0 0 0 1 0 0 1 1 1 1 0 0 1 1 0 0 0 0 1 1 0 1 1 0 0 0 0 0 0 0 0 1 0\n",
      " 1 0 1 0 1 0 0 0 0 1 0 0 0 0 0 0 0 0 1 0 1 0 0 0 0 0 0 0 0 0 1 0 0 0 0 0 0\n",
      " 0 0 0 0 1 0 0 1 1 0 1 1 0 0 0 0 0 0 0 0 0 0 1 0 1 0 0 0 1 1 0 0 1 1 0 0 1\n",
      " 1 0 0 0 0 0 0 0 1 0 0 1 1 0 0 0 0 1 1]\n",
      "[[244.  93.]\n",
      " [ 92.  70.]]\n",
      "[[0.7240356  0.27596438]\n",
      " [0.56790125 0.43209878]]\n"
     ]
    }
   ],
   "source": [
    "import collections\n",
    "def compute_counts(state_seq, nstates):\n",
    "    wseq = np.array(state_seq)\n",
    "    word_pairs = [pair for pair in zip(wseq[:-1], wseq[1:])]\n",
    "    counter_pairs = collections.Counter(word_pairs)\n",
    "    counts = np.zeros((nstates, nstates))\n",
    "    for (k,v) in counter_pairs.items():\n",
    "        counts[k[0], k[1]] = v\n",
    "    return counts\n",
    "\n",
    "def normalize_counts(counts):\n",
    "    ncounts = vmap(lambda v: normalize(v)[0], in_axes=0)(counts)\n",
    "    return ncounts\n",
    "\n",
    "init_dist = jnp.array([1.0, 0.0])\n",
    "trans_mat = jnp.array([[0.7, 0.3], [0.5, 0.5]])\n",
    "rng_key = jax.random.PRNGKey(0)\n",
    "seq_len = 500\n",
    "state_seq = markov_chain_sample(rng_key, init_dist, trans_mat, seq_len)\n",
    "print(state_seq)\n",
    "\n",
    "counts = compute_counts(state_seq, nstates=2)\n",
    "print(counts)\n",
    "\n",
    "trans_mat_empirical = normalize_counts(counts)\n",
    "print(trans_mat_empirical)\n",
    "\n",
    "assert jnp.allclose(trans_mat, trans_mat_empirical, atol=1e-1)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}