{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Import standard libraries\n",
    "\n",
    "import abc\n",
    "from dataclasses import dataclass\n",
    "import functools\n",
    "from functools import partial\n",
    "import itertools\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from typing import Any, Callable, NamedTuple, Optional, Union, Tuple\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",
    "import jax.random as jr\n",
    "\n",
    "\n",
    "\n",
    "import distrax\n",
    "import optax\n",
    "\n",
    "import jsl\n",
    "import ssm_jax\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import inspect\n",
    "import inspect as py_inspect\n",
    "import rich\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": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# meta-data does not work yet in VScode\n",
    "# https://github.com/microsoft/vscode-jupyter/issues/1121\n",
    "\n",
    "{\n",
    "    \"tags\": [\n",
    "        \"hide-cell\"\n",
    "    ]\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 distrax\n",
    "except:\n",
    "    %pip install --upgrade  distrax\n",
    "    import distrax\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": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "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 inspect\n",
    "import inspect as py_inspect\n",
    "import rich\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": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">class GaussianHMM<span style=\"font-weight: bold\">(</span>BaseHMM<span style=\"font-weight: bold\">)</span>:\n",
       "    def __init__<span style=\"font-weight: bold\">(</span>self,\n",
       "                 initial_probabilities,\n",
       "                 transition_matrix,\n",
       "                 emission_means,\n",
       "                 emission_covariance_matrices<span style=\"font-weight: bold\">)</span>:\n",
       "        <span style=\"color: #008000; text-decoration-color: #008000\">\"\"</span>\"_summary_\n",
       "\n",
       "        Args:\n",
       "            initial_probabilities <span style=\"font-weight: bold\">(</span>_type_<span style=\"font-weight: bold\">)</span>: _description_\n",
       "            transition_matrix <span style=\"font-weight: bold\">(</span>_type_<span style=\"font-weight: bold\">)</span>: _description_\n",
       "            emission_means <span style=\"font-weight: bold\">(</span>_type_<span style=\"font-weight: bold\">)</span>: _description_\n",
       "            emission_covariance_matrices <span style=\"font-weight: bold\">(</span>_type_<span style=\"font-weight: bold\">)</span>: _description_\n",
       "        <span style=\"color: #008000; text-decoration-color: #008000\">\"\"</span>\"\n",
       "        super<span style=\"font-weight: bold\">()</span>.__init__<span style=\"font-weight: bold\">(</span>initial_probabilities,\n",
       "                         transition_matrix<span style=\"font-weight: bold\">)</span>\n",
       "\n",
       "        self._emission_distribution = tfd.MultivariateNormalFullCovariance<span style=\"font-weight: bold\">(</span>\n",
       "            emission_means, emission_covariance_matrices<span style=\"font-weight: bold\">)</span>\n",
       "\n",
       "    @classmethod\n",
       "    def random_initialization<span style=\"font-weight: bold\">(</span>cls, key, num_states, emission_dim<span style=\"font-weight: bold\">)</span>:\n",
       "        key1, key2, key3 = jr.split<span style=\"font-weight: bold\">(</span>key, <span style=\"color: #000080; text-decoration-color: #000080; font-weight: bold\">3</span><span style=\"font-weight: bold\">)</span>\n",
       "        initial_probs = jr.dirichlet<span style=\"font-weight: bold\">(</span>key1, jnp.ones<span style=\"font-weight: bold\">(</span>num_states<span style=\"font-weight: bold\">))</span>\n",
       "        transition_matrix = jr.dirichlet<span style=\"font-weight: bold\">(</span>key2, jnp.ones<span style=\"font-weight: bold\">(</span>num_states<span style=\"font-weight: bold\">)</span>, <span style=\"font-weight: bold\">(</span>num_states,<span style=\"font-weight: bold\">))</span>\n",
       "        emission_means = jr.normal<span style=\"font-weight: bold\">(</span>key3, <span style=\"font-weight: bold\">(</span>num_states, emission_dim<span style=\"font-weight: bold\">))</span>\n",
       "        emission_covs = jnp.tile<span style=\"font-weight: bold\">(</span>jnp.eye<span style=\"font-weight: bold\">(</span>emission_dim<span style=\"font-weight: bold\">)</span>, <span style=\"font-weight: bold\">(</span>num_states, <span style=\"color: #000080; text-decoration-color: #000080; font-weight: bold\">1</span>, <span style=\"color: #000080; text-decoration-color: #000080; font-weight: bold\">1</span><span style=\"font-weight: bold\">))</span>\n",
       "        return cls<span style=\"font-weight: bold\">(</span>initial_probs, transition_matrix, emission_means, emission_covs<span style=\"font-weight: bold\">)</span>\n",
       "\n",
       "    # Properties to get various parameters of the model\n",
       "    @property\n",
       "    def emission_distribution<span style=\"font-weight: bold\">(</span>self<span style=\"font-weight: bold\">)</span>:\n",
       "        return self._emission_distribution\n",
       "\n",
       "    @property\n",
       "    def emission_means<span style=\"font-weight: bold\">(</span>self<span style=\"font-weight: bold\">)</span>:\n",
       "        return self.emission_distribution.mean<span style=\"font-weight: bold\">()</span>\n",
       "\n",
       "    @property\n",
       "    def emission_covariance_matrices<span style=\"font-weight: bold\">(</span>self<span style=\"font-weight: bold\">)</span>:\n",
       "        return self.emission_distribution.covariance<span style=\"font-weight: bold\">()</span>\n",
       "\n",
       "    @property\n",
       "    def unconstrained_params<span style=\"font-weight: bold\">(</span>self<span style=\"font-weight: bold\">)</span>:\n",
       "        <span style=\"color: #008000; text-decoration-color: #008000\">\"\"</span>\"Helper property to get a PyTree of unconstrained parameters.\n",
       "        <span style=\"color: #008000; text-decoration-color: #008000\">\"\"</span>\"\n",
       "        return tfb.SoftmaxCentered<span style=\"font-weight: bold\">()</span>.inverse<span style=\"font-weight: bold\">(</span>self.initial_probabilities<span style=\"font-weight: bold\">)</span>, \\\n",
       "               tfb.SoftmaxCentered<span style=\"font-weight: bold\">()</span>.inverse<span style=\"font-weight: bold\">(</span>self.transition_matrix<span style=\"font-weight: bold\">)</span>, \\\n",
       "               self.emission_means, \\\n",
       "               PSDToRealBijector.forward<span style=\"font-weight: bold\">(</span>self.emission_covariance_matrices<span style=\"font-weight: bold\">)</span>\n",
       "\n",
       "    @classmethod\n",
       "    def from_unconstrained_params<span style=\"font-weight: bold\">(</span>cls, unconstrained_params, hypers<span style=\"font-weight: bold\">)</span>:\n",
       "        initial_probabilities = tfb.SoftmaxCentered<span style=\"font-weight: bold\">()</span>.forward<span style=\"font-weight: bold\">(</span>unconstrained_params<span style=\"font-weight: bold\">[</span><span style=\"color: #000080; text-decoration-color: #000080; font-weight: bold\">0</span><span style=\"font-weight: bold\">])</span>\n",
       "        transition_matrix = tfb.SoftmaxCentered<span style=\"font-weight: bold\">()</span>.forward<span style=\"font-weight: bold\">(</span>unconstrained_params<span style=\"font-weight: bold\">[</span><span style=\"color: #000080; text-decoration-color: #000080; font-weight: bold\">1</span><span style=\"font-weight: bold\">])</span>\n",
       "        emission_means = unconstrained_params<span style=\"font-weight: bold\">[</span><span style=\"color: #000080; text-decoration-color: #000080; font-weight: bold\">2</span><span style=\"font-weight: bold\">]</span>\n",
       "        emission_covs = PSDToRealBijector.inverse<span style=\"font-weight: bold\">(</span>unconstrained_params<span style=\"font-weight: bold\">[</span><span style=\"color: #000080; text-decoration-color: #000080; font-weight: bold\">3</span><span style=\"font-weight: bold\">])</span>\n",
       "        return cls<span style=\"font-weight: bold\">(</span>initial_probabilities, transition_matrix, emission_means, emission_covs, \n",
       "*hypers<span style=\"font-weight: bold\">)</span>\n",
       "\n",
       "</pre>\n"
      ],
      "text/plain": [
       "<rich.jupyter.JupyterRenderable at 0x7fdea0b54460>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import ssm_jax\n",
    "from ssm_jax.hmm.models import GaussianHMM\n",
    "\n",
    "print_source(GaussianHMM)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "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"
     ]
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">    def sample<span style=\"font-weight: bold\">(</span>self, key, num_timesteps<span style=\"font-weight: bold\">)</span>:\n",
       "        <span style=\"color: #008000; text-decoration-color: #008000\">\"\"</span>\"Sample a sequence of latent states and emissions.\n",
       "\n",
       "        Args:\n",
       "            key <span style=\"font-weight: bold\">(</span>_type_<span style=\"font-weight: bold\">)</span>: _description_\n",
       "            num_timesteps <span style=\"font-weight: bold\">(</span>_type_<span style=\"font-weight: bold\">)</span>: _description_\n",
       "        <span style=\"color: #008000; text-decoration-color: #008000\">\"\"</span>\"\n",
       "        def _step<span style=\"font-weight: bold\">(</span>state, key<span style=\"font-weight: bold\">)</span>:\n",
       "            key1, key2 = jr.split<span style=\"font-weight: bold\">(</span>key, <span style=\"color: #000080; text-decoration-color: #000080; font-weight: bold\">2</span><span style=\"font-weight: bold\">)</span>\n",
       "            emission = self.emission_distribution.sample<span style=\"font-weight: bold\">(</span><span style=\"color: #808000; text-decoration-color: #808000\">seed</span>=<span style=\"color: #800080; text-decoration-color: #800080\">key1</span><span style=\"font-weight: bold\">)</span>\n",
       "            next_state = self.transition_distribution.sample<span style=\"font-weight: bold\">(</span><span style=\"color: #808000; text-decoration-color: #808000\">seed</span>=<span style=\"color: #800080; text-decoration-color: #800080\">key2</span><span style=\"font-weight: bold\">)</span>\n",
       "            return next_state, <span style=\"font-weight: bold\">(</span>state, emission<span style=\"font-weight: bold\">)</span>\n",
       "\n",
       "        # Sample the initial state\n",
       "        key1, key = jr.split<span style=\"font-weight: bold\">(</span>key, <span style=\"color: #000080; text-decoration-color: #000080; font-weight: bold\">2</span><span style=\"font-weight: bold\">)</span>\n",
       "        initial_state = self.initial_distribution.sample<span style=\"font-weight: bold\">(</span><span style=\"color: #808000; text-decoration-color: #808000\">seed</span>=<span style=\"color: #800080; text-decoration-color: #800080\">key1</span><span style=\"font-weight: bold\">)</span>\n",
       "\n",
       "        # Sample the remaining emissions and states\n",
       "        keys = jr.split<span style=\"font-weight: bold\">(</span>key, num_timesteps<span style=\"font-weight: bold\">)</span>\n",
       "        _, <span style=\"font-weight: bold\">(</span>states, emissions<span style=\"font-weight: bold\">)</span> = lax.scan<span style=\"font-weight: bold\">(</span>_step, initial_state, keys<span style=\"font-weight: bold\">)</span>\n",
       "        return states, emissions\n",
       "\n",
       "</pre>\n"
      ],
      "text/plain": [
       "<rich.jupyter.JupyterRenderable at 0x7fde91c661c0>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Set dimensions\n",
    "num_states = 5\n",
    "emission_dim = 2\n",
    "\n",
    "# Specify parameters of the HMM\n",
    "initial_probs = jnp.ones(num_states) / num_states\n",
    "transition_matrix = 0.95 * jnp.eye(num_states) + 0.05 * jnp.roll(jnp.eye(num_states), 1, axis=1)\n",
    "emission_means = jnp.column_stack([\n",
    "    jnp.cos(jnp.linspace(0, 2 * jnp.pi, num_states+1))[:-1],\n",
    "    jnp.sin(jnp.linspace(0, 2 * jnp.pi, num_states+1))[:-1]\n",
    "])\n",
    "emission_covs = jnp.tile(0.1**2 * jnp.eye(emission_dim), (num_states, 1, 1))\n",
    "\n",
    "hmm = GaussianHMM(initial_probs,\n",
    "                       transition_matrix,\n",
    "                       emission_means, \n",
    "                       emission_covs)\n",
    "\n",
    "print_source(hmm.sample)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<distrax._src.utils.hmm.HMM object at 0x7fde82c856d0>\n"
     ]
    }
   ],
   "source": [
    "import distrax\n",
    "from distrax import HMM\n",
    "\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",
    "hmm = HMM(trans_dist=distrax.Categorical(probs=A),\n",
    "            init_dist=distrax.Categorical(probs=pi),\n",
    "            obs_dist=distrax.Categorical(probs=B))\n",
    "\n",
    "print(hmm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">  def sample<span style=\"font-weight: bold\">(</span>self,\n",
       "             *,\n",
       "             seed: chex.PRNGKey,\n",
       "             seq_len: chex.Array<span style=\"font-weight: bold\">)</span> -&gt; Tuple:\n",
       "    <span style=\"color: #008000; text-decoration-color: #008000\">\"\"</span>\"Sample from this HMM.\n",
       "\n",
       "    Samples an observation of given length according to this\n",
       "    Hidden Markov Model and gives the sequence of the hidden states\n",
       "    as well as the observation.\n",
       "\n",
       "    Args:\n",
       "      seed: Random key of shape <span style=\"font-weight: bold\">(</span><span style=\"color: #000080; text-decoration-color: #000080; font-weight: bold\">2</span>,<span style=\"font-weight: bold\">)</span> and dtype uint32.\n",
       "      seq_len: The length of the observation sequence.\n",
       "\n",
       "    Returns:\n",
       "      Tuple of hidden state sequence, and observation sequence.\n",
       "    <span style=\"color: #008000; text-decoration-color: #008000\">\"\"</span>\"\n",
       "    rng_key, rng_init = jax.random.split<span style=\"font-weight: bold\">(</span>seed<span style=\"font-weight: bold\">)</span>\n",
       "    initial_state = self._init_dist.sample<span style=\"font-weight: bold\">(</span><span style=\"color: #808000; text-decoration-color: #808000\">seed</span>=<span style=\"color: #800080; text-decoration-color: #800080\">rng_init</span><span style=\"font-weight: bold\">)</span>\n",
       "\n",
       "    def draw_state<span style=\"font-weight: bold\">(</span>prev_state, key<span style=\"font-weight: bold\">)</span>:\n",
       "      state = self._trans_dist.sample<span style=\"font-weight: bold\">(</span><span style=\"color: #808000; text-decoration-color: #808000\">seed</span>=<span style=\"color: #800080; text-decoration-color: #800080\">key</span><span style=\"font-weight: bold\">)</span>\n",
       "      return state, state\n",
       "\n",
       "    rng_state, rng_obs = jax.random.split<span style=\"font-weight: bold\">(</span>rng_key<span style=\"font-weight: bold\">)</span>\n",
       "    keys = jax.random.split<span style=\"font-weight: bold\">(</span>rng_state, seq_len - <span style=\"color: #000080; text-decoration-color: #000080; font-weight: bold\">1</span><span style=\"font-weight: bold\">)</span>\n",
       "    _, states = jax.lax.scan<span style=\"font-weight: bold\">(</span>draw_state, initial_state, keys<span style=\"font-weight: bold\">)</span>\n",
       "    states = jnp.append<span style=\"font-weight: bold\">(</span>initial_state, states<span style=\"font-weight: bold\">)</span>\n",
       "\n",
       "    def draw_obs<span style=\"font-weight: bold\">(</span>state, key<span style=\"font-weight: bold\">)</span>:\n",
       "      return self._obs_dist.sample<span style=\"font-weight: bold\">(</span><span style=\"color: #808000; text-decoration-color: #808000\">seed</span>=<span style=\"color: #800080; text-decoration-color: #800080\">key</span><span style=\"font-weight: bold\">)</span>\n",
       "\n",
       "    keys = jax.random.split<span style=\"font-weight: bold\">(</span>rng_obs, seq_len<span style=\"font-weight: bold\">)</span>\n",
       "    obs_seq = jax.vmap<span style=\"font-weight: bold\">(</span>draw_obs, <span style=\"color: #808000; text-decoration-color: #808000\">in_axes</span>=<span style=\"font-weight: bold\">(</span><span style=\"color: #000080; text-decoration-color: #000080; font-weight: bold\">0</span>, <span style=\"color: #000080; text-decoration-color: #000080; font-weight: bold\">0</span><span style=\"font-weight: bold\">))(</span>states, keys<span style=\"font-weight: bold\">)</span>\n",
       "\n",
       "    return states, obs_seq\n",
       "\n",
       "</pre>\n"
      ],
      "text/plain": [
       "<rich.jupyter.JupyterRenderable at 0x7fdea03643d0>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "print_source(hmm.sample)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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
}