{ "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 }