{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Model comparison\n",
    "\n",
    "To demonstrate the use of model comparison criteria in PyMC3, we implement the **8 schools** example from Section 5.5 of Gelman et al (2003), which attempts to infer the effects of coaching on SAT scores of students from 8 schools. Below, we fit a **pooled model**, which assumes a single fixed effect across all schools, and a **hierarchical model** that allows for a random effect that partially pools the data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import pymc3 as pm\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "plt.style.use('seaborn-darkgrid')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The data include the observed treatment effects and associated standard deviations in the 8 schools."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "J = 8\n",
    "y = np.array([28,  8, -3,  7, -1,  1, 18, 12])\n",
    "sigma = np.array([15, 10, 16, 11,  9, 11, 10, 18])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Pooled model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Auto-assigning NUTS sampler...\n",
      "Initializing NUTS using jitter+adapt_diag...\n",
      "Multiprocess sampling (2 chains in 2 jobs)\n",
      "NUTS: [mu]\n",
      "100%|██████████| 1500/1500 [00:01<00:00, 1003.81it/s]\n"
     ]
    }
   ],
   "source": [
    "with pm.Model() as pooled:\n",
    "    mu = pm.Normal('mu', 0, sd=1e6)\n",
    "    \n",
    "    obs = pm.Normal('obs', mu, sd=sigma, observed=y)\n",
    "    \n",
    "    trace_p = pm.sample(1000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 48596 characters]\n",
      "text/plain": [
       "<matplotlib.figure.Figure at 0x11750a240>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "pm.traceplot(trace_p);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Hierarchical model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Auto-assigning NUTS sampler...\n",
      "Initializing NUTS using jitter+adapt_diag...\n",
      "Multiprocess sampling (2 chains in 2 jobs)\n",
      "NUTS: [tau_log__, mu, eta]\n",
      "100%|██████████| 1500/1500 [00:03<00:00, 484.03it/s]\n"
     ]
    }
   ],
   "source": [
    "with pm.Model() as hierarchical:\n",
    "    \n",
    "    eta = pm.Normal('eta', 0, 1, shape=J)\n",
    "    mu = pm.Normal('mu', 0, sd=1e6)\n",
    "    tau = pm.HalfCauchy('tau', 5)\n",
    "    \n",
    "    theta = pm.Deterministic('theta', mu + tau*eta)\n",
    "    \n",
    "    obs = pm.Normal('obs', theta, sd=sigma, observed=y)\n",
    "    \n",
    "    trace_h = pm.sample(1000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 47128 characters]\n",
      "text/plain": [
       "<matplotlib.figure.Figure at 0x118a34ef0>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "pm.traceplot(trace_h, varnames=['mu']);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 13280 characters]\n",
      "text/plain": [
       "<matplotlib.figure.Figure at 0x10a0a33c8>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "pm.forestplot(trace_h, varnames=['theta']);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Widely-applicable Information Criterion (WAIC)\n",
    "\n",
    "WAIC (Watanabe 2010) is a fully Bayesian criterion for estimating out-of-sample expectation, using the computed log pointwise posterior predictive density (LPPD) and correcting for the effective number of parameters to adjust for overfitting."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "61.120811331601089"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pooled_waic = pm.waic(trace_p, pooled)\n",
    "    \n",
    "pooled_waic.WAIC"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "61.229456725964631"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "hierarchical_waic = pm.waic(trace_h, hierarchical)\n",
    "    \n",
    "hierarchical_waic.WAIC"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "PyMC3 includes two convenience functions to help compare WAIC for different models. The first of this functions is `compare`, this one computes WAIC (or LOO) from a set of traces and models and returns a DataFrame."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "hierarchical.name = 'hierarchical'\n",
    "pooled.name = 'pooled'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style>\n",
       "    .dataframe thead tr:only-child th {\n",
       "        text-align: right;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: left;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>WAIC</th>\n",
       "      <th>pWAIC</th>\n",
       "      <th>dWAIC</th>\n",
       "      <th>weight</th>\n",
       "      <th>SE</th>\n",
       "      <th>dSE</th>\n",
       "      <th>warning</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>pooled</th>\n",
       "      <td>61.12</td>\n",
       "      <td>0.67</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>2.21</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>hierarchical</th>\n",
       "      <td>61.23</td>\n",
       "      <td>0.94</td>\n",
       "      <td>0.11</td>\n",
       "      <td>0</td>\n",
       "      <td>2.01</td>\n",
       "      <td>0.27</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "               WAIC pWAIC dWAIC weight    SE   dSE warning\n",
       "pooled        61.12  0.67     0      1  2.21     0       0\n",
       "hierarchical  61.23  0.94  0.11      0  2.01  0.27       1"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_comp_WAIC = pm.compare({hierarchical: trace_h, pooled: trace_p})\n",
    "df_comp_WAIC"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have many columns so let check one by one the meaning of them:\n",
    "\n",
    "\n",
    "1. The first column clearly contains the values of WAIC. The DataFrame is always sorted from lowest to highest WAIC. The index reflects the order in which the models are passed to this function.\n",
    "\n",
    "2. The second column is the estimated effective number of parameters. In general, models with more parameters will be more flexible to fit data and at the same time could also lead to overfitting. Thus we can interpret pWAIC as a penalization term, intuitively we can also interpret it as measure of how flexible each model is in fitting the data. \n",
    "\n",
    "3. The third column is the relative difference between the value of WAIC for the top-ranked model and the value of WAIC for each model. For this reason we will always get a value of 0 for the first model.\n",
    "\n",
    "4. Sometimes when comparing models, we do not want to select the \"best\" model, instead we want to perform predictions by averaging along all the models (or at least several models). Ideally we would like to perform a weighted average, giving more weight to the model that seems to explain/predict the data better. There are many approaches to perform this task, one of them is to use Akaike weights based on the values of WAIC for each model. These weights can be loosely interpreted as the probability of each model (among the compared models) given the data. One caveat of this approach is that the weights are based on point estimates of WAIC (i.e. the uncertainty is ignored).\n",
    "\n",
    "5. The fifth column records the standard error for the WAIC computations. The standard error can be useful to assess the uncertainty of the WAIC estimates. Nevertheless, caution need to be taken because the estimation of the standard error assumes normality and hence could be problematic when the sample size is low.\n",
    "\n",
    "6. In the same way that we can compute the standard error for each value of WAIC, we can compute the standard error of the differences between two values of WAIC. Notice that both quantities are not necessarily the same, the reason is that the uncertainty about WAIC is correlated between models. This quantity is always 0 for the top-ranked model.\n",
    "\n",
    "7. Finally we have the last column named \"warning\". A value of 1 indicates that the computation of WAIC may not be reliable, this warning is based on an empirical determined cutoff value and need to be interpreted with caution. For more details you can read this [paper](https://arxiv.org/abs/1507.04544)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The second convenience function takes the output of `compare` and produces a summary plot in the style of the one used in the book [Statistical Rethinking](http://xcelab.net/rm/statistical-rethinking/) by Richard McElreath (check also [this port](https://github.com/aloctavodia/Statistical-Rethinking-with-Python-and-PyMC3) of the examples in the book to PyMC3)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 7184 characters]\n",
      "text/plain": [
       "<matplotlib.figure.Figure at 0x11a06a4e0>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "pm.compareplot(df_comp_WAIC);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The empty circle represents the values of WAIC and the black error bars associated with them are the values of the standard deviation of WAIC. \n",
    "\n",
    "The value of the lowest WAIC is also indicated with a vertical dashed grey line to ease comparison with other WAIC values.\n",
    "\n",
    "The filled black dots are the in-sample deviance of each model, which for WAIC is  2 pWAIC from the corresponding WAIC value.\n",
    "\n",
    "For all models except the top-ranked one we also get a triangle indicating the value of the difference of WAIC between that model and the top model and a grey errobar indicating the standard error of the differences between the top-ranked WAIC and WAIC for each model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Leave-one-out Cross-validation (LOO)\n",
    "\n",
    "LOO cross-validation is an estimate of the out-of-sample predictive fit. In cross-validation, the data are repeatedly partitioned into training and holdout sets, iteratively fitting the model with the former and evaluating the fit with the holdout data. Vehtari et al. (2016) introduced an efficient computation of LOO from MCMC samples, which are corrected using Pareto-smoothed importance sampling (PSIS) to provide an estimate of point-wise out-of-sample prediction accuracy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "61.159472594603194"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pooled_loo = pm.loo(trace_p, pooled)\n",
    "    \n",
    "pooled_loo.LOO"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/fonnescj/Repos/pymc3/pymc3/stats.py:293: UserWarning: Estimated shape parameter of Pareto distribution is\n",
      "        greater than 0.7 for one or more samples.\n",
      "        You should consider using a more robust model, this is because\n",
      "        importance sampling is less likely to work well if the marginal\n",
      "        posterior and LOO posterior are very different. This is more likely to\n",
      "        happen with a non-robust model and highly influential observations.\n",
      "  happen with a non-robust model and highly influential observations.\"\"\")\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "61.336245930247486"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "hierarchical_loo  = pm.loo(trace_h, hierarchical)\n",
    "    \n",
    "hierarchical_loo.LOO"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also use `compare` with LOO."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style>\n",
       "    .dataframe thead tr:only-child th {\n",
       "        text-align: right;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: left;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>LOO</th>\n",
       "      <th>pLOO</th>\n",
       "      <th>dLOO</th>\n",
       "      <th>weight</th>\n",
       "      <th>SE</th>\n",
       "      <th>dSE</th>\n",
       "      <th>warning</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>pooled</th>\n",
       "      <td>61.16</td>\n",
       "      <td>0.69</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>2.21</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>hierarchical</th>\n",
       "      <td>61.34</td>\n",
       "      <td>0.99</td>\n",
       "      <td>0.18</td>\n",
       "      <td>0</td>\n",
       "      <td>2.03</td>\n",
       "      <td>0.25</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                LOO  pLOO  dLOO weight    SE   dSE warning\n",
       "pooled        61.16  0.69     0      1  2.21     0       0\n",
       "hierarchical  61.34  0.99  0.18      0  2.03  0.25       1"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_comp_LOO = pm.compare({hierarchical: trace_h, pooled: trace_p}, ic='LOO')\n",
    "df_comp_LOO"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The columns return the equivalent values for LOO, notice that in this example we get two warnings. Also notice that the order of the models is not the same as the one for WAIC.\n",
    "\n",
    "We can also plot the results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 7172 characters]\n",
      "text/plain": [
       "<matplotlib.figure.Figure at 0x119ff3d30>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "pm.compareplot(df_comp_LOO);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Interpretation\n",
    "\n",
    "Though we might expect the hierarchical model to outperform a complete pooling model, there is little to choose between the models in this case, giving that both models gives very similar values of the information criteria. This is more clearly appreciated when we take into account the uncertainty (in terms of standard errors) of WAIC and LOO.\n",
    "\n",
    "## Reference\n",
    "\n",
    "[Gelman, A., Hwang, J., & Vehtari, A. (2014). Understanding predictive information criteria for Bayesian models. Statistics and Computing, 24(6), 997–1016.](http://doi.org/10.1007/s11222-013-9416-2)\n",
    "\n",
    "[Vehtari, A, Gelman, A, Gabry, J. (2016). Practical Bayesian model evaluation using leave-one-out cross-validation and WAIC. Statistics and Computing](http://link.springer.com/article/10.1007/s11222-016-9696-4)"
   ]
  }
 ],
 "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.6.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
