{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Expectation maximization\n",
    "\n",
    "$$ \\begin{eqnarray} \n",
    " z_i &\\sim & Multinomial(\\phi) \\\\\n",
    " x_i|z_i & \\sim & \\mathcal{N}(\\mu_{z_i}, \\sigma_{z_i})) \\\\\n",
    "\\end{eqnarray}$$ \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 38680 characters]\n",
      "text/plain": [
       "<Figure size 1152x432 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from scipy import stats\n",
    "\n",
    "np.random.seed(654)\n",
    "# Draw samples from two Gaussian w.p. z_i ~ Bernoulli(phi)\n",
    "generative_m = np.array([stats.norm(-1, 1), stats.norm(5, 1.8)])\n",
    "z = stats.bernoulli(0.9).rvs(100)\n",
    "x_i = np.array([g.rvs() for g in generative_m[z]])\n",
    "\n",
    "# plot generated data and the latent distributions\n",
    "x = np.linspace(-5, 12, 150)\n",
    "plt.figure(figsize=(16, 6))\n",
    "plt.plot(x, generative_m[0].pdf(x))\n",
    "plt.plot(x, generative_m[1].pdf(x))\n",
    "plt.plot(x, generative_m[0].pdf(x) + generative_m[1].pdf(x), lw=1, ls='-.', color='black')\n",
    "plt.fill_betweenx(generative_m[0].pdf(x), x, alpha=0.1)\n",
    "plt.fill_betweenx(generative_m[1].pdf(x), x, alpha=0.1)\n",
    "plt.vlines(x_i, 0, 0.01, color=np.array(['C0', 'C1'])[z])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Log likelihood\n",
    "Let $\\theta = \\{ \\mu, \\sigma \\}$.\n",
    "$$\\ell(\\theta) = \\sum_{i=1}^n \\log P(x_i; \\theta)  $$\n",
    "\n",
    "Or in vector form:\n",
    "\n",
    "$$\\ell(\\theta) = \\log P(x; \\theta)  $$\n",
    "\n",
    "This log likelihood function cannot be solved directly as it not includes the latent variable $z$.\n",
    "\n",
    "$$ \\ell(\\theta) = \\sum_{k=1}^K \\log P(x | z_k; \\theta) P(z_k) $$\n",
    "\n",
    "## Continuous form\n",
    "\n",
    "$$ \\begin{eqnarray}\n",
    "\\ell(\\theta) & = & \\int_z \\log P(x | \\theta, z) P(z) dz \\\\\n",
    "\\log P(x; \\theta)& = & \\int_z \\log P(x | \\theta, z) P(z) dz\n",
    "\\end{eqnarray}$$\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Jensens inequality:**\n",
    "*Only for concave functions*\n",
    "\n",
    "let $f$ be a concave function.\n",
    "\n",
    "$$ E[f(X)] \\le f(E[X]) $$\n",
    "\n",
    "**Jensens equality:**\n",
    "$$ E[f(X)] = f(E[X]) \\iff X = E[X] $$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evidence Lower BOund\n",
    "\n",
    "\n",
    "Multiply the **evidence** w/ an arbitrary distribution over $Z, \\quad \\frac{Q(z)}{Q(z)} = 1$ \n",
    "\n",
    "$$ \\begin{eqnarray}\n",
    "\\log P(x; \\theta)& = & \\int \\log Q(z) \\frac{ P(x | \\theta, z) P(z)}{Q(z)} dz \\\\\n",
    "& = & \\log E_{z \\sim Q}[ \\frac{ P(x | \\theta, z) P(z)}{Q(z)}] \\\\\n",
    "\\end{eqnarray}$$\n",
    "\n",
    "As the log function is a concave function, we can apply Jensen's inequality.\n",
    "\n",
    "$$ \\begin{eqnarray}\n",
    "\\log P(x; \\theta)& = & \\log E_{z \\sim Q}[ \\frac{ P(x | \\theta, z) P(z)}{Q(z)}] \\\\\n",
    "& \\ge &  E_{z \\sim Q}[ \\log \\frac{ P(x | \\theta, z) P(z)}{Q(z)}]\n",
    "\\end{eqnarray}$$\n",
    "\n",
    "\n",
    "\n",
    "Resulting in a lower bound on the evidence called the **ELBO**:\n",
    "\n",
    "$$ \\begin{eqnarray}\n",
    "\\log P(x; \\theta)& \\ge &  E_{z \\sim Q}[ \\log \\frac{ P(x | \\theta, z) P(z)}{Q(z)}]\n",
    "\\end{eqnarray}$$\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Now what?\n",
    "\n",
    "Jensen's inequality holds to **equaltiy** if and only if $X \\iff E[X]$. This will be the case if we choose a distribution that leads to:\n",
    "\n",
    "$$ \\begin{eqnarray}\n",
    "\\frac{P(x|z) P(z)}{Q(z)} &=& \\frac{P(x, z)}{Q(z)} = c\n",
    "\\end{eqnarray}$$\n",
    "\n",
    "Wich means that \n",
    "\n",
    "$$ P(x, z; \\theta) \\propto Q(z) $$\n",
    "\n",
    "A next requirement is that $\\int_z Q(z) dz = 1$ (requirement of probability distribution). And because $Q(z)$ is a probability distribution, it must integrate to one; $\\int_z Q(z) dz = 1$. So ideally we want to take joint distribution $P(x, z; \\theta)$ and transform it so that it is proportional to itself, but also sums to 1 over all values of $z_i$. It turns out we can find that by normalizing the joint distribution.\n",
    "\n",
    "$$ \\begin{eqnarray}\n",
    "\\int_z \\frac{P(x, z; \\theta)}{\\int_z P(x, z; \\theta) dz} dz &=& 1 \\\\\n",
    "\\end{eqnarray}$$\n",
    "\n",
    "This will lead to:\n",
    "\n",
    "$$ \\begin{eqnarray}\n",
    "Q(z) &=& \\frac{P(x, z; \\theta)}{\\int_z P(x, z; \\theta) dz}  \\\\\n",
    "&=& \\frac{P(x, z; \\theta)}{\\int_z P(x; \\theta)} \\\\\n",
    "&=& P(z|x; \\theta)\n",
    "\\end{eqnarray}$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Gaussian mixture model\n",
    "\n",
    "### Expectation\n",
    "Given a Gaussian mixture model, observed data points $x$ and current parameters $\\theta_t$ find $Q(z; \\theta_t) = P(z|x; \\theta_t)$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([ 4.28244273,  4.66972656,  3.35275075,  5.65085616,  3.74751768,\n",
       "        9.75240814,  5.15756521,  4.80213993,  5.90145058,  4.84288191,\n",
       "        4.71273852, -0.63758867, -1.38915329,  6.71697406,  7.00596187,\n",
       "        5.15422491,  5.09580485,  5.20293865,  6.76574344, -1.51645305,\n",
       "        2.65304672,  6.81892055,  2.77103106,  5.06418788,  5.72517362,\n",
       "        3.54313443,  5.67858554, -1.96777689,  4.35169594,  5.51107055,\n",
       "        2.57224619,  4.97843073,  4.32989217,  5.51861638,  4.51797044,\n",
       "        4.52825271,  6.93340533,  4.94273792,  2.54969132,  2.81541702,\n",
       "        4.81172184,  5.19576022,  6.11967179,  3.5235398 , -2.24692387,\n",
       "        4.23428727,  7.15230773,  6.50022463,  8.60608829,  4.27073679,\n",
       "        7.75388195,  6.79259227, -2.37595108,  4.26260727,  4.09584051,\n",
       "        0.26993383,  4.86431429,  8.38977374,  6.48961063,  4.09783105,\n",
       "        4.40498267,  3.27433241,  7.44164281,  7.62286463,  3.21797707,\n",
       "        4.42538064,  7.06135769,  3.94733606,  1.91385988,  5.23226903,\n",
       "        3.61125312, -1.13133337,  5.71568321,  5.91236947,  5.48316563,\n",
       "        3.75523264,  2.69128758,  8.02274291, -0.53004188,  5.35222504,\n",
       "       -0.27283765,  4.15793385,  5.33299762, -1.99908258,  7.04351757,\n",
       "        6.14452194,  2.62636128,  4.55251749, -3.38101307,  4.27776483,\n",
       "        4.23396713,  3.75797097,  2.34067487,  4.27915132,  6.13345271,\n",
       "        3.97351015, -0.17274158,  4.57300419,  6.06374396,  3.45567689])"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Data points x\n",
    "x_i"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "import torch.distributions as dist\n",
    "import torch.nn.functional as F"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/miniconda3/lib/python3.7/site-packages/ipykernel_launcher.py:1: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).\n",
      "  \"\"\"Entry point for launching an IPython kernel.\n"
     ]
    }
   ],
   "source": [
    "x_i = torch.tensor(x_i, dtype=torch.float)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<BarContainer object of 100 artists>"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 9676 characters]\n",
      "text/plain": [
       "<Figure size 1152x576 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "torch.manual_seed(1)\n",
    "\n",
    "# Number of clusters\n",
    "k = torch.tensor([2])\n",
    "\n",
    "# Current values for theta\n",
    "mu = dist.Uniform(x_i.min(), x_i.max()).sample(k) \n",
    "std = torch.ones(k)\n",
    "phi = torch.ones(k)\n",
    "\n",
    "def expectation(x, mu, std, phi, k):\n",
    "    # P(z|x)\n",
    "    qz = torch.zeros((k, x.shape[0]))\n",
    "    phi = F.softmax(phi, dim=-1)\n",
    "    \n",
    "    for z in range(k):\n",
    "        # Define the joint: P(z, x)\n",
    "        qz[z] = (dist.Normal(mu[z], std[z]).log_prob(x) + torch.log(phi[z])).exp() \n",
    "    \n",
    "    # normalize zo that marginalizing z would lead to p = 1\n",
    "    # make conditional by dividing: 1/ P(x)\n",
    "    qz = qz / qz.sum(0)\n",
    "    \n",
    "    return qz\n",
    "\n",
    "qz = expectation(x_i, mu, std, phi, k)    \n",
    "\n",
    "plt.figure(figsize=(16, 8)) \n",
    "plt.title(r\"$Q(z) = P(z|x; \\theta)$\")\n",
    "plt.bar(x_i, qz[0], width=0.01)\n",
    "plt.bar(x_i, qz[1], width=0.01)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Maximization\n",
    "\n",
    "Optimize $\\theta$ with the $Q(z; \\theta_t)$ found in the expectation step.\n",
    "$$ \\text{ELBO} =  \\sum_{z_i} Q(z_i) \\log ( \\frac{p(x_i| z_i; \\theta_{t + 1}) p(z_i)}  {Q(z_i)}) $$\n",
    "\n",
    "$$ \\underset{\\theta_{t + 1}}{\\text{argmax ELBO}} $$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/miniconda3/lib/python3.7/site-packages/ipykernel_launcher.py:5: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).\n",
      "  \"\"\"\n",
      "/opt/miniconda3/lib/python3.7/site-packages/ipykernel_launcher.py:6: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).\n",
      "  \n",
      "/opt/miniconda3/lib/python3.7/site-packages/ipykernel_launcher.py:7: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).\n",
      "  import sys\n"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(0)\n",
    "k = torch.tensor([2])\n",
    "mu = dist.Uniform(x_i.min(), x_i.max()).sample(k) \n",
    "\n",
    "mu = torch.tensor(mu, dtype=torch.float, requires_grad=True)\n",
    "std = torch.tensor(torch.ones(k), requires_grad=True)\n",
    "phi = torch.tensor(torch.ones(k), requires_grad=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(-941.1580, grad_fn=<AddBackward0>)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def log_likelihood(mu, std, phi, x):\n",
    "    return torch.log(phi) + dist.Normal(mu, std).log_prob(x)\n",
    "\n",
    "def get_elbo(mu, std, phi):\n",
    "    elbo = 0\n",
    "    for k in range(2):\n",
    "        z_i = torch.tensor(1.)  # prior on z\n",
    "        mu_i = mu[k]\n",
    "        std_i = std[k]\n",
    "        phi_i = F.softmax(phi, dim=-1)[k]\n",
    "        \n",
    "        elbo += (qz[k] * (log_likelihood(mu_i, std_i, phi_i, x_i) + torch.log(z_i) - torch.log(qz[k]))).sum()\n",
    "\n",
    "    return elbo\n",
    "\n",
    "elbo = get_elbo(mu, std, phi)\n",
    "elbo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "optim = torch.optim.Adam([mu, std, phi])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "neg. ELBO: 313.53\tmu: tensor([3.0438, 6.6150])\tstd: tensor([1.0943, 1.0870])\tphi: tensor([0.5431, 0.4569])\n",
      "neg. ELBO: 230.57\tmu: tensor([1.7261, 5.3456])\tstd: tensor([1.9472, 1.4770])\tphi: tensor([0.2626, 0.7374])\n",
      "neg. ELBO: 224.19\tmu: tensor([0.4331, 5.1358])\tstd: tensor([2.2290, 1.5297])\tphi: tensor([0.1794, 0.8206])\n",
      "neg. ELBO: 222.64\tmu: tensor([-0.5704,  5.0750])\tstd: tensor([2.0617, 1.5510])\tphi: tensor([0.1491, 0.8509])\n",
      "neg. ELBO: 219.15\tmu: tensor([-1.3539,  5.0196])\tstd: tensor([1.0274, 1.5799])\tphi: tensor([0.1279, 0.8721])\n",
      "neg. ELBO: 219.14\tmu: tensor([-1.3579,  5.0186])\tstd: tensor([1.0029, 1.5811])\tphi: tensor([0.1277, 0.8723])\n",
      "neg. ELBO: 219.14\tmu: tensor([-1.3579,  5.0186])\tstd: tensor([1.0029, 1.5811])\tphi: tensor([0.1277, 0.8723])\n",
      "neg. ELBO: 219.14\tmu: tensor([-1.3579,  5.0186])\tstd: tensor([1.0029, 1.5811])\tphi: tensor([0.1277, 0.8723])\n"
     ]
    }
   ],
   "source": [
    "lr = 0.01\n",
    "for i in range(150):\n",
    "    \n",
    "    # expectation step\n",
    "    with torch.no_grad():\n",
    "        qz = expectation(x_i, mu, std, phi, k)\n",
    "    \n",
    "    # maximization step\n",
    "    for j in range(100):\n",
    "        negative_elbo = -get_elbo(mu, std, phi)\n",
    "        negative_elbo.backward()\n",
    "        \n",
    "        optim.step()\n",
    "        optim.zero_grad()\n",
    "    \n",
    "    if i % 20 == 0:\n",
    "        print(\"neg. ELBO: {:.2f}\"\n",
    "              \"\\tmu: {}\"\n",
    "              \"\\tstd: {}\"\n",
    "              \"\\tphi: {}\"\n",
    "             \"\".format(negative_elbo.item(), mu.data, std.data, F.softmax(phi, dim=-1).data))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x7f3a8b76cb38>]"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 53748 characters]\n",
      "text/plain": [
       "<Figure size 1152x432 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "fitted_m = [stats.norm(mu, std) for mu, std in zip(mu.data, std.data)]\n",
    "\n",
    "plt.figure(figsize=(16, 6))\n",
    "plt.vlines(x_i, 0, 0.01, color=np.array(['C0', 'C1'])[z])\n",
    "plt.plot(x, fitted_m[0].pdf(x))\n",
    "plt.plot(x, fitted_m[1].pdf(x))\n",
    "plt.plot(x, generative_m[0].pdf(x), color='black', lw=1, ls='-.')\n",
    "plt.plot(x, generative_m[1].pdf(x), color='black', lw=1, ls='-.')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## GMM\n",
    "\n",
    "### Expectation step\n",
    "$$\\begin{eqnarray}\n",
    "w_{ij} := p(z_i = j|x_i; \\theta) \\\\\n",
    "\\mu_j := \\frac{\\sum_{i=1}^nw_{ij} x_i} {\\sum_{i=1}^nw_{ij}}\n",
    "\\end{eqnarray}$$\n",
    "\n",
    "\n",
    "### Maximization step\n",
    "$$\\begin{eqnarray}\n",
    "\\mu_j := \\frac{\\sum_{i=1}^nw_{ij} x_i} {\\sum_{i=1}^nw_{ij}} \\\\\n",
    "\\sigma_j :=  \\sqrt{\\frac{\\sum_{i=1}^nw_{ij}(x_i - \\mu_j)^2} {\\sum_{i=1}^nw_{ij}}}\n",
    "\\end{eqnarray}$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x7f9520175f60>]"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 53488 characters]\n",
      "text/plain": [
       "<Figure size 1152x432 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "class EM:\n",
    "    def __init__(self, k):\n",
    "        self.k = k\n",
    "        self.mu = None\n",
    "        self.std = np.ones(k)\n",
    "        self.w_ij = None\n",
    "    [1] https://www.ritchievink.com/blog/2019/05/24/algorithm-breakdown-expectation-maximization/\n",
    "        self.phi = np.ones(k) / k\n",
    "\n",
    "    def expectation_step(self, x):\n",
    "        for z_i in range(self.k):\n",
    "            self.w_ij[z_i] = stats.norm(self.mu[z_i], self.std[z_i]).pdf(x) * self.phi[z_i]\n",
    "        \n",
    "        # normalize zo that marginalizing z would lead to p = 1\n",
    "        self.w_ij /= self.w_ij.sum(0)\n",
    "\n",
    "    def maximization_step(self, x):\n",
    "        self.phi = self.w_ij.mean(1)\n",
    "        self.std = ((self.w_ij * (x - self.mu[:, None])**2).sum(1) / self.w_ij.sum(1))**0.5\n",
    "        self.mu = (self.w_ij * x).sum(1) / self.w_ij.sum(1)\n",
    "\n",
    "    def fit(self, x):\n",
    "        self.mu = np.random.uniform(x.min(), x.max(), size=self.k)\n",
    "        self.w_ij = np.zeros((self.k, x.shape[0]))\n",
    "\n",
    "        last_mu = np.ones(self.k) * np.inf\n",
    "        while ~np.all(np.isclose(self.mu, last_mu)):\n",
    "            last_mu = self.mu\n",
    "            self.expectation_step(x)\n",
    "            self.maximization_step(x)\n",
    "\n",
    "m = EM(2)\n",
    "m.fit(np.array(x_i))\n",
    "\n",
    "fitted_m = [stats.norm(mu, std) for mu, std in zip(m.mu, m.std)]\n",
    "\n",
    "plt.figure(figsize=(16, 6))\n",
    "plt.vlines(x_i, 0, 0.01, color=np.array(['C0', 'C1'])[z])\n",
    "plt.plot(x, fitted_m[0].pdf(x))\n",
    "plt.plot(x, fitted_m[1].pdf(x))\n",
    "plt.plot(x, generative_m[0].pdf(x), color='black', lw=1, ls='-.')\n",
    "plt.plot(x, generative_m[1].pdf(x), color='black', lw=1, ls='-.')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# References\n",
    "[1] https://www.ritchievink.com/blog/2019/05/24/algorithm-breakdown-expectation-maximization/"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
