{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Create some convenience routines for plotting\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "def compute_sigma_level(trace1, trace2, nbins=20):\n",
    "    \"\"\"From a set of traces, bin by number of standard deviations\"\"\"\n",
    "    L, xbins, ybins = np.histogram2d(trace1, trace2, nbins)\n",
    "    L[L == 0] = 1E-16\n",
    "    logL = np.log(L)\n",
    "\n",
    "    shape = L.shape\n",
    "    L = L.ravel()\n",
    "\n",
    "    # obtain the indices to sort and unsort the flattened array\n",
    "    i_sort = np.argsort(L)[::-1]\n",
    "    i_unsort = np.argsort(i_sort)\n",
    "\n",
    "    L_cumsum = L[i_sort].cumsum()\n",
    "    L_cumsum /= L_cumsum[-1]\n",
    "\n",
    "    xbins = 0.5 * (xbins[1:] + xbins[:-1])\n",
    "    ybins = 0.5 * (ybins[1:] + ybins[:-1])\n",
    "\n",
    "    return xbins, ybins, L_cumsum[i_unsort].reshape(shape)\n",
    "\n",
    "\n",
    "def plot_MCMC_trace(ax,trace,model,scatter=False, **kwargs):\n",
    "    \"\"\"Plot traces and contours\"\"\"\n",
    "    if model==\"gaussian\":\n",
    "        xbins, ybins, sigma = compute_sigma_level(trace[1], trace[2])\n",
    "        ax.contour(xbins, ybins, sigma.T, levels=[0.683, 0.955], **kwargs)\n",
    "        if scatter:\n",
    "            ax.plot(trace[1], trace[2], ',k', alpha=0.1)\n",
    "        ax.set_xlabel(r'$\\mu$')\n",
    "        ax.set_ylabel(r'$\\sigma$')\n",
    "    elif model==\"sin\":\n",
    "        xbins, ybins, sigma = compute_sigma_level(trace[0], trace[1])\n",
    "        ax.contour(xbins, ybins, sigma.T, levels=[0.683, 0.955], **kwargs)\n",
    "        if scatter:\n",
    "            ax.plot(trace[0], trace[1], ',k', alpha=0.1)\n",
    "        ax.set_xlabel(\"a\")        \n",
    "    elif (model==\"linear\"):\n",
    "        xbins, ybins, sigma = compute_sigma_level(trace[0], trace[1])\n",
    "        ax.contour(xbins, ybins, sigma.T, levels=[0.683, 0.955], **kwargs)\n",
    "        if scatter:\n",
    "            ax.plot(trace[0], trace[1], ',k', alpha=0.1)\n",
    "        ax.set_xlabel(\"a\")\n",
    "        ax.set_ylabel(\"b\")\n",
    "    elif model==\"power law\":\n",
    "        xbins, ybins, sigma = compute_sigma_level(trace[0], trace[1])\n",
    "        ax.contour(xbins, ybins, sigma.T, levels=[0.683, 0.955], **kwargs)\n",
    "        if scatter:\n",
    "            ax.plot(trace[0], trace[1], ',k', alpha=0.1)   \n",
    "        ax.set_xlabel(\"a\")\n",
    "        ax.set_ylabel(\"b\")\n",
    "    elif model==\"power log\": # fit was done linearly, plots linearly, but params given as power law\n",
    "        xbins, ybins, sigma = compute_sigma_level(trace[0], trace[1])\n",
    "        ax.contour(np.exp(xbins), ybins, sigma.T, levels=[0.683, 0.955], **kwargs)\n",
    "        if scatter:\n",
    "            ax.plot(np.exp(trace[0]), trace[1], ',k', alpha=0.1) \n",
    "        ax.set_xlabel(\"a\")\n",
    "        ax.set_ylabel(\"b\")\n",
    "    \n",
    "    \n",
    "    ax.set_title(\"Traces and 68%, 95% contours\")\n",
    "\n",
    "\n",
    "def plot_MCMC_model(ax,hist, trace, model,units,bins,bar):\n",
    "    \"\"\"Plot the linear model and 2sigma contours\"\"\"\n",
    "\n",
    "    if model ==\"gaussian\":\n",
    "        ax.hist(hist,normed=True,bins=bins)\n",
    "        a,mean,stddev = trace[0],trace[1],trace[2]\n",
    "        xfit = bins\n",
    "        yfit=a[:, None]*np.exp(-(xfit-mean[:, None])**2/(2*stddev[:, None]**2))\n",
    "        mu = yfit.mean(0)\n",
    "        sig = 2 * yfit.std(0)\n",
    "        ax.plot(xfit, mu, '-k')\n",
    "        ax.fill_between(xfit, mu - sig, mu + sig, color='lightgray')\n",
    "    elif model==\"sin\":\n",
    "        ax.hist(hist,normed=True,bins=bins)\n",
    "        a,b = trace[0],trace[1]\n",
    "        xfit = bins\n",
    "        yfit= (a[:, None])*np.sin(xfit*b[:, None])\n",
    "        mu = yfit.mean(0)\n",
    "        sig = 2 * yfit.std(0)\n",
    "        ax.plot(xfit, mu, '-k')\n",
    "        ax.fill_between(xfit, mu - sig, mu + sig, color='lightgray')\n",
    "    elif model==\"linear\":\n",
    "        ax.hist(hist,normed=True,bins=bins)     \n",
    "        a,b = trace[0],trace[1]\n",
    "        xfit = bins\n",
    "        yfit=b[:, None]*xfit + a[:, None]\n",
    "        mu = yfit.mean(0)\n",
    "        sig = 2 * yfit.std(0)\n",
    "        ax.plot(xfit, mu, '-k')\n",
    "        ax.fill_between(xfit, mu - sig, mu + sig, color='lightgray')   \n",
    "    elif (model==\"power law\") and (bar==False):\n",
    "        ax.hist(hist,normed=True,bins=bins)\n",
    "        a,b = trace[0],trace[1]\n",
    "        xfit = bins\n",
    "        yfit= a[:, None]*xfit**b[:, None]\n",
    "        mu = yfit.mean(0)\n",
    "        sig = 2 * yfit.std(0)\n",
    "        ax.plot(xfit, mu, '-k')\n",
    "        ax.fill_between(xfit, mu - sig, mu + sig, color='lightgray')   \n",
    "    elif (model==\"power law\") and (bar==True):\n",
    "        ax.bar(bins, hist, width=.02)     \n",
    "        a,b = trace[0],trace[1]\n",
    "        xfit = bins\n",
    "        yfit= a[:, None]*xfit**b[:, None]\n",
    "        mu = yfit.mean(0)\n",
    "        sig = 2 * yfit.std(0)\n",
    "        ax.plot(xfit, mu, '-k')\n",
    "        ax.fill_between(xfit, mu - sig, mu + sig, color='lightgray')   \n",
    "    elif model==\"power log\": # fit was done linearly, plots linearly, but params given as power law\n",
    "        ax.hist(hist,normed=True,bins=bins)\n",
    "        a,b = trace[0],trace[1]\n",
    "        xfit = bins\n",
    "        yfit= np.exp(a[:, None])*xfit**b[:, None]\n",
    "        mu = yfit.mean(0)\n",
    "        sig = 2 * yfit.std(0)\n",
    "        ax.plot(xfit, mu, '-k')\n",
    "        ax.fill_between(xfit, mu - sig, mu + sig, color='lightgray')      \n",
    "        ax.set_xscale(\"log\", nonposx='clip')\n",
    "        ax.set_yscale(\"log\", nonposy='clip') \n",
    "    ax.set_xlabel(units[0])\n",
    "    ax.set_ylabel(units[1])\n",
    "    ax.set_title(\"Probability distribution with fit function\")\n",
    "\n",
    "# IF POWER LAW: PASS IN THE HIST IN THE POWER LAW FORM\n",
    "def plot_MCMC_results(hist, trace,model,units,title,bins,bar=False,colors='k'):\n",
    "    \"\"\"Plot both the trace and the model together\"\"\"\n",
    "    \n",
    "    fig, ax = plt.subplots(1, 2, figsize=(20, 8))\n",
    "        \n",
    "    plot_MCMC_trace(ax[0],trace, model, True, colors=colors)\n",
    "    plot_MCMC_model(ax[1],hist, trace,model,units,bins,bar)\n",
    "    plt.suptitle(title,fontsize=20)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}