HMM inference complete.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "experienced-victim",
   "metadata": {},
   "source": [
    "# Inference for Hidden Markov Models"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "animated-devices",
   "metadata": {},
   "source": [
    "Probabilistic Modelling and Reasoning 2023-2024. University of Edinburgh.\n",
    "\n",
    "In this note, we will study and implement inference algorithms for a hidden Markov model (HMM).\n",
    "\n",
    "The notes in this tutorial are based on the PMR lecture slides, which you can find [here](https://opencourse.inf.ed.ac.uk/pmr/).\n",
    "\n",
    "## Table of Content\n",
    "\n",
    "* Inference basics\n",
    "* Forward-backward\n",
    "* Viterbi-backtracking\n",
    "* Forward-filtering backward-sampling"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "toxic-secondary",
   "metadata": {},
   "source": [
    "# Inference basics"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "hairy-phenomenon",
   "metadata": {},
   "source": [
    "In slides on \"Exact Inference for Hidden Markov Models\", the following inference problems were introduced:\n",
    "* Filtering (inferring the present): $p(h_t | v_{1:t})$\n",
    "* Smoothing (infering the past):  $p(h_t | v_{1:T}), t < T$, with $T$ being the lenght of the observed sequence\n",
    "* Most likely hidden path: $h^*_{1:T} = \\text{argmax}_{h_{1:T}} p(h_{1:T} | v_{1:T})$\n",
    "* Posterior Sampling: $h_{1:T} \\sim p(h_{1:T} | v_{1:T})$\n",
    "\n",
    "In this notebook, we will implement the following inference algorithms for a discrete-valued HMM model specified in terms of probability tables:\n",
    "* The alpha-recursion and beta-recursion algorithms as `forward` and `backward` functions, respectively.\n",
    "* The Viterbi and Viterbi-backtracking algorithms.\n",
    "* The forward-filtering backward-sampling algorithm.\n",
    "\n",
    "These inference operations will allow us to solve the previously mentioned inference problems."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "economic-formation",
   "metadata": {},
   "source": [
    "To start, recall the simple language example from the previous tutorial, where sentences were generated by templates\n",
    "* Hidden template: Subject -> Verb -> Object\n",
    "* Observed sentence: Jack loves dogs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "completed-desert",
   "metadata": {},
   "source": [
    "Suppose we just observe the sentence but do not know the template. The inference problems then translate to:\n",
    "\n",
    "* Filtering and smoothing, i.e. having observed the sentence \"Jack loves dogs\", what is the probability of:\n",
    "    * the first latent state  $p(h_1 | v_{1:3} = \\text{\"Jack loves dogs\"})$\n",
    "    * the second latent state $p(h_2 | v_{1:3} = \\text{\"Jack loves dogs\"})$\n",
    "    * the third latent state  $p(h_3 | v_{1:3} = \\text{\"Jack loves dogs\"})$\n",
    "* Partition function, the marginal probability $p(v_{1:3} = \\text{\"Jack loves dogs\"})$.\n",
    "* The most likely hidden path $h^*_{1:3}$:\n",
    "    * Having observed the sentence \"Jack loves dogs\", what is the most probabe latent template $h^*_{1:3}$ that generates this sentence, e.g. is it [Subject -> Verb -> Object]?\n",
    "* Other hidden paths (i.e. templates) that are likely to generate \"Jack loves dogs\" sentence. Are there other hidden paths than [Subject -> Verb -> Object] that could also generate \"Jack loves dogs\"?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "opened-maria",
   "metadata": {},
   "source": [
    "## Working in the log domain\n",
    "\n",
    "Before diving into the details of the inference algorithms, we present some background on numerics. In brief, computers are more accurate when computing \"good numbers\" (not too small or too large in magnitude, say 5, 10, 100) than \"bad numbers\" (too small in magnitude like $10^{-10}$ or too large like $10^{10}$, try adding these two numbers in python yourself and see the result, also see how computers [represent floating point numbers](https://en.wikipedia.org/wiki/Single-precision_floating-point_format#IEEE_754_single-precision_binary_floating-point_format:_binary32)). To retain sufficient numerical precision we would prefer that all our inference algorithms worked with \"good numbers\".\n",
    "\n",
    "Directly working with the raw probabilities may result in very small numbers. Consider a simple Markov chain transitioning 10 steps with each step having a transition probability of 0.1. Then, the probability for the full chain is $0.1^{10}$, which is a very small number, in fact, too small for computers to operate on accurately.\n",
    "\n",
    "So we need to find a smarter way to represent this number. We do so by transforming it into the log-probability domain (natural logarithm) and perform any subsequent operations in the log domain only, e.g.: \n",
    "\n",
    "$\\log 0.1^{10} = 10 \\log 0.1 = -23.0258 \\ldots$\n",
    "\n",
    "Notice that we can represent $-23.0258$ with significantly more precision than $0.1^{10}$ because it is not too small or too large.\n",
    "\n",
    "In the following sections, we will work with the log-probabilities to implement the algorithms introduced in the lecture."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fossil-incentive",
   "metadata": {},
   "source": [
    "# Forward-backward algorithm for computing the partition function and performing filtering and smoothing"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "noted-duplicate",
   "metadata": {},
   "source": [
    "## Forward algorithm: implementation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "approximate-tanzania",
   "metadata": {},
   "source": [
    "We start with the forward algorithm that computes the alpha-recursion. The core recursion is just one line (see \"Exact Inference for Hidden Markov Models\" slides):\n",
    "\n",
    "$\\alpha_t(j) = \\sum_i \\alpha_{t-1}(i)\\phi_t(i, j) $\n",
    "\n",
    "$\\phi_t(i, j) = p(h_t = j| h_{t - 1} = i)p(v_t | h_t = j)$\n",
    "\n",
    "But we need to be more specific and ask the following questions before implementing the code:\n",
    "* What is the input to the algorithm?\n",
    "* What are the initial values? \n",
    "* What is the output of the algorithm?\n",
    "\n",
    "Note that these are the core questions that you should **always ask before implementing any algorithms.**"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "focused-intersection",
   "metadata": {},
   "source": [
    "Here are the answers to the above questions:\n",
    "\n",
    "Inputs:\n",
    "\n",
    "* Initial probability $\\pi(i)$: suppose there are $N$ latent states in total, then $\\pi$ is a vector of length $N$ where $\\pi(i)$ denotes the probability of state $i$ being the initial state, i.e. $p(h_1 = i)$.\n",
    "* Transition probability $P$: again, there are $N$ latent states in total, then $P$ is a $N \\times N$ shaped matrix where $P(i, j)$ denotes the probability of the transition from state $i$ to state $j$, i.e. $p(h_{t+1} = j \\mid h_{t}=i)$.\n",
    "* Emission probability $E$: suppose there are $V$ observed states in total, then $E$ is a $V \\times N$ shaped matrix where $E(i, j)$ denotes the probability of hidden state $j$ emitting the observation $i$, i.e. $p(v_t=i \\mid h_t=j)$.\n",
    "* Observed sequence $v$: this is a vector of length $T$, where $v_t \\in \\{1, 2, ..., V\\}$.\n",
    "\n",
    "**It is extremely important to be clear about the shape of these vectors and matrices.** This will often help prevent bugs early in the implementation.\n",
    " \n",
    "\n",
    "Initial values:\n",
    "\n",
    "* $\\alpha_1(i) = p(h_1 = i)p(v_1 | h_1 = i)= \\pi(i) E(v_1, i)$\n",
    "* What is the shape of $\\alpha_1$?\n",
    "* Answer: $\\alpha_1$ is a vector of length $N$ (think about why).\n",
    "\n",
    "Output:\n",
    "* $\\alpha_{1:T}$: A matrix of shape $T \\times N$\n",
    "* $Z_t$: a real valued number, the marginal probability $p(v)$ of the given observation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5dc25cc6",
   "metadata": {},
   "source": [
    "We now implement the forward algorithm below. Compare it with equations provided above and the description in \"Exact Inference for Hidden Markov Models\", also think about the questions provided as comments in the code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "worthy-group",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "from scipy.special import logsumexp\n",
    "\n",
    "def forward(log_initial, log_transition, log_emission, v):\n",
    "    \"\"\"\n",
    "    Args:\n",
    "        initial: a vector of length N\n",
    "        transition: a matrix of shape N * N\n",
    "        emission: a matrix of shape V * N\n",
    "        v: observed sequence, a vector of length T\n",
    "        \n",
    "    Returns:\n",
    "        log_alpha: a matrix of shape T * N, log_alpha[t, n] denote the sum of all possible latent state sequences ending at step t, state n\n",
    "        log_Z: a float number\n",
    "    \"\"\"\n",
    "    T = len(v)\n",
    "    N = log_emission.shape[1]\n",
    "    \n",
    "    log_alpha = np.zeros((T, N)) \n",
    "    log_alpha[0] = log_emission[v[0]] + log_initial # Q: why this is an addition, not a multiplication?\n",
    "    \n",
    "    for t in range(1, T):\n",
    "        log_emission_t = log_emission[v[t]]\n",
    "        log_alpha[t] = logsumexp(log_alpha[t - 1].reshape(N, 1) +\n",
    "                                 log_transition + \n",
    "                                 log_emission_t.reshape(1, N),\n",
    "                                 axis=0) # Q: what does the `logsumexp` function do here? \n",
    "                                         # Q: what does the function `reshape` do here? \n",
    "\n",
    "    log_Z = logsumexp(log_alpha[T - 1], axis=0)  # Q: what does the `logsumexp` function do here? \n",
    "    return log_alpha, log_Z"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e772ed49",
   "metadata": {},
   "source": [
    "Here are answers to the in-line questions:\n",
    "\n",
    "* Multiplications in the probability domain $p(v \\mid h) \\pi(h)$ become additions in the log-domain $\\log \\left( p(v \\mid h) \\pi(h) \\right) = \\log p(v \\mid h) + \\log \\pi(h)$ (see e.g. [wiki](https://en.wikipedia.org/wiki/Logarithm#Product,_quotient,_power,_and_root)).\n",
    "* The `log_alpha[t] = ...` lines correspond to the equation $\\sum_i \\alpha_{t-1}(i)\\phi_t(i, j)$ computed in the log-domain:\n",
    "  $$\n",
    "  \\log \\sum_i \\alpha_{t-1}(i)\\phi_t(i, j) = \\log \\sum_i \\exp\\left( \\log \\alpha_{t-1}(i) + \\log \\phi_t(i, j)   \\right)\n",
    "  $$\n",
    "  Hence the `logsumexp`, which computes the operation in a numerically more stable way. Here is a short [blog post](https://gregorygundersen.com/blog/2020/02/09/log-sum-exp/) that explains the numerical stability advantage and implementation of `logsumexp` in more detail.\n",
    "* The `reshape` operation on `log_alpha` makes sure that the sum operation in $\\log \\alpha_{t-1}(i) + \\log \\phi_t(i, j)$ is performed on the index $i$.\n",
    "* The `reshape` operation on `log_emission_t` makes sure that the sum operation in $\\log p(h_t = j \\mid h_{t - 1} = i) + \\log p(v_t \\mid h_t = j)$ is performed on the index $j$.\n",
    "* The `log_Z = ...` line corresponds to computing the partition function $\\log p(v_{1:T}) = \\log \\sum_{h_T} \\exp \\log \\alpha_T(h_T) = \\log \\sum_{h_T} \\alpha_T(h_T)$."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "measured-reservation",
   "metadata": {},
   "source": [
    "## Forward algorithm: verifying the implementation "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "guided-science",
   "metadata": {},
   "source": [
    "How would you verify that the above implementation is correct? \n",
    "\n",
    "Perhaps the simplest way would be hand-calculating the probabilities on a simple example problem and see if the algorithm's output matches the hand-calculated values.\n",
    "\n",
    "Firstly we define the parameters of the simple HMM: i.e. initial, transition, and emission probabilities."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "patent-fraction",
   "metadata": {},
   "outputs": [],
   "source": [
    "log_initial = np.log(np.array([0.3, 0.7]))\n",
    "log_transition = np.log(np.array([[0.2, 0.8], \n",
    "                                  [0.6, 0.4]]))\n",
    "log_emission = np.log(np.array([[0.6, 0.7], \n",
    "                                [0.4, 0.3]]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "closed-legislation",
   "metadata": {},
   "source": [
    "This is a simple HMM with only two hidden states and two types of observations. Now we enumerate all observations of length 2:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "received-recovery",
   "metadata": {},
   "outputs": [],
   "source": [
    "v0 = [0, 0]\n",
    "v1 = [0, 1]\n",
    "v2 = [1, 0]\n",
    "v3 = [1, 1]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "mighty-imaging",
   "metadata": {},
   "source": [
    "Then we call the forward algorithm to calculate the probabilities of these observations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "periodic-cliff",
   "metadata": {},
   "outputs": [],
   "source": [
    "log_Zs = []\n",
    "for v in [v0, v1, v2, v3]:\n",
    "    log_alpha, log_Z = forward(log_initial, log_transition, log_emission, v)\n",
    "    log_Zs.append(log_Z)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "gothic-petersburg",
   "metadata": {},
   "source": [
    "And print out all the probabilities from the algorithm. Note that we exponentiate the log-probabilities back to the raw probability domain:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "second-reward",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.4359999999999999, 0.23399999999999996, 0.21600000000000003, 0.114]\n"
     ]
    }
   ],
   "source": [
    "probs = [np.exp(log_Z) for log_Z in log_Zs]\n",
    "print(probs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "controlling-frank",
   "metadata": {},
   "source": [
    "Your task: calculate the probabilities for the observations (say calculate the last one [1, 1]) by hand, and verify the calculated probability matches the code output)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "78f2b88b",
   "metadata": {},
   "source": [
    "Answer: \n",
    "\n",
    "$$\n",
    "\\begin{align*}\n",
    "p(v = [1,1]) &= \\sum_{h_1, h_2} p(v_1=1 \\mid h_1) p(h_1) p(h_2 \\mid h_1) p(v_2=1 \\mid h_2)\\\\\n",
    "&= \\left. p(v_1=1 \\mid h_1) p(h_1) p(h_2 \\mid h_1) p(v_2=1 \\mid h_2) \\right|_{h_1 = 0, h_2 = 0}\\\\\n",
    "&\\phantom{=} + \\left. p(v_1=1 \\mid h_1) p(h_1) p(h_2 \\mid h_1) p(v_2=1 \\mid h_2) \\right|_{h_1 = 0, h_2 = 1}\\\\\n",
    "&\\phantom{=} + \\left. p(v_1=1 \\mid h_1) p(h_1) p(h_2 \\mid h_1) p(v_2=1 \\mid h_2) \\right|_{h_1 = 1, h_2 = 1}\\\\\n",
    "&\\phantom{=} + \\left. p(v_1=1 \\mid h_1) p(h_1) p(h_2 \\mid h_1) p(v_2=1 \\mid h_2) \\right|_{h_1 = 1, h_2 = 0}\\\\\n",
    "&= 0.4*0.3*0.2*0.4\\\\\n",
    "&\\phantom{=} + 0.4*0.3*0.8*0.3\\\\\n",
    "&\\phantom{=} + 0.3*0.7*0.4*0.3\\\\\n",
    "&\\phantom{=} + 0.3*0.7*0.6*0.4\\\\\n",
    "&= 0.114\n",
    "\\end{align*}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "designing-chinese",
   "metadata": {},
   "source": [
    "Additionally, a good sanity check practice is to sum them over and see if the result equals to one (Q: why the results should equal to one?)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "blank-elder",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.9999999999999999"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.sum(probs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "casual-speaking",
   "metadata": {},
   "source": [
    "It's OK if the resulting number is not exactly 1.0. It happens due to the rounding errors when working with floating point numbers. \n",
    "\n",
    "Yet, if the error is more than ~$10^{-5}$, you should  check if there are bugs. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bcf7c047",
   "metadata": {},
   "source": [
    "## Forward algorithm: simple language model example\n",
    "\n",
    "Let us now revisit the simple language model example from the previous tutorial. \n",
    "\n",
    "We first specify an HMM model with the same initial, transition, and emission probability tables as before.\n",
    "Then we compute the probability of the sentences, e.g. $p(v = \\text{\"Jack loves dogs\"})$, under the specified model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "d35be7cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "id2state = {0: 'Subject', 1: 'Adjective', 2: 'Adverb', 3: 'Verb', 4: 'Object', 5: '<EOS>'}\n",
    "state2id = {id2state[s]: s for s in id2state}\n",
    "word2id = {'I': 0, 'He': 1, 'Jack': 2, 'Mary': 3, 'likes': 4, 'loves': 5, 'hates': 6, 'really': 7, 'extremely': 8, 'pretty': 9, 'cute': 10, 'adorable': 11, 'cats': 12, 'dogs': 13, '.': 14}\n",
    "id2word = {word2id[w]: w for w in word2id}\n",
    "\n",
    "\n",
    "log_initial_lang = np.log([0.7, 0.3, 0., 0., 0., 0.])\n",
    "log_transition_lang = np.log(np.array([[0., 0., 0.3, 0.7, 0., 0.], \n",
    "                                       [0.4, 0.1, 0., 0., 0.5, 0.], \n",
    "                                       [0., 0.3, 0., 0.7, 0., 0.],\n",
    "                                       [0., 0.3, 0.2, 0., 0.5, 0.],\n",
    "                                       [0., 0., 0., 0., 0., 1.],\n",
    "                                       [0., 0., 0., 0., 0., 1.],\n",
    "                                      ]))\n",
    "log_emission_lang = np.log(np.array([\n",
    "    [0.2, 0.2, 0.2, 0.2, 0, 0, 0, 0, 0, 0, 0, 0, 0.1, 0.1, 0.],\n",
    "    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0.5, 0.25, 0.25, 0, 0, 0.],\n",
    "    [0, 0, 0, 0, 0, 0, 0, 0.25, 0.25, 0.5, 0, 0, 0, 0, 0],\n",
    "    [0, 0, 0, 0, 0.3, 0.4, 0.3, 0, 0, 0, 0, 0, 0, 0, 0],\n",
    "    [0, 0, 0.2, 0.2, 0, 0, 0, 0, 0, 0, 0, 0, 0.3, 0.3, 0],\n",
    "    [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1.]\n",
    "    ])).T"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4505fa3d",
   "metadata": {},
   "source": [
    "What are the probabilities of the following sentences under our model?\n",
    "* \"Jack loves dogs.\"\n",
    "* \"Mary likes adorable cats.\"\n",
    "* \"like dogs Jack.\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "52da3d88",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Jack loves dogs . p=0.005879999999999998\n",
      "Mary likes adorable cats . p=0.0003307499999999996\n",
      "likes dogs Jack . p=0.0\n"
     ]
    }
   ],
   "source": [
    "v = [\"Jack\", \"loves\", \"dogs\", \".\"]\n",
    "v_id = [word2id[word] for word in v]\n",
    "log_alpha, log_Z = forward(log_initial_lang, log_transition_lang, log_emission_lang, v=v_id)\n",
    "print(\" \".join(v), f'p={np.exp(log_Z)}')\n",
    "\n",
    "v = [\"Mary\", \"likes\", \"adorable\", \"cats\", \".\"]\n",
    "v_id = [word2id[word] for word in v]\n",
    "log_alpha, log_Z = forward(log_initial_lang, log_transition_lang, log_emission_lang, v=v_id)\n",
    "print(\" \".join(v), f'p={np.exp(log_Z)}')\n",
    "\n",
    "v = [\"likes\", \"dogs\", \"Jack\", \".\"]\n",
    "v_id = [word2id[word] for word in v]\n",
    "log_alpha, log_Z = forward(log_initial_lang, log_transition_lang, log_emission_lang, v=v_id)\n",
    "print(\" \".join(v), f'p={np.exp(log_Z)}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "146de915",
   "metadata": {},
   "source": [
    "Note that our HMM has a 0.0 probability to generate an incorrect sentence \"likes dogs Jack.\", while the probabilities are nonzero for the grammatically correct sentences."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f589a138",
   "metadata": {},
   "source": [
    "We can also use the forward algorithm to perform filtering: to infer the probability of a hidden state $h_t$ given the observed states $v_{1:t}$: $p(h_t \\mid v_{1:t})$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "6cac2618",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "p(h_1 | v[1:1] = \"cute\")=\n",
      " {'Subject': 0.0, 'Adjective': 1.0, 'Adverb': 0.0, 'Verb': 0.0, 'Object': 0.0, '<EOS>': 0.0}\n",
      "p(h_2 | v[1:2] = \"cute Jack\")=\n",
      " {'Subject': 0.44, 'Adjective': 0.0, 'Adverb': 0.0, 'Verb': 0.0, 'Object': 0.56, '<EOS>': 0.0}\n",
      "p(h_3 | v[1:3] = \"cute Jack loves\")=\n",
      " {'Subject': 0.0, 'Adjective': 0.0, 'Adverb': 0.0, 'Verb': 1.0, 'Object': 0.0, '<EOS>': 0.0}\n",
      "p(h_4 | v[1:4] = \"cute Jack loves dogs\")=\n",
      " {'Subject': 0.0, 'Adjective': 0.0, 'Adverb': 0.0, 'Verb': 0.0, 'Object': 1.0, '<EOS>': 0.0}\n",
      "p(h_5 | v[1:5] = \"cute Jack loves dogs .\")=\n",
      " {'Subject': 0.0, 'Adjective': 0.0, 'Adverb': 0.0, 'Verb': 0.0, 'Object': 0.0, '<EOS>': 1.0}\n"
     ]
    }
   ],
   "source": [
    "v = [\"cute\", \"Jack\", \"loves\", \"dogs\", \".\"]\n",
    "v_id = [word2id[word] for word in v]\n",
    "\n",
    "for t in range(1, len(v)+1):\n",
    "    log_alpha, log_Z = forward(log_initial_lang, log_transition_lang, log_emission_lang, v=v_id[:t])\n",
    "    p = {id2state[i]: round(p,2) for i, p in enumerate(np.exp(log_alpha[-1] - log_Z))}\n",
    "    sentence = \" \".join(v[:t])\n",
    "    print(f'p(h_{t} | v[1:{t}] = \"{sentence}\")=\\n {p}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a892cdd3",
   "metadata": {},
   "source": [
    "Note the resulting probabilities for $p(h_2 | v_{1:2} = \\text{\"cute Jack\"})$ are not certain whether the hidden state of $h_2$ is an object or a subject. Think about why this may be the case. We provide an answer in the smoothing section below."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "expensive-romania",
   "metadata": {},
   "source": [
    "## Backward algorithm: implementation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "returning-instrumentation",
   "metadata": {},
   "source": [
    "Compare the following implementation with slides on \"Exact Inference for Hidden Markov Models\", and with the implementation of the forward algorithm."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "national-above",
   "metadata": {},
   "outputs": [],
   "source": [
    "def backward(log_initial, log_transition, log_emission, v):\n",
    "    \"\"\"\n",
    "    Args:\n",
    "        initial: a vector of length N\n",
    "        transition: a matrix of shape N * N\n",
    "        emission: a matrix of shape V * N\n",
    "        v: observed sequence, a vector of length t\n",
    "\n",
    "    Returns:\n",
    "        log_beta: a matrix of shape T * N\n",
    "        log_Z: a float number\n",
    "    \"\"\"\n",
    "    T = len(v)\n",
    "    N = log_emission.shape[1]\n",
    "    \n",
    "    log_beta = np.zeros((T, N)) # Q: Why is it initialised to zero?\n",
    "    \n",
    "    for t in range(T - 2, -1, -1):\n",
    "        log_emission_t = log_emission[v[t + 1]]\n",
    "        log_beta[t] = logsumexp(log_beta[t + 1].reshape(1, N) + \n",
    "                                log_transition + \n",
    "                                log_emission_t.reshape(1, N), \n",
    "                                axis=1)\n",
    "    log_Z = logsumexp(log_beta[0] + log_emission[v[0]] + log_initial) # Q: why is it equal to the log-likelihood? \n",
    "\n",
    "    return log_beta, log_Z"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "efb4f506",
   "metadata": {},
   "source": [
    "Answers:\n",
    "\n",
    "* In the backward algorithm $\\beta_T(h_s) = 1$, hence in the log domain $\\log \\beta_T(h_s) = 0$.\n",
    "* Recall from slides that $\\beta_s(h_{s}) = p(v_{s+1:T} \\mid h_{s})$.  \n",
    "  Hence, $p(v_{1:T}) = \\sum_{h_1} \\beta_1(h_1) p(v_1 \\mid h_1) p(h_1) = \\sum_{h_1} p(v_{2:T} \\mid h_{1}) p(v_1 \\mid h_1) p(h_1)$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "precious-negotiation",
   "metadata": {},
   "source": [
    "NB: Be careful about the boundaries, i.e., the initial values `log_beta[0]` and the final output `log_Z` . \n",
    "\n",
    "Compare this with the above forward algorithm. What are the differences before and after the loop? What are the differences inside the loop?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "representative-legislation",
   "metadata": {},
   "source": [
    "## Backward algorithm: verifying the implementation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "sacred-summary",
   "metadata": {},
   "source": [
    "We firstly check the partition functions produced by the backward algorithm, which should be the same as the forward algorithm:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "diverse-freight",
   "metadata": {},
   "outputs": [],
   "source": [
    "log_Zs = []\n",
    "for v in [v0, v1, v2, v3]:\n",
    "    log_beta, log_Z = backward(log_initial, log_transition, log_emission, v)\n",
    "    log_Zs.append(log_Z)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "contrary-aaron",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.4359999999999999, 0.234, 0.21599999999999991, 0.114] 0.9999999999999999\n"
     ]
    }
   ],
   "source": [
    "probs = [np.exp(log_Z) for log_Z in log_Zs]\n",
    "print(probs, np.sum(probs))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "active-sally",
   "metadata": {},
   "source": [
    "Note that the above numbers are the same as the output of the forward algorithm, hence giving us confidence that the implementation is correct."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "laughing-interview",
   "metadata": {},
   "source": [
    "## Combining forward and backward (alpha and beta) to get all marginals"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "marine-meditation",
   "metadata": {},
   "source": [
    "Finally, we combine the alpha and beta variables to get the marginal probabilities $p(h_t|v_{1:T})$ for each step $t$. The equations are given in \"Exact Inference for Hidden Markov Models\":\n",
    "\n",
    "$p(h_t | v_{1:T}) = \\frac{1}{Z}\\alpha(h_t)\\beta(h_t)$\n",
    "\n",
    "Again, we compute this in the log domain. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "structural-aircraft",
   "metadata": {},
   "outputs": [],
   "source": [
    "v = [0, 1, 0, 1] # an example observation sequence. Try other sequences yourself.\n",
    "\n",
    "log_alpha, log_Z = forward(log_initial, log_transition, log_emission, v)\n",
    "log_beta, log_Z = backward(log_initial, log_transition, log_emission, v)\n",
    "log_marginal = log_alpha + log_beta - log_Z"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "russian-condition",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[0.23806971, 0.76193029],\n",
       "       [0.58981233, 0.41018767],\n",
       "       [0.31313673, 0.68686327],\n",
       "       [0.53619303, 0.46380697]])"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.exp(log_marginal)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "final-reply",
   "metadata": {},
   "source": [
    "Note that the shape of output `log_marginal` is $T \\times N$ where `log_marginal[t, i]` represents the marginal probability $\\log p(h_t = i|v_{1:T})$. In the next tutorial we will use this marginal probability in the EM algorithm for learning HMM parameters from data.\n",
    "\n",
    "We can also further confirm that the distribution corresponds to a conditional probability table by summing the rows and checking that they are equal to 1:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "d40f8c7b",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert np.allclose(np.exp(log_marginal).sum(-1), 1.)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "79ba480a",
   "metadata": {},
   "source": [
    "## Smoothing: simple language model example\n",
    "\n",
    "We can now perform smoothing, i.e. compute $p(h_t | v_{1:T})$ for $t < T$, on the simple language model.\n",
    "\n",
    "Let us compute the marginals $p(h_t | v_{1:T})$ for the sentence \"cute Jack loves dogs.\" from the previous example. Compare the marginals with the probabilities $p(h_t | v_{1:t})$ that we obtained from filtering, why are they different?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "08400729",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "p(h_1 | v[1:T] = \"cute Jack loves dogs .\")=\n",
      " {'Subject': 0.0, 'Adjective': 1.0, 'Adverb': 0.0, 'Verb': 0.0, 'Object': 0.0, '<EOS>': 0.0}\n",
      "p(h_2 | v[1:T] = \"cute Jack loves dogs .\")=\n",
      " {'Subject': 1.0, 'Adjective': 0.0, 'Adverb': 0.0, 'Verb': 0.0, 'Object': 0.0, '<EOS>': 0.0}\n",
      "p(h_3 | v[1:T] = \"cute Jack loves dogs .\")=\n",
      " {'Subject': 0.0, 'Adjective': 0.0, 'Adverb': 0.0, 'Verb': 1.0, 'Object': 0.0, '<EOS>': 0.0}\n",
      "p(h_4 | v[1:T] = \"cute Jack loves dogs .\")=\n",
      " {'Subject': 0.0, 'Adjective': 0.0, 'Adverb': 0.0, 'Verb': 0.0, 'Object': 1.0, '<EOS>': 0.0}\n",
      "p(h_5 | v[1:T] = \"cute Jack loves dogs .\")=\n",
      " {'Subject': 0.0, 'Adjective': 0.0, 'Adverb': 0.0, 'Verb': 0.0, 'Object': 0.0, '<EOS>': 1.0}\n"
     ]
    }
   ],
   "source": [
    "v = [\"cute\", \"Jack\", \"loves\", \"dogs\", \".\"]\n",
    "v_id = [word2id[word] for word in v]\n",
    "\n",
    "log_alpha, log_Z = forward(log_initial_lang, log_transition_lang, log_emission_lang, v=v_id)\n",
    "log_beta, log_Z = backward(log_initial_lang, log_transition_lang, log_emission_lang, v=v_id)\n",
    "\n",
    "log_marginal = log_alpha + log_beta - log_Z\n",
    "marginal = np.exp(log_marginal)\n",
    "\n",
    "sentence = \" \".join(v)\n",
    "for t in range(1, len(v)+1):\n",
    "    p = {id2state[i]: round(p,2) for i, p in enumerate(marginal[t-1])}\n",
    "    print(f'p(h_{t} | v[1:T] = \"{sentence}\")=\\n {p}')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d637a62f",
   "metadata": {},
   "source": [
    "Answer:\n",
    "\n",
    "The difference is in $p(h_2 \\mid v_{1:T})$ (smoothing) and $p(h_2 \\mid v_{1:2})$ (filtering). Since filtering does not use \"future\" observations it couldn't tell whether the word \"Jack\" was the object or the subject of the sentence. On the other hand, by incorporating \"future\" observations (i.e. the following words \"... loves dogs.\"), smoothing can correctly infer that \"Jack\" is the subject of the sentence."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "spiritual-dialogue",
   "metadata": {},
   "source": [
    "# Viterbi-backtracking for computing the most probable latent sequence"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "visible-maria",
   "metadata": {},
   "source": [
    "We now switch our attention to another important problem, given an observed sequence, what is the probability of the most probable latent sequence:\n",
    "\n",
    "$\\max_{h_{1:T}} p(h_{1:T} | v_{1:T})$\n",
    "\n",
    "We can compute this using the Viterbi algorithm from Tutorial 5. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "systematic-paris",
   "metadata": {},
   "source": [
    "The implementation of the Viterbi algorithm (in the non log-domain) resembles that of the forward algorithm: it replaces the `sum` in the forward algorithm to a `max`, as follows:\n",
    "\n",
    "* Forward recursion: $\\alpha_t(j) = \\sum_i \\alpha_{t-1}(i)\\phi_t(i, j) $\n",
    "* Viterbi recursion: $s_t(j) = \\max_i s_{t-1}(i)\\phi_t(i, j) $\n",
    "\n",
    "This is due to the more general relationship between sum-product and max-product message passing. Here we change the intermediate variables from $\\alpha$ to $s$ just for clarity. \n",
    "\n",
    "Following the max-sum algorithm (as used in Tutorial 5), we implement the Viterbi algorithm in the log-domain below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "literary-vocabulary",
   "metadata": {},
   "outputs": [],
   "source": [
    "def viterbi(log_initial, log_transition, log_emission, v):\n",
    "    \"\"\"\n",
    "    Args:\n",
    "        initial: a vector of length N\n",
    "        transition: a matrix of shape N * N\n",
    "        emission: a matrix of shape V * N\n",
    "        v: observed sequence, a vector of length T\n",
    "        \n",
    "    Returns:\n",
    "        max_log_p: a float number, maximum probability of the latent sequence\n",
    "    \"\"\"\n",
    "    T = len(v)\n",
    "    N = log_emission.shape[1]\n",
    "    \n",
    "    log_s = np.zeros((T, N)) \n",
    "    log_s[0] = log_emission[v[0]] + log_initial \n",
    "    \n",
    "    for t in range(1, T):\n",
    "        log_s[t] = np.max(log_s[t-1].reshape(N, 1) + \n",
    "                          log_transition + \n",
    "                          log_emission[v[t]].reshape(1, N), \n",
    "                          axis=0) # NOTE: change logsumexp to max here, compare this to Forward\n",
    "\n",
    "    max_log_p_h = np.max(log_s[T - 1]) # NOTE: change the logsumexp to a max function here\n",
    "    return max_log_p_h"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "portable-belfast",
   "metadata": {},
   "source": [
    "However, the above implementation only gives us the $\\max_{h_{1:T}} \\log p(h_{1:T} | v_{1:T})$, i.e., the maximum of $\\log p(h_{1:T} | v_{1:T})$. \n",
    "\n",
    "Most often we would also like to know the actual latent sequence that corresponds to the maximum probability, i.e. $\\text{argmax}_{h_{1:T}} p(h_{1:T} | v_{1:T})$. We obtain this sequence by tracking back-pointers during the Viterbi algorithm and then backtracking after reaching the last step:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "burning-render",
   "metadata": {},
   "outputs": [],
   "source": [
    "def viterbi_backtracking(log_initial, log_transition, log_emission, v):\n",
    "    \"\"\"\n",
    "    Args:\n",
    "        initial: a vector of length N\n",
    "        transition: a matrix of shape N * N\n",
    "        emission: a matrix of shape V * N\n",
    "        v: observed sequence, a vector of length T\n",
    "        \n",
    "    Returns:\n",
    "        max_log_p: a float number, maximum probability of the latent sequence\n",
    "        max_h: max probability latent sequence, a vector of length T\n",
    "    \"\"\"\n",
    "    T = len(v)\n",
    "    N = log_emission.shape[1]\n",
    "    \n",
    "    log_s = np.zeros((T, N)) \n",
    "    max_ptr = np.zeros((T, N))\n",
    "    log_s[0] = log_emission[v[0]] + log_initial \n",
    "    \n",
    "    for t in range(1, T):\n",
    "        log_phi_t = log_transition + log_emission[v[t]].reshape(1, N)\n",
    "        log_s_phi_t = log_s[t-1].reshape(N, 1) + log_phi_t\n",
    "        log_s[t] = np.max(log_s_phi_t, axis=0) \n",
    "        max_ptr[t] = np.argmax(log_s_phi_t, axis=0) # record the back pointer\n",
    "\n",
    "    max_log_p_h = np.max(log_s[T - 1])\n",
    "    \n",
    "    # backtracking start\n",
    "    max_h = np.zeros(T).astype(int)\n",
    "    max_h[T - 1] = np.argmax(log_s[T - 1])\n",
    "    for t in range(T - 2, -1, -1):\n",
    "        max_h[t] = max_ptr[t + 1, max_h[t + 1]]\n",
    "    return max_log_p_h, max_h"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "moderate-macintosh",
   "metadata": {},
   "source": [
    "It is important to understand what happens during the backtracking procedure. Specifically, think about the following questions:\n",
    "* What does `max_ptr[t][i]` mean? \n",
    "* What does the line `max_ptr[t + 1, max_h[t + 1]]` in the backtracking loop do? \n",
    "\n",
    "You may want to draw a figure to help yourself understand the backtracking process."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "690bbc50",
   "metadata": {},
   "source": [
    "Answer:\n",
    "\n",
    "* `max_ptr[t, i]` stores a look-up table that maps a state $h_{t+1}=i$ at timestep $t+1$ to the state index of $h_t$ that maximises $p(h_{1:t} | v_{1:t}, h_{t+1})$.\n",
    "* In the backtracking loop it starts with the state $h_T$ at the final step $T$ that maximises $p(h_{1:T} | v_{1:T})$ and looks up the  index of the previous state $h_{t}$ that maximises $p(h_{1:t} | v_{1:t}, h_{t+1})$ from the look-up table at $t=T-1$. It continues until the first state in the Markov chain $t=1$ is reached.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "architectural-zimbabwe",
   "metadata": {},
   "source": [
    "## Verifying the Viterbi implementation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "disabled-berlin",
   "metadata": {},
   "source": [
    "Let us run the Viterbi-backtracking code reusing the above simple HMM (2 latent states and 2 observed states)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "administrative-swedish",
   "metadata": {},
   "outputs": [],
   "source": [
    "log_initial = np.log(np.array([0.3, 0.7]))\n",
    "log_transition = np.log(np.array([[0.2, 0.8], \n",
    "                                  [0.6, 0.4]]))\n",
    "log_emission = np.log(np.array([[0.6, 0.7], \n",
    "                                [0.4, 0.3]]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "golden-trademark",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.17639999999999997 0.4045871559633028 [1 0]\n"
     ]
    }
   ],
   "source": [
    "v = [0, 0]\n",
    "_, log_Z = forward(log_initial, log_transition, log_emission, v)\n",
    "max_log_p_h, max_h = viterbi_backtracking(log_initial, log_transition, log_emission, v)\n",
    "print(np.exp(max_log_p_h), np.exp(max_log_p_h - log_Z), max_h)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bibliographic-webster",
   "metadata": {},
   "source": [
    "This means that for the observed sequence v = [0, 0], the most probable latent sequence is h = [1, 0], and its corresponding probability is $p(h | v) \\approx 0.4045$, and the joint probability $p(h, v) \\approx 0.1763$.\n",
    "\n",
    "Note that even though we performed the computations in the log-domain, i.e. we computed $h^*_{t:T} = \\text{argmax}_{h_{1:T}} \\log p(h_{1:T} | v_{1:T})$, the obtained latent sequence $h^*_{t:T}$ is the same as the one computed in the raw probability domain  $h^*_{t:T} = \\text{argmax}_{h_{1:T}} p(h_{1:T} | v_{1:T})$, due to the $\\log$ being a strictly increasing function.\n",
    "\n",
    "You can verify that the implementation is correct by calculating the probability for all latent sequences (i.e., calculate the probabilities for $h = $[0, 0], [0, 1], [1, 0] or [1, 1]) by hand, and comparing the hand-calculated probability to the code output.\n",
    "\n",
    "Note that the posterior probability $p(h | v)$ above is obtained by the Bayes rule: \n",
    "\n",
    "$p(h | v) = \\frac{p(h, v)}{p(v)}$\n",
    "\n",
    "where the probability of the observation $p(v)$ (the denominator) is obtained using the forward algorithm."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "89ac5fc6",
   "metadata": {},
   "source": [
    "## Viterbi-backtracking: simple language model example\n",
    "\n",
    "We now illustrate the Viterbi-backtracking on the simple language model to infer the most likely template of a sentence.\n",
    "\n",
    "Let us consider the sentence \"Mary likes pretty cute dogs\". Note that we excluded puctuation in this example such that the final word may be either a subject or an object. (In our grammar only an object can directly transition into the \\<EOS\\> state, hence it immediately gives away that the previous word is an object)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "d98078f8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Subject Verb Adverb Adjective Object\n"
     ]
    }
   ],
   "source": [
    "v = [\"Mary\", \"likes\", \"pretty\", \"cute\", \"dogs\"]\n",
    "v_id = [word2id[word] for word in v]\n",
    "\n",
    "_, max_h = viterbi_backtracking(log_initial_lang, log_transition_lang, log_emission_lang, v=v_id)\n",
    "print(' '.join([id2state[h] for h in max_h]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ea52125",
   "metadata": {},
   "source": [
    "Hence, the most likely hidden path under the given HMM model is [Subject Verb Adverb Adjective Object]. We will later compare this with other possible hidden paths by sampling."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "intelligent-orange",
   "metadata": {},
   "source": [
    "# Forward-filtering backward-sampling"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "technological-lesbian",
   "metadata": {},
   "source": [
    "We now consider the sampling of latent states $h$ given an observation $v$. Firstly, it is important to differentiate sampling from Viterbi-backtracking:\n",
    "* Viterbi-backtracking: $h_{1:t}^* = \\text{argmax } p(h_{1:T} | v_{1:t})$\n",
    "    * There is only one $h_{1:T}^*$ that has the largest probability (under the assumption of a unique global maximum).\n",
    "* Sampling: $h_{1:t} \\sim p(h_{1:T} | v_{1:t})$\n",
    "    * In sampling we attempt to generate several latent sequences $h_{1:T}$, not just the most probable.\n",
    "    * Note that the sequence with $h_{1:T}^*$ still has the largest probability to get sampled.\n",
    "\n",
    "\n",
    "It is also helpful to view sampling as a random function (as opposed to a deterministic function): \n",
    "* Viterbi-backtracking is a deterministic function since it always gives the same output.\n",
    "* Sampling is a random function since each time it may give us different outputs.\n",
    "\n",
    "Sampling from the posterior $p(h_{1:T} | v_{1:T})$ can be done with backward-sampling:\n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "    h_T &\\sim p(h_T | v_{1:T}) \\\\ \n",
    "    h_{t-1} &\\sim p(h_{t-1} | h_t, v_{1:T}) \n",
    "\\end{align}\n",
    "$$\n",
    "\n",
    "The problem is to calculate $p(h_{t-1} | h_t, v_{1:T})$, which is: \n",
    "\n",
    "$$\n",
    "\\begin{align}\n",
    "p(h_{t - 1}| h_t, v_{1:T}) &= \\frac{p(h_{t - 1}, h_t | v_{1:T})}{p(h_t | v_{1:T})} \\\\ \n",
    "    &= \\frac{\\alpha(h_{t-1}) p(h_t | h_{t-1})p(v_t | h_t) \\beta(h_t) \\frac{1}{Z}}{\\alpha(h_t)\\beta(h_t) \\frac{1}{Z}} \\\\ \n",
    "    &= \\frac{\\alpha(h_{t-1})  p(h_t | h_{t-1}) p(v_t | h_t)}{\\alpha(h_t)}\n",
    "\\end{align}\n",
    "$$ \n",
    "where the first line is obtained by the Bayes' rule, the second line plugs in the definitions of $p(h_{t - 1}, h_t | v_{1:T})$ and $p(h_t | v_{1:T})$, and the final line is obtained by eliminating $\\beta(h_t)$ and $Z$. (See Tutorial 5 for details.)\n",
    "\n",
    "Since the probabilities above are expressed as a function of the transition, emission, and the alpha variables, we will use the forward algorithm to obtain them. This is why the sampling algorithm is called Forward-Filtering Backward-Sampling: because it calls the forward algorithm first and then samples in a backward fashion."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "boring-vessel",
   "metadata": {},
   "outputs": [],
   "source": [
    "def sampling(log_initial, log_transition, log_emission, v):\n",
    "    \"\"\"\n",
    "    Args:\n",
    "        initial: a vector of length N\n",
    "        transition: a matrix of shape N * N\n",
    "        emission: a matrix of shape V * N\n",
    "        v: observed sequence, a vector of length T\n",
    "        \n",
    "    Returns:\n",
    "        h: sampled sequence, a vector of length T\n",
    "    \"\"\"\n",
    "    log_alpha, log_Z = forward(log_initial, log_transition, log_emission, v)\n",
    "    \n",
    "    T = len(v)\n",
    "    N = log_emission.shape[1]\n",
    "    h = np.zeros(T).astype(int)\n",
    "    p_h = np.zeros(T)\n",
    "    \n",
    "    p = np.exp(log_alpha[T - 1] - log_Z)\n",
    "    h[T - 1] = np.random.choice(N, p=p)\n",
    "    p_h[T - 1] = p[h[T - 1]]\n",
    "    for t in range(T - 2, -1, -1):\n",
    "        log_p = log_alpha[t].reshape(N, 1) +\\\n",
    "            log_transition +\\\n",
    "            log_emission[v[t + 1]].reshape(1, N) -\\\n",
    "            log_alpha[t + 1].reshape(1, N) # Q: why p is calculated this way?\n",
    "        p = np.exp(log_p[:, h[t + 1]])\n",
    "        h[t] = np.random.choice(N, p=p)\n",
    "        p_h[t] = p[h[t]]\n",
    "    return p_h, h"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3ef5e1a3",
   "metadata": {},
   "source": [
    "Answer:\n",
    "\n",
    "The computation of the equations above are performed in the log domain, hence\n",
    "$$\n",
    "\\begin{align*}\n",
    "\\log p(h_{t - 1}| h_t, v_{1:T}) &= \\log \\left(\\frac{\\alpha(h_{t-1}) p(h_t | h_{t-1}) p(v_t | h_t)}{\\alpha(h_t)} \\right)\\\\\n",
    "&= \\log \\alpha(h_{t-1}) + \\log p(h_t | h_{t-1}) + \\log p(v_t | h_t) - \\log \\alpha(h_t)\n",
    "\\end{align*}\n",
    "$$\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "prepared-microwave",
   "metadata": {},
   "source": [
    "## Forward-filtering backward-sampling: verifying the implementation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "stretch-loading",
   "metadata": {},
   "source": [
    "Again, we use the previous simple example HMM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "boolean-lotus",
   "metadata": {},
   "outputs": [],
   "source": [
    "log_initial = np.log(np.array([0.3, 0.7]))\n",
    "log_transition = np.log(np.array([[0.2, 0.8], \n",
    "                                  [0.6, 0.4]]))\n",
    "log_emission = np.log(np.array([[0.6, 0.7], \n",
    "                                [0.4, 0.3]]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "strategic-recommendation",
   "metadata": {},
   "source": [
    "Then we call the sampling function to see its output. Try running it multiple times to see that the outputs can be different."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "personal-ontario",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.3146788990825688 [1 1]\n"
     ]
    }
   ],
   "source": [
    "v = [0, 0]\n",
    "p_h, h = sampling(log_initial, log_transition, log_emission, v)\n",
    "print(np.prod(p_h), h)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "falling-speed",
   "metadata": {},
   "source": [
    "Finally, we verify our implementation by the law of large numbers: if we call the sampling function many times, then the frequency of each sampled latent sequence will be approximately equal to their probability"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "alike-redhead",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "latent 00 freq 0.0487\n",
      "latent 01 freq 0.2310\n",
      "latent 10 freq 0.4036\n",
      "latent 11 freq 0.3167\n"
     ]
    }
   ],
   "source": [
    "h_freq = {'00': 0, '01': 0, '10': 0, '11': 0}\n",
    "N = 10000\n",
    "for _ in range(N):\n",
    "    p_h, h = sampling(log_initial, log_transition, log_emission, v)\n",
    "    h_freq[f'{h[0]}{h[1]}'] += 1\n",
    "\n",
    "for h in h_freq:\n",
    "    print('latent %s freq %.4f' % (h, h_freq[h] / N))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "right-cincinnati",
   "metadata": {},
   "source": [
    "Note that the frequency of the latent sequence [1, 0] returned from the simulation is close to 0.4045, i.e., approximately same as the true probability returned by the Viterbi algorithm. You may increase the number `N=100000`, which should return an even closer result."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "491c5829",
   "metadata": {},
   "source": [
    "## Forward-filtering backward-sampling: simple language model example\n",
    "\n",
    "Let us now consider sampling the hidden paths (sentence templates) under our simple language HMM model, i.e. $h_{1:T} \\sim p(h_{1:T} \\mid v_{1:T})$. We use the same sentence as in the Viterbi-backtracking example \"Mary likes pretty cute dogs\". Note that running the below code generates different sentence templates, including the one obtained in the Viterbi-backtracking example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "bb06457a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Subject Verb Adjective Adjective Object\n",
      "0.263157894736842\n"
     ]
    }
   ],
   "source": [
    "v = [\"Mary\", \"likes\", \"pretty\", \"cute\", \"dogs\"]\n",
    "v_id = [word2id[word] for word in v]\n",
    "\n",
    "p_h, h = sampling(log_initial_lang, log_transition_lang, log_emission_lang, v=v_id)\n",
    "print(' '.join([id2state[h] for h in h]))\n",
    "print(np.prod(p_h))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7de9c4b",
   "metadata": {},
   "source": [
    "If you run the above code cell multiple times you may generate a template [Subject Verb Adjective Adjective Object], which is the correct template for this sentence. Note that it does not correspond to the result from Viterbi-backtracking example, which is due to our choice of the HMM model parameters, which make the template [Subject Verb Adverb Adjective Object] more likely for the sentence."
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "254022097a015c41d59f82ff72f67c12fa97a3ff9a429f4c69c3538babda9c68"
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.12"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}