{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "{\n",
    "    \"tags\": [\n",
    "        \"hide-cell\"\n",
    "    ]\n",
    "}\n",
    "\n",
    "\n",
    "### 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 jsl\n",
    "import ssm_jax\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(sec:inference)=\n",
    "# States estimation (inference)\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "Given the sequence of observations, and a known model,\n",
    "one of the main tasks with SSMs\n",
    "to perform posterior inference,\n",
    "about the hidden states; this is also called\n",
    "state estimation.\n",
    "At each time step $t$,\n",
    "there are multiple forms of posterior we may be interested in computing,\n",
    "including the following:\n",
    "- the filtering distribution\n",
    "$p(\\hidden_t|\\obs_{1:t})$\n",
    "- the smoothing distribution\n",
    "$p(\\hidden_t|\\obs_{1:T})$ (note that this conditions on future data $T>t$)\n",
    "- the fixed-lag smoothing distribution\n",
    "$p(\\hidden_{t-\\ell}|\\obs_{1:t})$ (note that this\n",
    "infers $\\ell$ steps in the past given data up to the present).\n",
    "\n",
    "We may also want to compute the\n",
    "predictive distribution $h$ steps into the future:\n",
    "\\begin{align}\n",
    "p(\\obs_{t+h}|\\obs_{1:t})\n",
    "= \\sum_{\\hidden_{t+h}} p(\\obs_{t+h}|\\hidden_{t+h}) p(\\hidden_{t+h}|\\obs_{1:t})\n",
    "\\end{align}\n",
    "where the hidden state predictive distribution is\n",
    "\\begin{align}\n",
    "p(\\hidden_{t+h}|\\obs_{1:t})\n",
    "&= \\sum_{\\hidden_{t:t+h-1}}\n",
    " p(\\hidden_t|\\obs_{1:t}) \n",
    " p(\\hidden_{t+1}|\\hidden_{t})\n",
    " p(\\hidden_{t+2}|\\hidden_{t+1})\n",
    "\\cdots\n",
    " p(\\hidden_{t+h}|\\hidden_{t+h-1})\n",
    "\\end{align}\n",
    "See \n",
    "{numref}`fig:dbn-inference` for a summary of these distributions.\n",
    "\n",
    "```{figure} /figures/inference-problems-tikz.png\n",
    ":scale: 30%\n",
    ":name: fig:dbn-inference\n",
    "\n",
    "Illustration of the different kinds of inference in an SSM.\n",
    " The main kinds of inference for state-space models.\n",
    "    The shaded region is the interval for which we have data.\n",
    "    The arrow represents the time step at which we want to perform inference.\n",
    "    $t$ is the current time,  $T$ is the sequence length,\n",
    "$\\ell$ is the lag and $h$ is the prediction horizon.\n",
    "```\n",
    "\n",
    "In addition  to comuting posterior marginals,\n",
    "we may want to compute the most probable hidden sequence,\n",
    "i.e., the joint MAP estimate\n",
    "```{math}\n",
    "\\arg \\max_{\\hidden_{1:T}} p(\\hidden_{1:T}|\\obs_{1:T})\n",
    "```\n",
    "or sample sequences from the posterior\n",
    "```{math}\n",
    "\\hidden_{1:T} \\sim p(\\hidden_{1:T}|\\obs_{1:T})\n",
    "```\n",
    "\n",
    "Algorithms for all these task are discussed in the following chapters,\n",
    "since the details depend on the form of the SSM.\n",
    "\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(sec:casino-inference)=\n",
    "## Example: inference in the casino HMM\n",
    "\n",
    "\n",
    "We now illustrate filtering, smoothing and MAP decoding applied\n",
    "to the casino HMM from {ref}`sec:casino` and [](sec:casino). \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "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 = np.array([0.5, 0.5])\n",
    "\n",
    "(nstates, nobs) = np.shape(B)\n",
    "\n",
    "import distrax\n",
    "from distrax import HMM\n",
    "\n",
    "\n",
    "hmm = HMM(trans_dist=distrax.Categorical(probs=A),\n",
    "            init_dist=distrax.Categorical(probs=pi),\n",
    "            obs_dist=distrax.Categorical(probs=B))\n",
    "\n",
    "\n",
    "seed = 314\n",
    "n_samples = 300\n",
    "z_hist, x_hist = hmm.sample(seed=PRNGKey(seed), seq_len=n_samples)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/anaconda3/lib/python3.8/site-packages/jax/_src/numpy/lax_numpy.py:4457: UserWarning: Explicitly requested dtype <class 'jax.numpy.int64'> requested in astype is not available, and will be truncated to dtype int32. To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. See https://github.com/google/jax#current-gotchas for more.\n",
      "  lax_internal._check_user_dtype_supported(dtype, \"astype\")\n"
     ]
    }
   ],
   "source": [
    "# Call inference engine\n",
    "\n",
    "filtered_dist, _, smoothed_dist, loglik = hmm.forward_backward(x_hist)\n",
    "map_path = hmm.viterbi(x_hist)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Find the span of timesteps that the    simulated systems turns to be in state 1\n",
    "def find_dishonest_intervals(z_hist):\n",
    "    spans = []\n",
    "    x_init = 0\n",
    "    for t, _ in enumerate(z_hist[:-1]):\n",
    "        if z_hist[t + 1] == 0 and z_hist[t] == 1:\n",
    "            x_end = t\n",
    "            spans.append((x_init, x_end))\n",
    "        elif z_hist[t + 1] == 1 and z_hist[t] == 0:\n",
    "            x_init = t + 1\n",
    "    return spans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot posterior\n",
    "def plot_inference(inference_values, z_hist, ax, state=1, map_estimate=False):\n",
    "    n_samples = len(inference_values)\n",
    "    xspan = np.arange(1, n_samples + 1)\n",
    "    spans = find_dishonest_intervals(z_hist)\n",
    "    if map_estimate:\n",
    "        ax.step(xspan, inference_values, where=\"post\")\n",
    "    else:\n",
    "        ax.plot(xspan, inference_values[:, state])\n",
    "\n",
    "    for span in spans:\n",
    "        ax.axvspan(*span, alpha=0.5, facecolor=\"tab:gray\", edgecolor=\"none\")\n",
    "    ax.set_xlim(1, n_samples)\n",
    "    # ax.set_ylim(0, 1)\n",
    "    ax.set_ylim(-0.1, 1.1)\n",
    "    ax.set_xlabel(\"Observation number\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(0.5, 1.0, 'Filtered')"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 40804 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "filenames": {
       "image/png": "/Users/kpmurphy/github/ssm-book/_build/jupyter_execute/chapters/ssm/inference_7_1.png"
      },
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    " # Filtering\n",
    "fig, ax = plt.subplots()\n",
    "plot_inference(filtered_dist, z_hist, ax)\n",
    "ax.set_ylabel(\"p(loaded)\")\n",
    "ax.set_title(\"Filtered\")\n",
    " \n",
    "\n",
    "  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(0.5, 1.0, 'Smoothed')"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 33688 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "filenames": {
       "image/png": "/Users/kpmurphy/github/ssm-book/_build/jupyter_execute/chapters/ssm/inference_8_1.png"
      },
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Smoothing\n",
    "fig, ax = plt.subplots()\n",
    "plot_inference(smoothed_dist, z_hist, ax)\n",
    "ax.set_ylabel(\"p(loaded)\")\n",
    "ax.set_title(\"Smoothed\")\n",
    "\n",
    "\n",
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(0.5, 1.0, 'Viterbi')"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 8672 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "filenames": {
       "image/png": "/Users/kpmurphy/github/ssm-book/_build/jupyter_execute/chapters/ssm/inference_9_1.png"
      },
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# MAP estimation\n",
    "fig, ax = plt.subplots()\n",
    "plot_inference(map_path, z_hist, ax, map_estimate=True)\n",
    "ax.set_ylabel(\"MAP state\")\n",
    "ax.set_title(\"Viterbi\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO: posterior samples\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example: inference in the tracking LG-SSM\n",
    "\n",
    "We now illustrate filtering, smoothing and MAP decoding applied\n",
    "to the 2d tracking HMM from {ref}`sec:tracking-lds`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "key = jax.random.PRNGKey(314)\n",
    "timesteps = 15\n",
    "delta = 1.0\n",
    "A = jnp.array([\n",
    "    [1, 0, delta, 0],\n",
    "    [0, 1, 0, delta],\n",
    "    [0, 0, 1, 0],\n",
    "    [0, 0, 0, 1]\n",
    "])\n",
    "\n",
    "C = jnp.array([\n",
    "    [1, 0, 0, 0],\n",
    "    [0, 1, 0, 0]\n",
    "])\n",
    "\n",
    "state_size, _ = A.shape\n",
    "observation_size, _ = C.shape\n",
    "\n",
    "Q = jnp.eye(state_size) * 0.001\n",
    "R = jnp.eye(observation_size) * 1.0\n",
    "mu0 = jnp.array([8, 10, 1, 0]).astype(float)\n",
    "Sigma0 = jnp.eye(state_size) * 1.0\n",
    "\n",
    "from jsl.lds.kalman_filter import LDS, smooth, filter\n",
    "\n",
    "lds = LDS(A, C, Q, R, mu0, Sigma0)\n",
    "\n",
    "\n",
    "z_hist, x_hist = lds.sample(key, timesteps)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "from jsl.demos.plot_utils import plot_ellipse\n",
    "\n",
    "def plot_tracking_values(observed, filtered, cov_hist, signal_label, ax):\n",
    "    timesteps, _ = observed.shape\n",
    "    ax.plot(observed[:, 0], observed[:, 1], marker=\"o\", linewidth=0,\n",
    "            markerfacecolor=\"none\", markeredgewidth=2, markersize=8, label=\"observed\", c=\"tab:green\")\n",
    "    ax.plot(*filtered[:, :2].T, label=signal_label, c=\"tab:red\", marker=\"x\", linewidth=2)\n",
    "    for t in range(0, timesteps, 1):\n",
    "        covn = cov_hist[t][:2, :2]\n",
    "        plot_ellipse(covn, filtered[t, :2], ax, n_std=2.0, plot_center=False)\n",
    "    ax.axis(\"equal\")\n",
    "    ax.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "L2-filter: 3.2481\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 37540 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "filenames": {
       "image/png": "/Users/kpmurphy/github/ssm-book/_build/jupyter_execute/chapters/ssm/inference_14_1.png"
      },
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Filtering\n",
    "\n",
    "mu_hist, Sigma_hist, mu_cond_hist, Sigma_cond_hist = filter(lds, x_hist)\n",
    "l2_filter = jnp.linalg.norm(z_hist[:, :2] - mu_hist[:, :2], 2)\n",
    "print(f\"L2-filter: {l2_filter:0.4f}\")\n",
    "fig_filtered, axs = plt.subplots()\n",
    "plot_tracking_values(x_hist, mu_hist, Sigma_hist, \"filtered\", axs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "L2-smooth: 2.0450\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 26772 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "filenames": {
       "image/png": "/Users/kpmurphy/github/ssm-book/_build/jupyter_execute/chapters/ssm/inference_15_1.png"
      },
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Smoothing\n",
    "\n",
    "mu_hist_smooth, Sigma_hist_smooth = smooth(lds, mu_hist, Sigma_hist, mu_cond_hist, Sigma_cond_hist)\n",
    "l2_smooth = jnp.linalg.norm(z_hist[:, :2] - mu_hist_smooth[:, :2], 2)\n",
    "print(f\"L2-smooth: {l2_smooth:0.4f}\")\n",
    "fig_smoothed, axs = plt.subplots()\n",
    "plot_tracking_values(x_hist, mu_hist_smooth, Sigma_hist_smooth, \"smoothed\", axs)"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "6407c60499271029b671b4ff687c4ed4626355c45fd34c44476827f4be42c4d7"
  },
  "kernelspec": {
   "display_name": "Python 3.9.2 ('spyder-dev')",
   "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
}