notebooks/Henbun_structure2.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# More about Henbun\n",
    "\n",
    "In this notebook, we describe more details about Henbun through Variational method."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Libraries to be used."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import tensorflow as tf\n",
    "import Henbun as hb"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Toy data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "X = np.linspace(0,6,40).reshape(-1,1)\n",
    "Y = 0.5*X + np.random.randn(40,1)*0.3 + 0.4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=(5,3))\n",
    "plt.plot(X, Y, 'o', label='data')\n",
    "plt.plot(X, 0.5*X+0.4, '--k', label='true')\n",
    "plt.legend(loc='best')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Variational linear model.\n",
    "\n",
    "Here, we construct a linear model within Bayesian framework.\n",
    "\n",
    "The likelihood is assumed as Gaussian,\n",
    "$$\n",
    "p(y_i | a, b, x_i) = \\mathcal{N}(y_i | a + b*x_i, c)\n",
    "$$\n",
    "where $\\sigma$ is the variance parameter.\n",
    "\n",
    "We assume weak prior for $a$, $b$ and $\\sigma$,\n",
    "$$\n",
    "p(a) = \\mathcal{N}(a|0,1) \\\\\n",
    "p(b) = \\mathcal{N}(b|0,1) \\\\\n",
    "p(c) = \\mathcal{N}(c|0,1) \\\\\n",
    "$$\n",
    "\n",
    "In this model, we approximate the posterior distribution by independent Gaussian, i.e.\n",
    "$$\n",
    "p(a|\\mathbf{x},\\mathbf{y}) \\sim q(a) = \\mathcal{N}(a| \\mu_a, \\sigma_a) \\\\\n",
    "p(b|\\mathbf{x},\\mathbf{y}) \\sim q(b) = \\mathcal{N}(b| \\mu_b, \\sigma_b) \\\\\n",
    "p(c|\\mathbf{x},\\mathbf{y}) \\sim q(c) = \\mathcal{N}(c| \\mu_c, \\sigma_c) \\\\\n",
    "$$\n",
    "where $\\mu_a, \\mu_b, \\mu_c$ and $\\sigma_a, \\sigma_b, \\sigma_c$ are the variational parameters to be optimized.\n",
    "\n",
    "By this approximation, the free energy (evidence lower bound: ELBO) to be maximized is\n",
    "$$\n",
    "\\mathrm{ELBO}(a, b, c) = \n",
    "\\int q(a)q(b)q(c) \\log p(\\mathbf{y}|\\mathbf{x}, a, b, c)\\; \\mathrm{d}a\\;\\mathrm{d}b\\;\\mathrm{d}c \\\\\n",
    "- \\mathrm{KL}[q(a)||p(a)]  \\\\\n",
    "- \\mathrm{KL}[q(b)||p(b)]  \\\\\n",
    "- \\mathrm{KL}[q(c)||p(c)]\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Construct the variational model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Construct a variational linear model\n",
    "class VariationalLinearModel(hb.model.Model):\n",
    "    def setUp(self):\n",
    "        # data should be stored in hb.param.Data class\n",
    "        self.X = hb.param.Data(X)\n",
    "        self.Y = hb.param.Data(Y)\n",
    "        \n",
    "        # Variational parameters \n",
    "        self.a = hb.variationals.Normal(shape=[1])\n",
    "        self.b = hb.variationals.Normal(shape=[1])\n",
    "        self.c = hb.variationals.Normal(shape=[1])\n",
    "        \n",
    "        # To access the posterior samples, \n",
    "        # we define op that returns samples of the fit.\n",
    "        # As described in Henbun_structure.ipynb, tf_mode is used here.\n",
    "        with self.tf_mode():\n",
    "            self.fit = self.a + self.b*self.X\n",
    "        \n",
    "    @hb.model.AutoOptimize()\n",
    "    def elbo(self):\n",
    "        \"\"\"\n",
    "        This method returns the sum of log-likelihood and log-prior.\n",
    "        \"\"\"\n",
    "        # log likelihood\n",
    "        log_lik = hb.densities.gaussian(self.Y, self.fit, tf.exp(self.c))\n",
    "        # KL() methods automatically gather KL for all the variational parameters.\n",
    "        KL = self.KL()\n",
    "        return tf.reduce_sum(log_lik) - KL"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "vlinear_model = VariationalLinearModel()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## What values are inside `hb.variationals.Normal`\n",
    "\n",
    "hb.variationals.Normal is a variational parameter with a Normal prior.  \n",
    "This class has two hb.param.Variables, `q_mu` and `q_sqrt`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "vlinear_model.initialize()\n",
    "print(vlinear_model.a.q_mu.value, vlinear_model.a.q_sqrt.value)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`q_mu` corresponds to $\\mu_a$ and `q_sqrt` corresponds to $\\sqrt{\\log \\sigma_a}$.\n",
    "\n",
    "Henbun will seek the best variational parameters (`q_mu` and `q_sqrt`) that maximizes ELBO."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Compilation and optimization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "vlinear_model.elbo().compile()\n",
    "vlinear_model.elbo().optimize(maxiter=50000)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# a\n",
    "print(vlinear_model.a.q_mu.value, vlinear_model.a.q_sqrt.value)\n",
    "# b\n",
    "print(vlinear_model.b.q_mu.value, vlinear_model.b.q_sqrt.value)\n",
    "# c\n",
    "print(vlinear_model.c.q_mu.value, vlinear_model.c.q_sqrt.value)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# plot the samples from the posterior\n",
    "plt.figure(figsize=(5,3))\n",
    "plt.plot(X, Y, 'o', label='data')\n",
    "plt.plot(X, 0.5*X+0.4, '--k', label='true')\n",
    "for _ in range(50):\n",
    "    plt.plot(X, vlinear_model.run(vlinear_model.fit), '-r', alpha=0.1)\n",
    "plt.plot()\n",
    "plt.legend(loc='best')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Note\n",
    "\n",
    "Note that we approximate the independent posterior for a, b, c, although the true posterior should have correlation.\n",
    "\n",
    "Henbun also provides 'fullrank' variational approximation.  \n",
    "See docstring for `hb.variationals.Variational`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# How does variational model work.\n",
    "\n",
    "## Variational expectation\n",
    "The main difficulty for calculating ELBO is the variational-expectation term\n",
    "$$\n",
    "\\int q(a)q(b)q(c) \\log p(\\mathbf{y}|\\mathbf{x}, a, b, c)\\; \\mathrm{d}a\\;\\mathrm{d}b\\;\\mathrm{d}c\n",
    "$$\n",
    "where in a, b, c are usually multidimensional (in our case, it is 3-dimensional integration).\n",
    "\n",
    "We approximate this integral by Monte-Carlo method,\n",
    "$$\n",
    "\\mathrm{ELBO} \\sim \\log p(\\mathbf{y}|\\mathbf{x}, a^\\mathrm{s}, b^\\mathrm{s}, c^\\mathrm{s})\n",
    "$$\n",
    "where $a^\\mathrm{s}, b^\\mathrm{s}, c^\\mathrm{s}$ are the samples from $q(a), q(b), q(c)$ respectively.\n",
    "\n",
    "Since this equation gives unbiased approximation of the True ELBO,  \n",
    "we can optimize this by stochastic optimization.\n",
    "\n",
    "## kullback leibler divergence\n",
    "The KL term is also approximated by Monte-Carlo method,\n",
    "$$\n",
    "\\mathrm{KL}[q(a)||p(a)]\n",
    "= \\int q(a) \\log\\frac{q(a)}{p(a)} \\mathrm{d}a\n",
    "\\sim a^\\mathrm{s} \\log\\frac{q(a^\\mathrm{s})}{p(a^\\mathrm{s})}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": []
  }
 ],
 "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.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}