intro_to_generative_ML_with_Pyro.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Introduction to Generative ML with Pyro"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "import pyro\n",
    "\n",
    "pyro.set_rng_seed(101)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Our goal is to understand causal modeling within the context of generative machine learning. We just examined one generative machine learning framework called Bayesian networks (BNs) and how we can use BNs as causal models.\n",
    "\n",
    "**Bayesian Networks (BNs)** Framework that defines a probabilistic generative model of the world in terms of a directed acyclic graph.\n",
    "\n",
    "**causal Bayesian networks:** Bayesian networks where the direction of edges in the DAG represent causality.\n",
    "\n",
    "Bayesian networks provide a general-purpose framework for representing a causal data generating story for how the world works.\n",
    "\n",
    "Now we will introduce probabilistic programming, a framework that is more expressive than Bayesian networks.\n",
    "\n",
    "### What is a probabilistic programming language?\n",
    "\n",
    "\"A probabilistic programming language (PPL) is a programming language designed to describe probabilistic models and then perform inference in those models. PPLs are closely related to graphical models and Bayesian networks but are more expressive and flexible. Probabilistic programming represents an attempt to \"Unify general purpose programming\" with probabilistic modeling.\"\n",
    "\n",
    "-Wikipedia\n",
    "\n",
    "A PPL is a domain-specific programming language for that lets you write a data generating story as a program. As with a causal Bayesian network, you can write your program in a way that orders the steps of its execution according to cause and effect.\n",
    "\n",
    "### How exactly do Bayesian networks and probabilistic programming differ?\n",
    "\n",
    "**Representation of relationships between variables**.  BNs restricted to representing the relationships between variables in terms of conditional probability distributions (CPDs) factored according to a DAG.  Frameworks typically limit you to a small set of parametric CPDs (e.g., Gaussian, multinomial).\n",
    "\n",
    "Just as computer programs are more expressive than flow charts, PPLs let you represent relations any way you like so long as you can represent them in code.  PPL relationships can include control flow and recursion.  In causal models, we will see that this allows you to be more specific about mechanism than you can with CPDs.\n",
    "\n",
    "**DAG vs. open world models**.  BNs restrict the representation of the joint distribution to a DAG.  This constraint enables you to reason easily about the joint distribution through graph-theoretic operations like d-separation.  PPLs need not be constrained to a DAG.  For example (using an imaginary Python PPL package):\n",
    "\n",
    "    X = Bernoulli(p)\n",
    "    if X == 1:\n",
    "        Y = Gaussian(0, 1)\n",
    "\n",
    "In a DAG, you have a fixed set of variables, i.e. a \"closed world\".  In the above model, the variable Y is only instantiated if X==1.  Y may or may not exist depending on how the generative process unfolds.  For a more extreme example, consider this:\n",
    "\n",
    "    X = Poisson(λ)\n",
    "    Y = [Gaussian(0, 1)]\n",
    "    for i in range(1, X):\n",
    "        Y[i] = Gaussian(Y[i-1], 1))\n",
    "\n",
    "Here you have the total number of Y variables itself being a random variable X.  Further, the mean of the ith Y is a random variable given by the i-1th Y.  You can't do that with a Bayes net!  Unfortunately, we can't reason about this as directly as we can with a DAG.  For example, recall that with the DAG, we had a convenient algorithm called `CPDAG` that converts the DAG to a partially directed acyclic graph structure called a PDAG that provides a compact representation of all the DAGs in an equivalence class.  How might we define an equivalence class on this program?  Certainly, enumerating all programs with an equivalent representation of the joint distribution would be very difficult even with constraints on the length of the program.  In general, enumerating all programs of minimal description that provide equivalent representations of a joint distribution is an NP-hard problem.\n",
    "\n",
    "**Inference**  When you have a DAG and a constrained set of parametric CPDs, as well as constraints on the kind of inference,  queries the user can make, you can implement some inference algorithms in your BN framework that will generally work in a reasonable amount of time.\n",
    "\n",
    "PPLs are more flexible than BNs, but the trade-off s that getting inference to work is harder.   PPL's develop several abstractions for inference and leave it to the user to apply them, requiring the user to become something of an expert in inference algorithms.  PPL developers make design decisions to make inference easier for the user, though this often sacrifices some flexibility.  One emergent pattern is to build PPLs on tensor-based frameworks like Tensorflow and PyTorch.  Tensor-based PPLs allow a data scientist with experience building deep learning models to rely on that experience when doing inference.\n",
    "\n",
    "![Image](fig/inference.png)\n",
    "\n",
    "[$$\\texttt{Kevin Smith - Tutorial: Probabilistic Programming}$$](https://www.youtube.com/watch?v=9SEIYh5BCjc)\n",
    "\n",
    "## Introduction to Pyro\n",
    "\n",
    "Pyro is a universal probabilistic programming language (PPL) written in Python and supported by PyTorch on the backend. Pyro enables flexible and expressive deep probabilistic modeling, unifying the best of modern deep learning and Bayesian modeling.\n",
    "\n",
    "Our purpose of this class, pyro has \"do\"-operator that allows intervention and counterfactual inference in these probabilistic models.\n",
    "\n",
    "### Stochastic Functions\n",
    "\n",
    "The basic unit of probabilistic programs is the stochastic function. A stochastic function is an arbitrary Python callable that combines two ingredients:\n",
    "\n",
    "- deterministic Python code; and\n",
    "- primitive stochastic functions that call a random number generator\n",
    "\n",
    "For this course, we will consider these stochastic functions as **models**. Stochastic functions can be used to represent simplified or abstract descriptions of a data-generating process.\n",
    "\n",
    "### Primitive stochastic functions\n",
    "\n",
    "We call them distributions. We can explicitly compute the probability of the outputs given the inputs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loc = 0.   # mean zero\n",
    "scale = 1. # unit variance\n",
    "normal = torch.distributions.Normal(loc, scale) # create a normal distribution object\n",
    "x = normal.rsample() # draw a sample from N(0,1)\n",
    "print(\"sample: \", x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Pyro simplifies this process of sampling from distributions. It uses `pyro.sample()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = pyro.sample(\"my_sample\", pyro.distributions.Normal(loc, scale))\n",
    "print(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Just like a direct call to `torch.distributions.Normal().rsample()`, this returns a sample from the unit normal distribution. The crucial difference is that this sample is named. Pyro’s backend uses these names to uniquely identify sample statements and change their behavior at runtime depending on how the enclosing stochastic function is being used. This is how Pyro can implement the various manipulations that underlie inference algorithms."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's write a simple `weather` model.\n",
    "\n",
    "### A simple model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pyro.distributions as dist\n",
    "\n",
    "def weather():\n",
    "    cloudy = pyro.sample('cloudy', dist.Bernoulli(0.3))\n",
    "    cloudy = 'cloudy' if cloudy.item() == 1.0 else 'sunny'\n",
    "    mean_temp = {'cloudy': 55.0, 'sunny': 75.0}[cloudy]\n",
    "    scale_temp = {'cloudy': 10.0, 'sunny': 15.0}[cloudy]\n",
    "    temp = pyro.sample('temp', dist.Normal(mean_temp, scale_temp))\n",
    "    return cloudy, temp.item()\n",
    "\n",
    "for _ in range(3):\n",
    "    print(weather())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First two lines introduce a binary variable `cloudy`, which is given by a draw from the Bernoulli distribution with a parameter of $0.3$. The Bernoulli distribution returns either $0$ or $1$, line `2` converts that into a string. So, So according to this model, $30%$ of the time it’s cloudy and $70%$ of the time it’s sunny.\n",
    "\n",
    "In line `4` and `5`, we initialize mean and scale of the temperature for both values. We then sample, the temperature from a Normal distribution and return that along with `cloudy` variable.\n",
    "\n",
    "\n",
    "We can build complex model by modularizing and reusing the concepts into functions and use them as programmers use functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def ice_cream_sales():\n",
    "    cloudy, temp = weather()\n",
    "    expected_sales = 200. if cloudy == 'sunny' and temp > 80.0 else 50.\n",
    "    ice_cream = pyro.sample('ice_cream', pyro.distributions.Normal(expected_sales, 10.0))\n",
    "    return ice_cream"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Inference\n",
    "\n",
    "As we discussed earlier, the reason we use PPLs is because they can easily go backwards and reason about cause given the observed effect. There are myriad of inference algorithms available in pyro. Let's try it on an even simpler model.\n",
    "\n",
    "$$weight \\mid guess \\sim \\mathcal{N}(guess, 1)$$\n",
    "\n",
    "$$measurement \\mid guess, weight \\sim \\mathcal{N}(weight, 0.75)$$\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def scale(guess):\n",
    "    weight = pyro.sample(\"weight\", dist.Normal(guess, 1.0))\n",
    "    measurement = pyro.sample(\"measurement\", dist.Normal(weight, 0.75))\n",
    "    return measurement\n",
    "\n",
    "scale(14.)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Suppose we observe that the `measurement` of an object was $14$ lbs. What would have we guessed if we tried to guess it's `weight` first?\n",
    "\n",
    "This question is answered in two steps.\n",
    "\n",
    "1. Condition the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "conditioned_scale = pyro.condition(scale,{\"measurement\": torch.tensor(14.)})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "2. Set the prior and infer the posterior. We will use the Hamiltonian Monte Carlo (HMC) inference algorithm with the no U-turn sampler (NUTS), like you would in a language like Stan."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyro.infer import NUTS, MCMC, EmpiricalMarginal\n",
    "\n",
    "guess_prior = torch.tensor(13.)\n",
    "\n",
    "nuts_kernel = NUTS(conditioned_scale, adapt_step_size=True)\n",
    "posterior = MCMC(nuts_kernel, \n",
    "                 num_samples=1000, \n",
    "                 warmup_steps=300).run(guess_prior)\n",
    "marginal = EmpiricalMarginal(posterior, \"weight\")\n",
    "\n",
    "plt.hist([marginal().item() for _ in range(1000)],)\n",
    "plt.title(\"P(weight | measurement = 14)\")\n",
    "plt.xlabel(\"Weight\")\n",
    "plt.ylabel(\"#\")\n",
    "print(posterior)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Shapes in distribution:\n",
    "\n",
    "We know that PyTorch tensors have single `shape` attribute, `Distribution`s have two shape attributes with special meaning. \n",
    "* `.batch_shape`: Indices over `.batch_shape` denote conditionally independent random variables,\n",
    "* `.event_shape`: indices over `.event_shape` denote dependent random variables (ie one draw from a distribution).\n",
    "\n",
    "These two combine to define the total shape of a sample. Thus the total shape of `.log_prob()` of distribution is `.batch_shape`.\n",
    "\n",
    "Also, `Distribution.sample()` also has a `sample_shape` attribute that indexes over independent and identically distributed(iid) random variables.\n",
    "\n",
    "```\n",
    "      |      iid     | independent | dependent\n",
    "------+--------------+-------------+------------\n",
    "shape = sample_shape + batch_shape + event_shape\n",
    "```\n",
    "\n",
    "To learn more see [broadcasting tensors in PyTorch](https://pytorch.org/docs/master/notes/broadcasting.html).\n",
    "\n",
    "### Examples\n",
    "\n",
    "One way to introduce batch_shape is use `expand`.\n",
    "\n",
    "```{python}\n",
    "d = dist.MultivariateNormal(torch.zeros(3), torch.eye(3, 3)).expand([5]) # expand - 3 of these Multivariate Normal Dists\n",
    "print(\"batch_shape: \", d.batch_shape)\n",
    "print(\"event_shape: \", d.event_shape)\n",
    "\n",
    "#x = d.sample(torch.Size([5]))\n",
    "x = d.sample()\n",
    "print(\"x shape: \", x.shape)          # == sample_shape + batch_shape + event_shape\n",
    "print(\"d.log_prob(x) shape:\", d.log_prob(x).shape)  # == batch_shape\n",
    "```\n",
    "\n",
    "The other way is using `plate` context manager.\n",
    "\n",
    "Pyro models can use the context manager `pyro.plate` to declare that certain batch dimensions are independent. Inference algorithms can then take advantage of this independence to e.g. construct lower variance gradient estimators or to enumerate in linear space rather than exponential space.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with pyro.plate(\"x_axis\", 5):\n",
    "    d = dist.MultivariateNormal(torch.zeros(3), torch.eye(3, 3))\n",
    "    x = pyro.sample(\"x\", d)\n",
    "x.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In fact, we can also nest `plates`. The only thing we need to care about is, which dimensions are independent. Pyro automatically manages this but sometimes we need to explicitely specify the dimensions. Once we specify that, we can leverage PyTorch's CUDA enabled capabilities to run inference on GPUs.\n",
    "\n",
    "```\n",
    "with pyro.plate(\"x_axis\", 320):\n",
    "    # within this context, batch dimension -1 is independent\n",
    "    with pyro.plate(\"y_axis\", 200):\n",
    "        # within this context, batch dimensions -2 and -1 are independent\n",
    "```\n",
    "\n",
    "Note that we always count from the right by using negative indices like $-2$, $-1$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Gaussian Mixture Model\n",
    "\n",
    "![Image](fig/gmm.png)\n",
    "\n",
    "[$$\\texttt{Blei - Build, Compute, Critique, Repeat:Data Analysis with Latent Variable Models}$$](http://www.cs.columbia.edu/~blei/papers/Blei2014b.pdf)\n",
    "\n",
    "Here is an example of how to implement a Gaussian mixture model in Pyro."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import print_function\n",
    "import os\n",
    "from collections import defaultdict\n",
    "import numpy as np\n",
    "import scipy.stats\n",
    "import torch\n",
    "from torch.distributions import constraints\n",
    "\n",
    "from pyro import poutine\n",
    "from pyro.contrib.autoguide import AutoDelta\n",
    "from pyro.optim import Adam\n",
    "from pyro.infer import SVI, TraceEnum_ELBO, config_enumerate, infer_discrete\n",
    "\n",
    "from matplotlib import pyplot\n",
    "# %matplotlib inline\n",
    "\n",
    "pyro.enable_validation(True)\n",
    "\n",
    "data = torch.tensor([0., 1., 10., 11., 12.])\n",
    "\n",
    "K = 2  # Fixed number of components.\n",
    "\n",
    "@config_enumerate\n",
    "def model(data):\n",
    "    # Global variables.\n",
    "    weights = pyro.sample('weights', dist.Dirichlet(0.5 * torch.ones(K)))\n",
    "    scale = pyro.sample('scale', dist.LogNormal(0., 2.))\n",
    "    with pyro.plate('components', K):\n",
    "        locs = pyro.sample('locs', dist.Normal(0., 10.))\n",
    "    with pyro.plate('data', len(data)):\n",
    "        # Local variables.\n",
    "        assignment = pyro.sample('assignment', dist.Categorical(weights))\n",
    "        pyro.sample('obs', dist.Normal(locs[assignment], scale), obs=data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Review of Approximate Inference\n",
    "\n",
    "We have variables $Z$s (cluster assignments) and $X$s (data points) in our mixture model, where $X$ is observed and $Z$ is latent (unobserved). As we saw earlier, a generative model entails a joint distribution\n",
    "$$p(Z,X)$$\n",
    "\n",
    "Inference of unknown can be achieved through conditioning on the observations.\n",
    "\n",
    "$$p(Z \\mid X) = \\frac{p(Z, X)}{p(X)}$$\n",
    "\n",
    "For non-trivial interesting problems, the integral for the denominator(marginal) is not tractable.\n",
    "$$p(X) = \\int dZp(X \\mid Z)p(Z)$$ \n",
    "\n",
    "So we have to directly approximate $p(Z \\mid X)$. The most popular way of doing this is with *stochastic variational inference* because it can make use of the loss function optimization infrastructure in PyTorch.\n",
    "\n",
    "### Variational Inference:\n",
    "\n",
    "We can't compute $p(Z \\mid X)$ directly, so let's approximate with some other distribution $q(Z; \\nu)$ over Z that is tractable (for example, Gaussions or other exponential family).\n",
    "\n",
    "![Image](fig/VI.png)\n",
    "\n",
    "[$$\\texttt{David Blei - Variational Inference (NeurIPS 2016 Tutorial)}$$](https://www.youtube.com/watch?v=ogdv_6dbvVQ)\n",
    "\n",
    "Since q is tractable, we can play with it's parameter $\\nu$ such that it reaches as close to $p(Z\\mid X)$ as possible. More precisely, we want to minimize the KL divergence between $q$ and $p$. With this trick, we just turned an **inference** problem to an **optimization** problem!\n",
    "\n",
    "$$\n",
    "\\begin{align*}\n",
    "KL(q(Z;\\nu) \\mid\\mid p(Z\\mid X)) &= -\\int dZ\\ q(Z) \\log\\frac{P(Z\\mid X)}{q(Z)}\\\\\n",
    "&= -\\int dZ\\ q(Z) \\log \n",
    "\\frac{\\frac{p(Z,X)}{p(X)}}{q(Z)}\\\\\n",
    "&= -\\int dZ\\ q(Z) \\log \n",
    "\\frac{p(Z,X)}{p(X)q(Z)}\\\\\n",
    "&= -\\int dZ\\ q(Z) \n",
    "\\left[\n",
    "\\log \n",
    "\\frac{p(Z,X)}{q(Z)}\n",
    "- \\log p(X)\n",
    "\\right]\\\\\n",
    "&= - \\int dZ\\ \\log \\frac{p(Z,X)}{q(Z)} + \\underbrace{\\int dZ\\ q(Z)}_{\\text{=1}}\\log p(X)\\\\\n",
    "&= - \\int dZ\\ \\log \\frac{p(Z,X)}{q(Z)} + \\log p(X)\\\\\n",
    "\\log p(X) &= KL(q(Z;\\nu)\\mid\\mid p(Z\\mid X) + \\underbrace{\\int dZ\\ q(Z;\\nu) \\log \\frac{p(Z,X)}{q(Z;\\nu)}}_{\\mathcal{L}}\\\\\n",
    "\\end{align*}\n",
    "$$\n",
    "\n",
    "Note that we already observed $X$ and we conditioned the model to get $p(Z \\mid X)$. But given $X$, $\\log p(X)$ is constant! So, minimizing KL is equivalent to maximizing $\\mathcal{L}$.\n",
    "\n",
    "**How do you maximize $\\mathcal{L}$**? Take $\\nabla_{\\nu} \\mathcal{L}$.\n",
    "\n",
    "$\\mathcal{L}$ is called **variational lower bound**. It is often called ELBO.\n",
    "\n",
    "**Stochastic Variational Inference** scales variational inference to massive data. Just like in stochastic variational inference, you subsample the data and update the posterior!\n",
    "\n",
    "<!-- #region -->\n",
    "### Stochastic Optimization\n",
    "\n",
    "In stochastic optimization, we replace the gradient with cheaper noisy estimate which is guranteed to converge to a local optimum.\n",
    "\n",
    "$$\\nu_{t+1} = \\nu_t + \\rho_t \\hat{\\nabla}_{\\nu} \\mathcal{L}(\\nu_t)$$\n",
    "\n",
    "Requirements:\n",
    "\n",
    "* Unbiased gradients, i.e.\n",
    "$$\\mathbb{E}[\\hat{\\nabla}_{\\nu} \\mathcal{L}(\\nu_t)] = \\nabla_{\\nu}\\mathcal{L}(\\nu)$$\n",
    "\n",
    "* Step-size sequence $\\rho_t$ that follows Robbins-Monro conditions.\n",
    "\n",
    "\n",
    "Stochastic variational inference takes inspiration from stochastic optimization and natural graidient. We follow the same procedure as stochastic gradient descent.\n",
    "<!-- #endregion -->\n",
    "\n",
    "### A Rough Stochastic variational inference algorithm:\n",
    "\n",
    "* Initialize $q$ with some $\\nu$\n",
    "* Until Converge:\n",
    "    * subsample from Data:\n",
    "        * compute gradient $\\hat{\\nabla_{\\nu}}\\mathcal{L}_{\\nu_t}$\n",
    "        * update global parameter $\\nu_{t+1} = \\nu_t + \\rho_t \\hat{\\nabla_{\\nu}}\\mathcal{L}_{\\nu_t}$\n",
    "* Return $q(Z;\\nu)$\n",
    "\n",
    "### Training a MAP estimator\n",
    "\n",
    "Let's start by learning model parameters `weights`, `locs`, and `scale` given priors and data. We will use `AutoDelta` guide function. Our model will learn global mixture weights, the location of each mixture component, and a shared scale that is common to both components.\n",
    "\n",
    "During inference, `TraceEnum_ELBO` will marginalize out the assignments of datapoints to clusters.\n",
    "\n",
    "`max_plate_nesting` lets Pyro know that we’re using the rightmost dimension plate and that Pyro can use any other dimension for parallelization.\n",
    "\n",
    "## Some other Pyro vocabulary\n",
    "\n",
    "* poutine - Beneath the built-in inference algorithms, Pyro has a library of composable effect handlers for creating new inference algorithms and working with probabilistic programs. Pyro’s inference algorithms are all built by applying these handlers to stochastic functions.\n",
    "    * poutine.block - blocks pyro premitives. By default, it blocks everything.\n",
    "\n",
    "* param - Parameters in Pyro are basically thin wrappers around PyTorch Tensors that carry unique names. As such Parameters are the primary stateful objects in Pyro. Users typically interact with parameters via the Pyro primitive `pyro.param`. Parameters play a central role in stochastic variational inference, where they are used to represent point estimates for the parameters in parameterized families of models and guides.\n",
    "\n",
    "* param_store - Global store for parameters in Pyro. This is basically a key-value store."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "global_guide = AutoDelta(poutine.block(model, expose=['weights', 'locs', 'scale']))\n",
    "\n",
    "optim = pyro.optim.Adam({'lr': 0.1, 'betas': [0.8, 0.99]})\n",
    "elbo = TraceEnum_ELBO(max_plate_nesting=1)\n",
    "svi = SVI(model, global_guide, optim, loss=elbo)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def initialize(seed):\n",
    "    pyro.set_rng_seed(seed)\n",
    "    pyro.clear_param_store()\n",
    "    # Initialize weights to uniform.\n",
    "    pyro.param('auto_weights', 0.5 * torch.ones(K), constraint=constraints.simplex)\n",
    "    # Assume half of the data variance is due to intra-component noise.\n",
    "    pyro.param('auto_scale', (data.var() / 2).sqrt(), constraint=constraints.positive)\n",
    "    # Initialize means from a subsample of data.\n",
    "    pyro.param('auto_locs', data[torch.multinomial(torch.ones(len(data)) / len(data), K)]);\n",
    "    loss = svi.loss(model, global_guide, data)\n",
    "    return loss\n",
    "\n",
    "# Choose the best among 100 random initializations.\n",
    "loss, seed = min((initialize(seed), seed) for seed in range(100))\n",
    "initialize(seed)\n",
    "print('seed = {}, initial_loss = {}'.format(seed, loss))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Register hooks to monitor gradient norms.\n",
    "gradient_norms = defaultdict(list)\n",
    "for name, value in pyro.get_param_store().named_parameters():\n",
    "    value.register_hook(lambda g, name=name: gradient_norms[name].append(g.norm().item()))\n",
    "losses = []\n",
    "for i in range(200):\n",
    "    loss = svi.step(data)\n",
    "    losses.append(loss)\n",
    "    print('.' if i % 100 else '\\n', end='')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(10,3), dpi=100).set_facecolor('white')\n",
    "plt.plot(losses)\n",
    "plt.xlabel('iters')\n",
    "plt.ylabel('loss')\n",
    "plt.yscale('log')\n",
    "plt.title('Convergence of SVI');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "map_estimates = global_guide(data)\n",
    "weights = map_estimates['weights']\n",
    "locs = map_estimates['locs']\n",
    "scale = map_estimates['scale']\n",
    "print('weights = {}'.format(weights.data.numpy()))\n",
    "print('locs = {}'.format(locs.data.numpy()))\n",
    "print('scale = {}'.format(scale.data.numpy()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X = np.arange(-3,15,0.1)\n",
    "Y1 = weights[0].item() * scipy.stats.norm.pdf((X - locs[0].item()) / scale.item())\n",
    "Y2 = weights[1].item() * scipy.stats.norm.pdf((X - locs[1].item()) / scale.item())\n",
    "\n",
    "plt.figure(figsize=(10, 4), dpi=100).set_facecolor('white')\n",
    "plt.plot(X, Y1, 'r-')\n",
    "plt.plot(X, Y2, 'b-')\n",
    "plt.plot(X, Y1 + Y2, 'k--')\n",
    "plt.plot(data.data.numpy(), np.zeros(len(data)), 'k*')\n",
    "plt.title('Density of two-component mixture model')\n",
    "plt.ylabel('probability density')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "language": "python",
   "name": "venv"
  },
  "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}