statrethink_numpyro_w07.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "7DzCWnAGlmVG"
   },
   "source": [
    "# Rethinking Statistics course in NumPyro - Week 7"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "nlScIa9GlmVG"
   },
   "source": [
    "Lecture 13: Monsters & Mixtures (Poisson GLMs, survival, zero-inflation)\n",
    "\n",
    "- [Video](https://www.youtube.com/watch?v=p7g-CgGCS34)\n",
    "- [Slides](https://speakerdeck.com/rmcelreath/l13-statistical-rethinking-winter-2019)\n",
    "\n",
    "Lecture 14: Ordered Categories, Left & Right\n",
    "\n",
    "- [Video](https://www.youtube.com/watch?v=zA3Jxv8LOrA)\n",
    "- [Slides](https://speakerdeck.com/rmcelreath/l14-statistical-rethinking-winter-2019)\n",
    "\n",
    "[Proposed problems](https://github.com/gbosquechacon/statrethinking_winter2019/blob/master/homework/week07.pdf) and [solutions in R](https://github.com/gbosquechacon/statrethinking_winter2019/blob/master/homework/week07_solutions.pdf) for the exercises of the week."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "executionInfo": {
     "elapsed": 4156,
     "status": "ok",
     "timestamp": 1607275920095,
     "user": {
      "displayName": "Andrés Suárez",
      "photoUrl": "https://lh5.googleusercontent.com/-s0kzcIwylzA/AAAAAAAAAAI/AAAAAAAAQXA/v8Sc6WgQy7c/s64/photo.jpg",
      "userId": "06409440331868776168"
     },
     "user_tz": -60
    },
    "id": "mTdV4mj8of07"
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "\n",
    "from jax import random\n",
    "import jax.numpy as jnp\n",
    "import numpyro\n",
    "import numpyro.distributions as dist\n",
    "from numpyro.distributions.transforms import OrderedTransform\n",
    "from numpyro.infer import NUTS, MCMC, Predictive"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "executionInfo": {
     "elapsed": 4150,
     "status": "ok",
     "timestamp": 1607275920096,
     "user": {
      "displayName": "Andrés Suárez",
      "photoUrl": "https://lh5.googleusercontent.com/-s0kzcIwylzA/AAAAAAAAAAI/AAAAAAAAQXA/v8Sc6WgQy7c/s64/photo.jpg",
      "userId": "06409440331868776168"
     },
     "user_tz": -60
    },
    "id": "73XDim7BopOZ",
    "outputId": "6291d964-d15c-404c-a450-3c4462513ecc"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Last updated: Sat Mar 27 2021\n",
      "\n",
      "Python implementation: CPython\n",
      "Python version       : 3.9.1\n",
      "IPython version      : 7.21.0\n",
      "\n",
      "numpyro: 0.6.0\n",
      "pandas : 1.2.3\n",
      "jax    : 0.2.10\n",
      "numpy  : 1.20.1\n",
      "\n",
      "Watermark: 2.2.0\n",
      "\n"
     ]
    }
   ],
   "source": [
    "%load_ext watermark\n",
    "%watermark -n -u -v -iv -w"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "executionInfo": {
     "elapsed": 3859,
     "status": "ok",
     "timestamp": 1607275920098,
     "user": {
      "displayName": "Andrés Suárez",
      "photoUrl": "https://lh5.googleusercontent.com/-s0kzcIwylzA/AAAAAAAAAAI/AAAAAAAAQXA/v8Sc6WgQy7c/s64/photo.jpg",
      "userId": "06409440331868776168"
     },
     "user_tz": -60
    },
    "id": "In86umhVov6K"
   },
   "outputs": [],
   "source": [
    "rng_key = random.PRNGKey(0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4hyJfeczlmVG"
   },
   "source": [
    "## Exercise 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "JusJHQKZlmVG"
   },
   "source": [
    "> In the Trolley data, `Trolley`, we saw how education level (modeled as an ordered category) is associated with responses. Is this association causal? One plausible confound is that education is also associated with age, through a causal process: People are older when they finish school than when they begin it.\n",
    "\n",
    "> Reconsider the `Trolley` data in this light. Draw a DAG that represents hypothetical causal relationships among response, education, and age. Which statical model or models do you need to evaluate the causal influence of education on responses? Fit these models to the trolley data. What do you conclude about the causal relationships among these three variables?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "sLRHSDJ_lmVG"
   },
   "source": [
    "This is my DAG:\n",
    "\n",
    "<img src=\"./fig/w7_img1.png\" width=\"30%\">\n",
    "\n",
    "Age could influence both response `R` and education `E`. It could influence `R`, because people at different ages could have different attitudes. Age could influence education, because the longer you have lived, the more education you could have completed (up to a point). It's like the age causing marriage example from earlier in the course.\n",
    "\n",
    "To evaluate the causal influence of `E` on `R`, we need to block the back-door from `E` through `A` to `R`. So we need a model that conditions on both `E` and `A`. Then the estimate for `E` should be the causal influence of `E`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 142
    },
    "executionInfo": {
     "elapsed": 1085,
     "status": "ok",
     "timestamp": 1607275920099,
     "user": {
      "displayName": "Andrés Suárez",
      "photoUrl": "https://lh5.googleusercontent.com/-s0kzcIwylzA/AAAAAAAAAAI/AAAAAAAAQXA/v8Sc6WgQy7c/s64/photo.jpg",
      "userId": "06409440331868776168"
     },
     "user_tz": -60
    },
    "id": "yofca_3zlmVG",
    "outputId": "e1609462-d9af-436d-b932-7c2d1cef8724"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>case</th>\n",
       "      <th>response</th>\n",
       "      <th>order</th>\n",
       "      <th>id</th>\n",
       "      <th>age</th>\n",
       "      <th>male</th>\n",
       "      <th>edu</th>\n",
       "      <th>action</th>\n",
       "      <th>intention</th>\n",
       "      <th>contact</th>\n",
       "      <th>story</th>\n",
       "      <th>action2</th>\n",
       "      <th>edu_cat</th>\n",
       "      <th>age_std</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>9927</th>\n",
       "      <td>ilshi</td>\n",
       "      <td>7</td>\n",
       "      <td>7</td>\n",
       "      <td>98;299</td>\n",
       "      <td>66</td>\n",
       "      <td>1</td>\n",
       "      <td>Graduate Degree</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>shi</td>\n",
       "      <td>0</td>\n",
       "      <td>7</td>\n",
       "      <td>2.003041</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9928</th>\n",
       "      <td>ilswi</td>\n",
       "      <td>2</td>\n",
       "      <td>18</td>\n",
       "      <td>98;299</td>\n",
       "      <td>66</td>\n",
       "      <td>1</td>\n",
       "      <td>Graduate Degree</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>swi</td>\n",
       "      <td>0</td>\n",
       "      <td>7</td>\n",
       "      <td>2.003041</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9929</th>\n",
       "      <td>nfrub</td>\n",
       "      <td>2</td>\n",
       "      <td>17</td>\n",
       "      <td>98;299</td>\n",
       "      <td>66</td>\n",
       "      <td>1</td>\n",
       "      <td>Graduate Degree</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>rub</td>\n",
       "      <td>1</td>\n",
       "      <td>7</td>\n",
       "      <td>2.003041</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "       case  response  order      id  age  male              edu  action  \\\n",
       "9927  ilshi         7      7  98;299   66     1  Graduate Degree       0   \n",
       "9928  ilswi         2     18  98;299   66     1  Graduate Degree       0   \n",
       "9929  nfrub         2     17  98;299   66     1  Graduate Degree       1   \n",
       "\n",
       "      intention  contact story  action2  edu_cat   age_std  \n",
       "9927          1        0   shi        0        7  2.003041  \n",
       "9928          1        0   swi        0        7  2.003041  \n",
       "9929          0        0   rub        1        7  2.003041  "
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "d = pd.read_csv('./dat/Trolley.csv', header=0, sep=';')\n",
    "elvl = d['edu'].unique()\n",
    "idx = [7 , 0 , 6 , 4 , 2 , 1, 3, 5]\n",
    "cat = pd.Categorical(d.edu, categories=list(elvl[idx]), ordered=True)\n",
    "d['edu_cat'] = pd.factorize(cat, sort=True)[0].astype('int')\n",
    "d['age_std'] = (d.age - d.age.mean())/d.age.std()\n",
    "d.tail(3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "executionInfo": {
     "elapsed": 1380,
     "status": "ok",
     "timestamp": 1607275920676,
     "user": {
      "displayName": "Andrés Suárez",
      "photoUrl": "https://lh5.googleusercontent.com/-s0kzcIwylzA/AAAAAAAAAAI/AAAAAAAAQXA/v8Sc6WgQy7c/s64/photo.jpg",
      "userId": "06409440331868776168"
     },
     "user_tz": -60
    },
    "id": "QNHlv_diw5om"
   },
   "outputs": [],
   "source": [
    "# https://fehiepsi.github.io/rethinking-numpyro/12-monsters-and-mixtures.html\n",
    "\n",
    "def model(action, contact, intention, age_std, edu_cat, response=None):\n",
    "    # priors\n",
    "    cutpoints = numpyro.sample('cutpoints', \n",
    "                               dist.TransformedDistribution( \n",
    "                                   dist.Normal(0, 1.5).expand([6]), \n",
    "                                   OrderedTransform()\n",
    "                                   ),\n",
    "                               )\n",
    "    norm = lambda label: numpyro.sample(label, dist.Normal(0,0.5))\n",
    "    bA = norm('bA')\n",
    "    bC = norm('bC')\n",
    "    bI = norm('bI')\n",
    "    bIA = norm('bIA')\n",
    "    bIC = norm('bIC')\n",
    "    bAge = norm('bAge')\n",
    "    bEdu = norm('bEdu')\n",
    "    edu_coefs = numpyro.sample('edu_coefs', dist.Dirichlet(jnp.repeat(2,7)))\n",
    "    # likelihood\n",
    "    delta_j = jnp.pad(edu_coefs, (1, 0))\n",
    "    delta_E = jnp.sum(jnp.where(jnp.arange(8) <= edu_cat[..., None], delta_j, 0), -1)\n",
    "    BI = bI + bIA*action + bIC*contact\n",
    "    phi = bA*action + BI*intention + bC*contact + bAge*age_std + bEdu*delta_E\n",
    "    response_hat = numpyro.sample('response_hat', dist.OrderedLogistic(phi, cutpoints), obs=response-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "sample: 100%|██████████| 1000/1000 [01:14<00:00, 13.41it/s, 31 steps of size 1.41e-01. acc. prob=0.94]\n",
      "sample: 100%|██████████| 1000/1000 [00:53<00:00, 18.57it/s, 31 steps of size 1.74e-01. acc. prob=0.92]\n",
      "sample: 100%|██████████| 1000/1000 [00:48<00:00, 20.64it/s, 31 steps of size 1.42e-01. acc. prob=0.94]\n",
      "sample: 100%|██████████| 1000/1000 [00:53<00:00, 18.66it/s, 63 steps of size 9.42e-02. acc. prob=0.95]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "                  mean       std    median      5.0%     95.0%     n_eff     r_hat\n",
      "          bA     -0.48      0.05     -0.48     -0.57     -0.39   1383.00      1.00\n",
      "        bAge     -0.10      0.02     -0.10     -0.13     -0.06    582.99      1.00\n",
      "          bC     -0.35      0.07     -0.35     -0.47     -0.24   1340.60      1.00\n",
      "        bEdu      0.21      0.14      0.23      0.04      0.41    235.69      1.01\n",
      "          bI     -0.30      0.06     -0.30     -0.40     -0.20   1188.45      1.00\n",
      "         bIA     -0.43      0.08     -0.43     -0.55     -0.30   1289.67      1.00\n",
      "         bIC     -1.24      0.10     -1.24     -1.39     -1.07   1182.44      1.00\n",
      "cutpoints[0]     -2.51      0.11     -2.49     -2.65     -2.32    279.62      1.01\n",
      "cutpoints[1]     -1.81      0.11     -1.80     -1.94     -1.63    264.82      1.01\n",
      "cutpoints[2]     -1.21      0.11     -1.20     -1.34     -1.03    267.12      1.01\n",
      "cutpoints[3]     -0.17      0.11     -0.16     -0.32     -0.01    267.94      1.02\n",
      "cutpoints[4]      0.50      0.11      0.51      0.36      0.67    255.45      1.02\n",
      "cutpoints[5]      1.40      0.11      1.42      1.25      1.57    251.75      1.02\n",
      "edu_coefs[0]      0.11      0.08      0.10      0.01      0.22   1165.66      1.00\n",
      "edu_coefs[1]      0.12      0.08      0.11      0.01      0.23   1762.44      1.00\n",
      "edu_coefs[2]      0.09      0.06      0.08      0.00      0.18   1069.67      1.00\n",
      "edu_coefs[3]      0.07      0.06      0.06      0.00      0.14    528.74      1.01\n",
      "edu_coefs[4]      0.42      0.16      0.44      0.17      0.72    292.82      1.01\n",
      "edu_coefs[5]      0.08      0.06      0.07      0.00      0.16   2250.29      1.00\n",
      "edu_coefs[6]      0.10      0.07      0.09      0.00      0.19   1320.81      1.00\n",
      "\n",
      "Number of divergences: 0\n"
     ]
    }
   ],
   "source": [
    "kernel = NUTS(model)\n",
    "mcmc = MCMC(kernel, num_warmup=500, num_samples=500, num_chains=4, chain_method='sequential')\n",
    "dat = {k:v.to_numpy() for k,v in d[['action', 'contact', 'intention', 'age_std', 'edu_cat', 'response']].items()}\n",
    "mcmc.run(rng_key, **dat)\n",
    "mcmc.print_summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "KH14kdI3lmVH"
   },
   "source": [
    "You may recall from the chapter that education has a negative effect in the model without age. Now that we include age, education has a positive influence (with some overlap with zero). So age has indeed soaked up some of the previous influence assigned to education. The back-door may be real.\n",
    "\n",
    "I'd summarize this model, assuming this DAG is true, as saying that age causes people to give slightly lower responses. This could be a cohort effect, and not a causal influence of age. Either way, it is small. Education seems to cause higher responses (more approval). This suggests that education trains people to see some or all of the features A, I, C as more permissible. A model that interacted education with each might shed more light on things. Remember: A DAG doesn't say whether you need an interaction effect or not. That is a separate problem."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "WuHiFDrKlmVH"
   },
   "source": [
    "## Exercise 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "E5Ag3kdTlmVH"
   },
   "source": [
    "> Consider one more variable in the Trolley data: Gender. Suppose that gender might influence education as well as response directly. Draw the DAG now that includes response, education, age, and gender.\n",
    "\n",
    "> Using only the DAG, is it possible that the inferences from Problem 1 are confounded by gender? If so, define any additional models you need to infer the causal influence of education on response. What do you conclude?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5meG1bJflmVH"
   },
   "source": [
    "This is my DAG:\n",
    "\n",
    "<img src=\"./fig/w7_img2.png\" width=\"30%\">\n",
    "\n",
    "Here's the model we need, which includes education, age, and gender (female dummy variable):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 142
    },
    "executionInfo": {
     "elapsed": 883,
     "status": "ok",
     "timestamp": 1607276542874,
     "user": {
      "displayName": "Andrés Suárez",
      "photoUrl": "https://lh5.googleusercontent.com/-s0kzcIwylzA/AAAAAAAAAAI/AAAAAAAAQXA/v8Sc6WgQy7c/s64/photo.jpg",
      "userId": "06409440331868776168"
     },
     "user_tz": -60
    },
    "id": "QzKDJctfJAIy",
    "outputId": "a1ceab9d-dded-4cea-d3a2-059913d5ccba"
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>case</th>\n",
       "      <th>response</th>\n",
       "      <th>order</th>\n",
       "      <th>id</th>\n",
       "      <th>age</th>\n",
       "      <th>male</th>\n",
       "      <th>edu</th>\n",
       "      <th>action</th>\n",
       "      <th>intention</th>\n",
       "      <th>contact</th>\n",
       "      <th>story</th>\n",
       "      <th>action2</th>\n",
       "      <th>edu_cat</th>\n",
       "      <th>age_std</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>9927</th>\n",
       "      <td>ilshi</td>\n",
       "      <td>7</td>\n",
       "      <td>7</td>\n",
       "      <td>98;299</td>\n",
       "      <td>66</td>\n",
       "      <td>1</td>\n",
       "      <td>Graduate Degree</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>shi</td>\n",
       "      <td>0</td>\n",
       "      <td>7</td>\n",
       "      <td>2.003041</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9928</th>\n",
       "      <td>ilswi</td>\n",
       "      <td>2</td>\n",
       "      <td>18</td>\n",
       "      <td>98;299</td>\n",
       "      <td>66</td>\n",
       "      <td>1</td>\n",
       "      <td>Graduate Degree</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>swi</td>\n",
       "      <td>0</td>\n",
       "      <td>7</td>\n",
       "      <td>2.003041</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9929</th>\n",
       "      <td>nfrub</td>\n",
       "      <td>2</td>\n",
       "      <td>17</td>\n",
       "      <td>98;299</td>\n",
       "      <td>66</td>\n",
       "      <td>1</td>\n",
       "      <td>Graduate Degree</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>rub</td>\n",
       "      <td>1</td>\n",
       "      <td>7</td>\n",
       "      <td>2.003041</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "       case  response  order      id  age  male              edu  action  \\\n",
       "9927  ilshi         7      7  98;299   66     1  Graduate Degree       0   \n",
       "9928  ilswi         2     18  98;299   66     1  Graduate Degree       0   \n",
       "9929  nfrub         2     17  98;299   66     1  Graduate Degree       1   \n",
       "\n",
       "      intention  contact story  action2  edu_cat   age_std  \n",
       "9927          1        0   shi        0        7  2.003041  \n",
       "9928          1        0   swi        0        7  2.003041  \n",
       "9929          0        0   rub        1        7  2.003041  "
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "d = pd.read_csv('./dat/Trolley.csv', header=0, sep=';')\n",
    "elvl = d['edu'].unique()\n",
    "idx = [7 , 0 , 6 , 4 , 2 , 1, 3, 5]\n",
    "cat = pd.Categorical(d.edu, categories=list(elvl[idx]), ordered=True)\n",
    "d['edu_cat'] = pd.factorize(cat, sort=True)[0].astype('int')\n",
    "d['age_std'] = (d.age - d.age.mean())/d.age.std()\n",
    "d.tail(3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "executionInfo": {
     "elapsed": 1269,
     "status": "ok",
     "timestamp": 1607276543509,
     "user": {
      "displayName": "Andrés Suárez",
      "photoUrl": "https://lh5.googleusercontent.com/-s0kzcIwylzA/AAAAAAAAAAI/AAAAAAAAQXA/v8Sc6WgQy7c/s64/photo.jpg",
      "userId": "06409440331868776168"
     },
     "user_tz": -60
    },
    "id": "T1PU90J0JFOH"
   },
   "outputs": [],
   "source": [
    "def model(action, contact, intention, age_std, edu_cat, male, response=None):\n",
    "    # priors\n",
    "    cutpoints = numpyro.sample('cutpoints', \n",
    "                               dist.TransformedDistribution( \n",
    "                                   dist.Normal(0, 1.5).expand([6]), \n",
    "                                   OrderedTransform()\n",
    "                                   ),\n",
    "                               )\n",
    "    norm = lambda label: numpyro.sample(label, dist.Normal(0,0.5))\n",
    "    bA = norm('bA')\n",
    "    bC = norm('bC')\n",
    "    bI = norm('bI')\n",
    "    bIA = norm('bIA')\n",
    "    bIC = norm('bIC')\n",
    "    bAge = norm('bAge')\n",
    "    bEdu = norm('bEdu')\n",
    "    edu_coefs = numpyro.sample('edu_coefs', dist.Dirichlet(jnp.repeat(2,7)))\n",
    "    bMale = norm('bMale')\n",
    "    # likelihood\n",
    "    delta_j = jnp.pad(edu_coefs, (1, 0))\n",
    "    delta_E = jnp.sum(jnp.where(jnp.arange(8) <= edu_cat[..., None], delta_j, 0), -1)\n",
    "    BI = bI + bIA*action + bIC*contact\n",
    "    phi = bA*action + BI*intention + bC*contact + bAge*age_std + bEdu*delta_E + bMale*male\n",
    "    response_hat = numpyro.sample('response_hat', dist.OrderedLogistic(phi, cutpoints), obs=response-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "sample: 100%|██████████| 1000/1000 [01:41<00:00,  9.86it/s, 63 steps of size 8.29e-02. acc. prob=0.96]\n",
      "sample: 100%|██████████| 1000/1000 [01:11<00:00, 14.00it/s, 31 steps of size 8.61e-02. acc. prob=0.95]\n",
      "sample: 100%|██████████| 1000/1000 [01:08<00:00, 14.54it/s, 31 steps of size 8.76e-02. acc. prob=0.95]\n",
      "sample: 100%|██████████| 1000/1000 [01:18<00:00, 12.67it/s, 63 steps of size 8.40e-02. acc. prob=0.95]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "                  mean       std    median      5.0%     95.0%     n_eff     r_hat\n",
      "          bA     -0.48      0.05     -0.48     -0.57     -0.39   1275.96      1.00\n",
      "        bAge     -0.07      0.02     -0.07     -0.11     -0.03    934.27      1.00\n",
      "          bC     -0.35      0.07     -0.35     -0.47     -0.25   1279.55      1.00\n",
      "        bEdu     -0.00      0.18      0.01     -0.28      0.27    469.78      1.01\n",
      "          bI     -0.29      0.06     -0.29     -0.39     -0.20   1195.19      1.00\n",
      "         bIA     -0.44      0.08     -0.44     -0.58     -0.31   1220.66      1.00\n",
      "         bIC     -1.26      0.10     -1.26     -1.42     -1.10   1207.35      1.00\n",
      "       bMale      0.57      0.04      0.57      0.51      0.63   1696.12      1.00\n",
      "cutpoints[0]     -2.37      0.14     -2.36     -2.58     -2.15    523.24      1.01\n",
      "cutpoints[1]     -1.67      0.14     -1.66     -1.91     -1.47    508.45      1.01\n",
      "cutpoints[2]     -1.07      0.14     -1.06     -1.29     -0.86    509.96      1.01\n",
      "cutpoints[3]     -0.01      0.14     -0.00     -0.22      0.21    514.76      1.01\n",
      "cutpoints[4]      0.68      0.14      0.69      0.47      0.90    513.91      1.01\n",
      "cutpoints[5]      1.60      0.14      1.61      1.37      1.81    510.99      1.01\n",
      "edu_coefs[0]      0.16      0.10      0.14      0.01      0.30   1522.06      1.00\n",
      "edu_coefs[1]      0.14      0.09      0.12      0.01      0.26   2125.46      1.00\n",
      "edu_coefs[2]      0.14      0.09      0.12      0.01      0.26   1712.43      1.00\n",
      "edu_coefs[3]      0.14      0.10      0.11      0.00      0.29    996.18      1.00\n",
      "edu_coefs[4]      0.19      0.15      0.13      0.00      0.42    504.14      1.01\n",
      "edu_coefs[5]      0.12      0.08      0.10      0.01      0.24   1688.08      1.00\n",
      "edu_coefs[6]      0.13      0.08      0.11      0.01      0.25   1877.20      1.00\n",
      "\n",
      "Number of divergences: 0\n"
     ]
    }
   ],
   "source": [
    "kernel = NUTS(model)\n",
    "mcmc = MCMC(kernel, num_warmup=500, num_samples=500, num_chains=4, chain_method='sequential')\n",
    "dat = {k:v.to_numpy() for k,v in d[['action', 'contact', 'intention', 'age_std', 'edu_cat', 'male', 'response']].items()}\n",
    "mcmc.run(rng_key, **dat)\n",
    "mcmc.print_summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "i1XxGWaDlmVI"
   },
   "source": [
    "Age is still negative (and weak), while education is right near zero and straddles both sides. Gender seems to have accounted for all of the previous influenced assigned to education. It looks like female respondents gave lower average responses, indicating less approval.\n",
    "\n",
    "It would be worth figuring out how gender is associated with education in this sample. It could be true for example that some education levels under-sampled men or women, and this leads to another kind of confound. Consider for example if older men are less likely to respond, so the sample becomes increasingly female with age. Then education level will also be increasingly female with age. Since the sample is\n",
    "not a representative sample of the population, there are probably some biases of this sort."
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "rethink_stats_numpyro_w07.ipynb",
   "provenance": []
  },
  "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.9.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}