{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\nEmcee and the Model Interface\n=============================\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import corner\nimport matplotlib.pyplot as plt\nimport numpy as np\n\nimport lmfit"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Set up a double-exponential function and create a Model\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "def double_exp(x, a1, t1, a2, t2):\n    return a1*np.exp(-x/t1) + a2*np.exp(-(x-0.1) / t2)\n\n\nmodel = lmfit.Model(double_exp)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Generate some fake data from the model with added noise\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "truths = (3.0, 2.0, -5.0, 10.0)\nx = np.linspace(1, 10, 250)\nnp.random.seed(0)\ny = double_exp(x, *truths)+0.1*np.random.randn(x.size)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Create model parameters and give them initial values\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "p = model.make_params(a1=4, t1=3, a2=4, t2=3)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Fit the model using a traditional minimizer, and show the output:\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "result = model.fit(data=y, params=p, x=x, method='Nelder', nan_policy='omit')\n\nlmfit.report_fit(result)\nresult.plot()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Calculate parameter covariance using emcee:\n\n - start the walkers out at the best-fit values\n - set is_weighted to False to estimate the noise weights\n - set some sensible priors on the uncertainty to keep the MCMC in check\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "emcee_kws = dict(steps=1000, burn=300, thin=20, is_weighted=False,\n                 progress=False)\nemcee_params = result.params.copy()\nemcee_params.add('__lnsigma', value=np.log(0.1), min=np.log(0.001), max=np.log(2.0))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "run the MCMC algorithm and show the results:\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "result_emcee = model.fit(data=y, x=x, params=emcee_params, method='emcee',\n                         nan_policy='omit', fit_kws=emcee_kws)\n\nlmfit.report_fit(result_emcee)\n\nax = plt.plot(x, model.eval(params=result.params, x=x), label='Nelder', zorder=100)\nresult_emcee.plot_fit(ax=ax, data_kws=dict(color='gray', markersize=2))\nplt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "check the acceptance fraction to see whether emcee performed well\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "plt.plot(result_emcee.acceptance_fraction)\nplt.xlabel('walker')\nplt.ylabel('acceptance fraction')\nplt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "try to compute the autocorrelation time\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "if hasattr(result_emcee, \"acor\"):\n    print(\"Autocorrelation time for the parameters:\")\n    print(\"----------------------------------------\")\n    for i, p in enumerate(result.params):\n        print(p, result.acor[i])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Plot the parameter covariances returned by emcee using corner\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "emcee_corner = corner.corner(result_emcee.flatchain, labels=result_emcee.var_names,\n                             truths=list(result_emcee.params.valuesdict().values()))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "print(\"\\nmedian of posterior probability distribution\")\nprint('--------------------------------------------')\nlmfit.report_fit(result_emcee.params)\n\n# find the maximum likelihood solution\nhighest_prob = np.argmax(result_emcee.lnprob)\nhp_loc = np.unravel_index(highest_prob, result_emcee.lnprob.shape)\nmle_soln = result_emcee.chain[hp_loc]\nprint(\"\\nMaximum likelihood Estimation\")\nprint('-----------------------------')\nfor ix, param in enumerate(emcee_params):\n    print(param + ': ' + str(mle_soln[ix]))\n\nquantiles = np.percentile(result_emcee.flatchain['t1'], [2.28, 15.9, 50, 84.2, 97.7])\nprint(\"\\n\\n1 sigma spread\", 0.5 * (quantiles[3] - quantiles[1]))\nprint(\"2 sigma spread\", 0.5 * (quantiles[4] - quantiles[0]))"
      ]
    }
  ],
  "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.7.6"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}