neural_networks/fit_dist_with_pyro/fit_beta_distribution.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fitting a Beta Distribution with Pyro\n",
    "\n",
    "> Needs re-run\n",
    "\n",
    "Here we assume we are flipping a slightly biased coin.\n",
    "We think the probability of a heads is close to 0.5, but we are not sure.\n",
    "We want to fit a beta distribution to the random observed data.\n",
    "\n",
    "References:\n",
    "  * [http://pyro.ai/examples/svi_part_i.html](http://pyro.ai/examples/svi_part_i.html)\n",
    "\n",
    "Import the required libraries:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "from torch.distributions import constraints\n",
    "import pyro\n",
    "import pyro.infer\n",
    "import pyro.optim\n",
    "import pyro.distributions as dist\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "plt.style.use(\"seaborn-whitegrid\")\n",
    "torch.manual_seed(0)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generate observed data\n",
    "We use the distribution module of pytorch to generate random data from Bernoulli trials with a known probability of success $$P(p)=0.4$$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate data from actual distribution\n",
    "true_dist = dist.Bernoulli(0.4)\n",
    "n = 100\n",
    "data = true_dist.sample(sample_shape=(n, 1))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prior distribution\n",
    "We propose the probability of a heads comes from a beta distribution.\n",
    "We assume that the probability is close to 0.5 but with some error.\n",
    "This is characterised in the prior distribution as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prior = dist.Beta(10, 10)\n",
    "\n",
    "plt.figure(num=None, figsize=(10, 6), dpi=80)\n",
    "x_range = np.linspace(0, 1, num=100)\n",
    "\n",
    "y_values = torch.exp(prior.log_prob(torch.tensor(x_range)))\n",
    "plt.plot(x_range, y_values, label=\"prior\")\n",
    "\n",
    "plt.title(\"PDF\")\n",
    "plt.legend()\n",
    "plt.savefig(\"images/beta_prior_pdf.png\")\n",
    "plt.show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![](images/beta_prior_pdf.png)\n",
    "\n",
    "## Analytical Posterior\n",
    "Using the random data we have generated, we can calculate the posterior distribution.\n",
    "In the case of a beta distribution - the posterior has an analytical solution,\n",
    "based on conjugacy:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Analytical posterior\n",
    "posterior = dist.Beta(\n",
    "    prior.concentration1 + data.sum(),\n",
    "    prior.concentration0 + len(data) - data.sum(),\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Variational inference\n",
    "We can solve the same problem with variational inference using `pyro`.\n",
    "We setup the model to sample from a Bernoulli trial,\n",
    "where the probability of a heads comes from a beta distribution.\n",
    "The model is conditioned to give the generated data when it is sampled from."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def data_model(params):\n",
    "    # returns a Bernoulli trial outcome\n",
    "    beta = pyro.sample(\"beta_dist\", dist.Beta(params[0], params[1]))\n",
    "    return pyro.sample(\"data_dist\", dist.Bernoulli(beta))\n",
    "\n",
    "\n",
    "conditioned_data_model = pyro.condition(data_model, data={\"data_dist\": data})"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The guide function creates a pyro beta distribution object given a set of parameters,\n",
    "which we will track."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def guide(params):\n",
    "    # returns the Bernoulli probablility\n",
    "    alpha = pyro.param(\n",
    "        \"alpha\", torch.tensor(params[0]), constraint=constraints.positive\n",
    "    )\n",
    "    beta = pyro.param(\"beta\", torch.tensor(params[1]), constraint=constraints.positive)\n",
    "    return pyro.sample(\"beta_dist\", dist.Beta(alpha, beta))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We iterate over the above functions, starting from our prior distribution.\n",
    "Each step we converge towards an ideal posterior form of the guide."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "svi = pyro.infer.SVI(\n",
    "    model=conditioned_data_model,\n",
    "    guide=guide,\n",
    "    optim=pyro.optim.SGD({\"lr\": 0.001, \"momentum\": 0.8}),\n",
    "    loss=pyro.infer.Trace_ELBO(),\n",
    ")\n",
    "\n",
    "params_prior = [prior.concentration1, prior.concentration0]\n",
    "\n",
    "# Iterate over all the data and store results\n",
    "losses, alpha, beta = [], [], []\n",
    "pyro.clear_param_store()\n",
    "\n",
    "num_steps = 3000\n",
    "for t in range(num_steps):\n",
    "    losses.append(svi.step(params_prior))\n",
    "    alpha.append(pyro.param(\"alpha\").item())\n",
    "    beta.append(pyro.param(\"beta\").item())\n",
    "\n",
    "posterior_vi = dist.Beta(alpha[-1], beta[-1])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We plot the trajectories of the parameters to show they have converged sufficiently:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(num=None, figsize=(10, 6), dpi=80)\n",
    "plt.plot(alpha, label=\"alpha\")\n",
    "plt.plot(beta, label=\"beta\")\n",
    "plt.title(\"Parameter trajectories\")\n",
    "plt.xlabel(\"Iteration\")\n",
    "plt.legend()\n",
    "plt.savefig(\"images/beta_trajectories.png\")\n",
    "plt.show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![](images/beta_trajectories.png)\n",
    "\n",
    "## Comparing distributions\n",
    "We can compare the variational inference distribution to the analytical posterior."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(num=None, figsize=(10, 6), dpi=80)\n",
    "x_range = np.linspace(0, 1, num=100)\n",
    "\n",
    "y_values = torch.exp(prior.log_prob(torch.tensor(x_range)))\n",
    "plt.plot(x_range, y_values, label=\"prior\")\n",
    "\n",
    "y_values = torch.exp(posterior.log_prob(torch.tensor(x_range)))\n",
    "plt.plot(x_range, y_values, label=\"posterior\")\n",
    "\n",
    "y_values = torch.exp(posterior_vi.log_prob(torch.tensor(x_range)))\n",
    "plt.plot(x_range, y_values, label=\"posterior_vi\")\n",
    "\n",
    "plt.title(\"PDF\")\n",
    "plt.legend()\n",
    "plt.savefig(\"images/beta_pdfs.png\")\n",
    "plt.show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![](images/beta_pdfs.png)\n",
    "\n",
    "The estimated posterior from variational inference is very similar to the analytical posterior.\n",
    "\n",
    "## Extra: More Data\n",
    "If we generate much more data and repeat the process,\n",
    "we can get a tighter posterior:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 10000\n",
    "data_more = true_dist.sample(sample_shape=(n, 1))\n",
    "conditioned_data_model_more = pyro.condition(data_model, data={\"data_dist\": data_more})\n",
    "\n",
    "svi = pyro.infer.SVI(\n",
    "    model=conditioned_data_model_more,\n",
    "    guide=guide,\n",
    "    optim=pyro.optim.SGD({\"lr\": 0.0001, \"momentum\": 0.8}),\n",
    "    loss=pyro.infer.Trace_ELBO(),\n",
    ")\n",
    "\n",
    "# Iterate over all the data and store results\n",
    "losses, alpha, beta = [], [], []\n",
    "pyro.clear_param_store()\n",
    "\n",
    "num_steps = 3000\n",
    "for t in range(num_steps):\n",
    "    losses.append(svi.step(params_prior))\n",
    "    alpha.append(pyro.param(\"alpha\").item())\n",
    "    beta.append(pyro.param(\"beta\").item())\n",
    "\n",
    "posterior_vim = dist.Beta(alpha[-1], beta[-1])\n",
    "\n",
    "plt.figure(num=None, figsize=(10, 6), dpi=80)\n",
    "x_range = np.linspace(0, 1, num=1000)\n",
    "\n",
    "y_values = torch.exp(prior.log_prob(torch.tensor(x_range)))\n",
    "plt.plot(x_range, y_values, label=\"prior\")\n",
    "\n",
    "y_values = torch.exp(posterior.log_prob(torch.tensor(x_range)))\n",
    "plt.plot(x_range, y_values, label=\"posterior\")\n",
    "\n",
    "y_values = torch.exp(posterior_vi.log_prob(torch.tensor(x_range)))\n",
    "plt.plot(x_range, y_values, label=\"posterior_vi\")\n",
    "\n",
    "y_values = torch.exp(posterior_vim.log_prob(torch.tensor(x_range)))\n",
    "plt.plot(x_range, y_values, label=\"posterior_vi_more_data\")\n",
    "\n",
    "plt.title(\"PDF\")\n",
    "plt.legend()\n",
    "plt.savefig(\"images/beta_pdfs_more.png\")\n",
    "plt.show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![](images/beta_pdfs_more.png)"
   ]
  }
 ],
 "metadata": {
  "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
  },
  "orig_nbformat": 2
 },
 "nbformat": 4,
 "nbformat_minor": 2
}