Solutions-03.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Solutions-03"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This contains one possible solution for the three-parameter linear model with intrinsic scatter.\n",
    "\n",
    "Recall that we are working with this data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import seaborn; seaborn.set() # for plot formatting\n",
    "\n",
    "def make_data_scatter(intercept, slope, scatter,\n",
    "                      N=20, dy=2, rseed=42):\n",
    "    rand = np.random.RandomState(rseed)\n",
    "    x = 100 * rand.rand(20)\n",
    "    y = intercept + slope * x\n",
    "    y += np.sqrt(dy ** 2 + scatter ** 2) * rand.randn(20)\n",
    "    return x, y, dy * np.ones_like(x)\n",
    "\n",
    "\n",
    "# (intercept, slope, intrinsic scatter)\n",
    "theta = (25, 0.5, 3.0)\n",
    "x, y, dy = make_data_scatter(*theta)\n",
    "plt.errorbar(x, y, dy, fmt='o');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define the prior, likelihood, and posterior\n",
    "\n",
    "The likelihood for this model looks very similar to what we used above, except that the intrinsic scatter is added *in quadrature* to the measurement error.\n",
    "If $\\varepsilon_i$ is the measurement error on the point $(x_i, y_i)$, and $\\sigma$ is the intrinsic scatter, then the likelihood should look like this:\n",
    "\n",
    "$$\n",
    "P(x_i,y_i\\mid\\theta) = \\frac{1}{\\sqrt{2\\pi(\\varepsilon_i^2 + \\sigma^2)}} \\exp\\left(\\frac{-\\left[y_i - y(x_i;\\theta)\\right]^2}{2(\\varepsilon_i^2 + \\sigma^2)}\\right)\n",
    "$$\n",
    "\n",
    "For the prior, you can use either a flat or symmetric prior on the slope and intercept, but on the intrinsic scatter $\\sigma$ it is best to use a scale-invariant Jeffreys Prior:\n",
    "\n",
    "$$\n",
    "P(\\sigma)\\propto\\sigma^{-1}\n",
    "$$\n",
    "\n",
    "As discussed before, this has the nice feature that the resulting posterior will not depend on the units of measurement."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Define functions to compute the log-prior, log-likelihood, and log-posterior\n",
    "\n",
    "# theta = [intercept, slope, scatter]\n",
    "\n",
    "def log_prior(theta):\n",
    "    if theta[2] <= 0 or np.any(np.abs(theta[:2]) > 1000):\n",
    "        return -np.inf  # log(0)\n",
    "    else:\n",
    "        # Jeffreys Prior\n",
    "        return -np.log(theta[2])\n",
    "    \n",
    "def log_likelihood(theta, x, y, dy):\n",
    "    y_model = theta[0] + theta[1] * x\n",
    "    S = dy ** 2 + theta[2] ** 2\n",
    "    return -0.5 * np.sum(np.log(2 * np.pi * S) +\n",
    "                         (y - y_model) ** 2 / S)\n",
    "\n",
    "def log_posterior(theta, x, y, dy):\n",
    "    return log_prior(theta) + log_likelihood(theta, x, y, dy)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Sampling from the Posterior"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Using emcee, create and initialize a sampler and draw 200 samples from the posterior.\n",
    "# Remember to think about what starting guesses should you use!\n",
    "# You can use the above as a template\n",
    "\n",
    "import emcee\n",
    "\n",
    "ndim = 3  # number of parameters in the model\n",
    "nwalkers = 50  # number of MCMC walkers\n",
    "\n",
    "# initialize walkers\n",
    "starting_guesses = np.random.randn(nwalkers, ndim)\n",
    "starting_guesses[:, 2] = np.random.rand(nwalkers)\n",
    "\n",
    "sampler = emcee.EnsembleSampler(nwalkers, ndim, log_posterior,\n",
    "                                args=[x, y, dy])\n",
    "pos, prob, state = sampler.run_mcmc(starting_guesses, 200)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualizing the Chains"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Plot the three chains as above\n",
    "\n",
    "fig, ax = plt.subplots(3, sharex=True)\n",
    "for i in range(3):\n",
    "    ax[i].plot(sampler.chain[:, :, i].T, '-k', alpha=0.2);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Restarting and getting a clean sample"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Are your chains stabilized? Reset them and get a clean sample\n",
    "\n",
    "sampler.reset()\n",
    "pos, prob, state = sampler.run_mcmc(pos, 1000)\n",
    "\n",
    "fig, ax = plt.subplots(3, sharex=True)\n",
    "for i in range(3):\n",
    "    ax[i].plot(sampler.chain[:, :, i].T, '-k', alpha=0.2);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualizing the results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Use corner.py to visualize the three-dimensional posterior\n",
    "\n",
    "import corner\n",
    "corner.corner(sampler.flatchain, truths=theta,\n",
    "              labels=['intercept', 'slope', 'scatter']);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And visualizing the model over the data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Next plot ~100 of the samples as models over the data to get an idea of the fit\n",
    "\n",
    "chain = sampler.flatchain\n",
    "\n",
    "plt.errorbar(x, y, dy, fmt='o');\n",
    "\n",
    "thetas = [chain[i] for i in np.random.choice(chain.shape[0], 100)]\n",
    "\n",
    "xfit = np.linspace(0, 100)\n",
    "for i in range(100):\n",
    "    theta = thetas[i]\n",
    "    plt.plot(xfit, theta[0] + theta[1] * xfit,\n",
    "             color='black', alpha=0.05);"
   ]
  }
 ],
 "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.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}