code.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "<div class=\"title\"> Modern Bayesian methods: principles and practice</div>\n",
    "\n",
    "<hr/>\n",
    "\n",
    "## Vinayak Rao, Purdue University\n",
    "### Jan 4, 2021\n",
    "\n",
    "Material at $\\mathtt{github.com/varao/ds3\\_bayesian}$\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import scipy as sp\n",
    "import pystan\n",
    "import arviz as az\n",
    "from matplotlib import pyplot as plt\n",
    "import numdifftools as nd    # Optional"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [],
   "source": [
    "# Data from https://github.com/ImperialCollegeLondon/covid19model\n",
    "covid_agg = pd.read_csv('covid_agg.csv')\n",
    "covid_agg.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "covid_agg['rate'] = covid_agg.cases/covid_agg.popData2018\n",
    "covid_agg.rate.hist()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "### Modeling the infection rate\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### A brief intro to Stan (from the Stan website)\n",
    "\n",
    "\n",
    "\n",
    "Stan is a state-of-the-art platform for statistical modeling and high-performance statistical computation.\n",
    "\n",
    "Users specify log density functions in Stan’s probabilistic programming language and get:\n",
    "\n",
    "+ full Bayesian statistical inference with MCMC sampling (NUTS, HMC)\n",
    "\n",
    "+ approximate Bayesian inference with variational inference (ADVI)\n",
    "\n",
    "+ penalized maximum likelihood estimation with optimization (L-BFGS)\n",
    "\n",
    "Stan interfaces with the most popular data analysis languages (R, Python, shell, MATLAB, Julia, Stata) \n",
    "\n",
    "Stan User guide: https://mc-stan.org/docs/2_25/stan-users-guide/index.html\n",
    "\n",
    "Stan reference manual: https://mc-stan.org/docs/2_25/reference-manual/index.html\n",
    "\n",
    "PyStan: https://pystan.readthedocs.io/en/latest/api.html\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [],
   "source": [
    "model1_code = \"\"\"\n",
    "data {\n",
    "    int<lower=0> N; // number of countries\n",
    "    vector[N] x; // proportion of cases\n",
    "}\n",
    "\n",
    "parameters {\n",
    "    real<lower=0> a;\n",
    "    real<lower=0> b;\n",
    "\n",
    "}\n",
    "\n",
    "model {\n",
    "    a ~ exponential(.0001);\n",
    "    b ~ exponential(.0001);\n",
    "    for(i in 1:N) {\n",
    "      x[i] ~ beta(a,b);\n",
    "    }\n",
    "}\n",
    "\n",
    "\"\"\"\n",
    "sm1 = pystan.StanModel(model_code=model1_code)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [],
   "source": [
    "model1_data = {'N': len(covid_agg.rate), 'x': covid_agg.rate}\n",
    "fit = sm.sampling(data=model1_data, iter=5000, chains=1)\n",
    "az.plot_density(fit);\n",
    "\n",
    "print(fit);   #pd.DataFrame(fit.extract())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "### Exercise\n",
    "\n",
    "+ Is this a good fit to the data? How would you tell?\n",
    "\n",
    "+ What is the distribution of first observation?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "Instead of the infection rate, let us model the log-odds\n",
    "\n",
    "For a probability of success $p$, the log odds are $\\log(\\frac{p}{1-p})$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "covid_agg['logodds'] = np.log(covid_agg.rate/(1-covid_agg.rate))\n",
    "covid_agg.logodds.hist(bins=20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [],
   "source": [
    "model2_code = \"\"\"\n",
    "data {\n",
    "    int<lower=0> N; // number of countries\n",
    "    real<lower=0> sigma; \n",
    "    vector[N] x; // proportion of cases\n",
    "}\n",
    "\n",
    "parameters {\n",
    "    real mu;\n",
    "}\n",
    "\n",
    "model {\n",
    "    mu ~ normal(0, 10);\n",
    "    for(i in 1:N) {\n",
    "      x[i] ~ normal(mu,sigma);\n",
    "    }\n",
    "}\n",
    "    \n",
    "generated quantities {\n",
    "    real mn;\n",
    "    real mx; \n",
    "    real std;\n",
    "    {\n",
    "      vector[N] x_rep;\n",
    "      for(j in 1:N) {\n",
    "        x_rep[j] <- normal_rng(mu, sigma);\n",
    "      }\n",
    "      std = sd(x_rep);\n",
    "      mn = min(x_rep);\n",
    "      mx = max(x_rep);\n",
    "    }\n",
    "}\n",
    "\"\"\"\n",
    "sm2 = pystan.StanModel(model_code=model2_code)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [],
   "source": [
    "model2_data = {'N': len(covid_agg.logodds), 'sigma':11, 'x': covid_agg.logodds}\n",
    "fit = sm2.sampling(data=model2_data, iter=5000, chains=1)\n",
    "\n",
    "az.plot_density(fit);\n",
    "\n",
    "print(fit);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### Model checking\n",
    "\n",
    "Is the previous model a good fit of the data?\n",
    "\n",
    "How can you quantify this?\n",
    "\n",
    "+ Cross-validation\n",
    "\n",
    "+ Posterior predictive checks\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [],
   "source": [
    "model3_code = \"\"\"\n",
    "data {\n",
    "    int<lower=0> N; // number of countries\n",
    "    vector[N] x; // proportion of cases\n",
    "}\n",
    "\n",
    "parameters {\n",
    "    real mu;\n",
    "    real<lower=0> sigma; \n",
    "\n",
    "}\n",
    "\n",
    "model {\n",
    "    mu ~ normal(0, 10);\n",
    "    sigma ~ gamma(.01,.01);\n",
    "    for(i in 1:N) {\n",
    "      x[i] ~ normal(mu,sigma);\n",
    "    }\n",
    "}\n",
    "\"\"\"\n",
    "sm3 = pystan.StanModel(model_code=model3_code)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [],
   "source": [
    "model3_data = {'N': len(covid_agg.logodds), 'x': covid_agg.logodds}\n",
    "\n",
    "fit = sm3.sampling(data=model3_data, iter=5000, warmup=100, chains=1)\n",
    "az.plot_density(fit)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "### Exercises\n",
    "+ Can we place a prior on the prior mean?\n",
    "\n",
    "+ What properties of the data might the previous model fail to capture?\n",
    "\n",
    "+ Modify the code to quantify this"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### Hierarchical Bayes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [],
   "source": [
    "model_hier_code = \"\"\"\n",
    "data {\n",
    "    int<lower=0> N; // number of countries\n",
    "    int<lower = 1> C; // number of continents\n",
    "\n",
    "    vector[N] x; // log-odds of proportion of cases\n",
    "    int<lower=1, upper=C> cont[N];  // continent\n",
    "}\n",
    "\n",
    "parameters {\n",
    "    real mu0;\n",
    "    real<lower=0> sigma0; \n",
    "    \n",
    "    vector[C] mu_cont;\n",
    "    real<lower=0> sigma; \n",
    "}\n",
    "\n",
    "model {\n",
    "    mu0 ~ normal([0,0], 10);\n",
    "    sigma0 ~ gamma(.01,.01);\n",
    "    sigma ~ gamma(.01,.01);\n",
    "\n",
    "    for(i in 1:C) {\n",
    "      mu_cont[i] ~ normal(mu0,sigma0);\n",
    "    }\n",
    "    x ~ normal(mu_cont[cont],sigma);\n",
    "}\n",
    "\"\"\"\n",
    "sm_hier = pystan.StanModel(model_code=model_hier_code)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [],
   "source": [
    "model_hier_data = {'N': len(covid_agg.logodds), 'C': max(covid_agg.cont), 'x': covid_agg.logodds, \n",
    "               'cont': covid_agg.cont}\n",
    "fit = sm_hier.sampling(data=model_hier_data, iter=5000, warmup=1000, chains=1)\n",
    "az.plot_density(fit)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### Exercise\n",
    "\n",
    "+ What are different ways you can extend the hierarchical model above to allow more flexibility?\n",
    "+ How do you expect the parameter estimates to differ from the situation where each group is modeled independently?\n",
    "+ Write Stan code for the latter case\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "### Markov chain Monte Carlo (MCMC)\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "### The Titanic dataset\n",
    "\n",
    "\n",
    "\n",
    "Around 800 measurements including:\n",
    "\n",
    "<ul>\n",
    "<li> survival ($y$): Survival (0 = No; 1 = Yes)\n",
    "<li> pclass ($x_1$): Passenger Class (1 = 1st; 2 = 2nd; 3 = 3rd)\n",
    "<li> sex ($x_2$): Sex (1 = female, 0 = male)\n",
    "<li> age ($x_3$): Age\n",
    "<li> ticket ($x_3$): Ticket Number\n",
    "<li> fare ($x_4$): Passenger Fare\n",
    "</ul>\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "titanic = pd.read_csv('titanic.csv')\n",
    "titanic.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [],
   "source": [
    "model_code = \"\"\"\n",
    "\n",
    "data {\n",
    "  int<lower=0> N; // number of obs\n",
    "  int<lower=0> P; // number of predictors\n",
    "  int<lower=0,upper=1> y[N]; // outcomes\n",
    "  matrix[N, P] x; // predictor variables\n",
    "}\n",
    "parameters {\n",
    "  vector[P] theta; // theta coefficients\n",
    "}\n",
    "model {\n",
    "  vector[N] mu; \n",
    "  theta ~ normal(0, 100);   // cauchy(0,10) is what Gelman et al recommend\n",
    "  mu <- x*theta;\n",
    "  for (n in 1:N) mu[n] <- Phi(mu[n]);\n",
    "  y ~ bernoulli(mu);\n",
    "}\n",
    "\"\"\"\n",
    "sm_probit = pystan.StanModel(model_code=model_code)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [],
   "source": [
    "titanic_data = {'N': len(titanic.survived), 'P': titanic.shape[1]-1, \n",
    "                'x': titanic.drop('survived',axis=1), 'y': titanic.survived}\n",
    "fit = sm_probit.sampling(data=titanic_data, iter=5000, warmup=1000, chains=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "az.plot_density(fit)\n",
    "print(fit)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [],
   "source": [
    "def probit_loglik(theta, x,y):\n",
    "    mu = x @ theta\n",
    "    \n",
    "    a = y*sp.stats.norm(0,1).logcdf(mu)\n",
    "    b =  (1-y)*np.log(1-sp.stats.norm(0,1).cdf(mu))\n",
    "    return sum(a+b)-(theta.T @ theta)/(2*100)\n",
    "\n",
    "def probit_mh(niter, M, data):\n",
    "    N, P, x, y = data['N'], data['P'], data['x'], data['y']\n",
    "    y = y[:,np.newaxis]\n",
    "    \n",
    "    theta = np.zeros([P, niter])\n",
    "    \n",
    "    for i in np.arange(1, niter):\n",
    "        prop = np.random.multivariate_normal(theta[:,i-1], M)\n",
    "        prop = prop[:,np.newaxis]\n",
    "        #print(theta[i,].T)\n",
    "\n",
    "        #print(probit_loglik(theta[i,].T,x,y))\n",
    "\n",
    "        if np.log(np.random.uniform()) < (probit_loglik(prop,x,y) - probit_loglik(theta[:,[i-i]],x,y)):\n",
    "            theta[:,[i]] = prop\n",
    "        else:\n",
    "            theta[:,[i]] = theta[:,[i-1]]\n",
    "    return theta.T\n",
    "    \n",
    "M = np.identity(titanic_data['P'])\n",
    "M = np.linalg.inv(iM)\n",
    "rslt = probit_mh(1000,M*5,titanic_data)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [],
   "source": [
    "rslt.mean(0), rslt.std(0)\n",
    "#jnk = rslt[1,]\n",
    "#jnk[:,np.newaxis]\n",
    "plt.plot(rslt[:,0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "outputs": [],
   "source": [
    "[az.ess(rslt[:,[i]].T) for i in range(6) ]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "source": [
    "### Exercise\n",
    "\n",
    "+ Play around with the M above. What gives you higher effective sample size (ESS)? How does it compare to Stan's results?\n",
    "+ How would you debug your sampler to make sure it is working?\n",
    "+ Does a high ESS mean a better sampler?\n",
    "+ Does a high acceptance rate mean a better sampler?\n",
    "+ Can the proposal variance also depend on the current location?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "source": [
    "### A brief look at Bayesian nonparametrics\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [],
   "source": [
    "def my_kernel(X,s,l):\n",
    "    n = len(X)\n",
    "    p = 2\n",
    "    D = np.zeros([n,n])\n",
    "\n",
    "    for i in range(n):\n",
    "        D[i,] = np.abs((X[i] - X)) ** p\n",
    "    \n",
    "    return s*np.exp(-0.5*D/l) + 1e-5*np.eye(n)\n",
    "\n",
    "X = np.arange(0,10,step=.01)\n",
    "cv = my_kernel(X,3,1)\n",
    "\n",
    "plt.plot(X, np.random.multivariate_normal(X*0,cv))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "__Exercise__: play around with the parameters s,l,p of the the kernel. Note down your conclusions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "subslide"
    }
   },
   "outputs": [],
   "source": [
    "model_code = \"\"\"\n",
    "data {\n",
    "  int<lower=1> N;\n",
    "  real x[N];\n",
    "  vector[N] y;\n",
    "}\n",
    "\n",
    "parameters {\n",
    "  real<lower=0> rho;\n",
    "  real<lower=0> alpha;\n",
    "  real<lower=0> sigma;\n",
    "  vector[N] eta;\n",
    "}\n",
    "\n",
    "transformed parameters {\n",
    "  vector[N] f;\n",
    "  {\n",
    "    matrix[N, N] L_K;\n",
    "    matrix[N, N] K = cov_exp_quad(x, alpha, rho);\n",
    "\n",
    "    // diagonal elements\n",
    "    for (n in 1:N)\n",
    "      K[n, n] = K[n, n] + 1e-4;\n",
    "\n",
    "    L_K = cholesky_decompose(K);\n",
    "    f = L_K * eta;\n",
    "  }\n",
    "}\n",
    "\n",
    "model {\n",
    "  real sq_sigma = square(sigma);\n",
    "\n",
    "\n",
    "  eta ~ std_normal();\n",
    "\n",
    "  rho ~ inv_gamma(5, 5);\n",
    "  alpha ~ std_normal();\n",
    "  sigma ~ std_normal();\n",
    "\n",
    "  y ~ normal(f, sq_sigma);\n",
    "}\n",
    "\"\"\"\n",
    "\n",
    "sm_gp = pystan.StanModel(model_code=model_code)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_gp_data = {'N': len(covid_agg.rate), 'y': np.log(covid_agg.cases), 'x': np.log(covid_agg.popData2018), \n",
    "               'cont': covid_agg.cont}\n",
    "fit = sm_gp.sampling(data=model_gp_data, iter=1000,  chains=1)\n",
    "#az.plot_density(fit)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tmp = fit.extract('f')['f']\n",
    "mn = tmp.mean(axis=0)\n",
    "sd = tmp.std(axis=0)\n",
    "\n",
    "rslt = pd.DataFrame({'x': model_gp_data['x'], 'mn':mn, 'uci': mn+3*sd, 'lci': mn-3*sd})\n",
    "rslt =  rslt.sort_values(by='x')\n",
    "\n",
    "plt.plot(rslt.x,rslt.mn)\n",
    "plt.plot(rslt.x,rslt.uci)\n",
    "plt.plot(rslt.x,rslt.lci)\n",
    "plt.scatter(rslt.x, model_gp_data['y'])\n"
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Slideshow",
  "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}