neural_networks/fit_dist_with_pyro/fit_pyro_distribution.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fitting a Distribution with Pyro\n",
    "\n",
    "> Needs re-run\n",
    "\n",
    "In this simple example we will fit a Gaussian distribution to random data from a gaussian with some known mean and standard deviation.\n",
    "We want to estimate a distribution that best fits the data using variational inference with Pyro.\n",
    "\n",
    "References:\n",
    "  * [https://pyro.ai/examples/intro_part_ii.html](https://pyro.ai/examples/intro_part_ii.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",
    "from scipy.stats import norm\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "plt.style.use(\"seaborn-whitegrid\")\n",
    "np.random.seed(0)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generate observed data\n",
    "We use `numpy` to generate random data from a Gaussian distribution with a known mean and standard deviation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 1000\n",
    "x = np.random.randn(n, 1)\n",
    "std = 4\n",
    "mu = 2\n",
    "x = x * std + mu\n",
    "\n",
    "print(f\"Shape: {x.shape}\")\n",
    "print(f\"Mean: {np.mean(x)}\")\n",
    "print(f\"Standard deviation: {np.std(x)}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```\n",
    "Shape: (1000, 1)\n",
    "Mean: 1.8189731700392184\n",
    "Standard deviation: 3.9481326346761034\n",
    "```"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we have generated the data we need to use it to estimate the original distribution parameters.\n",
    "\n",
    "The traditional approach, when using a Gaussian distribution, would be to simply take the mean and standard deviation, as we did above. The following approach, however, will give us a full distribution rather than a point estimate, and it should generalise to many different problems.\n",
    "\n",
    "## Variational inference\n",
    "This is solving a bayesian posterior estimation problem with gradient descent.\n",
    "\n",
    "We specify the type of distribution the posterior should follow. We then tune the parameters to minimise the difference between the estimated posterior and the actual posterior. The difference measure used is the KL divergence. (KL divergence is not symmetrical, not sure why its the way round it is)\n",
    "\n",
    "$$D_{\\mathrm{KL}}(Q \\parallel P) \\triangleq \\sum_\\mathbf{Z}  Q(\\mathbf{Z}) \\log \\frac{Q(\\mathbf{Z})}{P(\\mathbf{Z}\\mid \\mathbf{X})}$$\n",
    "\n",
    "It starts with a class of approximating distributions. Then it finds the best approximation to the posterior distribution. It minimise the Kullback-Leibler divergence between our approximate distribution and the posterior. This is equivalent to maximising the evidence lower bound (ELBO). This requires calculating the joint distribution, rather than the true posterior.\n",
    "\n",
    "This is now an optimisation problem which can be solved with gradient descent algorithms.\n",
    "\n",
    "This is different to MCMC. With MCMC we get a numerical approximation to the exact posterior using a set of samples, Variational Bayes provides a locally-optimal, exact analytical solution to an approximation of the posterior.\n",
    "References:\n",
    "* [https://www.youtube.com/watch?v=3KGZDC3-_iY](https://www.youtube.com/watch?v=3KGZDC3-_iY)\n",
    "* [http://pyro.ai/examples/bayesian_regression.html](http://pyro.ai/examples/bayesian_regression.html)\n",
    "* [https://en.wikipedia.org/wiki/Variational_Bayesian_methods](https://en.wikipedia.org/wiki/Variational_Bayesian_methods)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Pyro approach\n",
    "`pyro` is a probabilistic library that sits on top of `pytorch` that enables variational inference.\n",
    "To solve this problem in pyro we need a few different components.\n",
    "\n",
    "We need a model of our data generating function."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Data generating function\n",
    "We assume a Guassian distribution as the model to generate our random data.\n",
    "This function takes parameters for our distributions and generates a random sample from the resulting distribution.\n",
    "Our model consists of a Gaussian distribution which has two priors: mean and standard deviation.\n",
    "These parameters come from distributions themselves.\n",
    "The mean is taken from another Gaussian distribution. The standard deviation comes from a Gamma distribution.\n",
    "\n",
    "This is represented as:\n",
    "\n",
    "$$x\\sim\\mathcal{N}\\left(\\mu,\\sigma^{2}\\right)$$\n",
    "\n",
    "$$\\mu\\sim\\mathcal{N}\\left(\\mu_{\\mu},\\mu_{\\sigma^{2}}\\right)$$\n",
    "\n",
    "$$\\sigma\\sim\\mathrm{Gamma}\\left(\\alpha,\\beta\\right)$$\n",
    "\n",
    "The parameters for these two distributions ($\\mu_{\\mu},\\mu_{\\sigma}, \\alpha,\\beta$) are the function inputs:\n",
    "```\n",
    "    params: [mu_prior + std_prior]\n",
    "        mu_prior - Gaussian - mu, std\n",
    "        std_prior - Gamma - a, b\n",
    "```\n",
    "We then condition the function, so that the samples produced from `data_dist` are enforced to match those from our original random data, `x`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def data_model(params):\n",
    "    mu_dist = pyro.sample(\"mu_dist\", dist.Normal(params[0], params[1]))\n",
    "    std_dist = pyro.sample(\"std_dist\", dist.Gamma(np.abs(params[2]), np.abs(params[3])))\n",
    "    return pyro.sample(\"data_dist\", dist.Normal(mu_dist, std_dist))\n",
    "\n",
    "\n",
    "conditioned_data_model = pyro.condition(\n",
    "    data_model, data={\"data_dist\": torch.tensor(x.flatten())}\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Guide function\n",
    "The guide function represents the family of distribution we want to consider as our posterior distribution, therefore it should be an approximation of the model posterior distribution. In this case we assume a Guassian distribution as the approximating class for the posterior distribution. This is an ideal case as we know the original data came from a Gaussian, in practice this would be based on domain knowledge.\n",
    "\n",
    "The guide has two requirements:\n",
    "* The guide function must take the same parameters as the generating model.\n",
    "* The data seen from the model must be valid outputs from the guide function.\n",
    "\n",
    "\n",
    "These functions are built with pyro primatives so that they can be used with gradient descent to optimise the KL divergence.\n",
    "The function params are in the same form as the above data generating model.\n",
    "The `pyro.param` statements recall the named parameters from the pyro param store. If no parameter exists with that name it will use the `param[.]` value passed to it, this happens on the first call only.\n",
    "We use the constraint property to ensure the distribution parameters are correctly $>0$.\n",
    "\n",
    "We use the `torch.abs` calls to ensure the distribution parameters are correctly $>0$.\n",
    "\n",
    "We make both `mu_dist` and `std_dist` as separate objects in order to optimise the mean and standard deviation of our data separately."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def parametrised_guide(params):\n",
    "    mu_mu = pyro.param(\"mu_mu\", torch.tensor(params[0]))\n",
    "    mu_std = pyro.param(\n",
    "        \"mu_std\", torch.tensor(params[1]), constraint=constraints.positive\n",
    "    )\n",
    "    std_a = pyro.param(\n",
    "        \"std_a\", torch.tensor(params[2]), constraint=constraints.positive\n",
    "    )\n",
    "    std_b = pyro.param(\n",
    "        \"std_b\", torch.tensor(params[3]), constraint=constraints.positive\n",
    "    )\n",
    "\n",
    "    mu_dist = pyro.sample(\"mu_dist\", dist.Normal(mu_mu, mu_std))\n",
    "    std_dist = pyro.sample(\"std_dist\", dist.Gamma(std_a, std_b))\n",
    "    return pyro.sample(\"data_dist\", dist.Normal(mu_dist, std_dist))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Setup variational inference descent\n",
    "This is setup via the object `pyro.infer.SVI()` using the functions we generated above.\n",
    "\n",
    "We use stochastic gradient descent. This is parameterised by the learning rate and momentum. These values were picked by trial and error so that it converges well.\n",
    "\n",
    "The loss function to optimise is the evidence lower bound."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "svi = pyro.infer.SVI(\n",
    "    model=conditioned_data_model,\n",
    "    guide=parametrised_guide,\n",
    "    optim=pyro.optim.SGD({\"lr\": 0.00001, \"momentum\": 0.8}),\n",
    "    loss=pyro.infer.Trace_ELBO(),\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Prior initialisation\n",
    "We choose uninformed priors for the mean (Gaussian) and standard deviation (Gamma) prior distributions.\n",
    "This suggests we want to learn from the data without assuming any significant previous knowledge.\n",
    "\n",
    "### Gradient descent\n",
    "Starting with our priors we iterate over our data. Each iteration we step the gradient descent optimiser.\n",
    "This should push our estimated posterior distribution closer to the actual posterior from the data each time.\n",
    "At each step we store the parameters so we can inspect them afterwards."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mu_prior = [0.0, 10.0]  # Gaussian - mu, std\n",
    "std_prior = [1.0, 0.1]  # Gamma - a, b\n",
    "params_prior = mu_prior + std_prior\n",
    "\n",
    "# Iterate over all the data\n",
    "losses, mu_mu, mu_std, std_a, std_b = [], [], [], [], []\n",
    "pyro.clear_param_store()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_steps = 5000\n",
    "for t in range(num_steps):\n",
    "    losses.append(svi.step(params_prior))\n",
    "    mu_mu.append(pyro.param(\"mu_mu\").item())\n",
    "    mu_std.append(pyro.param(\"mu_std\").item())\n",
    "    std_a.append(pyro.param(\"std_a\").item())\n",
    "    std_b.append(pyro.param(\"std_b\").item())"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Results\n",
    "The loss function has reduced with time:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convergence of the loss function\n",
    "plt.figure(num=None, figsize=(10, 6), dpi=80)\n",
    "plt.plot(losses)\n",
    "plt.title(\"ELBO\")\n",
    "plt.xlabel(\"Iteration\")\n",
    "plt.ylabel(\"loss\")\n",
    "plt.savefig(\"images/elbo.png\")\n",
    "plt.show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![](images/elbo.png)\n",
    "\n",
    "We can also see how the distribution parameters have converged:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(num=None, figsize=(10, 6), dpi=80)\n",
    "plt.subplot(2, 2, 1)\n",
    "plt.plot(mu_mu)\n",
    "plt.ylabel(\"mu_mu\")\n",
    "\n",
    "plt.subplot(2, 2, 2)\n",
    "plt.ylabel(\"mu_std\")\n",
    "plt.plot(mu_std)\n",
    "\n",
    "plt.subplot(2, 2, 3)\n",
    "plt.ylabel(\"std_a\")\n",
    "plt.plot(std_a)\n",
    "\n",
    "plt.subplot(2, 2, 4)\n",
    "plt.ylabel(\"std_b\")\n",
    "plt.plot(std_b)\n",
    "plt.savefig(\"images/params.png\")\n",
    "plt.show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![](images/params.png)\n",
    "\n",
    "The parameters for the mean distribution have converged well.\n",
    "The parameters of the standard deviation distribution have behaved differently.\n",
    "\n",
    "First, we can show the PDF of the mean distribution comparing the prior and posteriors:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot mean distributions\n",
    "plt.figure(num=None, figsize=(10, 6), dpi=80)\n",
    "mu_prior_dist = norm(loc=mu_prior[0], scale=mu_prior[1])\n",
    "x_range = np.linspace(mu_prior_dist.ppf(0.01), mu_prior_dist.ppf(0.99), num=100)\n",
    "y_values = mu_prior_dist.pdf(x_range)\n",
    "plt.plot(x_range, y_values, label=\"prior\")\n",
    "\n",
    "mu_post_dist = norm(loc=mu_mu[-1], scale=mu_std[-1])\n",
    "x_range = np.linspace(mu_post_dist.ppf(0.01), mu_post_dist.ppf(0.99), num=100)\n",
    "y_values = mu_post_dist.pdf(x_range)\n",
    "plt.plot(x_range, y_values, label=\"posterior\")\n",
    "\n",
    "plt.xlabel(\"x\")\n",
    "plt.ylabel(\"prob(x)\")\n",
    "plt.title(\"Mean PDF\")\n",
    "plt.legend()\n",
    "plt.savefig(\"images/mean_dist.png\")\n",
    "plt.show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![](images/mean_dist.png)\n",
    "\n",
    "The prior is mostly flat, the posterior on the other hand is very sharp.\n",
    "It is very confident that the actual mean is around 2, which would be correct.\n",
    "The standard deviation of our mean distribution, `mu_std[-1]=0.20`, is on a similar scale to what we would expect from the sample error of the mean: `np.std(x) / np.sqrt(n) = 0.12`.\n",
    "\n",
    "Similarly we look at the distribution of the standard deviation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot std distributions\n",
    "plt.figure(num=None, figsize=(10, 6), dpi=80)\n",
    "x_range = np.linspace(0, 10, num=100)\n",
    "\n",
    "std_prior_dist = dist.Gamma(std_prior[0], std_prior[1])\n",
    "y_values = torch.exp(std_prior_dist.log_prob(x_range))\n",
    "plt.plot(x_range, y_values, label=\"prior\")\n",
    "\n",
    "std_post_dist = dist.Gamma(std_a[-1], std_b[-1])\n",
    "y_values = torch.exp(std_post_dist.log_prob(x_range))\n",
    "plt.plot(x_range, y_values, label=\"posterior\")\n",
    "\n",
    "plt.title(\"Standard Deviation PDF\")\n",
    "plt.legend()\n",
    "plt.savefig(\"images/std_dist.png\")\n",
    "plt.show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![](images/std_dist.png)\n",
    "\n",
    "The prior is similarly mostly flat. The posterior has a peak around 4.0 which is a match to the true value of 4 and the sample standard deviation of 3.9.\n",
    "\n",
    "To look into the non-converging parameters let's look at distribution at different points in its training:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(num=None, figsize=(10, 6), dpi=80)\n",
    "x_range = np.linspace(0, 10, num=100)\n",
    "\n",
    "for idx in [500, 1000, 2000, 3000, 4000, 4999]:\n",
    "    std_post_dist = dist.Gamma(std_a[idx], std_b[idx])\n",
    "    y_values = torch.exp(std_post_dist.log_prob(x_range))\n",
    "    plt.plot(x_range, y_values, label=idx)\n",
    "\n",
    "plt.title(\"Standard Deviation PDF\")\n",
    "plt.legend()\n",
    "plt.savefig(\"images/std_dist_idx.png\")\n",
    "plt.show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![](images/std_dist_idx.png)\n",
    "\n",
    "The distribution is converging towards the correct value from about 500 iterations.\n",
    "The parameters kept changing in the same direction.\n",
    "As the distributions are converging this suggests that the two parameters\n",
    "$\\alpha$, $\\beta$ are some what correlated, allowing both to change to improve our loss function.\n",
    "This can cause the optimisation to struggle or take longer.\n",
    "I will not pursue this much further though, as the distribution has converged well.\n",
    "\n",
    "The data distribution can be plotted over the original data to see a goodness of fit:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(num=None, figsize=(10, 6), dpi=80)\n",
    "plt.hist(x, density=True)\n",
    "\n",
    "# plot prior\n",
    "prior_mu = mu_prior[0]\n",
    "prior_std = std_prior[0] / std_prior[1]  # distribution mean\n",
    "prior_dist = norm(loc=prior_mu, scale=prior_std)\n",
    "x_range = np.linspace(prior_dist.ppf(0.01), prior_dist.ppf(0.99), num=100)\n",
    "y_values = prior_dist.pdf(x_range)\n",
    "plt.plot(x_range, y_values, label=\"prior\")\n",
    "\n",
    "# plot posterior\n",
    "post_mu = mu_mu[-1]\n",
    "post_std = std_a[-1] / std_b[-1]\n",
    "post_dist = norm(loc=post_mu, scale=post_std)\n",
    "x_range = np.linspace(post_dist.ppf(0.01), post_dist.ppf(0.99), num=100)\n",
    "y_values = post_dist.pdf(x_range)\n",
    "plt.plot(x_range, y_values, label=\"post\")\n",
    "\n",
    "plt.legend()\n",
    "plt.title(\"Data histogram\")\n",
    "plt.savefig(\"images/data_dist.png\")\n",
    "plt.show()\n",
    "\n",
    "print(post_mu)\n",
    "print(post_std)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![](images/data_dist.png)\n",
    "\n",
    "The posterior (green line) fits the data histogram well as we would expect.\n",
    "The values of the posterior distribution (1.82, 4.01) are similar to those from the sample estimates (1.82, 3.95).\n",
    "However in the posterior case we have our confidence around those values rather than just point estimates."
   ]
  }
 ],
 "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
}