{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import numpy as np\n",
    "import torch\n",
    "from torch import nn\n",
    "from sklearn import datasets\n",
    "from torch.utils.data import DataLoader, TensorDataset\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "import torch.distributions as dist\n",
    "from dataclasses import dataclass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.collections.PathCollection at 0x7f11cfba83c8>"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 14344 characters]\n",
      "text/plain": [
       "<Figure size 1152x432 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "w0 = 0.125\n",
    "b0 = 5.\n",
    "x_range = [-20, 60]\n",
    "\n",
    "def load_dataset(n=150, n_tst=150):\n",
    "    np.random.seed(43)\n",
    "    \n",
    "    def s(x):\n",
    "        g = (x - x_range[0]) / (x_range[1] - x_range[0])\n",
    "        return 3 * (0.25 + g**2.)\n",
    "    \n",
    "    x = (x_range[1] - x_range[0]) * np.random.rand(n) + x_range[0]\n",
    "    eps = np.random.randn(n) * s(x)\n",
    "    y = (w0 * x * (1. + np.sin(x)) + b0) + eps\n",
    "    y = (y - y.mean()) / y.std()\n",
    "    idx = np.argsort(x)\n",
    "    x = x[idx]\n",
    "    y = y[idx]\n",
    "    return y[:, None], x[:, None]\n",
    "\n",
    "Y, X = load_dataset()\n",
    "\n",
    "plt.figure(figsize=(16, 6))\n",
    "plt.scatter(X, Y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "X = torch.tensor(X, dtype=torch.float)\n",
    "Y = torch.tensor(Y, dtype=torch.float)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MaximumLikelihood(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.out = nn.Sequential(\n",
    "            nn.Linear(1, 20),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(20, 1)\n",
    "        )\n",
    "    def forward(self, x):\n",
    "        return self.out(x)\n",
    "    \n",
    "epochs = 200\n",
    "m = MaximumLikelihood()\n",
    "optim = torch.optim.Adam(m.parameters(), lr=0.01)\n",
    "\n",
    "for epoch in range(epochs):\n",
    "    optim.zero_grad()\n",
    "    y_pred = m(X)\n",
    "    loss = (0.5 * (y_pred - Y)**2).mean()\n",
    "    loss.backward()\n",
    "    optim.step()\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 23616 characters]\n",
      "text/plain": [
       "<Figure size 1152x432 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "with torch.no_grad():\n",
    "    plt.figure(figsize=(16, 6))\n",
    "    plt.plot(X, m(X))\n",
    "    plt.scatter(X, Y, alpha=0.5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(0.5151)\n"
     ]
    }
   ],
   "source": [
    "with torch.no_grad():\n",
    "    print(torch.abs(m(X) - Y).mean())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$ \\text{ELBO}  = E_{\\theta \\sim Q}[\\log P(D|\\theta)] + D_{KL}(Q(\\theta)||P(\\theta)) $$\n",
    "\n",
    "$$ \\text{ELBO}  = E_{\\theta \\sim Q}[\\log P(D|\\theta)] + \\int Q(\\theta) \\log \\frac{P(\\theta)}{Q(\\theta)}d\\theta $$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def ll_gaussian(y, mu, log_var):\n",
    "    sigma = torch.exp(0.5 * log_var)\n",
    "    return -0.5 * torch.log(2 * np.pi * sigma**2) - (1 / (2 * sigma**2))* (y-mu)**2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "class VI(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.q_mu = nn.Sequential(\n",
    "            nn.Linear(1, 20),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(20, 10),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(10, 1)\n",
    "        )\n",
    "        self.q_log_var = nn.Sequential(\n",
    "            nn.Linear(1, 20),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(20, 10),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(10, 1)\n",
    "        )\n",
    "        \n",
    "    def reparameterize(self, mu, log_var):\n",
    "        # std can not be negative, thats why we use log variance\n",
    "        sigma = torch.exp(0.5 * log_var) + 1e-5\n",
    "        eps = torch.randn_like(sigma)\n",
    "        return mu + sigma * eps\n",
    "    \n",
    "    def forward(self, x):\n",
    "        mu = self.q_mu(x)\n",
    "        log_var = self.q_log_var(x)\n",
    "        return self.reparameterize(mu, log_var), mu, log_var"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def elbo(y_pred, y, mu, log_var):\n",
    "    sigma = log_var.exp()**0.5\n",
    "    # likelihood of observing y given Variational mu and sigma\n",
    "    likelihood = dist.Normal(mu, sigma).log_prob(y)\n",
    "    \n",
    "    # prior probability of y_pred\n",
    "    log_prior = dist.Normal(0, 1).log_prob(y_pred)\n",
    "    \n",
    "    # variational probability of y_pred\n",
    "    log_p_q = dist.Normal(mu, sigma).log_prob(y_pred) \n",
    "    \n",
    "    # by taking the mean we approximate the expectation\n",
    "    return (likelihood + log_prior - log_p_q).mean()\n",
    "\n",
    "def det_loss(y_pred, y, mu, log_var):\n",
    "    return -elbo(y_pred, y, mu, log_var)\n",
    "\n",
    "    \n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "epochs = 1500\n",
    "\n",
    "m = VI()\n",
    "optim = torch.optim.Adam(m.parameters(), lr=0.005)\n",
    "\n",
    "for epoch in range(epochs):\n",
    "    optim.zero_grad()\n",
    "    y_pred, mu, log_var = m(X)\n",
    "    loss = det_loss(y_pred, Y, mu, log_var)\n",
    "    loss.backward()\n",
    "    optim.step()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(0.5138)\n"
     ]
    }
   ],
   "source": [
    "with torch.no_grad():\n",
    "    print(torch.abs(m(X)[1] - Y).mean())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# draw samples from Q(theta)\n",
    "with torch.no_grad():\n",
    "    y_pred = torch.cat([m(X)[0] for _ in range(1000)], dim=1)\n",
    "    \n",
    "# Get some quantiles\n",
    "q1, mu, q2 = np.quantile(y_pred, [0.05, 0.5, 0.95], axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.collections.PolyCollection at 0x7f11cc2a3860>"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 37920 characters]\n",
      "text/plain": [
       "<Figure size 1152x432 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.figure(figsize=(16, 6))\n",
    "plt.scatter(X, Y)\n",
    "plt.plot(X, mu)\n",
    "plt.fill_between(X.flatten(), q1, q2, alpha=0.2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "def det_loss(y, y_pred, mu, log_var):    \n",
    "    reconstruction_error = (0.5 * (y - y_pred)**2).sum()\n",
    "    kl_divergence = (-0.5 * torch.sum(1 + log_var - mu**2 - log_var.exp()))\n",
    "\n",
    "    return reconstruction_error + kl_divergence"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$ Q(w) = \\mathcal{N}(\\mu_{\\theta}, \\sigma^2_{\\theta}) $ "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LinearVariational(nn.Module):\n",
    "    \"\"\"\n",
    "    Mean field approximation of nn.Linear\n",
    "    \"\"\"\n",
    "    def __init__(self, in_features, out_features, parent, n_batches, bias=True):\n",
    "        super().__init__()\n",
    "        self.in_features = in_features\n",
    "        self.out_features = out_features\n",
    "        self.include_bias = bias        \n",
    "        self.parent = parent\n",
    "        self.n_batches = n_batches\n",
    "        \n",
    "        if getattr(parent, 'accumulated_kl_div', None) is None:\n",
    "            parent.accumulated_kl_div = 0\n",
    "            \n",
    "        # Initialize the variational parameters.\n",
    "        # 𝑄(𝑤)=N(𝜇_𝜃,𝜎2_𝜃)\n",
    "        # Do some random initialization with 𝜎=0.01\n",
    "        self.w_mu = nn.Parameter(\n",
    "            torch.FloatTensor(in_features, out_features).normal_(mean=0, std=0.001)\n",
    "        )\n",
    "        # proxy for variance\n",
    "        # log(1 + exp(ρ))◦ eps\n",
    "        self.w_p = nn.Parameter(\n",
    "            torch.FloatTensor(in_features, out_features).normal_(mean=-2.5, std=0.001)\n",
    "        )\n",
    "        if self.include_bias:\n",
    "            self.b_mu = nn.Parameter(\n",
    "                torch.zeros(out_features)\n",
    "            )\n",
    "            # proxy for variance\n",
    "            self.b_p = nn.Parameter(\n",
    "                torch.zeros(out_features)\n",
    "            )\n",
    "        \n",
    "    def reparameterize(self, mu, p):\n",
    "        sigma = torch.log(1 + torch.exp(p)) \n",
    "        eps = torch.randn_like(sigma)\n",
    "        return mu + (eps * sigma)\n",
    "    \n",
    "    def kl_divergence(self, z, mu_theta, p_theta, prior_sd=1):\n",
    "        log_prior = dist.Normal(0, prior_sd).log_prob(z) \n",
    "        log_p_q = dist.Normal(mu_theta, torch.log(1 + torch.exp(p_theta))).log_prob(z) \n",
    "        return (log_p_q - log_prior).sum() / self.n_batches\n",
    "                    \n",
    "\n",
    "    def forward(self, x):\n",
    "        w = self.reparameterize(self.w_mu, self.w_p)\n",
    "        \n",
    "        if self.include_bias:\n",
    "            b = self.reparameterize(self.b_mu, self.b_p)\n",
    "        else:\n",
    "            b = 0\n",
    "            \n",
    "        z = x @ w + b\n",
    "        \n",
    "        self.parent.accumulated_kl_div += self.kl_divergence(w, \n",
    "                                                             self.w_mu,\n",
    "                                                             self.w_p, \n",
    "                                                             )\n",
    "        if self.include_bias:\n",
    "            self.parent.accumulated_kl_div += self.kl_divergence(b, \n",
    "                                                                 self.b_mu, \n",
    "                                                                 self.b_p,\n",
    "                                                                 )\n",
    "        return z\n",
    "            \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.2807],\n",
       "        [0.2807],\n",
       "        [0.2828],\n",
       "        [0.2903],\n",
       "        [0.2968],\n",
       "        [0.3007],\n",
       "        [0.3134],\n",
       "        [0.3250],\n",
       "        [0.3283],\n",
       "        [0.3316]], grad_fn=<SliceBackward>)"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "@dataclass\n",
    "class KL:\n",
    "    accumulated_kl_div = 0\n",
    "\n",
    "class Model(nn.Module):\n",
    "    def __init__(self, in_size, hidden_size, out_size, n_batches):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.kl_loss = KL\n",
    "        \n",
    "        self.layers = nn.Sequential(\n",
    "            LinearVariational(in_size, hidden_size, self.kl_loss, n_batches),\n",
    "            nn.ReLU(),\n",
    "            LinearVariational(hidden_size, hidden_size, self.kl_loss, n_batches),\n",
    "            nn.ReLU(),\n",
    "            LinearVariational(hidden_size, out_size, self.kl_loss, n_batches)\n",
    "        )\n",
    "    \n",
    "    @property\n",
    "    def accumulated_kl_div(self):\n",
    "        return self.kl_loss.accumulated_kl_div\n",
    "    \n",
    "    def reset_kl_div(self):\n",
    "        self.kl_loss.accumulated_kl_div = 0\n",
    "            \n",
    "    def forward(self, x):\n",
    "        return self.layers(x)\n",
    "    \n",
    "m = Model(1, 20, 1, 1)\n",
    "m(X)[:10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 417,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(908.0202, grad_fn=<AddBackward0>)"
      ]
     },
     "execution_count": 417,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m.accumulated_kl_div"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "912758.8125\n",
      "5630.6767578125\n",
      "4918.69482421875\n",
      "4853.455078125\n",
      "4451.697265625\n",
      "4369.171875\n",
      "4522.48388671875\n",
      "4442.5087890625\n"
     ]
    }
   ],
   "source": [
    "epochs = 2000\n",
    "\n",
    "def det_loss(y, y_pred, model):\n",
    "    batch_size = y.shape[0]\n",
    "    reconstruction_error = -dist.Normal(y_pred, .1).log_prob(y).sum() \n",
    "    kl = model.accumulated_kl_div\n",
    "    model.reset_kl_div()\n",
    "    return (reconstruction_error + kl) \n",
    "    \n",
    "m = Model(1, 20, 1, n_batches=1)\n",
    "optim = torch.optim.Adam(m.parameters(), lr=0.01)\n",
    "\n",
    "for epoch in range(epochs):\n",
    "    optim.zero_grad()\n",
    "    y_pred = m(X)\n",
    "    loss = det_loss(y_pred, Y, m)\n",
    "    loss.backward()\n",
    "    optim.step()\n",
    "    \n",
    "    if epoch % 250 == 0:\n",
    "        print(loss.data.item())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.collections.PolyCollection at 0x7f1196a75a58>"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 28372 characters]\n",
      "text/plain": [
       "<Figure size 1152x432 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "with torch.no_grad():\n",
    "    trace = np.array([m(X).flatten().numpy() for _ in range(1000)]).T\n",
    "q_25, q_95 = np.quantile(trace, [0.05, 0.95], axis=1)\n",
    "plt.figure(figsize=(16, 6))\n",
    "plt.plot(X, trace.mean(1))\n",
    "plt.scatter(X, Y)\n",
    "plt.fill_between(X.flatten(), q_25, q_95, alpha=0.2)\n"
   ]
  }
 ],
 "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
}
