{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collecting jax[cpu]\n",
      "  Downloading jax-0.3.5.tar.gz (946 kB)\n",
      "\u001b[K     |████████████████████████████████| 946 kB 2.7 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting absl-py\n",
      "  Downloading absl_py-1.0.0-py3-none-any.whl (126 kB)\n",
      "\u001b[K     |████████████████████████████████| 126 kB 47.7 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting numpy>=1.19\n",
      "  Downloading numpy-1.22.3-cp38-cp38-macosx_10_14_x86_64.whl (17.6 MB)\n",
      "\u001b[K     |████████████████████████████████| 17.6 MB 47.5 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting opt_einsum\n",
      "  Using cached opt_einsum-3.3.0-py3-none-any.whl (65 kB)\n",
      "Collecting scipy>=1.2.1\n",
      "  Downloading scipy-1.8.0-cp38-cp38-macosx_12_0_universal2.macosx_10_9_x86_64.whl (55.3 MB)\n",
      "\u001b[K     |████████████████████████████████| 55.3 MB 73.1 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting typing_extensions\n",
      "  Using cached typing_extensions-4.1.1-py3-none-any.whl (26 kB)\n",
      "Collecting jaxlib==0.3.5\n",
      "  Downloading jaxlib-0.3.5-cp38-none-macosx_10_9_x86_64.whl (70.5 MB)\n",
      "\u001b[K     |████████████████████████████████| 70.5 MB 723 kB/s  eta 0:00:01\n",
      "\u001b[?25hCollecting flatbuffers<3.0,>=1.12\n",
      "  Using cached flatbuffers-2.0-py2.py3-none-any.whl (26 kB)\n",
      "Requirement already satisfied: six in /opt/anaconda3/envs/scripts/lib/python3.8/site-packages (from absl-py->jax[cpu]) (1.16.0)\n",
      "Building wheels for collected packages: jax\n",
      "  Building wheel for jax (setup.py) ... \u001b[?25ldone\n",
      "\u001b[?25h  Created wheel for jax: filename=jax-0.3.5-py3-none-any.whl size=1095861 sha256=6886baa70817bbac3b5797b3720dcb81e49097f61cee7d2e1255823ea32ccad8\n",
      "  Stored in directory: /Users/kpmurphy/Library/Caches/pip/wheels/05/30/aa/908988293721511b4b29e0aadf9b5d133d0f14f6c0a188e764\n",
      "Successfully built jax\n",
      "Installing collected packages: numpy, typing-extensions, scipy, opt-einsum, flatbuffers, absl-py, jaxlib, jax\n",
      "Successfully installed absl-py-1.0.0 flatbuffers-2.0 jax-0.3.5 jaxlib-0.3.5 numpy-1.22.3 opt-einsum-3.3.0 scipy-1.8.0 typing-extensions-4.1.1\n",
      "Note: you may need to restart the kernel to use updated packages.\n",
      "Collecting git+https://github.com/probml/jsl\n",
      "  Cloning https://github.com/probml/jsl to /private/var/folders/mn/vt7cgfsx6zs9vblhvbbk7pf8003xtr/T/pip-req-build-i8seqdiw\n",
      "  Running command git clone -q https://github.com/probml/jsl /private/var/folders/mn/vt7cgfsx6zs9vblhvbbk7pf8003xtr/T/pip-req-build-i8seqdiw\n",
      "Collecting chex\n",
      "  Downloading chex-0.1.2-py3-none-any.whl (72 kB)\n",
      "\u001b[K     |████████████████████████████████| 72 kB 1.3 MB/s eta 0:00:011\n",
      "\u001b[?25hCollecting dataclasses\n",
      "  Using cached dataclasses-0.6-py3-none-any.whl (14 kB)\n",
      "Requirement already satisfied: jaxlib in /opt/anaconda3/envs/scripts/lib/python3.8/site-packages (from jsl==0.0.0) (0.3.5)\n",
      "Requirement already satisfied: jax in /opt/anaconda3/envs/scripts/lib/python3.8/site-packages (from jsl==0.0.0) (0.3.5)\n",
      "Collecting matplotlib\n",
      "  Downloading matplotlib-3.5.1-cp38-cp38-macosx_10_9_x86_64.whl (7.3 MB)\n",
      "\u001b[K     |████████████████████████████████| 7.3 MB 3.9 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting tensorflow_probability\n",
      "  Using cached tensorflow_probability-0.16.0-py2.py3-none-any.whl (6.3 MB)\n",
      "Collecting dm-tree>=0.1.5\n",
      "  Using cached dm_tree-0.1.6-cp38-cp38-macosx_10_14_x86_64.whl (95 kB)\n",
      "Requirement already satisfied: absl-py>=0.9.0 in /opt/anaconda3/envs/scripts/lib/python3.8/site-packages (from chex->jsl==0.0.0) (1.0.0)\n",
      "Requirement already satisfied: numpy>=1.18.0 in /opt/anaconda3/envs/scripts/lib/python3.8/site-packages (from chex->jsl==0.0.0) (1.22.3)\n",
      "Collecting toolz>=0.9.0\n",
      "  Downloading toolz-0.11.2-py3-none-any.whl (55 kB)\n",
      "\u001b[K     |████████████████████████████████| 55 kB 11.3 MB/s eta 0:00:01\n",
      "\u001b[?25hRequirement already satisfied: six in /opt/anaconda3/envs/scripts/lib/python3.8/site-packages (from absl-py>=0.9.0->chex->jsl==0.0.0) (1.16.0)\n",
      "Requirement already satisfied: typing-extensions in /opt/anaconda3/envs/scripts/lib/python3.8/site-packages (from jax->jsl==0.0.0) (4.1.1)\n",
      "Requirement already satisfied: scipy>=1.2.1 in /opt/anaconda3/envs/scripts/lib/python3.8/site-packages (from jax->jsl==0.0.0) (1.8.0)\n",
      "Requirement already satisfied: opt-einsum in /opt/anaconda3/envs/scripts/lib/python3.8/site-packages (from jax->jsl==0.0.0) (3.3.0)\n",
      "Requirement already satisfied: flatbuffers<3.0,>=1.12 in /opt/anaconda3/envs/scripts/lib/python3.8/site-packages (from jaxlib->jsl==0.0.0) (2.0)\n",
      "Collecting cycler>=0.10\n",
      "  Downloading cycler-0.11.0-py3-none-any.whl (6.4 kB)\n",
      "Collecting kiwisolver>=1.0.1\n",
      "  Downloading kiwisolver-1.4.2-cp38-cp38-macosx_10_9_x86_64.whl (65 kB)\n",
      "\u001b[K     |████████████████████████████████| 65 kB 8.5 MB/s  eta 0:00:01\n",
      "\u001b[?25hCollecting fonttools>=4.22.0\n",
      "  Downloading fonttools-4.32.0-py3-none-any.whl (900 kB)\n",
      "\u001b[K     |████████████████████████████████| 900 kB 35.8 MB/s eta 0:00:01\n",
      "\u001b[?25hRequirement already satisfied: pyparsing>=2.2.1 in /opt/anaconda3/envs/scripts/lib/python3.8/site-packages (from matplotlib->jsl==0.0.0) (3.0.7)\n",
      "Requirement already satisfied: packaging>=20.0 in /opt/anaconda3/envs/scripts/lib/python3.8/site-packages (from matplotlib->jsl==0.0.0) (21.3)\n",
      "Requirement already satisfied: python-dateutil>=2.7 in /opt/anaconda3/envs/scripts/lib/python3.8/site-packages (from matplotlib->jsl==0.0.0) (2.8.2)\n",
      "Collecting pillow>=6.2.0\n",
      "  Downloading Pillow-9.1.0-cp38-cp38-macosx_10_9_x86_64.whl (3.1 MB)\n",
      "\u001b[K     |████████████████████████████████| 3.1 MB 76.6 MB/s eta 0:00:01\n",
      "\u001b[?25hCollecting gast>=0.3.2\n",
      "  Downloading gast-0.5.3-py3-none-any.whl (19 kB)\n",
      "Requirement already satisfied: decorator in /opt/anaconda3/envs/scripts/lib/python3.8/site-packages (from tensorflow_probability->jsl==0.0.0) (5.1.1)\n",
      "Collecting cloudpickle>=1.3\n",
      "  Downloading cloudpickle-2.0.0-py3-none-any.whl (25 kB)\n",
      "Building wheels for collected packages: jsl\n",
      "  Building wheel for jsl (setup.py) ... \u001b[?25ldone\n",
      "\u001b[?25h  Created wheel for jsl: filename=jsl-0.0.0-py3-none-any.whl size=77852 sha256=e7365293dc97e2b3e72bf42cc19db7d7e355abec312fc4d87961fa2044fa06f0\n",
      "  Stored in directory: /private/var/folders/mn/vt7cgfsx6zs9vblhvbbk7pf8003xtr/T/pip-ephem-wheel-cache-63vxzlng/wheels/ed/8b/bf/0105dc839fecf1fc8db14f7267a6ce5ee876324b58565b359f\n",
      "Successfully built jsl\n",
      "Installing collected packages: toolz, pillow, kiwisolver, gast, fonttools, dm-tree, cycler, cloudpickle, tensorflow-probability, matplotlib, dataclasses, chex, jsl\n",
      "Successfully installed chex-0.1.2 cloudpickle-2.0.0 cycler-0.11.0 dataclasses-0.6 dm-tree-0.1.6 fonttools-4.32.0 gast-0.5.3 jsl-0.0.0 kiwisolver-1.4.2 matplotlib-3.5.1 pillow-9.1.0 tensorflow-probability-0.16.0 toolz-0.11.2\n",
      "Note: you may need to restart the kernel to use updated packages.\n",
      "Collecting rich\n",
      "  Downloading rich-12.2.0-py3-none-any.whl (229 kB)\n",
      "\u001b[K     |████████████████████████████████| 229 kB 2.9 MB/s eta 0:00:01\n",
      "\u001b[?25hRequirement already satisfied: typing-extensions<5.0,>=4.0.0 in /opt/anaconda3/envs/scripts/lib/python3.8/site-packages (from rich) (4.1.1)\n",
      "Requirement already satisfied: pygments<3.0.0,>=2.6.0 in /opt/anaconda3/envs/scripts/lib/python3.8/site-packages (from rich) (2.11.2)\n",
      "Collecting commonmark<0.10.0,>=0.9.0\n",
      "  Using cached commonmark-0.9.1-py2.py3-none-any.whl (51 kB)\n",
      "Installing collected packages: commonmark, rich\n",
      "Successfully installed commonmark-0.9.1 rich-12.2.0\n",
      "Note: you may need to restart the kernel to use updated packages.\n"
     ]
    }
   ],
   "source": [
    "{\n",
    "    \"tags\": [\n",
    "        \"hide-cell\"\n",
    "    ]\n",
    "}\n",
    "\n",
    "### 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": 1,
   "metadata": {},
   "outputs": [
    {
     "ename": "ModuleNotFoundError",
     "evalue": "No module named 'rich'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mModuleNotFoundError\u001b[0m                       Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-1-80f0cb72052a>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     29\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0minspect\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     30\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0minspect\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0mpy_inspect\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 31\u001b[0;31m \u001b[0;32mfrom\u001b[0m \u001b[0mrich\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0minspect\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0mr_inspect\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     32\u001b[0m \u001b[0;32mfrom\u001b[0m \u001b[0mrich\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mprint\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0mr_print\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     33\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mModuleNotFoundError\u001b[0m: No module named 'rich'"
     ]
    }
   ],
   "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 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": [
    "(sec:ssm-intro)=\n",
    "# What are State Space Models?\n",
    "\n",
    "\n",
    "A state space model or SSM\n",
    "is a partially observed Markov model,\n",
    "in which the hidden state,  $z_t$,\n",
    "evolves over time according to a Markov process.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```{figure} /figures/SSM-AR-inputs.png\n",
    ":scale: 100%\n",
    ":name: ssm-ar\n",
    "\n",
    "Illustration of an SSM as a graphical model.\n",
    "```\n",
    "\n",
    "```{figure} /figures/SSM-simplified.png\n",
    ":scale: 100%\n",
    ":name: ssm-simplifed\n",
    "\n",
    "Illustration of a simplified SSM.\n",
    "```"
   ]
  },
  {
   "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",
    "\n",
    "\n",
    "There are 2 hidden states, each of which emit 6 possible observations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "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": "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.9.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}