{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import cupy as cp\n",
    "\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib inline\n",
    "%config InlineBackend.figure_format = 'retina'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "-"
    }
   },
   "source": [
    "# A Fundamental Property of Gaussians\n",
    "\n",
    "A multivariate Gaussian is nothing more than a generalization of the univariate Gaussian. \n",
    "\n",
    "We parameterize univariate Gaussians with a $\\mu$ and $\\sigma$, where $\\mu$ and $\\sigma$ are scalars.\n",
    "\n",
    "A bivariate Gaussian is two univariate Gaussians that may also share a relationship to one another. We can jointly model both Gaussians by modelling not just how they vary independently, but also how they vary with one another. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def np(a):\n",
    "    return cp.asnumpy(a)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mu1 = 0\n",
    "mu2 = 0\n",
    "sig11 = 3.0\n",
    "sig12 = -2.0\n",
    "sig21 = -2.0\n",
    "sig22 = 4.0\n",
    "\n",
    "mean = cp.array([mu1, mu2])\n",
    "cov = cp.array([[sig11, sig12], [sig21, sig22]])\n",
    "draws = cp.random.multivariate_normal(mean, cov, size=1000)\n",
    "plt.scatter(*np(draws.T))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "One of the fundamental properties of Gaussians is that if you have a Multivariate Gaussian (e.g. a joint distribution of 2 or more Gaussian random variables), if we condition on any subset of Gaussians, the joint distribution of the rest of the Gaussians can be found analytically. There's a formula, and it's expressed in code below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x1 = 0\n",
    "mu_2g1 = mu2 + sig21 * 1 / sig11 * (x1 - mu1)\n",
    "sig_2g1 = sig22 - sig21 * 1 / sig11 * sig12\n",
    "mu_2g1, sig_2g1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Go ahead and play with the slider below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ipywidgets import interact, FloatSlider, IntSlider\n",
    "\n",
    "@interact(x1=FloatSlider(min=-4, max=4, continuous_update=False))\n",
    "def plot_conditional(x1):\n",
    "    fig, axes = plt.subplots(nrows=1, ncols=2, figsize=(12, 6))\n",
    "    axes[0].scatter(*np(draws.T))\n",
    "    axes[0].vlines(x=x1, ymin=draws[:, 1].min(), ymax=draws[:, 1].max(), color='red')\n",
    "    axes[0].hlines(y=0, xmin=draws[:, 0].min(), xmax=draws[:, 0].max(), color='black')\n",
    "    \n",
    "    # Compute Conditional\n",
    "    mu_2g1 = mu2 + sig21 * 1 / sig11 * (x1 - mu1)\n",
    "    sig_2g1 = sig22 - sig21 * 1 / sig11 * sig12\n",
    "    x2_draws = cp.random.normal(mu_2g1, sig_2g1, size=10000)\n",
    "    axes[1].hist(np(x2_draws), bins=100, color='red')\n",
    "    axes[1].vlines(x=0, ymin=0, ymax=400, color='black')\n",
    "    axes[1].set_xlim(-10, 10)\n",
    "    axes[1].set_ylim(0, 400)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Implementing GP Prior\n",
    "\n",
    "When we use a GP, we're essentially modelling the **outputs** as being described by a joint Gaussian distribution. \n",
    "\n",
    "We would like to be able to specify the covariance matrix as a function of the distances between the inputs - regardless of whether the inputs are 1-D, 2-D, or more. That is the key to generalizing from 1D examples to the 2D examples commonly shown."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import seaborn as sns\n",
    "\n",
    "n = 50\n",
    "x_train = cp.linspace(-5, 5, n).reshape(-1, 1)\n",
    "\n",
    "# sns.heatmap(x_train - x_train.T, cmap='RdBu')\n",
    "\n",
    "def sq_exp(x1, x2):\n",
    "    \"\"\"\n",
    "    Squared exponential kernel.\n",
    "    \n",
    "    Assumes that x1 and x2 have the same shape.\n",
    "    \"\"\"\n",
    "    diff = x1 - x2.T\n",
    "    sqdiff = cp.power(diff, 2)\n",
    "    return cp.exp(-0.5 * sqdiff)\n",
    "\n",
    "\n",
    "\n",
    "sns.heatmap(np(sq_exp(x_train, x_train)), cmap='RdBu')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Draw from prior."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "K = sq_exp(x_train, x_train)\n",
    "eps = 1E-6 * cp.eye(n)  # \n",
    "L = cp.linalg.cholesky(K + eps)\n",
    "\n",
    "f_prior = cp.dot(L, cp.random.normal(size=(n, 10)))\n",
    "plt.plot(np(x_train), np(f_prior))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def true_function_1d(x):\n",
    "    x = x + 1E-10\n",
    "    return cp.sin(x)\n",
    "\n",
    "n = 200\n",
    "x_samp = cp.array([2, 18, -10, 10, -12, 12, -2, 5, -13, 6, -18, 8, -8, 0, 15]).reshape(-1, 1)\n",
    "f_samp = true_function_1d(x_samp)\n",
    "K_samp = sq_exp(x_samp, x_samp)\n",
    "eps = 1E-6 * cp.eye(len(x_samp))\n",
    "L_samp = cp.linalg.cholesky(K_samp + eps)\n",
    "\n",
    "x_s = cp.linspace(-20, 20, n).reshape(-1, 1)\n",
    "K_ss = sq_exp(x_s, x_s)\n",
    "K_s = sq_exp(x_samp, x_s)\n",
    "\n",
    "mu_cond = K_s.T @ cp.linalg.inv(K_samp) @ f_samp\n",
    "sig_cond = K_ss - K_s.T @ cp.linalg.inv(K_samp) @ K_s\n",
    "\n",
    "f_posterior = cp.random.multivariate_normal(mu_cond.flatten(), sig_cond, size=100)\n",
    "\n",
    "for f in f_posterior:\n",
    "    plt.plot(np(x_s), np(f), color='grey', alpha=0.1)\n",
    "plt.scatter(np(x_samp), np(f_samp))\n",
    "plt.plot(np(x_s), np(true_function_1d(x_s)))\n",
    "plt.plot(np(x_s), np(mu_cond.flatten()))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sig_cond.min()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sns.heatmap(np(sig_cond), cmap='RdBu')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can rewrite extend this code to apply in two dimensions. Let's say that our data lived on a grid, rather than on a single dimension. A periodic function is applied on a 2D grid."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def true_function(x1, x2):\n",
    "    return cp.sin(x1) + cp.sin(x2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Prior"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's sample a prior from a 2D plane."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sq_exp??"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import scipy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x1 = cp.array([[2, 2], [2, 1], [1, 2], [1, 1]])\n",
    "\n",
    "def sq_exp2d(x1, x2):\n",
    "    d = scipy.spatial.distance.cdist(np(x1), np(x2))\n",
    "    return cp.exp(-0.5 * cp.power(cp.array(d), 2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x1 = cp.linspace(-5, 5, 20)\n",
    "x2 = cp.linspace(-5, 5, 20)\n",
    "xx1, xx2 = cp.meshgrid(x1, x2, )\n",
    "z = true_function(xx1, xx1)\n",
    "h = plt.contourf(np(x1), np(x2), np(z))\n",
    "plt.gca().set_aspect('equal')\n",
    "plt.title('true function')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "true_function(xx1, xx2).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "\n",
    "fig = plt.figure()\n",
    "ax = fig.add_subplot(111, projection='3d')\n",
    "z = true_function(xx1, xx2)\n",
    "ax.plot_surface(np(xx1), np(xx2), np(z))\n",
    "ax.set_title('true function')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll simulate sampling 5 starting points."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_samp = cp.array([[0, 3], [1, 2], [1, -5], [-2, -2], [-2, 2], [2, -2], [3, 5]])\n",
    "\n",
    "f_samp = true_function(x_samp[:, 0], x_samp[:, 1])\n",
    "\n",
    "K_samp = sq_exp2d(x_samp, x_samp)\n",
    "\n",
    "eps = 1E-6 * cp.eye(len(x_samp))\n",
    "L_samp = cp.linalg.cholesky(K_samp + eps)\n",
    "\n",
    "n = 35\n",
    "x_points = cp.linspace(-5, 5, n)\n",
    "xx1, xx2 = cp.meshgrid(x_points, x_points)\n",
    "x_spts = cp.vstack([xx1.flatten(), xx2.flatten()])\n",
    "K_ss = sq_exp2d(x_spts.T, x_spts.T)\n",
    "K_s = sq_exp2d(x_samp, x_spts.T)\n",
    "\n",
    "mu_cond = K_s.T @ cp.linalg.inv(K_samp) @ f_samp.flatten()\n",
    "sig_cond = K_ss - K_s.T @ cp.linalg.inv(K_samp) @ K_s\n",
    "\n",
    "n_samps = 1000\n",
    "f_posterior = cp.random.multivariate_normal(mu_cond, sig_cond, size=n_samps)\n",
    "\n",
    "# for f in f_posterior:\n",
    "#     plt.plot(x_s, f, color='grey', alpha=0.1)\n",
    "# plt.scatter(x_samp, f_samp)\n",
    "# plt.plot(x_s, true_function(x_train))\n",
    "# plt.plot(x_s, mu_cond.flatten())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f_posterior.reshape(n_samps, n, n).max(axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mu_cond.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure()\n",
    "ax = fig.add_subplot(111, projection='3d')\n",
    "ax.plot_surface(np(xx1), np(xx2), np(mu_cond.reshape(n, n)))\n",
    "ax.set_title('mean')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lower, upper = cp.percentile(f_posterior.reshape(n_samps, n, n), [2.5, 97.5], axis=0)\n",
    "uncertainty = upper - lower"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "uncertainty.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(nrows=1, ncols=3, figsize=(15, 5), sharex=True, sharey=True)\n",
    "axes[0].contourf(np(xx1), np(xx2), np(uncertainty), levels=50, cmap='inferno')\n",
    "axes[0].set_title('95% interval size')\n",
    "axes[0].scatter(*np(x_samp.T), color='red')\n",
    "axes[0].set_aspect('equal')\n",
    "axes[0].set_ylim(-5, 5)\n",
    "axes[0].set_xlim(-5, 5)\n",
    "\n",
    "axes[1].contourf(np(xx1), np(xx2), np(true_function(xx1, xx2)), levels=50)\n",
    "axes[1].set_title('ground truth')\n",
    "axes[1].scatter(*np(x_samp.T), color='red')\n",
    "axes[1].set_aspect('equal')\n",
    "\n",
    "axes[2].contourf(np(xx1), np(xx2), np(mu_cond.reshape(n, n)), levels=50)\n",
    "axes[2].set_title('mean')\n",
    "axes[2].scatter(*np(x_samp.T), color='red')\n",
    "axes[2].set_aspect('equal')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the plot sabove, red dots mark where we have sampled points on the 2D grid. \n",
    "\n",
    "The left plot shows the size of the 95% prediction interval at each point on the grid. We can see that we have the smallest uncertainty where we have sampled.\n",
    "\n",
    "The middle plot shows ground truth, and the values where we have sampled data.\n",
    "\n",
    "The right plot shows the value of the mean where we have sampled. It is evident that where we do not have better data, the function evaluations default to the mean function."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Parting Thoughts\n",
    "\n",
    "The key ingredient of a GP: A Kernel that can model \"distance\" of some kind between every pair of inputs. Thus, it isn't the number of dimensions that is limiting; it is the number of *data points that have been sampled* that is limiting! (Inversion of the matrix only depends on the data that we are conditioning on, and that is of order $O(n^3)$.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x2_new, x1_new = cp.where(uncertainty == uncertainty.max())  # row, then column, i.e. x2 then x1.\n",
    "xx1_s, xx2_s = cp.meshgrid(x_points, x_points)\n",
    "xx1_s.flatten()[x1_new], xx2_s.flatten()[x2_new]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "bayesian",
   "language": "python",
   "name": "bayesian"
  },
  "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.6.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}