just_the_code/roch_mmids_chap_prob_notebook.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d24687ee",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "remove-cell",
     "colab-keep"
    ]
   },
   "source": [
    "***\n",
    "\n",
    "*Course:* [Math 535](https://people.math.wisc.edu/~roch/mmids/) - Mathematical Methods in Data Science (MMiDS)  \n",
    "*Chapter:* 7-Probabilistic models   \n",
    "*Author:* [Sebastien Roch](https://people.math.wisc.edu/~roch/), Department of Mathematics, University of Wisconsin-Madison  \n",
    "*Updated:* April 12, 2024   \n",
    "*Copyright:* &copy; 2024 Sebastien Roch\n",
    "\n",
    "***"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a46e1a88",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "# You will need the files:\n",
    "#     * mmids.py\n",
    "# from https://github.com/MMiDS-textbook/MMiDS-textbook.github.io/tree/main/utils\n",
    "#\n",
    "# IF RUNNING ON GOOGLE COLAB (RECOMMENDED):\n",
    "# \"Upload to session storage\" from the Files tab on the left\n",
    "# Alternative instructions: https://colab.research.google.com/notebooks/io.ipynb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3d90a01",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "# PYTHON 3\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "import numpy as np\n",
    "from numpy import linalg as LA\n",
    "import matplotlib.pyplot as plt\n",
    "import pandas as pd\n",
    "import networkx as nx\n",
    "import torch\n",
    "from torchvision import datasets, transforms\n",
    "from torch.utils.data import DataLoader, TensorDataset\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torch.nn.functional as F\n",
    "import mmids\n",
    "seed = 535\n",
    "rng = np.random.default_rng(seed)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9653f05d",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "## Motivating example: tracking location"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6e6e4674",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "Suppose we let loose a cyborg corgi in a large park. We would like to know where it is at all time. For this purpose, it has an implanted location device that sends a signal to a tracking app.   "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f194b649",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "Here is an example of the data we might have. The red dots are recorded locations at regular time intervals. The dotted line helps keep track of the time order of the recordings. (We will explain later how this dataset is generated.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21eaea76",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "ss = 4\n",
    "os = 2\n",
    "F = np.array([[1., 0., 1., 0.],[0., 1., 0., 1.],[0., 0., 1., 0.],[0., 0., 0., 1.]]) \n",
    "H = np.array([[1., 0., 0., 0.],[0., 1, 0., 0.]])\n",
    "Q = 0.1 * np.diag(np.ones(ss))\n",
    "R = 10 * np.diag(np.ones(os))\n",
    "init_mu = np.array([0., 0., 1., 1.])\n",
    "init_Sig = 1 * np.diag(np.ones(ss))\n",
    "T = 50\n",
    "x, y = mmids.lgSamplePath(ss, os, F, H, Q, R, init_mu, init_Sig, T)\n",
    "plt.plot(y[0,:], y[1,:], marker='o', c='r', linestyle='dotted')\n",
    "plt.xlim((np.min(y[0,:])-5, np.max(y[0,:])+5)) \n",
    "plt.ylim((np.min(y[1,:])-5, np.max(y[1,:])+5))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c672557d",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "By convention, we start at $(0,0)$. Notice how squiggly the trajectory is. One issue might be that the times  at which the location is recorded are too far between. But, in fact, there is another issue: the tracking device is *inaccurate*. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa4051b4",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "To get a better estimate of the true trajectory, it is natural to try to model the noise in the measurement as well as the dynamics itself. Probabilistic models are perfectly suited for this. \n",
    "\n",
    "In this chapter, we will encounter of variety of such models and show how to take advantage of them to estimate unknown states (or parameters). In particular, conditional independence will play a key role.\n",
    "\n",
    "We will come back to location tracking later in the chapter."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f3cc5549",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "$\\newcommand{\\P}{\\mathbb{P}}$ $\\newcommand{\\E}{\\mathbb{E}}$ $\\newcommand{\\S}{\\mathcal{S}}$ $\\newcommand{\\X}{\\mathcal{X}}$ $\\newcommand{\\var}{\\mathrm{Var}}$ $\\newcommand{\\btheta}{\\boldsymbol{\\theta}}$ $\\newcommand{\\bbeta}{\\boldsymbol{\\beta}}$ $\\newcommand{\\bphi}{\\boldsymbol{\\phi}}$ $\\newcommand{\\bpi}{\\boldsymbol{\\pi}}$ $\\newcommand{\\bmu}{\\boldsymbol{\\mu}}$ $\\newcommand{\\bSigma}{\\boldsymbol{\\Sigma}}$ $\\newcommand{\\balpha}{\\boldsymbol{\\alpha}}$ $\\newcommand{\\indep}{\\perp\\!\\!\\!\\perp}$ $\\newcommand{\\bLambda}{\\boldsymbol{\\Lambda}}$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ade49aea",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "## Background: introduction to parametric families and maximum likelihood estimation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a7a9c857",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "### Exponential family"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "832e0597",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "**EXAMPLE:** **(Multivariate Gaussian)** A multivariate Gaussian vector $\\mathbf{X} = (X_1,\\ldots,X_d)$ on $\\mathbb{R}^d$ with mean $\\bmu \\in \\mathbb{R}^d$ and positive definite covariance matrix $\\bSigma \\in \\mathbb{R}^{d \\times d}$ has probability density function \n",
    "\n",
    "$$\n",
    "f_{\\bmu, \\bSigma}(\\mathbf{x})\n",
    "= \\frac{1}{(2\\pi)^{d/2} \\,|\\bSigma|^{1/2}}\n",
    "\\exp\\left(-\\frac{1}{2}(\\mathbf{x} - \\bmu)^T \\bSigma^{-1} (\\mathbf{x} - \\bmu)\\right).\n",
    "$$\n",
    "\n",
    "We use the notation $\\mathbf{X} \\sim N_d(\\bmu, \\bSigma)$.\n",
    "\n",
    "It can be shown that indeed the mean is\n",
    "\n",
    "$$\n",
    "\\E[\\mathbf{X}]\n",
    "= \\bmu\n",
    "$$\n",
    "\n",
    "and the covariance matrix is\n",
    "\n",
    "$$\n",
    "\\E[(\\mathbf{X} - \\bmu)(\\mathbf{X} - \\bmu)^T]\n",
    "= \\E[\\mathbf{X} \\mathbf{X}^T] - \\bmu \\bmu^T\n",
    "= \\bSigma.\n",
    "$$\n",
    "\n",
    "See, e.g., [[Bis](https://www.microsoft.com/en-us/research/uploads/prod/2006/01/Bishop-Pattern-Recognition-and-Machine-Learning-2006.pdf), Secion 2.3]."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d397b805",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "The following code, which plots the density in the bivariate case, was adapted from [gauss_plot_2d.ipynb](https://github.com/probml/pyprobml/blob/master/notebooks/book1/03/gauss_plot_2d.ipynb) by ChatGPT."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e2854c84",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "**LEARNING BY CHATTING:** Ask your favorite AI chatbot to explain the code! Try different covariance matrices. $\\ddagger$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16fc5155",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "from scipy.stats import multivariate_normal\n",
    "\n",
    "def gaussian_pdf(X, Y, mean, cov):\n",
    "    xy = np.stack([X.flatten(), Y.flatten()], axis=-1)\n",
    "    return multivariate_normal.pdf(\n",
    "        xy, mean=mean, cov=cov).reshape(X.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f670b5b6",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "def make_surface_plot(X, Y, Z):\n",
    "    fig = plt.figure()\n",
    "    ax = fig.add_subplot(111, projection='3d')\n",
    "    surf = ax.plot_surface(\n",
    "        X, Y, Z, cmap=plt.cm.viridis, antialiased=False)\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "50e6fdc6",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "We plot the density for mean $(0,0)$ with two different covariance matrices:\n",
    "\n",
    "$$\n",
    "\\bSigma_1 = \\begin{bmatrix} 1.0 & 0 \\\\ 0 & 1.0 \\end{bmatrix}\n",
    "\\quad \\text{and} \\quad \n",
    "\\bSigma_2 = \\begin{bmatrix} \\sigma_1^2 & \\rho \\sigma_1 \\sigma_2 \\\\ \\rho \\sigma_1 \\sigma_2 & \\sigma_2^2 \\end{bmatrix}\n",
    "$$\n",
    "\n",
    "where $\\sigma_1 = 1.5$, $\\sigma_2 = 0.5$ and $\\rho = -0.75$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "869696f7",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "start_point = 5\n",
    "stop_point = 5\n",
    "num_samples = 100\n",
    "points = np.linspace(-start_point, stop_point, num_samples)\n",
    "X, Y = np.meshgrid(points, points)\n",
    "\n",
    "mean = np.array([0., 0.])\n",
    "cov = np.array([[1., 0.], [0., 1.]])\n",
    "make_surface_plot(X, Y, gaussian_pdf(X, Y, mean, cov))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e01a3d41",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "mean = np.array([0., 0.])\n",
    "cov = np.array([[1.5 ** 2., -0.75 * 1.5 * 0.5], \n",
    "                 [-0.75 * 1.5 * 0.5, 0.5 ** 2.]])\n",
    "make_surface_plot(X, Y, gaussian_pdf(X, Y, mean, cov))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88c47d5a",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "Rewriting the density as\n",
    "\n",
    "$$\n",
    "f_{\\bmu, \\bSigma}(\\mathbf{x})\n",
    "= \\frac{e^{-(1/2) \\bmu^T \\bSigma^{-1} \\bmu}}{(2\\pi)^{d/2} \\,|\\bSigma|^{1/2}}\n",
    "\\exp\\left(- \\mathbf{x}^T \\bSigma^{-1}\\bmu \n",
    "- \\frac{1}{2} \\mathrm{tr}\\left(\\mathbf{x} \\mathbf{x}^T \\bSigma^{-1}\\right)\\right)\n",
    "$$\n",
    "\n",
    "where we used the symmetric nature of $\\bSigma^{-1}$ in the first term of the exponential and the previous trace identity in the second term. The expression in parentheses is linear in the entries of $\\mathbf{x}$ and $\\mathbf{x} \\mathbf{x}^T$, which can then be taken as  sufficient statistics (formally, using [vectorization](https://en.wikipedia.org/wiki/Vectorization_%28mathematics%29)). Indeed note that\n",
    "\n",
    "$$\n",
    "\\mathbf{x}^T \\bSigma^{-1}\\bmu \n",
    "= \\sum_{i=1}^d x_i (\\bSigma^{-1}\\bmu)_i\n",
    "$$\n",
    "\n",
    "and\n",
    "\n",
    "$$\n",
    "\\mathrm{tr}\\left(\\mathbf{x} \\mathbf{x}^T \\bSigma^{-1}\\right)\n",
    "= \\sum_{i = 1}^d \\left(\\sum_{j=1}^d (\\mathbf{x} \\mathbf{x}^T)_{i,j} (\\bSigma^{-1})_{j,i}\\right)\n",
    "= \\sum_{i = 1}^d \\sum_{j=1}^d x_i x_j (\\bSigma^{-1})_{j,i}.\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1b09bd4",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "So we can take \n",
    "\n",
    "$$\n",
    "\\bphi(\\mathbf{x})\n",
    "= (x_1,\\ldots,x_d, x_1 x_1, \\ldots, x_d x_1, x_1 x_2, \\ldots, x_d x_2, \\ldots, x_1 x_d, \\ldots, x_d x_d)\n",
    "$$\n",
    "\n",
    "\\begin{align*}\n",
    "\\btheta\n",
    "&= \\bigg(-(\\bSigma^{-1}\\bmu)_1,\\ldots,-(\\bSigma^{-1}\\bmu)_d,\\\\\n",
    "&\\qquad - \\frac{1}{2}(\\bSigma^{-1})_{1,1}, \\ldots, - \\frac{1}{2}(\\bSigma^{-1})_{1,d},\\\\\n",
    "&\\qquad - \\frac{1}{2}(\\bSigma^{-1})_{2,1}, \\ldots, - \\frac{1}{2}(\\bSigma^{-1})_{2,d},\\\\ \n",
    "&\\qquad \\ldots, - \\frac{1}{2}(\\bSigma^{-1})_{d,1}, \\ldots,- \\frac{1}{2}(\\bSigma^{-1})_{d,d}\\bigg)\n",
    "\\end{align*}\n",
    "\n",
    "and $h (\\mathbf{x}) \\equiv 1$. Expressing $Z(\\btheta)$ explicitly is  not straightforward. But note that $\\btheta$ includes all entries of $\\bSigma^{-1}$, from which $\\bSigma$ can be computed (e.g., from [Cramer's rule](https://en.wikipedia.org/wiki/Cramer%27s_rule#Finding_inverse_matrix)), and in turn from which $\\bmu$ can be extracted out of the entries of $\\bSigma^{-1}\\bmu$ in $\\btheta$. So the normalizing factor $\\frac{(2\\pi)^{d/2} \\,|\\bSigma|^{1/2}}{e^{-(1/2) \\bmu^T \\bSigma^{-1} \\bmu}}$ can in principle be expressed in terms of $\\btheta$.\n",
    "\n",
    "This shows that the multivariate normal is an exponential family. \n",
    "\n",
    "The matrix $\\bLambda = \\bSigma^{-1}$ is also known as the precision matrix."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4c7557eb",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "Alternatively, let $\\mathbf{Z}$ be a standard Normal $d$-vector,\n",
    "let $\\bmu \\in \\mathbb{R}^d$ and let $\\bSigma \\in \\mathbb{R}^{d \\times d}$ be positive definite. Then the transformed random variable $\\mathbf{X} = \\bmu + \\bSigma \\mathbf{Z}$ is a multivariate Gaussian with mean $\\bmu$ and covariance matrix $\\bSigma$. This can be proved using the [change of variables formula](https://en.wikipedia.org/wiki/Probability_density_function#Function_of_random_variables_and_change_of_variables_in_the_probability_density_function) (try it!). $\\lhd$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5dfd25a2",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "**NUMERICAL CORNER:** In Numpy, as we have seen before, the module [numpy.random](https://numpy.org/doc/stable/reference/random/index.html) provides a way to sample from a variety of standard distributions. We first initialize the [pseudorandom number generator](https://en.wikipedia.org/wiki/Pseudorandom_number_generator) with a [random seed](https://en.wikipedia.org/wiki/Random_seed). Recall that it allows the results to be reproducible: using the same seed produces the same results again."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1889210e",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "seed = 535\n",
    "rng = np.random.default_rng(535)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "628d3965",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "Here's are lists of available [probability distributions](https://numpy.org/doc/stable/reference/random/generator.html#distributions)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "602e5115",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "p = 0.1 # probability of success\n",
    "N = 5 # number of samples\n",
    "rng.binomial(1, p, size=N) # Bernoulli is special case of binomial with 1 trial"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7958143d",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "Here are a few other examples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14a4d99f",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "p = [0.1, 0.2, 0.7]\n",
    "n = 100\n",
    "rng.multinomial(n, p, size=N)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ac7fb41",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "mu = np.array([0.1, -0.3])\n",
    "sig = np.array([[2., 0.],[0., 3.]])\n",
    "rng.multivariate_normal(mu, sig, size=N)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cef537a6",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "$\\unlhd$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b445ed2",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "### Parameter estimation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cbb1809e",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "### Parameter estimation for exponential families"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "59148cd3",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "### Generalized linear models"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f9b0ffa",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "$\\newcommand{\\P}{\\mathbb{P}}$ $\\newcommand{\\E}{\\mathbb{E}}$ $\\newcommand{\\S}{\\mathcal{S}}$ $\\newcommand{\\var}{\\mathrm{Var}}$ $\\newcommand{\\btheta}{\\boldsymbol{\\theta}}$ $\\newcommand{\\bphi}{\\boldsymbol{\\phi}}$ $\\newcommand{\\bpi}{\\boldsymbol{\\pi}}$ $\\newcommand{\\bmu}{\\boldsymbol{\\mu}}$ $\\newcommand{\\blambda}{\\boldsymbol{\\lambda}}$ $\\newcommand{\\bSigma}{\\boldsymbol{\\Sigma}}$ $\\newcommand{\\balpha}{\\boldsymbol{\\alpha}}$ $\\newcommand{\\indep}{\\perp\\!\\!\\!\\perp}$ $\\newcommand{\\bp}{\\mathbf{p}}$ $\\newcommand{\\bx}{\\mathbf{x}}$ $\\newcommand{\\bX}{\\mathbf{X}}$ $\\newcommand{\\by}{\\mathbf{y}}$ $\\newcommand{\\bY}{\\mathbf{Y}}$ $\\newcommand{\\bz}{\\mathbf{z}}$ $\\newcommand{\\bZ}{\\mathbf{Z}}$ $\\newcommand{\\bw}{\\mathbf{w}}$ $\\newcommand{\\bW}{\\mathbf{W}}$ $\\newcommand{\\pa}{\\mathrm{pa}}$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6bbfe170",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "## Modeling more complex dependencies"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c10a3e4",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "### Technique 1: Using conditional independence"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "86ddeeb2",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "### Technique 2: Marginalizing out an unobserved random variable"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "497a064a",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "**EXAMPLE:** **(Gaussian mixture model)** For $i=1,\\ldots,K$, let $\\bmu_i$ and $\\bSigma_i$ be the mean and covariance matrix of a multivariate Gaussian. Let $\\bpi \\in \\Delta_K$. A Gaussian Mixture Model (GMM) is obtained as follows: take $Y \\sim \\mathrm{Cat}(\\bpi)$ and\n",
    "\n",
    "$$\n",
    "\\bX|\\{Y=i\\} \\sim N_d(\\bmu_i, \\bSigma_i).\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b940397e",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "Its probability density function (PDF) takes the form\n",
    "\n",
    "$$\n",
    "f_\\bX(\\bx)\n",
    "= \\sum_{i=1}^K \\pi_i \\frac{1}{(2\\pi)^{d/2} \\,|\\bSigma_i|^{1/2}}\n",
    "\\exp\\left(-\\frac{1}{2}(\\mathbf{x} - \\bmu_i)^T \\bSigma_i^{-1} (\\bx - \\bmu_i)\\right).\n",
    "$$\n",
    "\n",
    "This is illustrated below."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "170ea468",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "We plot the density for means $\\bmu_1 = (-2,-2)$ and $\\bmu_2 = (2,2)$ and covariance matrices\n",
    "\n",
    "$$\n",
    "\\bSigma_1 = \\begin{bmatrix} 1.0 & 0 \\\\ 0 & 1.0 \\end{bmatrix}\n",
    "\\quad \\text{and} \\quad \n",
    "\\bSigma_2 = \\begin{bmatrix} \\sigma_1^2 & \\rho \\sigma_1 \\sigma_2 \\\\ \\rho \\sigma_1 \\sigma_2 & \\sigma_2^2 \\end{bmatrix}\n",
    "$$\n",
    "\n",
    "where $\\sigma_1 = 1.5$, $\\sigma_2 = 0.5$ and $\\rho = -0.75$. The mixing weights are $\\pi_1 = 0.25$ and $\\pi_2 = 0.75$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01bb2ec6",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "from scipy.stats import multivariate_normal\n",
    "\n",
    "def gmm2_pdf(X, Y, mean1, cov1, pi1, mean2, cov2, pi2):\n",
    "    xy = np.stack([X.flatten(), Y.flatten()], axis=-1)\n",
    "    Z1 = multivariate_normal.pdf(\n",
    "        xy, mean=mean1, cov=cov1).reshape(X.shape) \n",
    "    Z2 = multivariate_normal.pdf(\n",
    "        xy, mean=mean2, cov=cov2).reshape(X.shape) \n",
    "    return pi1 * Z1 + pi2 * Z2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92bc9355",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "start_point = 6\n",
    "stop_point = 6\n",
    "num_samples = 100\n",
    "points = np.linspace(-start_point, stop_point, num_samples)\n",
    "X, Y = np.meshgrid(points, points)\n",
    "\n",
    "mean1 = np.array([-2., -2.])\n",
    "cov1 = np.array([[1., 0.], [0., 1.]])\n",
    "pi1 = 0.5\n",
    "mean2 = np.array([2., 2.])\n",
    "cov2 = np.array([[1.5 ** 2., -0.75 * 1.5 * 0.5], \n",
    "                 [-0.75 * 1.5 * 0.5, 0.5 ** 2.]])\n",
    "pi2 = 0.5\n",
    "Z = gmm2_pdf(X, Y, mean1, cov1, pi1, mean2, cov2, pi2)\n",
    "mmids.make_surface_plot(X, Y, Z)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57ef5626",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "We could also consider continuous mixing distributions. Several standard approaches, e.g. probabilistic PCA, factor analysis and independent components analysis, use such mixtures. See [[Bis](https://www.microsoft.com/en-us/research/uploads/prod/2006/01/Bishop-Pattern-Recognition-and-Machine-Learning-2006.pdf), Chapter 12]. $\\lhd$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7b987e8",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "**NUMERICAL CORNER:** In Numpy, as we have seen before, the module [numpy.random](https://numpy.org/doc/stable/reference/random/index.html) also provides a way to sample from mixture models by using [numpy.random.Generator.choice](https://numpy.org/doc/stable/reference/random/generated/numpy.random.Generator.choice.html). This is better seen on an example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f375a11a",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "outputs": [],
   "source": [
    "seed = 535 # setting the seed\n",
    "rng = np.random.default_rng(seed) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4827670",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "outputs": [],
   "source": [
    "# constructs a mixture of three normal distributions,\n",
    "# with prior probabilities [0.2, 0.5, 0.3]\n",
    "mus = np.array([-2.0, 0.0, 3.0])\n",
    "sigmas = np.array([1.2, 1.0, 2.5])\n",
    "coeffs = np.array([0.2, 0.5, 0.3])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f788bb2",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "outputs": [],
   "source": [
    "N = 5 # number of samples\n",
    "which = rng.choice(3,N,p=coeffs)\n",
    "rng.normal(mus[which],sigmas[which])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f61bbe18",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "More generally, we consider mixtures of multivariate Gaussians. We chage the notation slightly to track Python's indexing. For $i=0,1$, we have a mean $\\bmu_i \\in \\mathbb{R}^d$ and a positive definite covariance matrix $\\bSigma_i \\in \\mathbb{R}^{d \\times d}$. We also have mixture weights $\\phi_0, \\phi_1 \\in (0,1)$ such that $\\phi_0 + \\phi_1 = 1$. Suppose we want to generate a total of $n$ samples."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f44aa7ca",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "For each sample $j=1,\\ldots, n$, independently from everything else:\n",
    "\n",
    "1. We first pick a component $i \\in \\{0,1\\}$ at random according to the mixture weights, that is, $i=0$ is chosen with probability $\\phi_0$ and $i=1$ is chosen with probability $\\phi_1$.\n",
    "\n",
    "2. We generate a sample $\\bX_j = (X_{j,1},\\ldots,X_{j,d})$ according to a multivariate Gaussian with mean $\\bmu_i$ and covariance $\\bSigma_i$. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe8b2925",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "This is straightforward to implement by using again [`numpy.random.Generator.choice`](https://numpy.org/doc/stable/reference/random/generated/numpy.random.Generator.choice.html) to choose the component of each sample and [`numpy.random.Generator.multivariate_normal`](https://numpy.org/doc/stable/reference/random/generated/numpy.random.Generator.multivariate_normal.html) to generate multivariate Gaussians. For convenience, we will stack the means and covariances into one array with a new dimension. So, for instance, the covariance matrices will now be in a 3d-array, that is, an array with $3$ indices. The first index corresponds to the component (here $0$ or $1$)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ad23a33",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "outputs": [],
   "source": [
    "d = 2\n",
    "sigma0 = np.outer(np.array([2., 2.]), np.array([2., 2.])) \n",
    "sigma0 += np.outer(np.array([-0.5, 0.5]), np.array([-0.5, 0.5]))\n",
    "sigma1 = 2 * np.identity(d)\n",
    "sigma = np.stack((sigma0,sigma1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56fd6eef",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "outputs": [],
   "source": [
    "print(sigma[0,:,:])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41d532bd",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "outputs": [],
   "source": [
    "print(sigma[1,:,:])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a63b801",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "Here is a visualization borrowed from the [TensorFlow tutorial](https://www.tensorflow.org/guide/tensor#basics)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1f528d59",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "**Figure:** Three matrices ([Source](https://www.tensorflow.org/guide/tensor#basics))\n",
    "\n",
    "![Three matrices](https://www.tensorflow.org/static/guide/images/tensor/3-axis_numpy.png)\n",
    "\n",
    "$\\bowtie$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4fe70f72",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "**Figure:** Three matrices stacked into a 3d-array ([Source](https://www.tensorflow.org/guide/tensor#basics))\n",
    "\n",
    "\n",
    "![Three matrices stacked into a tensor](https://www.tensorflow.org/static/guide/images/tensor/3-axis_front.png)\n",
    "\n",
    "$\\bowtie$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f3428ff",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "The code is the following. It returns an `d` by `n` array `X`, where each row is a sample from a 2-component Gaussian mixture. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d344077b",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "outputs": [],
   "source": [
    "def gmm2(d, n, phi0, phi1, mu0, sigma0, mu1, sigma1):\n",
    "    \n",
    "    # merge components into tensors\n",
    "    phi = np.stack((phi0, phi1))\n",
    "    mu = np.stack((mu0, mu1))\n",
    "    sigma = np.stack((sigma0,sigma1))\n",
    "    \n",
    "    # initialization\n",
    "    X = np.zeros((n,d))\n",
    "    \n",
    "    # choose components of each data point, then generate samples\n",
    "    component = rng.choice(2, size=n, p=phi)\n",
    "    for i in range(n):\n",
    "        X[i,:] = rng.multivariate_normal(\n",
    "            mu[component[i],:],\n",
    "            sigma[component[i],:,:])\n",
    "    \n",
    "    return X"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d1b06e96",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "Let us try it with following parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8f2e153",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "outputs": [],
   "source": [
    "d, n, w = 2, 1000, 5.\n",
    "phi0 = 0.9\n",
    "phi1 = 0.1\n",
    "mu0 = np.concatenate(([w], np.zeros(d-1)))\n",
    "mu1 = np.concatenate(([-w], np.zeros(d-1)))\n",
    "sigma0 = np.outer(np.array([2., 2.]), np.array([2., 2.])) \n",
    "sigma0 += np.outer(np.array([-0.5, 0.5]), np.array([-0.5, 0.5]))\n",
    "sigma1 = 2 *np.identity(d)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "946cc129",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "outputs": [],
   "source": [
    "X = gmm2(d, n, phi0, phi1, mu0, sigma0, mu1, sigma1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3856e6f9",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "plt.scatter(X[:,0], X[:,1], marker='x')\n",
    "plt.axis('equal')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1793bf46",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "$\\unlhd$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ddb02d70",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "### Example 1: Naive Bayes"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d004e8f",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "**NUMERICAL CORNER:** We implement the Naive Bayes model. We use [Laplace smoothing](https://en.wikipedia.org/wiki/Additive_smoothing) to avoid overfitting issues."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f9128b0",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "We use a simple example from [Towards Data Science](https://towardsdatascience.com/all-about-naive-bayes-8e13cef044cf):\n",
    "\n",
    "> **Example:** let’s say we have data on 1000 pieces of fruit. The fruit being a Banana, Orange or some other fruit and imagine we know 3 features of each fruit, whether it’s long or not, sweet or not and yellow or not, as displayed in the table below.\n",
    "\n",
    "![Table](./figs/1*B_uXqox7nHfwTa1HH4Fc_A.png)\n",
    "\n",
    "> [...] Which should provide enough evidence to predict the class of another fruit as it’s introduced."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "772f4f2d",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "We encode the data into a table, where the rows are the classes and the columns are the features. The entries are the corresponding $N_{k,m}$s. In addition we provide the vector $N_k$, which is the last column above, and the value $N$, which is the sum of the entries of $N_k$. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d33840ec",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "N_km = np.array([[400., 350., 450.],\n",
    "                 [0., 150., 300.],\n",
    "                 [100., 150., 50.]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c56c344",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "def nb_fit_table(N_km, alpha=1., beta=1.):\n",
    "    \n",
    "    K, M = N_km.shape\n",
    "    N_k = np.sum(N_km,axis=-1)\n",
    "    N = np.sum(N_k)\n",
    "    \n",
    "    # MLE for pi_k's\n",
    "    pi_k = (N_k+alpha) / (N+K*alpha)\n",
    "    \n",
    "    # MLE for p_km's\n",
    "    p_km = (N_km+beta) / (N_k[:,None]+2*beta)\n",
    "\n",
    "    return pi_k, p_km"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a524948f",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "We run it on our simple dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c863b873",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "pi_k, p_km = nb_fit_table(N_km)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8fb09e64",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "print(pi_k)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76e20ea3",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "print(p_km)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44be9fb2",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "Continuing on with our previous example:\n",
    "\n",
    "> So let’s say we’re given the features of a piece of fruit and we need to predict the class. If we’re told that the additional fruit is Long, Sweet and Yellow, we can classify it using the [prediction] formula and subbing in the values for each outcome, whether it’s a Banana, an Orange or Other Fruit. The one with the highest probability (score) being the winner."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c531f4c",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "The next function computes the negative logarithm of $\\pi_k \\prod_{m=1}^M p_{k,m}^{x_m} (1-p_{k,m})^{1-x_m}$, that is, the score of $k$, and outputs a $k$ achieving the minimum score."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71d76c71",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "def nb_predict(pi_k, p_km, x, label_set):\n",
    "   \n",
    "    K = len(pi_k)\n",
    "    \n",
    "    # Computing the score for each k\n",
    "    score_k = np.zeros(K)\n",
    "    for k in range(K):\n",
    "       \n",
    "        score_k[k] -= np.log(pi_k[k])\n",
    "        score_k[k] -= np.sum(x * np.log(p_km[k,:]) \n",
    "                               + (1 - x)*np.log(1 - p_km[k,:]))\n",
    "    \n",
    "    # Computing the minimum\n",
    "    argmin = np.argmin(score_k, axis=0)\n",
    "    minscr = np.max(score_k, axis=0)\n",
    "\n",
    "    return label_set[argmin]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec9ba3a2",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "We run it on our dataset with the additional fruit from the quote above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7a52e8e",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "label_set = ['Banana', 'Orange', 'Other']\n",
    "x = np.array([1., 1., 1.])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa2ecfb9",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "nb_predict(pi_k, p_km, x, label_set)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6e4b06db",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "$\\unlhd$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cc5ebc79",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "### Example 2: Mixtures of multivariate Bernoullis and the EM algorithm"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c0a255e",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "**NUMERICAL CORNER:** We implement the EM algorithm for mixtures of multivariate Bernoullis. For this purpose, we adapt our previous Naive Bayes routines. We also allow for the possibility of using Laplace smoothing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5429273",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "def responsibility(pi_k, p_km, x):\n",
    "   \n",
    "    K = len(pi_k)\n",
    "        \n",
    "    # Computing the score for each k\n",
    "    score_k = np.zeros(K)\n",
    "    for k in range(K):\n",
    "       \n",
    "        score_k[k] -= np.log(pi_k[k])\n",
    "        score_k[k] -= np.sum(x*np.log(p_km[k,:]) \n",
    "                             + (1 - x)*np.log(1 - p_km[k,:]))\n",
    "    \n",
    "    # Computing responsibilities for each k\n",
    "    r_k = np.exp(-score_k)/(np.sum(np.exp(-score_k)))\n",
    "        \n",
    "    return r_k"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "113f3285",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "def update_parameters(eta_km, eta_k, eta, alpha, beta):\n",
    "\n",
    "    K = len(eta_k)\n",
    "    \n",
    "    # MLE for pi_k's\n",
    "    pi_k = (eta_k+alpha) / (eta+K*alpha)\n",
    "    \n",
    "    # MLE for p_km's\n",
    "    p_km = (eta_km+beta) / (eta_k[:,None]+2*beta)\n",
    "\n",
    "    return pi_k, p_km"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "840439e0",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "We implement the E and M Step next."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b68ac7f",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "def em_bern(X, K, pi_0, p_0, maxiters = 10, alpha=0., beta=0.):\n",
    "    \n",
    "    n, M = X.shape\n",
    "    pi_k = pi_0\n",
    "    p_km = p_0\n",
    "    \n",
    "    for _ in range(maxiters):\n",
    "    \n",
    "        # E Step\n",
    "        r_ki = np.zeros((K,n))\n",
    "        for i in range(n):\n",
    "            r_ki[:,i] = responsibility(pi_k, p_km, X[i,:])\n",
    "        \n",
    "        # M Step     \n",
    "        eta_km = np.zeros((K,M))\n",
    "        eta_k = np.sum(r_ki, axis=-1)\n",
    "        eta = np.sum(eta_k)\n",
    "        for k in range(K):\n",
    "            for m in range(M):\n",
    "                eta_km[k,m] = np.sum(X[:,m] * r_ki[k,:]) \n",
    "        pi_k, p_km = update_parameters(\n",
    "            eta_km, eta_k, eta, alpha, beta)\n",
    "        \n",
    "    return pi_k, p_km   "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "29d3eb0e",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "We test the algorithm on a very simple dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "86a7fd3f",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "X = np.array([[1., 1., 1.],\n",
    "              [1., 1., 1.],\n",
    "              [1., 1., 1.],\n",
    "              [1., 0., 1.],\n",
    "              [0., 1., 1.],\n",
    "              [0., 0., 0.],\n",
    "              [0., 0., 0.],\n",
    "              [0., 0., 1.]])\n",
    "n, M = X.shape\n",
    "K = 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "649e17c1",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "seed = 535\n",
    "rng = np.random.default_rng(seed)\n",
    "pi_0 = np.ones(K)/K\n",
    "p_0 = rng.random((K,M))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62969e3b",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "pi_k, p_km = em_bern(\n",
    "    X, K, pi_0, p_0, maxiters=100, alpha=0.01, beta=0.01)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dbedd784",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "print(pi_k)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f58bd59",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "print(p_km)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "37d9689d",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "We compute the probability that the vector $(0, 0, 1)$ is in each cluster."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c6d8d51",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "x_test = np.array([0., 0., 1.])\n",
    "responsibility(pi_k, p_km, x_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f8f1d9c7",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "To give a more involved example, we return to the MNIST dataset. There are two common ways to write a $2$. Let's see if a mixture of multivariate Bernoullis can find them. We load the dataset and extract the images labelled $2$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "093d3b28",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "# Download and load the MNIST dataset\n",
    "mnist = datasets.MNIST(root='./data', \n",
    "                       train=True, \n",
    "                       download=True, \n",
    "                       transform=transforms.ToTensor())\n",
    "\n",
    "# Convert the dataset to a PyTorch DataLoader\n",
    "train_loader = torch.utils.data.DataLoader(mnist, \n",
    "                                           batch_size=len(mnist), \n",
    "                                           shuffle=False)\n",
    "\n",
    "# Extract images and labels from the DataLoader\n",
    "imgs, labels = next(iter(train_loader))\n",
    "imgs = imgs.squeeze().numpy()\n",
    "labels = labels.numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74eddfda",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Filter out images with label 2\n",
    "mask = labels == 2\n",
    "imgs2 = imgs[mask]\n",
    "labels2 = labels[mask]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "452e74a5",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "Next, we transform the images into vectors and convert into black and white by rounding."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "122b7a47",
   "metadata": {},
   "outputs": [],
   "source": [
    "X = np.round(imgs2.reshape(len(imgs2), -1))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9784ecaa",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "We can convert back as follows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67d8788a",
   "metadata": {
    "scrolled": true,
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "plt.imshow(X[0,:].reshape((28,28)))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e4ee557f",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "In this example, the probabilities involved are very small and the responsibilities are close to $0$ or $1$. We use a variant, called hard EM, which replaces responsibilities with the one-hot encoding of the largest responsibility."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aebd6302",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "def hard_responsibility(pi_k, p_km, x):\n",
    "   \n",
    "    K = len(pi_k)\n",
    "        \n",
    "    # Computing the score for each k\n",
    "    score_k = np.zeros(K)\n",
    "    for k in range(K):\n",
    "       \n",
    "        score_k[k] -= np.log(pi_k[k])\n",
    "        score_k[k] -= np.sum(x*np.log(p_km[k,:]) \n",
    "                             + (1 - x)*np.log(1 - p_km[k,:]))\n",
    "    \n",
    "    # Computing responsibilities for each k\n",
    "    argmin = np.argmin(score_k, axis=0)\n",
    "    r_k = np.zeros(K)\n",
    "    r_k[argmin] = 1\n",
    "\n",
    "    return r_k"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "094f445e",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "def hard_em_bern(X, K, pi_0, p_0, maxiters = 10, alpha=0., beta=0.):\n",
    "    \n",
    "    n, M = X.shape\n",
    "    pi_k = pi_0\n",
    "    p_km = p_0\n",
    "    \n",
    "    for _ in range(maxiters):\n",
    "    \n",
    "        # E Step\n",
    "        r_ki = np.zeros((K,n))\n",
    "        for i in range(n):\n",
    "            r_ki[:,i] = hard_responsibility(pi_k, p_km, X[i,:])\n",
    "        \n",
    "        # M Step     \n",
    "        eta_km = np.zeros((K,M))\n",
    "        eta_k = np.sum(r_ki, axis=-1)\n",
    "        eta = np.sum(eta_k)\n",
    "        for k in range(K):\n",
    "            for m in range(M):\n",
    "                eta_km[k,m] = np.sum(X[:,m] * r_ki[k,:]) \n",
    "        pi_k, p_km = update_parameters(\n",
    "            eta_km, eta_k, eta, alpha, beta)\n",
    "        \n",
    "    return pi_k, p_km   "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b4b88d3",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "We run the algorithm with $2$ clusters. You may have to run it a few times to get a meaningful clustering."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6278bf14",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "n, M = X.shape\n",
    "K = 2\n",
    "pi_0 = np.ones(K)/K\n",
    "p_0 = rng.random((K,M))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24db2a6c",
   "metadata": {
    "scrolled": false,
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "pi_k, p_km = hard_em_bern(\n",
    "    X, K, pi_0, p_0, maxiters=10, alpha=1., beta=1.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51611a42",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "print(pi_k)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ecf43726",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "We vizualize the two uncovered clusters by rendering their means as an image. Here is one of them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "625bafb9",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "plt.imshow(p_km[0,:].reshape((28,28)))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b612994",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "Here is the other one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cbd5955e",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "plt.imshow(p_km[1,:].reshape((28,28)))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43593bdd",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "$\\unlhd$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c541fb33",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "$\\newcommand{\\P}{\\mathbb{P}}$ $\\newcommand{\\E}{\\mathbb{E}}$ $\\newcommand{\\S}{\\mathcal{S}}$ $\\newcommand{\\var}{\\mathrm{Var}}$ $\\newcommand{\\bmu}{\\boldsymbol{\\mu}}$ $\\newcommand{\\bSigma}{\\boldsymbol{\\Sigma}}$ $\\newcommand{\\btheta}{\\boldsymbol{\\theta}}$ $\\newcommand{\\bpi}{\\boldsymbol{\\pi}}$ $\\newcommand{\\indep}{\\perp\\!\\!\\!\\perp}$ $\\newcommand{\\bp}{\\mathbf{p}}$ $\\newcommand{\\bx}{\\mathbf{x}}$ $\\newcommand{\\bX}{\\mathbf{X}}$ $\\newcommand{\\by}{\\mathbf{y}}$ $\\newcommand{\\bY}{\\mathbf{Y}}$ $\\newcommand{\\bz}{\\mathbf{z}}$ $\\newcommand{\\bZ}{\\mathbf{Z}}$ $\\newcommand{\\bw}{\\mathbf{w}}$ $\\newcommand{\\bW}{\\mathbf{W}}$ $\\newcommand{\\bv}{\\mathbf{v}}$ $\\newcommand{\\bV}{\\mathbf{V}}$ $\\newcommand{\\bLambda}{\\boldsymbol{\\Lambda}}$ $\\newcommand{\\bu}{\\mathbf{u}}$ $\\newcommand{\\bU}{\\mathbf{U}}$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8d168519",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "## Linear-Gaussian models and Kalman filtering"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36fedfe1",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "### Multivariate Gaussians: marginals and conditionals"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "01816f8a",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "### Kalman filter"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12199db4",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "### Location tracking"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96e8916b",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "We apply Kalman filtering to location tracking. Returning to our cyborg corgi example, we imagine that we get noisy observations about its successive positions in a park. (Think of GPS measurements.) We seek to get a better estimate of its location using the method above. See for example this [blog post](https://towardsdatascience.com/optimal-estimation-algorithms-kalman-and-particle-filters-be62dcb5e83) on location tracking."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c13246d",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "We model the true location as a linear-Gaussian system over the 2d position $(z_{1,t}, z_{2,t})_t$ and velocity $(\\dot{z}_{1,t}, \\dot{z}_{2,t})_t$ sampled at $\\Delta$ intervals of time. Formally,\n",
    "\n",
    "$$\n",
    "\\bX_t = (z_{1,t}, z_{2,t}, \\dot{z}_{1,t}, \\dot{z}_{2,t}),\n",
    "\\quad\n",
    "F = \\begin{pmatrix}\n",
    "1 & 0 & \\Delta & 0\\\\\n",
    "0 & 1 & 0 & \\Delta\\\\\n",
    "0 & 0 & 1 & 0\\\\\n",
    "0 & 0 & 0 & 1\n",
    "\\end{pmatrix},\n",
    "$$\n",
    "\n",
    "so the unobserved dynamics are\n",
    "\n",
    "$$\n",
    "\\begin{pmatrix}\n",
    "z_{1,t+1}\\\\ \n",
    "z_{2,t+1}\\\\ \n",
    "\\dot{z}_{1,t+1}\\\\ \n",
    "\\dot{z}_{2,t+1}\n",
    "\\end{pmatrix}\n",
    "= \\bX_{t+1}\n",
    "= F \\,\\bX_t + \\bW_t\n",
    "= \n",
    "\\begin{pmatrix}\n",
    "z_{1,t} + \\Delta \\dot{z}_{1,t} + W_{1,t}\\\\\n",
    "z_{2,t} + \\Delta \\dot{z}_{2,t} + W_{2,t}\\\\\n",
    "\\dot{z}_{1,t} + \\dot{W}_{1,t}\\\\\n",
    "\\dot{z}_{2,t} + \\dot{W}_{2,t}\n",
    "\\end{pmatrix}\n",
    "$$\n",
    "\n",
    "where the $\\bW_t = (W_{1,t}, W_{2,t}, \\dot{W}_{1,t}, \\dot{W}_{2,t}) \\sim N_{d_0}(\\mathbf{0}, Q)$ with $Q$ known.\n",
    "\n",
    "In words, the velocity is unchanged, up to Gaussian perturbation. The position changes proportionally to the velocity in the corresponding dimension."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52e47bfa",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "The observations $(\\tilde{z}_{1,t}, \\tilde{z}_{2,t})_t$ are modeled as\n",
    "\n",
    "$$\n",
    "\\bY_t = (\\tilde{z}_{1,t}, \\tilde{z}_{2,t}),\n",
    "\\quad \n",
    "H = \\begin{pmatrix}\n",
    "1 & 0 & 0 & 0\\\\\n",
    "0 & 1 & 0 & 0\\\\\n",
    "\\end{pmatrix}.\n",
    "$$\n",
    "\n",
    "so the observed process satisfies\n",
    "\n",
    "$$\n",
    "\\begin{pmatrix}\n",
    "\\tilde{z}_{1,t}\\\\ \n",
    "\\tilde{z}_{2,t}\n",
    "\\end{pmatrix}\n",
    "= \\bY_t\n",
    "= H\\,\\bX_t + \\bV_t\n",
    "= \n",
    "\\begin{pmatrix}\n",
    "\\tilde{z}_{1,t} + \\tilde{V}_{1,t}\\\\ \n",
    "\\tilde{z}_{2,t} + \\tilde{V}_{2,t}\n",
    "\\end{pmatrix}\n",
    "$$\n",
    "\n",
    "where the $\\bV_t = (\\tilde{V}_{1,t}, \\tilde{V}_{2,t}) \\sim N_d(\\mathbf{0}, R)$ with $R$ known. \n",
    "\n",
    "In words, we only observe the positions, up to Gaussian noise."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45a22342",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "**Implementing the Kalman filter** We implement the Kalman filter as described above with known covariance matrices. We take $\\Delta = 1$ for simplicity. The code is adapted from [[Mur](https://github.com/probml)]."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42286c86",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "We will test Kalman filtering on a simulated path drawn from the linear-Gaussian model above. The following function creates such a path and its noisy observations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31aa3023",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    }
   },
   "outputs": [],
   "source": [
    "seed = 535\n",
    "rng = np.random.default_rng(seed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15b6cb2c",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "def lgSamplePath(ss, os, F, H, Q, R, init_mu, init_Sig, T):\n",
    "    x = np.zeros((ss,T)) \n",
    "    y = np.zeros((os,T))\n",
    "\n",
    "    x[:,0] = rng.multivariate_normal(init_mu, init_Sig)\n",
    "    for t in range(1,T):\n",
    "        x[:,t] = rng.multivariate_normal(F @ x[:,t-1],Q)\n",
    "        y[:,t] = rng.multivariate_normal(H @ x[:,t],R)\n",
    "    \n",
    "    return x, y"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61a4c249",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "Here is an example. In the plots, the red dots are the noisy observations and the green crosses are the unobserved true path."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1050727e",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "ss = 4 # state size\n",
    "os = 2 # observation size\n",
    "F = np.array([[1., 0., 1., 0.],\n",
    "              [0., 1., 0., 1.],\n",
    "              [0., 0., 1., 0.],\n",
    "              [0., 0., 0., 1.]]) \n",
    "H = np.array([[1., 0., 0., 0.],\n",
    "              [0., 1, 0., 0.]])\n",
    "Q = 0.1 * np.diag(np.ones(ss))\n",
    "R = 10 * np.diag(np.ones(os))\n",
    "init_mu = np.array([0., 0., 1., 1.])\n",
    "init_Sig = 1 * np.diag(np.ones(ss))\n",
    "T = 50\n",
    "x, y = lgSamplePath(ss, os, F, H, Q, R, init_mu, init_Sig, T)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3523e6ee",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "In the next plot (and throughout this section), the red dots are the noisy observations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12862d6d",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "plt.plot(y[0,:], y[1,:], \n",
    "         marker='o', c='r', linestyle='dotted')\n",
    "plt.xlim((np.min(y[0,:])-5, np.max(y[0,:])+5)) \n",
    "plt.ylim((np.min(y[1,:])-5, np.max(y[1,:])+5))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b8ffa85",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "In the next plot (and throughout this section), the green crosses are the unobserved true path."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2739e4a4",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "plt.plot(x[0,:], x[1,:], \n",
    "         marker='x', c='g', linestyle='dashed')\n",
    "plt.xlim((np.min(x[0,:])-5, np.max(x[0,:])+5)) \n",
    "plt.ylim((np.min(x[1,:])-5, np.max(x[1,:])+5))\n",
    "plt.scatter(y[0,:], y[1,:], \n",
    "            c='r', alpha=0.5)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "afc1e86c",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "The following function implements the Kalman filter. The full recursion is broken up into several steps. We use [`numpy.linalg.inv`](https://numpy.org/doc/stable/reference/generated/numpy.linalg.inv.html) to compute the Kalman gain matrix. Below, `mu_pred` is $F \\bmu_{t-1}$ and `Sig_pred` is $P_{t-1} = F \\bSigma_{t-1} F^T + Q$, which are the mean vector and covariance matrix of $\\bX_{t}$ given $\\bY_{1:t-1}$ as computed in the *Predict* step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5150e3f9",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "def kalmanUpdate(ss, F, H, Q, R, y_t, mu_prev, Sig_prev):\n",
    "    mu_pred = F @ mu_prev\n",
    "    Sig_pred = F @ Sig_prev @ F.T + Q\n",
    "    e_t = y_t - H @ mu_pred\n",
    "    S = H @ Sig_pred @ H.T + R\n",
    "    Sinv = LA.inv(S)\n",
    "    K = Sig_pred @ H.T @ Sinv\n",
    "    mu_new = mu_pred + K @ e_t\n",
    "    Sig_new = (np.diag(np.ones(ss)) - K @ H) @ Sig_pred\n",
    "    return mu_new, Sig_new"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c30a2766",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "def kalmanFilter(ss, os, y, F, H, Q, R, init_mu, init_Sig, T):\n",
    "    mu = np.zeros((ss, T))\n",
    "    Sig = np.zeros((ss, ss, T))\n",
    "    mu[:,0] = init_mu\n",
    "    Sig[:,:,0] = init_Sig\n",
    "\n",
    "    for t in range(1,T):\n",
    "        mu[:,t], Sig[:,:,t] = kalmanUpdate(ss, F, H, Q, R, \n",
    "                                           y[:,t], mu[:,t-1], \n",
    "                                           Sig[:,:,t-1])\n",
    "\n",
    "    return mu, Sig"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bff45d62",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "We apply this to the location tracking example. The inferred states, or more precisely their estimated mean, are in blue. Note that we also inferred the velocity at each time point, but we are not plotting that information. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a52e1a4",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "init_mu = np.array([0., 0., 1., 1.])\n",
    "init_Sig = 1 * np.diag(np.ones(ss))\n",
    "mu, Sig = kalmanFilter(ss, os, y, F, H, Q, R, init_mu, init_Sig, T)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d74d590c",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "plt.scatter(x[0,:], x[1,:], \n",
    "            marker='x', c='g', alpha=0.5)\n",
    "plt.xlim((np.min(x[0,:])-5, np.max(x[0,:])+5)) \n",
    "plt.ylim((np.min(x[1,:])-5, np.max(x[1,:])+5))\n",
    "plt.scatter(y[0,:], y[1,:], \n",
    "            c='r', alpha=0.5)\n",
    "plt.plot(mu[0,:], mu[1,:], \n",
    "         c='b', marker='s', linewidth=1)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "05d4ef6b",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "To quantify the improvement in the inferred means compared to the observations, we compute the mean squared error in both cases."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1c1cfe7",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "dobs = x[0:1,:] - y[0:1,:]\n",
    "mse_obs = np.sqrt(np.sum(dobs**2))\n",
    "print(mse_obs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e0e32ef",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "dfilt = x[0:1,:] - mu[0:1,:]\n",
    "mse_filt = np.sqrt(np.sum(dfilt**2))\n",
    "print(mse_filt)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98b92894",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "We indeed observe a substantial reduction."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "31d0cc88",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "**Missing data** We can also allow for the possibility that some observations are missing. Imagine for instance losing GPS signal while going through a tunnel. The recursions above are still valid, with the only modification that the *Update* equations involving $\\bY_t$ are dropped at those times $t$ where there is no observation. In Numpy, we can use [`NaN`](https://numpy.org/doc/stable/reference/constants.html#numpy.nan) to indicate the lack of observation. (Alternatively, one can use the [numpy.ma](https://numpy.org/doc/stable/reference/maskedarray.generic.html) module.) "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "370a0166",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "We use a same sample path as above, but mask observations at times $t=10,\\ldots,20$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92f90933",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "ss = 4\n",
    "os = 2\n",
    "F = np.array([[1., 0., 1., 0.],\n",
    "              [0., 1., 0., 1.],\n",
    "              [0., 0., 1., 0.],\n",
    "              [0., 0., 0., 1.]]) \n",
    "H = np.array([[1., 0., 0., 0.],\n",
    "              [0., 1, 0., 0.]])\n",
    "Q = 0.01 * np.diag(np.ones(ss))\n",
    "R = 10 * np.diag(np.ones(os))\n",
    "init_mu = np.array([0., 0., 1., 1.])\n",
    "init_Sig = Q\n",
    "T = 30\n",
    "x, y = lgSamplePath(ss, os, F, H, Q, R, init_mu, init_Sig, T)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a6ce34f",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "for i in range(10,20):\n",
    "    y[0,i] = np.nan\n",
    "    y[1,i] = np.nan"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7e85ab56",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "Here is the sample we are aiming to infer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0bb6dc28",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "plt.plot(x[0,:], x[1,:], \n",
    "         marker='x', c='g', linestyle='dashed', alpha=0.5)\n",
    "plt.xlim((np.min(x[0,:])-5, np.max(x[0,:])+5)) \n",
    "plt.ylim((np.min(x[1,:])-5, np.max(x[1,:])+5))\n",
    "plt.plot(y[0,:], y[1,:], \n",
    "         c='r', marker='o', linestyle='dotted')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0a36f64c",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "We modify the recursion accordingly, that is, skip the *Update* step when there is no observation to use for the update."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ebb5e631",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "def kalmanUpdate(ss, F, H, Q, R, y_t, mu_prev, Sig_prev):\n",
    "    mu_pred = F @ mu_prev\n",
    "    Sig_pred = F @ Sig_prev @ F.T + Q\n",
    "    if np.isnan(y_t[0]) or np.isnan(y_t[1]):\n",
    "        return mu_pred, Sig_pred\n",
    "    else:\n",
    "        e_t = y_t - H @ mu_pred\n",
    "        S = H @ Sig_pred @ H.T + R\n",
    "        Sinv = LA.inv(S)\n",
    "        K = Sig_pred @ H.T @ Sinv\n",
    "        mu_new = mu_pred + K @ e_t\n",
    "        Sig_new = (np.diag(np.ones(ss)) - K @ H) @ Sig_pred\n",
    "        return mu_new, Sig_new"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed4f92ad",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "init_mu = np.array([0., 0., 1., 1.])\n",
    "init_Sig = 1 * np.diag(np.ones(ss))\n",
    "mu, Sig = kalmanFilter(ss, os, y, F, H, Q, R, init_mu, init_Sig, T)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1ea68ad",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "plt.scatter(x[0,:], x[1,:], \n",
    "         marker='x', c='g', alpha=0.5)\n",
    "plt.xlim((np.min(x[0,:])-5, np.max(x[0,:])+5)) \n",
    "plt.ylim((np.min(x[1,:])-5, np.max(x[1,:])+5))\n",
    "plt.scatter(y[0,:], y[1,:], \n",
    "            c='r', alpha=0.5)\n",
    "plt.plot(mu[0,:], mu[1,:], \n",
    "         marker='s', linewidth=1)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f91ded86",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "$\\newcommand{\\P}{\\mathbb{P}}$ $\\newcommand{\\E}{\\mathbb{E}}$ $\\newcommand{\\S}{\\mathcal{S}}$ $\\newcommand{\\var}{\\mathrm{Var}}$ $\\newcommand{\\bmu}{\\boldsymbol{\\mu}}$ $\\newcommand{\\bSigma}{\\boldsymbol{\\Sigma}}$ $\\newcommand{\\btheta}{\\boldsymbol{\\theta}}$ $\\newcommand{\\bpi}{\\boldsymbol{\\pi}}$ $\\newcommand{\\indep}{\\perp\\!\\!\\!\\perp}$ $\\newcommand{\\bp}{\\mathbf{p}}$ $\\newcommand{\\bx}{\\mathbf{x}}$ $\\newcommand{\\bX}{\\mathbf{X}}$ $\\newcommand{\\by}{\\mathbf{y}}$ $\\newcommand{\\bY}{\\mathbf{Y}}$ $\\newcommand{\\bz}{\\mathbf{z}}$ $\\newcommand{\\bZ}{\\mathbf{Z}}$ $\\newcommand{\\bw}{\\mathbf{w}}$ $\\newcommand{\\bW}{\\mathbf{W}}$ $\\newcommand{\\bv}{\\mathbf{v}}$ $\\newcommand{\\bV}{\\mathbf{V}}$ $\\newcommand{\\Z}{\\mathcal{Z}}$ $\\newcommand{\\bh}{\\mathbf{h}}$ $\\newcommand{\\bb}{\\mathbf{b}}$ $\\newcommand{\\bc}{\\mathbf{c}}$ $\\newcommand{\\cE}{\\mathcal{E}}$ $\\newcommand{\\cP}{\\mathcal{P}}$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "395caf56",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "## Gibbs sampling with application to generating images"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3f6abc80",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "### Markov chain Monte Carlo (MCMC)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b12b566a",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "**Sampling from simple distributions** When $\\bpi$ is a standard distribution or $\\S$ is relatively small, this can be done efficiently by using a random number generator, as we have done previously."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5eeb9e11",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "**NUMERICAL CORNER:** Recall how this works. We first initialize the random number generator and use a `seed` for reproducibility."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53510a1f",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "seed = 535\n",
    "rng = np.random.default_rng(seed)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b75b3511",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "To generate, say $1000$, samples from a multivariate normal, say with mean $(0, 0)$ and covariance $\\begin{pmatrix}5 & 0\\\\0 & 1\\end{pmatrix}$, we use [`numpy.random.Generator.multivariate_normal`](https://numpy.org/doc/stable/reference/random/generated/numpy.random.Generator.multivariate_normal.html#numpy.random.Generator.multivariate_normal) as follows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0dad9bb5",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "mean = np.array([0., 0.])\n",
    "cov = np.array([[5., 0.], [0., 1.]])\n",
    "x, y = rng.multivariate_normal(mean, cov, 1000).T"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d588b8c2",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "Plotting the result we get:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ac73f90",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "plt.scatter(x, y, marker='x')\n",
    "plt.axis('equal')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e511d472",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "Computing the mean of each component we get:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "617c1a54",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "print(np.mean(x))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "78bea025",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "print(np.mean(y))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43c3c8ce",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "This is somewhat close to the expected answer: $(0,0)$. \n",
    "\n",
    "Using a larger number of samples, say $10,000$, gives a better result."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7b4957d",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "x, y = rng.multivariate_normal(mean, cov, 10000).T\n",
    "print(np.mean(x))\n",
    "print(np.mean(y))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8ab410b",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "$\\unlhd$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5cf3e072",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "Sampling from an arbitrary distribution on a finite set is also straightforward -- as long as the set is not too big. This can be done using [`numpy.random.Generator.choice`](https://numpy.org/doc/stable/reference/random/generated/numpy.random.Generator.choice.html). Borrowing the example from the documentation, the following:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d87030e0",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "aa_milne_arr = ['pooh', 'rabbit', 'piglet', 'christopher']\n",
    "rng.choice(aa_milne_arr, 5, p=[0.5, 0.1, 0.1, 0.3])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "35857cf4",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "generates $5$ samples from the set $\\S = \\{\\tt{pooh}, \\tt{rabbit}, \\tt{piglet}, \\tt{christopher}\\}$ with respective probabilities $0.5, 0.1, 0.1, 0.3$."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f02eef2b",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "But this may not be practical when the state space $\\S$ is very large. As an example, later in this section, we will learn a \"realistic\" distribution of handwritten digits. We will do so using the [MNIST dataset](https://en.wikipedia.org/wiki/MNIST_database), which we have encountered previously."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f4468643",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "We load it from PyTorch and turn the grayscale images (more precisely, each pixel is an integer between $0$ and $255$) into a black-and-white images by rounding the pixels (after dividing by $255$)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c7c9815c",
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "# Download and load the MNIST dataset\n",
    "mnist = datasets.MNIST(root='./data', \n",
    "                       train=True, \n",
    "                       download=True, \n",
    "                       transform=transforms.ToTensor())\n",
    "\n",
    "# Convert the dataset to a PyTorch DataLoader\n",
    "train_loader = torch.utils.data.DataLoader(mnist, \n",
    "                                           batch_size=len(mnist), \n",
    "                                           shuffle=False)\n",
    "\n",
    "# Extract images and labels from the DataLoader\n",
    "imgs, labels = next(iter(train_loader))\n",
    "imgs = imgs.squeeze().numpy()\n",
    "labels = labels.numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82fc9ea5",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "imgs = np.round(imgs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f87833f",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "The first image is the following."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b0d097b",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "plt.imshow(imgs[0], cmap=plt.cm.gray_r)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a14ef60a",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "It is $28 \\times 28$, so the total number of pixels is $784$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "110c1461",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "nx_pixels, ny_pixels = imgs[0].shape\n",
    "nx_pixels, ny_pixels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9dfc0779",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "n_pixels = nx_pixels * ny_pixels\n",
    "n_pixels"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "009b4929",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "To specify the a distribution over all possible black and white images of this size, we need in principle to assign a probability to a very large number of states. Our space here is $\\S = \\{0,1\\}^{784}$, imagining that $0$ encodes white and $1$ encodes black and that we have ordered the pixels in some arbitrary way. How big is this space?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5877fe6b",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "Answer: $2^{784}$.\n",
    "\n",
    "Or in base $10$, we compute $\\log_{10}(2^{784})$, which is:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f36a99d4",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "784 * np.log(2) / np.log(10)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "54d5777d",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "So a little more than $10^{236}$. \n",
    "\n",
    "This is much too large to naively plug into `rng.choice`!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "24500937",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "**EXAMPLE:** Suppose $\\S = \\{1,\\cdots, n\\} = [n]$ for some positive integer $n$ and $\\bpi$ is proportional to a Poisson distribution with mean $\\lambda > 0$. That is, \n",
    "\n",
    "$$\n",
    "\\pi_i = C e^{-\\lambda} \\frac{\\lambda^i}{i!}, \\quad \\forall i \\in \\S\n",
    "$$\n",
    "\n",
    "for some constant $C$ chosen so that $\\sum_{i=1}^{n} \\pi_i = 1$. Recall that we do not need to determine $C$ as it is enough to know the target distribution up to a scaling factor by the previous remark. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3f498324",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "To apply Metropolis-Hastings, we need a proposal chain. Consider the following choice. For each $1 < i < n$, move to $i+1$ or $i-1$ with probability $1/2$ each. For $i=1$ (respectively $i = n$), move to $2$ (respectively $n-1$) with probability $1/2$, otherwise stay where you are. For instance, if $n = 4$, then \n",
    "\n",
    "$$\n",
    "Q\n",
    "=\n",
    "\\begin{pmatrix}\n",
    "1/2 & 1/2 & 0 & 0\\\\\n",
    "1/2 & 0 & 1/2 & 0\\\\\n",
    "0 & 1/2 & 0 & 1/2\\\\\n",
    "0 & 0 & 1/2 & 1/2\n",
    "\\end{pmatrix},\n",
    "$$\n",
    "\n",
    "which is indeed a stochastic matrix. It is also symmetric, so it does not enter into the acceptance probability by the previous remark."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48ff60d9",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "To compute the acceptance probability, we only need to consider pairs of adjacent integers as they are the only one that have non-zero probability under $Q$. Consider state $1 < i < n$. Observe that\n",
    "\n",
    "$$\n",
    "\\frac{\\pi_{i+1}}{\\pi_{i}}\n",
    "= \\frac{C e^{-\\lambda} \\lambda^{i+1}/(i+1)!}{C e^{-\\lambda} \\lambda^{i}/i!}\n",
    "= \\frac{\\lambda}{i+1}\n",
    "$$\n",
    "\n",
    "so a move to $i+1$ happens with probability\n",
    "\n",
    "$$\n",
    "\\frac{1}{2} \\min\\left\\{1, \\frac{\\lambda}{i+1}\\right\\},\n",
    "$$\n",
    "\n",
    "where the $1/2$ factor from the proposal distribution.\n",
    "Similarly, it can be checked (try it!) that a move to $i-1$ occurs with probability\n",
    "\n",
    "$$\n",
    "\\frac{1}{2} \\min\\left\\{1, \\frac{i}{\\lambda}\\right\\}.\n",
    "$$\n",
    "\n",
    "And we stay at $i$ with probability $1 - \\frac{1}{2} \\min\\left\\{1, \\frac{\\lambda}{i+1}\\right\\} - \\frac{1}{2} \\min\\left\\{1, \\frac{i}{\\lambda}\\right\\}$. (Why is this guaranteed to be a probability?) \n",
    "\n",
    "A similar formula applies to $i = 1, n$. (Try it!)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "56bd9f36",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "We are ready to apply Metropolis-Hastings.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9403bf9f",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "def mh_transition_poisson(lmbd, n):\n",
    "    P = np.zeros((n,n))\n",
    "    for idx in range(n):\n",
    "        i = idx + 1 # index starts at 0 rather than 1\n",
    "        if (i > 1 and i < n):\n",
    "            P[idx, idx+1] = (1/2) * np.min(np.array([1, lmbd/(i+1)]))\n",
    "            P[idx, idx-1] = (1/2) * np.min(np.array([1, i/lmbd]))\n",
    "            P[idx, idx] = 1 - P[idx, idx+1] - P[idx, idx-1]\n",
    "        elif i == 1:\n",
    "            P[idx, idx+1] = (1/2) * np.min(np.array([1, lmbd/(i+1)]))\n",
    "            P[idx, idx] = 1 - P[idx, idx+1]\n",
    "        elif i == n:\n",
    "            P[idx, idx-1] = (1/2) * np.min(np.array([1, i/lmbd]))\n",
    "            P[idx, idx] = 1 - P[idx, idx-1]\n",
    "    return P"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b32eb74",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "Take $\\lambda = 1$ and $n = 6$. We get the following transition matrix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb824cd6",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "lmbd = 1\n",
    "n = 6"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76f429e4",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "P = mh_transition_poisson(lmbd, n)\n",
    "print(P)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a3b5bdae",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "We use our simulator from a previous chapter. We start from the uniform distribution and take $100$ steps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96c5df2a",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "mu = np.ones(n) / n\n",
    "T = 100\n",
    "X = mmids.SamplePath(mu, P, T)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "825062a1",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "Our sample is the final state of the trajectory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "443851b3",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "X[T]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "26f30edd",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "We repeat $1000$ times and plot the frequencies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7e65749",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "N_samples = 1000 # number of repetitions\n",
    "\n",
    "freq_z = np.zeros(n) # init of frequencies sampled\n",
    "for i in range(N_samples):\n",
    "    X = mmids.SamplePath(mu, P, T)\n",
    "    freq_z[int(X[T])-1] += 1 # adjust for index starting at 0\n",
    "    \n",
    "freq_z = freq_z/N_samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0d49066",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "plt.bar(range(1,n+1),freq_z)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "86def90d",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "If we increase the parameter $\\lambda$ (which is not quite the mean; why?), we expect the sampled distribution to shift to the right. We must recompute the transition matrix first."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "321a4341",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "lmbd = 10\n",
    "P = mh_transition_poisson(lmbd, n)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1a84f76",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "freq_z = np.zeros(n) # init of frequencies sampled\n",
    "for i in range(N_samples):\n",
    "    X = mmids.SamplePath(mu, P, T)\n",
    "    freq_z[int(X[T])-1] += 1 # adjust for index starting at 0\n",
    "    \n",
    "freq_z = freq_z/N_samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a3bb783",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "plt.bar(range(1,n+1),freq_z)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7980f961",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "**TRY IT!** Redo the simulations, but this time implement a general Metropolis-Hastings algorithm rather than specifying the transition matrix directly. That is, implement the algorithm for an arbitrary $\\bpi$ and $Q$. Assume the state space is $[n]$. ([Open in Colab](https://colab.research.google.com/github/MMiDS-textbook/MMiDS-textbook.github.io/blob/main/just_the_code/roch_mmids_chap_prob_notebook.ipynb))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "671406fd",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "$\\lhd$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cbd14de3",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "### Gibbs sampling"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "83c48e29",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "**An example: restricted Boltzmann machines (RBM)** We implement the Gibbs sampler on a specific probabilistic model, a so-called restricted Boltzmann machine (RBM), and apply it to the generation of random images from a \"realistic\" distribution. For more on Boltzmann machines, including their restricted and deep versions, see [here](https://en.wikipedia.org/wiki/Boltzmann_machine). We will not describe them in great details here, but only use them as an example of a complex distribution."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c9c36eb",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "*Probabilistic model:* An RBM has $m$ visible units (i.e., observed variables) and $n$ hidden units (i.e., hidden variables). It is represented by a complete bipartite graph between the two. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "90cfb47d",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "Visible unit $i$ is associated a variable $v_i$ and hidden unit $j$ is associated a variable $h_j$. We define the corresponding vectors $\\bv = (v_1,\\ldots,v_m)$ and $\\bh = (h_1,\\ldots,h_n)$. For our purposes, it will suffice to assume that $\\bv \\in \\{0,1\\}^m$ and $\\bh \\in \\{0,1\\}^n$. These are referred to as binary units."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1f6a865b",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "The probabilistic model has a number of parameters. Each visible unit $i$ has an offset $b_i \\in \\mathbb{R}$ and each hidden unit $j$ has an offset $c_j \\in \\mathbb{R}$. We write $\\bb = (b_1,\\ldots,b_m)$ and $\\bc = (c_1,\\ldots,c_n)$ for the offset vectors. For each pair $(i,j)$ of visible and hidden units (or, put differently, for each edge in the complete bipartite graph), there is a weight $w_{i,j} \\in \\mathbb{R}$. We write $W = (w_{i,j})_{i,j=1}^{m,n}$ for the weight matrix."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "92da4947",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "To define the probability distribution, we need the so-called [energy](https://en.wikipedia.org/wiki/Energy-based_model) (as you may have guessed, this terminology comes from related models in [physics](https://en.wikipedia.org/wiki/Boltzmann_distribution)): for $\\bv \\in \\{0,1\\}^m$ and $\\bh \\in \\{0,1\\}^n$,\n",
    "\n",
    "\\begin{align*}\n",
    "\\cE(\\bv, \\bh)\n",
    "&= \n",
    "- \\bv^T W \\bh\n",
    "- \\bb^T \\bv\n",
    "- \\bc^T \\bh\\\\\n",
    "&= \n",
    "- \\sum_{i=1}^m \\sum_{j=1}^n w_{i,j} v_i h_j\n",
    "- \\sum_{i=1}^m b_i v_i\n",
    "- \\sum_{j=1}^n c_j h_j.\n",
    "\\end{align*}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "72c563bc",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "The joint distribution of $\\bv$ and $\\bh$ is\n",
    "\n",
    "$$\n",
    "\\pi(\\bv, \\bh)\n",
    "= \\frac{1}{Z} \\exp\\left(- \\cE(\\bv, \\bh)\\right),\n",
    "$$\n",
    "\n",
    "where $Z$, the [partition function](https://en.wikipedia.org/wiki/Partition_function_%28statistical_mechanics%29) (a function of $W,\\bb,\\bc$), ensures that $\\pi$ indeed sums to $1$."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf30110a",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "We will be interested in sampling from the marginal over visible units, that is,\n",
    "\n",
    "$$\n",
    "\\rho(\\bv)\n",
    "= \\sum_{\\bh \\in \\{0,1\\}^n} \\pi(\\bv, \\bh).\n",
    "$$\n",
    "\n",
    "When $m$ and/or $n$ are large, computing $\\rho$ or $\\pi$ explicitly -- or even numerically -- is impractical. \n",
    "\n",
    "We develop the Gibbs sampler for this model next."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "46c70e92",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "*Gibbs sampling:* We sample from the joint distribution $\\pi$ and observe only $\\bv$.\n",
    "\n",
    "We need to compute the conditional probabilities given every other variable. The sigmoid function, which we have encountered previously, \n",
    "\n",
    "$$\n",
    "\\sigma(x)\n",
    "=\n",
    "\\frac{1}{1 + e^{-x}}\n",
    "$$\n",
    "\n",
    "will once again make an appearance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e5fa49f",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "def sigmoid(x): \n",
    "    return 1/(1 + np.exp(-x))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d1942ed",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "grid = np.linspace(-5, 5, 100)\n",
    "plt.plot(grid,sigmoid(grid),'r')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98b53117",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "Fix a visible unit $i \\in [m]$. For a pair $(\\bv, \\bh)$, we denote by $(\\bv_{[i]}, \\bh)$ the same pair where coordinate $i$ of $\\bv$ is flipped. Given every other variable, i.e., $(\\bv_{-i},\\bh)$, and using a superscript $\\text{v}$ to indicate the probability of a visible unit, the conditional probability of $v_i$ is\n",
    "\n",
    "\\begin{align*}\n",
    "\\pi^{\\text{v}}_i(v_i|\\bv_{-i},\\bh)\n",
    "&= \\frac{\\pi(\\bv, \\bh)}{\\pi(\\bv, \\bh) + \\pi(\\bv_{[i]}, \\bh)}\\\\\n",
    "&= \\frac{\\frac{1}{Z} \\exp\\left(- \\cE(\\bv, \\bh)\\right)}{\\frac{1}{Z} \\exp\\left(- \\cE(\\bv, \\bh)\\right) + \\frac{1}{Z} \\exp\\left(- \\cE(\\bv_{[i]}, \\bh)\\right)}.\n",
    "\\end{align*}\n",
    "\n",
    "In this last ratio, the partition functions (the $Z$'s) cancel out. Moreover, all the terms in the exponentials *not depending* on the $i$-th visible unit actually factor out and cancel out as well -- they are identical in all three exponentials. Similarly, the terms in the exponentials *depending only on $\\bh$* also factor out and cancel out."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d09e7030",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "What we are left with is:\n",
    "\n",
    "\\begin{align*}\n",
    "&\\pi^{\\text{v}}_i(v_i|\\bv_{-i},\\bh)\\\\\n",
    "&= \\frac{\\exp\\left(\\sum_{j=1}^n w_{i,j} v_i h_j\n",
    "+ b_i v_i\\right)}\n",
    "{\\exp\\left(\\sum_{j=1}^n w_{i,j} v_i h_j\n",
    "+ b_i v_i\\right) \n",
    "+ \\exp\\left(\\sum_{j=1}^n w_{i,j} (1-v_i) h_j\n",
    "+ b_i (1-v_i)\\right)},\n",
    "\\end{align*}\n",
    "\n",
    "where we used the fact that flipping $v_i \\in \\{0,1\\}$ is the same as setting it to $1 - v_i$, a transformation which indeed sends $0$ to $1$ and $1$ to $0$."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac4bd717",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "This expression does not depend on $\\bv_{-i}$. In other words, the $i$-th visible unit is conditionally independent of all other visible units given the hidden units. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "df1492d9",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "We simplify the expression\n",
    "\n",
    "\\begin{align*}\n",
    "&\\pi^{\\text{v}}_i(v_i|\\bv_{-i},\\bh)\\\\\n",
    "&= \\frac{1}\n",
    "{1 \n",
    "+ \\exp\\left(\\sum_{j=1}^n w_{i,j} (1-2 v_i) h_j\n",
    "+ b_i (1- 2v_i)\\right)}\\\\\n",
    "&= \\sigma\\left(\\sum_{j=1}^n w_{i,j} (2 v_i-1) h_j\n",
    "+ b_i (2v_i-1)\\right).\n",
    "\\end{align*}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77427fb9",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "In particular, the conditional mean of the $i$-th visible unit given everything else is \n",
    "\n",
    "\\begin{align*}\n",
    "0 \\cdot \\pi^{\\text{v}}_i(0|\\bv_{-i},\\bh) + 1 \\cdot \\pi^{\\text{v}}_i(1|\\bv_{-i},\\bh)\n",
    "&= \\pi^{\\text{v}}_i(1|\\bv_{-i},\\bh)\\\\\n",
    "&= \\sigma\\left(\\sum_{j=1}^n w_{i,j} h_j\n",
    "+ b_i \\right)\\\\\n",
    "&= \\sigma\\left((W \\bh + \\bb)_i\n",
    "\\right)\n",
    "\\end{align*}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ab91d8df",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "Similarly for the conditional of the $j$-th hidden unit given everything else, we have\n",
    "\n",
    "\\begin{align*}\n",
    "&\\pi^{\\text{h}}_j(h_j|\\bv,\\bh_{-j})\\\\\n",
    "&= \\sigma\\left(\\sum_{i=1}^m w_{i,j} v_i (2h_j -1) \n",
    "+ c_j (2h_j -1)\\right).\n",
    "\\end{align*}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1754f516",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "The conditional mean given everything else is \n",
    "\n",
    "\\begin{align*}\n",
    "0 \\cdot \\pi^{\\text{h}}_j(0|\\bv,\\bh_{-j}) + 1 \\cdot \\pi^{\\text{h}}_j(1|\\bv,\\bh_{-j})\n",
    "&= \\pi^{\\text{h}}_j(1|\\bv,\\bh_{-j})\n",
    "= \\sigma\\left((W^T \\bv + \\bc)_j \\right).\n",
    "\\end{align*}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b77f265",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "And the $j$-th hidden unit is conditionally independent of all other hidden units given the visible units. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c12eae0a",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "**NUMERICAL CORNER:** We implement the Gibbs sampler for an RBM. Rather than updating the units at random, we use a block approach. Specifically, we update all hidden units independently, given the visible units; then we update all visible units independently, given the hidden units. In each case, this is warranted by the conditional independence structure revealed above. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23abfcb3",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "We first implement the conditional means using the formulas previously derived. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33c3833a",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "def rbm_mean_hidden(v, W, c):\n",
    "    return sigmoid(W @ v + c)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62958b12",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "def rbm_mean_visible(h, W, b):\n",
    "    return sigmoid(W.T @ h + b)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f5298d74",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "We next implement one step of the sampler, which consists in updating all hidden units, followed by updating all visible units. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9d0aa10",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "def rbm_gibbs_update(v, W, b, c):\n",
    "    p_hidden = rbm_mean_hidden(v, W, c)\n",
    "    h = rng.binomial(1, p_hidden, p_hidden.shape)\n",
    "    p_visible = rbm_mean_visible(h, W, b)\n",
    "    v = rng.binomial(1, p_visible, p_visible.shape)\n",
    "    return v"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b000f97",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "Finally, we repeat these steps `k` times. We only return the visible units `v`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a85456bf",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "def rbm_gibbs_sampling(k, v_0, W, b, c):\n",
    "    counter = 0\n",
    "    v = v_0\n",
    "    while counter < k:\n",
    "        v = rbm_gibbs_update(v, W, b, c)\n",
    "        counter += 1\n",
    "    return v"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5694b73f",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "Here `v_0` is the initial visible unit states. We do not need to initialize the hidden ones as this is done automatically in the first update step. In the next subsection, we will take the initial distribution of $\\bv$ to be independent Bernoullis with success probability $1/2$."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c82100de",
   "metadata": {
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "$\\unlhd$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "421e05b0",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "### Sampling handwritten digits with an RBM"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0555120",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "**Trainging the model** We first need to train the model on the data. We will not show how this is done here, but instead use [`sklearn.neural_network.BernoulliRBM`](https://scikit-learn.org/stable/modules/generated/sklearn.neural_network.BernoulliRBM.html). (Some details of how this training is done is provided [here](https://scikit-learn.org/stable/modules/neural_networks_unsupervised.html#stochastic-maximum-likelihood-learning).)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f43ce574",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "from sklearn.neural_network import BernoulliRBM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d11ecc5",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "rbm = BernoulliRBM(random_state=seed, verbose=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36d87919",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "To simplify the analysis and speed up the training, we only keep digits $0$, $1$ and $5$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4874aed8",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "# Filter out images with labels 0, 1, or 5\n",
    "mask = (labels == 0) | (labels == 1) | (labels == 5)\n",
    "imgs = imgs[mask]\n",
    "labels = labels[mask]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2715cc09",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "We flatten the images (which have already been \"rounded\" to black-and-white; see the first subsection)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d96d1cba",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "X = imgs.reshape(len(imgs), -1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a5ce883",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "We now fit the model. Choosing the hyperparameters of the training algorithm is tricky. The following seem to work reasonably well. (For a more systematic approach to tuning hyperparameters, see [here](https://scikit-learn.org/stable/modules/grid_search.html).)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1eb2d72",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "rbm.n_components = 100\n",
    "rbm.learning_rate = 0.02\n",
    "rbm.batch_size = 50\n",
    "rbm.n_iter = 20\n",
    "rbm.fit(X)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d5d64511",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "We plot the learned weight matrix using a script [adapted from here](https://scikit-learn.org/stable/auto_examples/neural_networks/plot_rbm_logistic_classification.html). Each image shows the weights associated to all visible units by one hidden unit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "107bfa19",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "def plot_imgs(z, n_imgs, nx_pixels, ny_pixels):\n",
    "    nx_imgs = np.floor(np.sqrt(n_imgs))\n",
    "    ny_imgs = np.ceil(np.sqrt(n_imgs))\n",
    "    plt.figure(figsize=(8, 8))\n",
    "    for i, comp in enumerate(z):\n",
    "        plt.subplot(int(nx_imgs), int(ny_imgs), i + 1)\n",
    "        plt.imshow(comp.reshape((nx_pixels, ny_pixels)), \n",
    "                   cmap=plt.cm.gray_r)\n",
    "        plt.xticks([])\n",
    "        plt.yticks([])\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a2078d9",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "plot_imgs(rbm.components_, rbm.n_components, \n",
    "          nx_pixels, ny_pixels)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf2fa2db",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "**Back to Gibbs sampling** We are ready to sample from the trained RBM. We extract the learned parameters from `rbm`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b488a91",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "W = rbm.components_\n",
    "W.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ecda1ef4",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "b = rbm.intercept_visible_\n",
    "b.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6283ffd3",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "c = rbm.intercept_hidden_\n",
    "c.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "190d19b8",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "To generate $25$ samples, we first generate $25$ independent initial states. We stack them into a matrix, where each row is a different flattened random noise image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e3b12c7",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "n_samples = 25 # number of samples\n",
    "z = rng.binomial(1, 0.5, (n_samples, n_pixels))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0fc057c",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "plot_imgs(z, n_samples, nx_pixels, ny_pixels)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c8096dd4",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "To process all samples simultaneously, we make a small change to the code. We [`numpy.reshape`](https://numpy.org/doc/stable/reference/generated/numpy.reshape.html)\n",
    "to make the offsets into column vectors, which are then automatically added to all columns of the resulting weighted sum. \n",
    "(This is known as [broadcasting](https://numpy.org/doc/stable/user/basics.broadcasting.html).)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5204e923",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "def rbm_mean_hidden(v, W, c):\n",
    "    return sigmoid(W @ v + c.reshape(len(c),1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1389e62",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "def rbm_mean_visible(h, W, b):\n",
    "    return sigmoid(W.T @ h + b.reshape(len(b),1))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e5b9fdb",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "We are now ready to run our Gibbs sampler. The outcome depends on the number of steps we take. For instance, after one step, the result is still very noisy -- although note that the fraction of white pixels is already more realistic!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d1aaa0d",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "v_0 = z.T\n",
    "gen_v = rbm_gibbs_sampling(1, v_0, W, b, c)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f29fefb",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "plot_imgs(gen_v.T, n_samples, nx_pixels, ny_pixels)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a6ba29d",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "After $10$ steps, we already see shadows of digits appearing. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3f2e300",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "v_0 = z.T\n",
    "gen_v = rbm_gibbs_sampling(10, v_0, W, b, c)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad992761",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "plot_imgs(gen_v.T, n_samples, nx_pixels, ny_pixels)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb15958f",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "After $100$ steps, the outcome is quite realistic. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1210adf",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "v_0 = z.T\n",
    "gen_v = rbm_gibbs_sampling(100, v_0, W, b, c)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b947205",
   "metadata": {
    "scrolled": false,
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "plot_imgs(gen_v.T, n_samples, nx_pixels, ny_pixels)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9e00490f",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "RBMs can be stacked into [deep probabilistic models](https://en.wikipedia.org/wiki/Boltzmann_machine#Deep_Boltzmann_machine), similarly to what we encountered with neural networks."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e9577a23",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "$\\newcommand{\\P}{\\mathbb{P}}$ $\\newcommand{\\E}{\\mathbb{E}}$ $\\newcommand{\\S}{\\mathcal{S}}$ $\\newcommand{\\var}{\\mathrm{Var}}$ $\\newcommand{\\bmu}{\\boldsymbol{\\mu}}$ $\\newcommand{\\bSigma}{\\boldsymbol{\\Sigma}}$ $\\newcommand{\\btheta}{\\boldsymbol{\\theta}}$ $\\newcommand{\\bpi}{\\boldsymbol{\\pi}}$ $\\newcommand{\\indep}{\\perp\\!\\!\\!\\perp}$ $\\newcommand{\\bp}{\\mathbf{p}}$ $\\newcommand{\\bx}{\\mathbf{x}}$ $\\newcommand{\\bX}{\\mathbf{X}}$ $\\newcommand{\\by}{\\mathbf{y}}$ $\\newcommand{\\bY}{\\mathbf{Y}}$ $\\newcommand{\\bz}{\\mathbf{z}}$ $\\newcommand{\\bZ}{\\mathbf{Z}}$ $\\newcommand{\\bw}{\\mathbf{w}}$ $\\newcommand{\\bW}{\\mathbf{W}}$ $\\newcommand{\\bv}{\\mathbf{v}}$ $\\newcommand{\\bV}{\\mathbf{V}}$ $\\newcommand{\\bfbeta}{\\boldsymbol{\\beta}}$ $\\newcommand{\\bflambda}{\\boldsymbol{\\lambda}}$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "509fc7df",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "## More advanced material"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d8220d08",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "### Cholesky decomposition"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "718c87cf",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "**NUMERICAL CORNER:** We implement the algorithm above. In our naive implementation, we assume that $B$ is positive definite, and therefore that all steps are well-defined."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6bce2ba2",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "def cholesky(B):\n",
    "    n = B.shape[0] \n",
    "    L = np.zeros((n, n))\n",
    "    for j in range(n):\n",
    "        L[j,0:j] = mmids.forwardsubs(L[0:j,0:j],B[j,0:j])\n",
    "        L[j,j] = np.sqrt(B[j,j] - LA.norm(L[j,0:j])**2)\n",
    "    return L "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b9d58b06",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "Here is a simple example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "011776bd",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "B = np.array([[2., 1.],[1., 2.]])\n",
    "print(B)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6803d4ff",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "L = cholesky(B)\n",
    "print(L)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c82e9a1b",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "We can check that it produces the right factorization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20e858a9",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "print(L @ L.T)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "641e7688",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "$\\unlhd$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b568a9d2",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "**NUMERICAL CORNER:** We implement this algorithm below. In our naive implementation, we assume that $A$ has full column rank, and therefore that all steps are well-defined."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d7698e1",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "def ls_by_chol(A, b):\n",
    "    L = cholesky(A.T @ A)\n",
    "    z = mmids.forwardsubs(L, A.T @ b)\n",
    "    return mmids.backsubs(L.T, z)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43fffb1c",
   "metadata": {
    "slideshow": {
     "slide_type": "skip"
    },
    "tags": [
     "colab-keep"
    ]
   },
   "source": [
    "$\\unlhd$"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 5
}