statrethink_numpyro_w09.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Rethinking Statistics course in NumPyro - Week 9"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Lecture 17: Adventures in Covariance\n",
    "\n",
    "- [Video](https://www.youtube.com/watch?v=yfXpjmWgyXU)\n",
    "- [Slides](https://speakerdeck.com/rmcelreath/l17-statistical-rethinking-winter-2019)\n",
    "\n",
    "Lecture 18: Slopes, Instruments and Social Relations\n",
    "\n",
    "- [Video](https://www.youtube.com/watch?v=e5cgiAGBKzI)\n",
    "- [Slides](https://speakerdeck.com/rmcelreath/l18-statistical-rethinking-winter-2019)\n",
    "\n",
    "[Proposed problems](https://github.com/gbosquechacon/statrethinking_winter2019/blob/master/homework/week09.pdf) and [solutions in R](https://github.com/gbosquechacon/statrethinking_winter2019/blob/master/homework/week09_solutions.pdf) for the exercises of the week."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "\n",
    "import jax.numpy as jnp\n",
    "from jax import random\n",
    "from jax.scipy.special import expit as logistic\n",
    "import numpyro\n",
    "from numpyro.infer import NUTS, MCMC, Predictive\n",
    "import numpyro.distributions as dist\n",
    "\n",
    "import seaborn as sns\n",
    "import arviz as az"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "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",
      "pandas : 1.2.3\n",
      "seaborn: 0.11.1\n",
      "numpyro: 0.6.0\n",
      "arviz  : 0.11.2\n",
      "numpy  : 1.20.1\n",
      "jax    : 0.2.10\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": {},
   "outputs": [],
   "source": [
    "rng_key = random.PRNGKey(0)\n",
    "sns.set_style('whitegrid')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> Revisit the Bangladesh fertility data, data(`bangladesh`). Fit a model with both varying intercepts by `district_id` and varying slopes of `urban` (as a 0/1 indicator variable) by `district_id`. You are still predicting `use.contraception`. Inspect the correlation between the intercepts and slopes. Can you interpret this correlation, in terms of what it tells you about the pattern of contraceptive use in the sample? It might help to plot the varying effect estimates for both the intercepts and slopes by district. Then you can visualize the correlation and maybe more easily think through what it means to have a particular correlation. Plotting predicted proportion of women using contraception, in each district, with urban women on one axis and rural on the other, might also help."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "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>woman</th>\n",
       "      <th>district</th>\n",
       "      <th>contra</th>\n",
       "      <th>living.children</th>\n",
       "      <th>age.centered</th>\n",
       "      <th>urban</th>\n",
       "      <th>did</th>\n",
       "      <th>children</th>\n",
       "      <th>age</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>1931</th>\n",
       "      <td>1932</td>\n",
       "      <td>61</td>\n",
       "      <td>0</td>\n",
       "      <td>4</td>\n",
       "      <td>14.44</td>\n",
       "      <td>0</td>\n",
       "      <td>59</td>\n",
       "      <td>1.084782</td>\n",
       "      <td>1.601843</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1932</th>\n",
       "      <td>1933</td>\n",
       "      <td>61</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>-13.56</td>\n",
       "      <td>0</td>\n",
       "      <td>59</td>\n",
       "      <td>-1.329451</td>\n",
       "      <td>-1.504697</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1933</th>\n",
       "      <td>1934</td>\n",
       "      <td>61</td>\n",
       "      <td>0</td>\n",
       "      <td>4</td>\n",
       "      <td>10.44</td>\n",
       "      <td>0</td>\n",
       "      <td>59</td>\n",
       "      <td>1.084782</td>\n",
       "      <td>1.158052</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "      woman  district  contra  living.children  age.centered  urban  did  \\\n",
       "1931   1932        61       0                4         14.44      0   59   \n",
       "1932   1933        61       0                1        -13.56      0   59   \n",
       "1933   1934        61       0                4         10.44      0   59   \n",
       "\n",
       "      children       age  \n",
       "1931  1.084782  1.601843  \n",
       "1932 -1.329451 -1.504697  \n",
       "1933  1.084782  1.158052  "
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "d = pd.read_csv('./dat/bangladesh.csv', header=0, sep=';')\n",
    "d['did'] = pd.factorize(d['district'])[0]\n",
    "normalize = lambda x: (x - x.mean())/x.std()\n",
    "d['children'] = normalize(d['living.children'])\n",
    "d['age'] = normalize(d['age.centered'])\n",
    "d.rename(columns={'use.contraception':'contra'}, inplace=True)\n",
    "d.tail(3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is a conventional varying slopes model, with a centered parameterization. No surprises. If you peek at the posterior distributions for the average effects, you'll see that the average slope is positive:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# https://fehiepsi.github.io/rethinking-numpyro/14-adventures-in-covariance.html\n",
    "\n",
    "def model(did, urban, contra=None):\n",
    "    # fixed priors\n",
    "    a_bar = numpyro.sample('a_bar', dist.Normal(0, 1))\n",
    "    b_bar = numpyro.sample('b_bar', dist.Normal(0, 0.5))\n",
    "    sigma_ab = numpyro.sample('sigma_ab', dist.Exponential(1), sample_shape=(2,))\n",
    "    rho = numpyro.sample('rho', dist.LKJ(2,2))\n",
    "    # adaptive priors\n",
    "    mu = jnp.stack([a_bar,b_bar])\n",
    "    cov = jnp.outer(sigma_ab, sigma_ab)*rho\n",
    "    i, j = len(d.did.unique()), len(d.urban.unique())\n",
    "    ab_did = numpyro.sample('ab_did', dist.MultivariateNormal(mu, cov), sample_shape=(i,j))\n",
    "    # likelihood\n",
    "    a_did, b_did = ab_did[:,:,0], ab_did[:,:,1]\n",
    "    theta = a_did[did,urban] + b_did[did,urban]*urban\n",
    "    contra_hat = numpyro.sample('contra_hat', dist.Bernoulli(logits=theta), obs=contra)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "sample: 100%|██████████| 1000/1000 [00:17<00:00, 57.73it/s, 63 steps of size 1.22e-01. acc. prob=0.80] \n",
      "sample: 100%|██████████| 1000/1000 [00:05<00:00, 167.03it/s, 31 steps of size 1.36e-01. acc. prob=0.82]\n",
      "sample: 100%|██████████| 1000/1000 [00:06<00:00, 153.30it/s, 63 steps of size 6.78e-02. acc. prob=0.79]\n",
      "sample: 100%|██████████| 1000/1000 [00:10<00:00, 93.79it/s, 127 steps of size 4.00e-02. acc. prob=0.79]\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[['did','urban','contra']].items()}\n",
    "mcmc.run(rng_key, **dat)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "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>mean</th>\n",
       "      <th>sd</th>\n",
       "      <th>hdi_5.5%</th>\n",
       "      <th>hdi_94.5%</th>\n",
       "      <th>mcse_mean</th>\n",
       "      <th>mcse_sd</th>\n",
       "      <th>ess_bulk</th>\n",
       "      <th>ess_tail</th>\n",
       "      <th>r_hat</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>a_bar</th>\n",
       "      <td>-0.68</td>\n",
       "      <td>0.10</td>\n",
       "      <td>-0.83</td>\n",
       "      <td>-0.52</td>\n",
       "      <td>0.01</td>\n",
       "      <td>0.00</td>\n",
       "      <td>318.0</td>\n",
       "      <td>717.0</td>\n",
       "      <td>1.02</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>b_bar</th>\n",
       "      <td>0.63</td>\n",
       "      <td>0.17</td>\n",
       "      <td>0.32</td>\n",
       "      <td>0.89</td>\n",
       "      <td>0.02</td>\n",
       "      <td>0.01</td>\n",
       "      <td>108.0</td>\n",
       "      <td>100.0</td>\n",
       "      <td>1.05</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>sigma_ab[0]</th>\n",
       "      <td>0.58</td>\n",
       "      <td>0.08</td>\n",
       "      <td>0.45</td>\n",
       "      <td>0.72</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>394.0</td>\n",
       "      <td>671.0</td>\n",
       "      <td>1.01</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>sigma_ab[1]</th>\n",
       "      <td>0.27</td>\n",
       "      <td>0.20</td>\n",
       "      <td>0.04</td>\n",
       "      <td>0.52</td>\n",
       "      <td>0.06</td>\n",
       "      <td>0.05</td>\n",
       "      <td>7.0</td>\n",
       "      <td>22.0</td>\n",
       "      <td>1.51</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "             mean    sd  hdi_5.5%  hdi_94.5%  mcse_mean  mcse_sd  ess_bulk  \\\n",
       "a_bar       -0.68  0.10     -0.83      -0.52       0.01     0.00     318.0   \n",
       "b_bar        0.63  0.17      0.32       0.89       0.02     0.01     108.0   \n",
       "sigma_ab[0]  0.58  0.08      0.45       0.72       0.00     0.00     394.0   \n",
       "sigma_ab[1]  0.27  0.20      0.04       0.52       0.06     0.05       7.0   \n",
       "\n",
       "             ess_tail  r_hat  \n",
       "a_bar           717.0   1.02  \n",
       "b_bar           100.0   1.05  \n",
       "sigma_ab[0]     671.0   1.01  \n",
       "sigma_ab[1]      22.0   1.51  "
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "samples_11 = az.from_numpyro(mcmc)\n",
    "az.summary(samples_11, var_names=['a_bar','b_bar', 'sigma_ab'], hdi_prob=.89).round(2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This implies that urban areas use contraception more. Not surprising. The correlation between the intercepts and slopes is quite negative. Let's plot the individual effects to appreciate this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "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>a</th>\n",
       "      <th>b</th>\n",
       "      <th>u0</th>\n",
       "      <th>u1</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>57</th>\n",
       "      <td>-0.917290</td>\n",
       "      <td>0.661106</td>\n",
       "      <td>0.285510</td>\n",
       "      <td>0.436302</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>58</th>\n",
       "      <td>-1.079067</td>\n",
       "      <td>0.634777</td>\n",
       "      <td>0.253683</td>\n",
       "      <td>0.390719</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>59</th>\n",
       "      <td>-1.093814</td>\n",
       "      <td>0.618861</td>\n",
       "      <td>0.250901</td>\n",
       "      <td>0.383445</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "           a         b        u0        u1\n",
       "57 -0.917290  0.661106  0.285510  0.436302\n",
       "58 -1.079067  0.634777  0.253683  0.390719\n",
       "59 -1.093814  0.618861  0.250901  0.383445"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# mean chain value\n",
    "ab = np.concatenate(samples_11.posterior.ab_did.values, axis=0)\n",
    "ab = ab.mean(axis=0)\n",
    "# aggregate rural and urban results\n",
    "a, b = ab[:,:,0], ab[:,:,1]\n",
    "a, b = a.mean(axis=1), b.mean(axis=1)\n",
    "# dataframe\n",
    "aux = pd.DataFrame({'a':a, 'b':b})\n",
    "aux['u0'] = logistic(aux.a)\n",
    "aux['u1'] = logistic(aux.a+aux.b)\n",
    "aux.tail(3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There's the negative correlation: districts with higher use outside urban areas (a values) have smaller slopes. Since the slope is the difference between urban and non-urban areas, you can see this as saying that districts with high use in rural areas have urban areas that aren't as different. On the outcome scale, what this ends up meaning is that urban places are much the same in all districts, but rural areas vary a lot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAX8AAAEECAYAAADAoTRlAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjMuNCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8QVMy6AAAACXBIWXMAAAsTAAALEwEAmpwYAAAoLklEQVR4nO3df3xT1f0/8FeTNKFtilB+iF+kCJ2FgnZpdX4UP1W3UgWUh/KrLR11uA/TwVAR3GDogEftJxYnjqHQ4dRWccNWVL4D0bkqUlZRZiHDYiiO0loL1EJLaRqa0N77/YNvA6FpmyY5yU3zev4DyUnOfTc5ed9zzzn33jBZlmUQEVFIUQU6ACIi8j8mfyKiEMTkT0QUgpj8iYhCEJM/EVEI0gQ6AHeYTCbodLpAh+HEZrMpLiZ3MfbAYOyBE6zxV1ZWQpIkJCQkePR+m80Gg8Hgsiwokr9Op/P4jxfFbDYrLiZ3MfbAYOyBE6zxL1y4EFarFfv37/fo/WazudsyDvsQEYWgoOj5ExGFoqeffho1NTVC6mbyJyJSqMmTJ/c4dOMNDvsQESmUyWQSlvzZ8yciUqglS5bAarVi5syZPq+byZ+IAACSJKP6TCvqz7Xh6oEDcN2QKKhUYYEOiwRh8iciSJKMDw+fwtJiE9ouSBgQrsIL6QZMmTiCO4B+imP+RITqM62OxA8AbRckLC02ofpMa4AjI1GE9PwlScKaNWtQWVkJrVaL3NxcjB49GsDFkw6MRqPjtSaTCRs3bsQdd9whIhS38ZCXQln9uTZH4u/UdkHC9y1tGDtMH6CoSCQhyb+kpAR2ux1FRUUwmUzIy8tDfn4+ACAhIQFbtmwBAHzwwQcYPny4IhI/D3kplF09cAAGhKucdgADwlUYHj0ggFGR0WhEdXW1kLqFDPuUl5cjJSUFAGAwGFBRUdHlNVarFS+++CKeeuopESH0CQ95KdRdNyQKL6QbMCD8Ykro7ABdNyQqwJGFtkmTJiEpKUlI3UJ6/haLBXr9pUNFtVqN9vZ2aDSXNrdt2zZMmTIFMTExvdZns9mErXUFgBMXolwe8h4/1Qjb6VqX72lraxMak0iMPTCUHnucTo2//iwRTec7MDhCjSipBZWVZwEoP/beBGv8Bw8ehN1uF1K3kOSv1+vR2nqp1yxJklPiB4AdO3Zgw4YNbtUn+sJuugaLy0PeMSNiMHZYrMv3BOuFogDGHiiMPXCCNf7OC7vNnz/fo/f7/cJuycnJKC0tBXBxQjc+Pt6pvKWlBXa7Hddcc42IzfcZD3k9J0kyqhos2HfsNKoaLJAkOdAhuS2YYyfylpCef1paGsrKypCZmQlZlmE0GlFQUIDY2Fikpqbi+PHjGDlypIhNe0SlCsOUiSMw/rEUfN/ShuHRXO3jjmCeKA/m2Il8QUjyV6lUyMnJcXouLi7O8f/ExERs2rRJxKY9plKFYewwfUgsa1Op1ahqsHi9rLW7ifLxj6Uo/nPsLvYJj6dAkuH2Z8MlwhSseIZviJEkGeYWHVZs2et1jzeY14a7in1wpBYHvj2Lle995dZnE2xHD9xR0eV4hm+IqT7TihXbv/bJstbOteGX621tuFLG2V3FPufmax2JH+j9swmmJcKdO6ppG/Zi7p+/wLQNe/Hh4VMB+fyV0gaCwfr167FixQohdTP5K5yvfyg99db7qq8T5UpKQK5ijx8e3afPxpefpWgidlSetE0ltYFgYDAYhK1S4rCPgokYVhge7fpMzmF6z87kHHd1NDb9NBlROg2ujtYhNqb7oYRAzBF0N9ThapJfltGns1yD6axYXw/Redo2g3meKBBKSkpQU1MjZAfAnr+CieitqVXA46nXO/V4H0+9Huo+toTOH/+9L+7Fzwu/xM9e24+vT7b0+B5/95R762V2TvLfOnYoxg7TY8zQvh3JBNMSYU+G6HriadsMpqMlJcjNzcXmzZuF1M2ev4KJmFA92dyGN/bV4H/+eyzCwgBZBt7YV4Ok2EG4bqj7dbrTg7uy133NVZd6ytdcNQAzk6+FWgVEhGsgSbLPJx/72svs65Lf3l5/5d+vUqu9+nu8mbDt3FFd2VP3dEfladsMpqOl/o7JX8FE/FCuHjgATVY7Nu7+j1d19vbjlyQZn1TW49B3zZBkQB0GxF8djfUZBjz7gRkZN8diwyffoO2ChJdLq4SskvEkQfV1yW93r3c1LJL3wARc7+FOztshQF+fy+Jp2/T1Tog812+HffrDigIRwwrXDYlC3gMTvK6zt2GEbxtb8U29BS+XVuGlT/6DzaVV+LbRiu/PtSFvZqIj8QPeD2dd/l3/u/Ys/lV9BlUNFseRRncxiuTqqGPF9q+7/Ru7a6+dz/+rutHrIcArh7m82dF62jY7d0K7HkvBWw//F3Y9lqLYpbH9Xb/s+Qfb+uvuiDjzWKUKQ0K0Dbu8rLO3Hlz9ORv++LFzgv/jx9/g97N/iH1VZ3w2nOXqu37sJ9ej6MtvsXxKAl7KSsLivx70ey+zL0cd3bXXuxOuxkfmeiwtNmFBylhFnVPhTdsMpRMqlaxfJv/+tKJAxA9F6ujwus7efvyt9naXycpqa4fUx1U1PXH1XW/45Bv8z3+PxdJiE95/NMXrHZ0n+jIs0l17LXr4VqfnlTZWziQu3ubNm3Hs2DEhdffLYZ/+tKJAycNXPQ0jjI6Jcjnk0mCxYce/67B2VqJPhrO6+67Dwpy/c9nLj62v34OrYZG8Bya4/Bu7+xtONl96/p3y7/DYT6736RAgKd+4ceMwZswYIXX3y55/f1lREMzDV53LJi+P/YnJ8fjr/hosn5KAuxOuxo0jr/K6R97dd925Zt/WLmHaBu8uZeHJ9+DqyOhC0wmXr+/ub7h8ddTJ5jZs+bwGD98xFkmjBmH0kCjFX56Bl5Pw3o4dO1BbW8t1/u4KpvXXPfF0LXVfe6kiji4un9jb+ov/QtHDt+Km0YNQMP8WTJk4AhqNyieTj66+68d+cj12HqrD/864Eav/VuH1xLKr72Hth2Z8VXe2x8/syiMjqaPD7b/hhXQDJl5zldPzTVY7xo8YiDvjh3s9YSsaz+T1jXXr1qGwsFBI3f2y599fLtHsyVLFvvZSRR5d+GNM+PLvuv5cGyK1alzokHD7D4bgi6ozqDlz3un1nkySXvk9XHPVAGTcHIuMlz93fGbGGTciOXaQyzOcO3vAJy5EQddg6dIWe2qvwdaOO//W6jOtqDx1DoMjtY7hq0CdzU2u9cvkDwT3ZFRnIz5/oQOPp/4AxV9+h5PNF8euexu+6utkd3+YHHf1Xe87dhrnbB0+Gf67clhmZvK1XZaqrnzvKzx8x1iMHzHQacfp7s61u/Yquh37MmF2t/Jqy+c1jh2AqNVJwTxEGij9ctgnmF1+uPzzwi+xubQKD9422jH+29vwVV8nu/vT5Pjlrh44AJ8fa8AL6QY8lvoDLP7JDzB6SATWzkrs8/DflcMyahVcfmaSjC7DSlUNFsVe+dPXQzPdrbyamXwtALHzbsF0hVWl6Lc9/2DlqhH/8eNv8PpDt2BYtK7XnllfJ7v7y+T4lWIHRyLzltFOPcH/nXEjpk5w7pW70+u9cvglIlyDl0urXE4yX967bW+XcPjkOUWtz7+cr4/6elp55Yt5t56+r2C+t0SgsOevMN01YhmyW5N8fZ3s7i+T41f6tsmKp7c7T/Y+9d5X+K754hxAX3u9l0/e3jjyKpeTzO8e+M5px3n4ZDOqGiwBO8u4N74+6uvurO+UHwz1+kze3r4vX1+4Tim2bNmCvLw8IXWz568w3vbEfX1xsmDT2Ts8Wt/SY0/Qm15v52c27tEUmE+dw9H6Fmz5vAZNVrvTjvNkcxuKv7y4Pr9zjmBAuAq5D9zg0c7V1xOavjrquzyuP2ffjKf/71eoOXPe0ZH40XUxXren3r6v/nrNoFGjRsFisQipm8lfYXzRiH11cbJg09k7XPuhGU/ePb7HxObtMIFKFYa44RcvAz3hmoGYFDeky47zmqsi0GS1Y8vnl66iqgoDxo+I7nMy9HZC09WOwxdtzVVca2clYuSgAYiJ6n2Y0l29fV/9rRPTqaioCHV1dULW+QtJ/pIkYc2aNaisrIRWq0Vubi5Gjx7tKN+zZw82btwIWZYxceJErF69GmFhwf0l+Up/bcT+UH2mFWs/vHjF0Oc/OtKlx315YvNVr7enHefEawYi94Eb8PT2Cmzc/R9Hr3/81QM9+ts8PVK5PEEPjtRizs3XIn54NBKuGYi7E6726vIXruJa/s4h7PLxajF3vq/+0om5XH5+PqxWK5YuXerzuoUk/5KSEtjtdhQVFcFkMiEvLw/5+fkAAIvFgt///vd44403EBMTgz//+c9oampCTEyMiFCCUn9sxP5Qf64N9yWOdCT8zh63WgWkjh+OG0cOciQ2b3u97gzBaDQqPPDDkbh+uB6nmtsQE6lGUuxQaDR9n2rz5kilM0EPjtQi+9bRXXaIUyaO8Lit+Wuitb8O6wSSkORfXl6OlJQUABfvQVlRUeEoO3jwIOLj47F27VrU1tZizpw5TPzkE1cPHOC0DPNkc5vjvgWT4oa4fWJVb/oyBKPRqPDDUYPxw1GA2WyGRjPco7+tu9tvunOk0pmgXZ2f4O05Hf5aLcYjYt8TkvwtFgv0+kuNSa1Wo729HRqNBk1NTfjiiy+wfft2REZG4qc//SkMBkOPFy+y2Wwwm80iQvVYW1ub4mJyV3+NXaVW46ZRg1wmI63UfRu6CoDNBlSedi+G8+EDsbT4UJckOuJniYi4cM6j2HuiUqtRdT4Cj6de77hM9oBwFfJmTMCFphMwn3Z92YhOA8IHYkC4ynGxO6eYLkg4fqoRttO1PdbRXewqtRp5D0zAiu1fX4rrAffi8lRfvy8geNu81WqFJElCYheS/PV6PVpbL51cIUkSNJqLmxo0aBBuvPFGDBs2DABw8803w2w295j8dTqdsDvYe8psNisuJncFOnZvVq30Fntcu4S1sxKx/J1DTr1yw9gRUKlG+iT+fcdOu0yidpUOyT3E5unnXtVgwRNb9mJwpNZp4tgQO8StW29KkowX0lWoPHXO5Y5xzIgYjB0W22MdPcV+vSQjMXaIonvkgW7znoqMjITVavU49p52GkKSf3JyMnbv3o1p06bBZDIhPj7eUTZx4kQcPXoUjY2NGDhwIP79738jPT1dRBikQKJPw9doVJie+H98csXQ7vj7xLjOYZvLh7GAi0NZ7iT/ziGTCddEY/SQKKx87yufjptzjkqcbdu24ejRo0LqFpL809LSUFZWhszMTMiyDKPRiIKCAsTGxiI1NRXLli3DggULAABTpkxx2jlQ/+aPawmJTkb+nnz0xc5GpQrDdUP1iI2JgmHUIEX30umSoUOHoqGhQUjdQpK/SqVCTk6O03NxcXGO/99777249957RWyaFK4/nIbv78lHX+5s2EsPLoWFhThx4kTwrPMn6k5/uZaQP5MoV7qErsLCQlitVqxcudLndfPaPuRX/fVaQqL1dMtMIk+w509+Feq9WN5whJSCyZ/8LlTHnVVqtVsrnbiDIH/gsA+Rn7Sqonq94QjvfUv+wuRP5CdN1o5er5/v7R2pJElGVYOlxxvLU/DYtWsX/vSnPwmpm8M+RH4SE6nudaWTN0theR/b/icyMhIRERFC6mbPn8iFvvag3Xl9pNTa60onb+5IxfvY9j+bNm3C1q1bhdTNnj/RFfrag3b39VJHR68rnbw5oas/nEBHzoqLi2G1WrucNOsLTP5EV+jrJSj68vreVjp5sxS2v5xAR/7BYR+iK/T1xua+vhG6pyd08QQ66gv2/Imu0NcetFJ63KF+Al1PVGo1qhosPHfiMuz5E12hrz1oJfW4eRmIriRJhrlFx3MnrsCeP9EV+tqDZo9b2arPtDruNAaIuYy4KJ9++qmwO5Ax+RO50NdLUITqJSuCAVdBucZhHyLq17w5dyLQnn/+ebz22mtC6mbyJ6J+7bohUch7YIIi5mT6aufOndizZ4+QujnsQ0T9mkoVhoRoG3ZxTsYJkz8R9XtSRwfnZK7AYR8iohAkpOcvSRLWrFmDyspKaLVa5ObmYvTo0Y7y3NxcHDhwAFFRF8fcNm3ahOjoaBGhEBEFrYiICHR0dAipW0jyLykpgd1uR1FREUwmE/Ly8pCfn+8oP3z4MF555RXExMSI2DwRUb/wwQcfCFvnL2TYp7y8HCkpKQAAg8GAiooKR5kkSaipqcGqVauQmZmJbdu2iQiBiIh6IKTnb7FYoNdfmlhRq9Vob2+HRqOB1WrFvHnz8NBDD6GjowMPPvggbrjhBowfP77b+mw2m7C9n6fa2toUF5O7GHtgMPbACdb48/Pz0d7ejkcffdTndQtJ/nq9Hq2tl9+XVIJGc3FTERERePDBBx13p7n11ltx5MiRHpO/TqdDQkKCiFA9ZjabFReTuxh7YDD2wAnW+A8dOgSr1epx7D3t8IQM+yQnJ6O0tBQAYDKZEB8f7yirrq7G3Llz0dHRgQsXLuDAgQOYOHGiiDCIiKgbQnr+aWlpKCsrQ2ZmJmRZhtFoREFBAWJjY5Gamor7778f6enpCA8Px/3334/rr79eRBhERNQNIclfpVJ1ue1YXFyc4/8LFizAggULRGyaiIjcwDN8iYgUasiQIdBqtULqZvInIlKod955J7jW+RMRkbIx+RMRKdRvf/tbvPDCC0Lq5rAPEZFC7du3D1arVUjd7PkTEYUgJn8iohDE5E9EFII45k9EpFDXXnstmpubhdTN5E9EpFBvvvkm1/kTEZHvsOdPRKRQS5YsQWNjI9544w2f183kT0SkUCaTiev8iYjId5j8iYhCEJM/EVEI4pg/EZFCxcfH4+zZs0LqZvInIlKol19+mev8iYjId9jzJyJSqIcffhhnz55FcXGxz+sW0vOXJAmrVq1CRkYGsrOzUVNT4/I1CxYswNatW0WEQEQU9I4ePYrq6mohdQtJ/iUlJbDb7SgqKsKyZcuQl5fX5TXr16/HuXPnRGyeiIh6IWTYp7y8HCkpKQAAg8GAiooKp/IPP/wQYWFhjtf0xmazCZv08FRbW5viYnIXYw8Mxh44wRq/1WqFJElCYheS/C0WC/R6veOxWq1Ge3s7NBoNjh49ip07d2LDhg3YuHGjW/XpdDokJCSICNVjZrNZcTG5i7EHBmMPnGCNPzIyElar1ePYe9ppCEn+er0era2tjseSJEGjubip7du3o76+Hj/72c9QV1eH8PBwjBw5EnfccYeIUIiIgpbBYEBjY6OQuoUk/+TkZOzevRvTpk2DyWRCfHy8o+w3v/mN4/8vvvgihg4dysRPROTC+vXrhQ1XCUn+aWlpKCsrQ2ZmJmRZhtFoREFBAWJjY5Gamipik0RE1AdCkr9KpUJOTo7Tc3FxcV1e9+ijj4rYPBFRvzBv3jw0Nzdjx44dPq+bJ3kRESnUd999x+v5ExGR7zD5ExGFICZ/IqIQxDF/IiKFuu2223D69GkhdTP5ExEp1LPPPquM6/k3NTUJCYKIiPzLreS/Z88epKam4qGHHsI999yDL774QnRcREQhb9asWXj88ceF1O3WsM9LL72Et99+GzExMWhoaMCvfvUrITcXICKiS86cORPYdf5RUVGIiYkBAAwbNgwRERFCgiEiIv/osef/wgsvAAA6OjrwyCOP4KabbsKhQ4eg1Wr9EhwREYnRY/IfM2aM078AeGE2IqJ+oMfkP2PGDH/FQUREV0hNTUVDQ4OQurnOn4hIoX73u98pY50/ERH1D+z5ExEp1NSpU2GxWLB3716f182ePxGRQp0/fx42m01I3Uz+REQhiMmfiCgEMfkTEYUgIRO+kiRhzZo1qKyshFarRW5uLkaPHu0o/8tf/oJ3330XYWFh+PnPf45p06aJCIOIKKjdd999qK+vF1K3kORfUlICu92OoqIimEwm5OXlIT8/HwDQ2NiIrVu34r333oPNZsO9996LqVOnIiwsTEQoRERB68knnxS2zl9I8i8vL0dKSgoAwGAwoKKiwlEWExOD7du3Q6PRoK6uDjqdrtfEb7PZhH0Anmpra1NcTO5i7IHB2AMnmOMXFbuQ5G+xWKDX6x2P1Wo12tvbodFc3JxGo8Gbb76JF198EdnZ2b3Wp9PpkJCQICJUj5nNZsXF5C7GHhiMPXCCNf677roLVqsV+/fv9+j9Pe00hEz46vV6tLa2Oh5LkuRI/J3mzZuHvXv34l//+hc+//xzEWEQEVE3hCT/5ORklJaWAgBMJhPi4+MdZVVVVVi8eDFkWUZ4eDi0Wi1UKi46IiLyJyHDPmlpaSgrK0NmZiZkWYbRaERBQQFiY2ORmpqK8ePHIyMjA2FhYUhJScEtt9wiIgwiIuqGkOSvUqmQk5Pj9FxcXJzj/4sXL8bixYtFbJqIiNzAC7sRESlUeno6Tp06JaRuJn8iIoVatGgRr+dPRBRqrFYrzp8/L6Ru9vyJiBRq2rRpXq3z7wl7/kREIYjJn4goBDH5ExGFICZ/IqIQxAlfIiKFmj9/Pk6cOCGkbiZ/IiKFmj9/Ptf5ExGFmtOnT6OpqUlI3ez5ExEp1OzZs7nOn4iIfIfJn4goBDH5ExGFICZ/IqIQxAlfIiKFWrhwIerq6oTUzeRPRKRQGRkZXOdPRBRqamtrcfLkSSF1s+dPRKRQ2dnZwtb5C0n+kiRhzZo1qKyshFarRW5uLkaPHu0oLywsxPvvvw8AuPPOO3kzdyIiPxMy7FNSUgK73Y6ioiIsW7YMeXl5jrLa2lr87W9/w1tvvYXi4mL885//xJEjR0SEQURE3RDS8y8vL0dKSgoAwGAwoKKiwlE2YsQIvPLKK1Cr1QCA9vZ26HS6Huuz2WzCJj081dbWpriY3MXYA4OxB06wxm+1WiFJkpDYhSR/i8UCvV7veKxWq9He3g6NRoPw8HDExMRAlmU899xzmDBhAsaMGdNjfTqdDgkJCSJC9ZjZbFZcTO5i7IHB2AMnWOOPjIyE1Wr1OPaedhpCkr9er0dra6vjsSRJ0Ggubcpms2HlypWIiorC6tWrRYRARBT0li1bhtraWiF1CxnzT05ORmlpKQDAZDIhPj7eUSbLMhYtWoRx48YhJyfHMfxDRETOpk+fjh//+MdC6hbS809LS0NZWRkyMzMhyzKMRiMKCgoQGxsLSZKwf/9+2O127N27FwCwdOlSJCUliQiFiChoVVZW4vjx40KGrIQkf5VKhZycHKfn4uLiHP//6quvRGyWiKhfeeSRR2C1WjFt2jSf180zfImIQhCTPxFRCGLyJyIKQUz+REQhiBd2IyJSqKeffho1NTVC6mbyJyJSqMmTJ/N6/kREocZkMglL/uz5ExEp1JIlS2C1WjFz5kyf182ePxFRCGLyJyIKQUz+REQhiMmfiCgEccKXiEihjEYjqqurhdTN5E9EpFCTJk3C4MGDhdTNYR8iIoX67LPPcPDgQSF1s+dPRKRQK1euhNVqRVZWls/rZs+fiCgEMfkTEYUgJn8iohAkJPlLkoRVq1YhIyMD2dnZLi9J2tjYiHvuuQc2m01ECERE1AMhE74lJSWw2+0oKiqCyWRCXl4e8vPzHeV79+7FunXr0NDQIGLzRET9wvr161FVVSWkbiE9//LycqSkpAAADAYDKioqnDeqUqGgoACDBg0SsXkion7BYDAgISFBSN1Cev4WiwV6vd7xWK1Wo729HRrNxc3dfvvtfarPZrMJu6a1p9ra2hQXk7sYe2Aw9sAJ1vg/++wzXLhwQUjdQpK/Xq9Ha2ur47EkSY7E7wmdTids7+cps9msuJjcxdgDg7EHTrDGv3DhQlitVvzyl7/06P097fCEDPskJyejtLQUwMU70cTHx4vYDBEReUhIzz8tLQ1lZWXIzMyELMswGo0oKChAbGwsUlNTRWySiIj6QEjyV6lUyMnJcXouLi6uy+s++eQTEZsnIqJe8CQvIqIQxAu7EREp1ObNm3Hs2DEhdTP5ExEp1Lhx4yBJkpC6OexDRKRQO3bswO7du4XUzeRPRKRQ69atQ2FhoZC6mfyJiEIQkz8RUQhi8iciCkFM/kREIYhLPYmIFGrLli345ptvhNTN5E9EpFCjRo2CxWIRUjeTPxGRQhUVFaGurk7I5ag55k9EpFD5+fl46623hNTN5E9EFIKY/ImIQhCTPxFRCGLyJyIKQVztQ0SkUNu2bcPRo0eF1M2ePxGRQg0dOhSDBw8WUjd7/kREClVYWIgTJ04Ezzp/SZKwatUqZGRkIDs7GzU1NU7lxcXFmDlzJtLT04XdqICIKNgVFhZi+/btQuoW0vMvKSmB3W5HUVERTCYT8vLykJ+fDwBoaGjAli1b8M4778BmsyErKwu33347tFqtiFCIiMgFIcm/vLwcKSkpAACDwYCKigpH2aFDh5CUlAStVgutVovY2FgcOXIEiYmJ3dZXWVmJhQsXOj2Xnp6ORYsWwWq1Ytq0aV3eM3/+fMyfPx+nT5/G7Nmzu5QvXLgQGRkZqK2tRXZ2dpfyZcuWYfr06aisrMQjjzzSpTw7OxsJCQkwmUxYsmRJl3Kj0YhJkybhs88+w8qVK7uUr1+/HgaDASUlJcjNze1SvnnzZowbNw47duzAunXrupRv2bIFo0aNQlFRkWPHerlt27Zh6NChKCws7HInIKvVik8//RSRkZHYtGkTiouLu7z/008/BQA8//zz2Llzp1NZREQEPvjgAwDAM888g48//tipfMiQIXjnnXcAAL/97W+xb98+p/Jrr70Wb775JgBgyZIlMJlMTuXx8fF4+eWXAQAPP/yw04SX1WrFpEmTsH79egDAvHnz8N133zm9/7bbbsOzzz4LAJg1axbOnDnjVJ6amorf/e53AICpU6fi/PnzTuX33XcfnnzySQDAXXfdhSt52vasVisiIyO9bntPP/00Jk+e7Ne21xk74F3bA4Bdu3b5ve11xu9N2wMu5jN/tj2TyQRJkhztsK9tz9X300lI8rdYLNDr9Y7HarUa7e3t0Gg0sFgsiI6OdpRFRUX1euEiSZJgtVqdnjt16hTMZjPOnz/fpQwATpw4AbPZjKamJpfldXV1MJvNOHnypMvy2tpamM1mHD9+3GX5hQsXYDabUVVV5bK8uroagwcPRnV1tcvyqqoq6HQ61NTUuCw/duwYJElCbW2ty/JvvvkGFosFdXV1LsuPHj2KhoYGnDhxoku5JEk4cuQIIiIicOrUKZfvN5vNAID6+vou5R0dHY7yhoaGLuVardZRfvr06S7lzc3NjvLGxsYu5WfPnnWUnz171qlckiQ0NjY6ypubm7u8//Tp047ylpaWLuUNDQ2OcovFApvN5lReX1/vKHf12Xja9jrbsbdtr6amxu9t7/LfoDdtD0BA2l5n/N60vc73+LPtSZIEWZYd9Xib95zIAhiNRvn99993PE5JSXH8v6SkRF69erXj8aJFi+RDhw71WN/XX3/t8xi9pcSY3MXYA4OxB06wxn/nnXfKP/rRjzx+f09/t5AJ3+TkZJSWlgIATCYT4uPjHWWJiYkoLy+HzWZDS0sLjh075lROREQX7dq1C3/605+E1C1k2CctLQ1lZWXIzMyELMswGo0oKChAbGwsUlNTkZ2djaysLMiyjCeeeAI6nU5EGEREQS0yMhIRERFC6haS/FUqFXJycpyei4uLc/w/PT0d6enpIjZNRNRvbNq0CadOneqST32BZ/gSESlUcXExPvzwQyF1M/kTEYUgJn8iohDE5E9EFIKY/ImIQlCYLMtyoIPojclk4nJQIqI+stlsMBgMLsuCIvkTEZFvcdiHiCgEMfkTEYUgJn8iohDE5E9EFIKY/ImIQhCTPxFRCGLy76N//OMfWLZsmcuywsJCzJkzB3PmzMFLL73k58h611PsxcXFmDlzJtLT07F7924/R9a9trY2PProo8jKysIvfvELNDY2dnnNs88+i9mzZyM9PR3l5eUBiNI1d2J/9913MWfOHMycORMbN24MQJSuuRM7AJw/fx7333+/4/4dSuBO7GvXrkVGRgZmzZrl8laS/iZJElatWoWMjAxkZ2ejpqbGqVzI79PjW8SEoGeeeUa+55575CVLlnQp+/bbb+UZM2bI7e3tsiRJckZGhmw2mwMQpWs9xf7999/L9913n2yz2eRz5845/q8Er732mrxhwwZZlmV5586d8jPPPONUbjab5Tlz5siSJMnHjx+XZ8yYEYgwXeot9pqaGnn27Nny+fPn5Y6ODvkPf/iDbLfbAxFqF73F3mnFihXy/fffL+/Zs8ef4fWot9j37dsnL1q0SJZlWbbZbPLkyZPls2fP+j3Oy/3973+Xly9fLsuyLB88eFD+5S9/6SgT9ftkz78PkpOTsWbNGpdlI0aMwCuvvAK1Wo2wsDC0t7cr6qzknmI/dOgQkpKSoNVqER0djdjYWBw5csS/AXajvLwcKSkpAIA77rijy83ghw8fjgEDBsBut8NisUCjEXKLCo/0Fvtnn32GG264AcuXL8e8efOQnJyM8PDwQITaRW+xA8Crr76KpKQkjB8/3t/h9ai32JOSkmA0Gh2POzo6At5uLo/ZYDCgoqLCUSbq96mcX4qCvP3223j99dednjMajZg2bRq++OILl+8JDw9HTEwMZFnGc889hwkTJmDMmDH+CNeJJ7FbLBZER0c7HkdFRcFisQiN0xVXsQ8ZMsQRW1RUFFpaWpzKNRoNVCoVpk6dipaWFjzzzDN+i/dynsTe1NSEL7/8Elu3boXNZkNWVhYMBgMGDhzot7gBz2Lft28fampqkJOTgwMHDvgt1it5ErtOp4NOp8OFCxewYsUKZGRkICoqym8xu2KxWKDX6x2P1Wo12tvbodFohP0+mfxd6By37yubzYaVK1ciKioKq1evFhBZ7zyJXa/Xo7W11fG4tbXVqbH5i6vYFy9e7IittbW1S2Lcvn07hg4dildffRWtra2OBDpixAi/xQ14FvugQYNwyy23QK/XQ6/XY+zYsaiurkZiYqLf4gY8i33btm2oq6tDdnY2qqqqcPjwYQwbNgwJCQl+ixvwLHYAaG5uxmOPPYZbbrkFjzzyiF9i7cmVv0FJkhxHI6J+nxz28RFZlrFo0SKMGzcOOTk5UKvVgQ7JbYmJiSgvL4fNZkNLSwuOHTuG+Pj4QIcF4OJw1Z49ewAApaWluOmmm5zKBw4ciMjISKjVakRFRUGr1cJqtQYi1C56iz05ORn79++HzWaD1WrFsWPHEBsbG4hQu+gt9nXr1uGtt97Cli1bkJKSgl//+td+T/zd6S32trY2zJ8/H7NmzcKvfvWrQITYRXJysmPS3GQyOf3+RP0+2fP3UueN6SVJwv79+2G327F3714AwNKlS5GUlBTgCLvXGXtqaiqys7ORlZUFWZbxxBNPKGa+Yu7cuVi+fDnmzp2L8PBwrFu3DgDw3HPPYcqUKZg+fToOHDiAzMxMdHR0YPr06Rg7dmyAo76ot9gTExMxa9YszJ0719F5GDRoUGCD/v/ciV2peov9wIEDqK2txdtvv423334bwMWh0VGjRgUs5rS0NJSVlSEzMxOyLMNoNAr/ffKqnkREIYjDPkREIYjJn4goBDH5ExGFICZ/IqIQxORPRBSCmPyJiEIQkz8RUQjiSV5EHrBYLHjqqafQ0tKC77//HllZWcjKygp0WERu40leRB44fPgw6urqcPfdd6O+vh7Z2dn46KOPAh0WkdvY8yfywNChQ/H666/jo48+gl6vR3t7e6BDIuoTjvkTeeC1116DwWDA888/jylTpoAH0BRs2PMn8sCPf/xj5ObmYteuXYiOjoZarYbdbodWqw10aERu4Zg/EVEI4rAPEVEIYvInIgpBTP5ERCGIyZ+IKAQx+RMRhSAmfyKiEMTkT0QUgv4fIl4tvg/rHBIAAAAASUVORK5CYII=\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "g = sns.scatterplot(data=aux, x='a', y='b')\n",
    "g.axvline(0, ls='--', color='black')\n",
    "g.axhline(0, ls='--', color='black');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The horizontal axis is probability of contraceptive use in rural area of a district. The vertical is the probability in urban area of same district. There urban areas above and below 0.5. All rural areas are below 0.5. However, for the same ditrict, the urban area is above the rural one. The negative correlation between the intercepts and slopes is necessary to encode this pattern."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAYQAAAEECAYAAAAoDUMLAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjMuNCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8QVMy6AAAACXBIWXMAAAsTAAALEwEAmpwYAAAqiklEQVR4nO3dfViUZd438O/MMDO8qYFYuiYYrqjh2oTWYx4L7Ybem+hmZqvgarqb4dph5hOR5WHmYfOgPEWHi29pbtLq6mJa7Ooa7YOVeCvZrTa14AAmgoaKvJgyDAwv1/X8gYyMzAxvczkzF9/PX8Q5c855OjS/Od9+p0IURRFERNTnKd3dACIi8gwMCEREBIABgYiIbmFAICIiAAwIRER0i4+7G9AbBoMBWq3W3c1wGYvFIqv+APLsEyDPfsmxT4D8+lVUVARBEDBmzJgePd9isUCn09kt8+qAoNVqe/yP4omMRqOs+gPIs0+APPslxz4B8uvXkiVLYDab8c033/To+Uaj0WEZp4yIiAiAl48QiIj6mlWrVqGsrEySuhkQiIi8yOTJk51O+/QGp4yIiLyIwWCQLCBwhEBE5EWWL18Os9mMZ555xuV1SxIQBEHAmjVrUFRUBI1GA71ej7CwMGv50aNHsXnzZoiiiMjISLz11luwWCxITk5GdXU1AgICkJqaiuDgYCmaR0QyIwgiSqvrUHGzAff198XwgQFQKhXubpbXkWTKKCcnB42NjcjMzERSUhLWr19vLTOZTHjnnXfw/vvv4+OPP8bQoUNx/fp17N27FxEREdizZw+efvppbNmyRYqmEZHMCIKI7IKriEs/hoQPTiIu/RiyC65CEJjIubskCQinT59GdHQ0AECn0yE/P99a9u233yIiIgKpqamYO3cuQkJCEBwcbPOcmJgY5OXlSdE0IpKZ0uo6vLLPgIYmAQDQ0CTglX0GlFbXubll3keSKSOTyYTAwEDrf6tUKjQ3N8PHxwfXr1/HyZMnkZWVBX9/f/z+97+HTqeDyWRCv379AAABAQGora3t9HUsFotkiyvu0NDQIKv+APLsEyDPfnlrny43BViDQZuGJgEXrtbAUnXJa/vliNlshiAIkvRJkoAQGBiIurrb0VkQBPj4tL7UPffcg1/84hcYNGgQAGDChAkwGo02z6mrq0P//v07fR2eVPZ8cuwTIM9+eWuftJUm+KqVNkHBV63EA4ODET4o1Gv75ciGDRtQWlra4z7d9ZPKUVFRyM3NBdC6RSoiIsJaFhkZieLiYtTU1KC5uRnfffcdfv7znyMqKgpHjx4FAOTm5mL8+PFSNI2IZGb4wAC8N1sHX3Xrx5mvWon3ZuswfGCAm1smjUmTJuHhhx+WpG5JRghTpkzB8ePHER8fD1EUkZKSgp07dyI0NBSxsbFISkrCokWLAABPPvkkIiIiMGzYMKxYsQIJCQlQq9VIS0uTomlEJDNKpQJPRg7G6GXRuFbbgHv7yXuX0YkTJ3o1QnBGkoCgVCqxdu1am9+NGDHC+vO0adMwbdo0m3I/Pz+kp6dL0RwikjmlUoHwQYEIHxTY+YO93MqVK2E2mzF37lyX182DaUTk0XjG4O5hQCAij9V2xqBtW2nb+sCTkYMZFCTAXEZEJDlBEFFSaULe+SqUVJq6fGiMZwzuLo4QiEhSvfmWX3Gzwe4Zg2u1DX1iveBuY0AgIkk5+pY/ell0px/q9/X3tXvG4N5+vpK22ZNt2LABJSUlktTNKSMikpSzb/mdccUZA6VK1aPpKk+l0+kkO2jHEQIRSao33/J7e8ZAEEQYa7V4fdcx2SxK5+TkoKysTJKgwBECEUmqt9/y284YTAwPQfigwG59kJdW1+H1rLOyWpTW6/XYtm2bJHVzhEBEknLnSWIuSncPAwIRSc5dJ4m5KN09nDIiItkaPjAA659+sM8kvustjhCISLaUSgXG9LPgcB9JfNdbDAhEJGtCS4usEt9t27YN58+fl6RuBgQi6jUmoLt7Ro0aBUEQOn9gDzAgEJFTnX3YMwHd3XXw4EFcunSJ5xCIqPe6k2iu7cM+Lv0YEj44ibj0Y8guuGrzHCagu7vS0tKQkZEhSd0MCER9SFc+4Nvryod9b1JTkGdhQCDqQ7r7bd7Zh33bSKO+qQUvx/4cQwbc3tvPvf7eiWsIRH2Isw/4AXYe7+hg1+D+vh3WDV6OHYm/5pXhurmRe/29FAMCUR/i7OSuxdLx8W15iO5cMG4R0GGk8ecj5/DRHx7FoH5a7jLyUgwIRH2Iow/44QMDUFTV8fGO8hCdvFBtd6QhQpTNfn9PtWvXLpw7d06SuhkQiPqQniSas5eHiDmC3GfYsGEwmUyS1M1FZaI+pjfppNu44uIa6pnMzEx89tlnktTNEQIRdVn7Q2qj7uuH7JejcfUmcwTdTVu3boXZbMYrr7zi8roZEIioS3giWf44ZUREXcITyfLHgEBEXcITyfLHgEBEXdK2s6g97iySF64hEMmIlGmonZ1hoLtn//79KC4ulqRuBgQimZB60bcnZxjI9UJCQlBZWSlJ3ZwyIpKJu7Ho64ozDNQ7GRkZ+PTTTyWpW5IRgiAIWLNmDYqKiqDRaKDX6xEWFmYt1+v1OHPmDAICWoeaW7ZsQUtLC37zm98gIiICADB58mQsWLBAiuYRyZKzRV+mk5CPjIwMmM1mrFy50uV1SxIQcnJy0NjYiMzMTBgMBqxfvx5bt261lhcUFGDHjh0IDg62/u7EiROYPn063nzzTSmaRCR7TCdBvSXJlNHp06cRHR0NANDpdMjPz7eWCYKAsrIyrF69GvHx8di/fz8AID8/HwUFBZg3bx6WLVuGa9euSdE0ItliOgnqLUlGCCaTCYGBt4eoKpUKzc3N8PHxgdlsxrx58/CHP/wBLS0teO655zB27FiEh4dj7NixmDRpEv75z39Cr9cjPT3d6etYLBYYjUYpuuAWDQ0NsuoPIM8+AZ7brxFaFfYsGIfr9S0I8lMhQKhFUdFPUKpUqFMG4Lq5BcH+KvgLdRBaWmye66l96i259ctsNkMQBEn6JElACAwMRF3d7YUsQRDg49P6Un5+fnjuuefg5+cHAJg4cSIKCwsxefJk6++mTJnSaTAAAK1WK8lF0+5iNBpl1R9Ann0CvKtfXd195E196g659cvf3x9ms7nHfXIWSCSZMoqKikJubi4AwGAwWBeKAaC0tBQJCQloaWlBU1MTzpw5g8jISKxatQqff/45ACAvLw+RkZFSNI2oz2HKCXk5fPgw3n//fUnqlmSEMGXKFBw/fhzx8fEQRREpKSnYuXMnQkNDERsbixkzZmD27NlQq9WYMWMGRo4ciaSkJKxcuRJ79+6Fn58f9Hq9FE0j6nO4+0he/P39rbMpriZJQFAqlVi7dq3N70aMGGH9edGiRVi0aJFN+bBhw7Br1y4pmkPUp3H3kbxs2bIFV69e7fAZ6wo8mEYkc9x9JC/79u1Ddna2JHUzdQWRzDHlBHUVAwKRDNlLcnfnvchEd2JAIJKJtiBQXWfB5Z8asOLA97zZjLqFawhEMtB21iAu/Ri+KqqyBgOA20yp6zhCIPIQvbnLoP1ZA4UC3GYqY1999ZVkJ68ZEIg8QG/vMrjzrAG3mVJPcMqIyAP09jRx++stD5z+EcueGMltpjL17rvv4sMPP5Skbo4QiDxAb08Tt7/e8sqNBmSeuojt8ydArVK4/CpNcq9Dhw7BbDZLUjcDApEH6O1pYp41IFfglBGRmwiCiJJKE/LOV0GpQK9PE/N6S+otjhCI3MDeIvKmuQ/jXy9Fo9LEb/jkHhwhELmBvUXkpXu+hUIBfsMnp/z8/KDVaiWpmyMEIjdgSmrqqc8++4znEIi82Z2HzoYMYEpq8jycMiKSWPu0EgkfnERc+jGcvVKLTXMftllE3jT3YYgikHe+CiWVJgiC6OaWkyd6++23sXXrVknq5giBSGKODp3966VoHL61TXRwf1+cvVKLaRuPMSEdOXXkyBHJziFwhEAkMUfrBZWmBus2UUFEl08qt9+uypEEuRJHCEQS68qhs64uMvc25xGRMxwhEEmssyssBUFEc4toLW9jb5G5tzmPiJzhCIFIYp2llSitrsOqf/wHy54YifQvzlm/+afOGtfhpDK3q9LAgQOh0WgkqZsBgeguaEsrYe9Du+JmA8qq67Hr6zI8/8twKBSAKAJD7/HtMA3U25xH5P0OHDjAcwhEnkSpUqGk0tSjy2zu1PYhf+VGAzZ/+QOA1g/5WVFDOzy2fVbT9msITG1NrsCAQOSAoxvMBEGEsVaL13e5Zotodz7kmdWU3njjDVRVVeGDDz5wed0MCER2ONvNU1pdh9ezznZY2B29LLpH8/jd/ZB3Nv1E8peXl8dzCER3k7PdPM4WdnuKqavJE3CEQGSHsw/99gu7Qwb44pmo+6FSAn5qHwiCyA9z8locIRDZ0f6O4jZtu3mGDwzA+qcfRNhAP8yfGIa//HcJ0o/8gDnb85BdcJUnh8lrMSAQ2eHsMJlSqcCYfhZsTHjYem4A6PohMaaeoN64//77cd9990lSN6eMiOzobKFXqVLhQpW524fEmHqCemv37t2SnUPgCIHIAWcLvT8p+qGk0tSldBPtMfUEeTJJRgiCIGDNmjUoKiqCRqOBXq9HWFiYtVyv1+PMmTMICGjdZ71lyxY0NTXh1VdfRUNDA+69916sW7cOfn5+UjSPqNeu1TZh36kfO6Sb0D891ukhMaaeoN5avnw5ampq8Ne//tXldUsSEHJyctDY2IjMzEwYDAasX7/e5kKHgoIC7NixA8HBwdbf6fV6TJ8+Hc888wy2b9+OzMxMLFy40OnrFBUVYcmSJTa/mz17Nl588UWYzWbExcV1eM7ChQuxcOFCVFVV4dlnn+1QvmTJEsyZMweXLl3C/PnzO5QnJSXht7/9LYqKirB48eIO5atWrcLkyZNhMBiwfPnyDuUpKSmYNGkSTpw4gZUrV9qUmc1mbN++HTqdDjk5OdDr9R2ev23bNowaNQoHDx5EWlpah/Jdu3Zh2LBhyMzMtHuJxv79+xESEoKMjAxkZGR0KD98+DD8/f2xZcsW7Nu3r0P5V199BQB49913cejQIZsyPz8/fPbZZwBaL/Foy9vu7+8PoDUHy4EDBwC0Hq7Jy8uzef7999+P3bt3A2j9ozcYDDblERER2L59OwAgMTERxcXFNuU6nQ4bNmwAAMybNw8//vijTfljjz2GdevWAQBmzZqF6upqm/LY2Fi8+eabAICpU6eivr7epnz69Ol49dVXAQCrF8+G8UotCj5QYFCgFioFEBkzFaMHP4aGhnqHf3sx056FT1Mtfvz4/1h/r1QosPzYAPzvZUvd9rcHAMuWLcOYMWNk87fXxmw2Y9iwYbL528vIyIAgCPjVr35lLe/O554zkgSE06dPIzo6GkDrP1R+fr61TBAElJWVYfXq1dYP5WeffRanT5+2/pHHxMTgvffe6zQgCILQ4YDG1atXYTQaUV9fb/fwxuXLl2E0GnH9+nW75eXl5TAajbhy5Yrd8kuXLsFoNOLChQt2y8vKymA0GlFSUmK3vLS0FEFBQSgtLe1QLggCSkpKoNVqUVZWZvf558+fhyAIuHTpkt3yc+fOwWQyoby83G55cXExKisrcfnyZbvlhYWF8PPzw9WrV+2Wt81dVlRUdChvaWmxlldWVsJsNtu8RxqNxlpeVVXV4fk3btywltfU1HQo/+mnn6zlP/30U4fympoaa/mNGzc6lFdVVVnLa2trO5RXVlZay00mEywWi015RUWFtVwpNOGBkABcqKpD+U/1UCoU+LW/GtqGahQWdnxtoPVvr+n6ZayOG4U/7VdAEEUoFQqEh/hDaG50698eADQ2NsJoNMrmb6+NIAiora2Vzd+eIAgQRdGmju587gUFBXUotxIlsHLlSvGrr76y/vfjjz8uNjU1iaIoirW1teLmzZtFs9ks1tbWijNnzhSNRqM4efJksb6+XhRFUbx48aIYHx/f6eucPXtWiua7jdz6I4qe06eWFkE8f61WPPFDpXj+Wq3Y0iL06vFnz54Vm5paRMPFGjH7P5dFw8UasamppVttyTvftbbcLZ7yXrma3Pr1+OOPi4888kiPn+/s30OSEUJgYCDq6m4vkgmCAB+f1pfy8/PDc889Z10fmDhxIgoLC63P8fX1RV1dHfr37y9F06gP6u7OnuZmASdKqnGqrAaCCBz8rhwrnhzT4fE+Pko8NCwIDw2zfS17+Y/aY+oJ8lSS7DKKiopCbm4uAMBgMCAiIsJaVlpaioSEBLS0tKCpqQlnzpxBZGQkoqKicPToUQBAbm4uxo8fL0XTqA/qzs4eQRDxr/wrSNx1CulHfsCOYyWYMyEUqdnGLp0vyC64irj0Y0j44CTi0o/xoBq5XEREBIYPHy5J3ZKMEKZMmYLjx48jPj4eoigiJSUFO3fuRGhoKGJjYzFjxgzMnj0barUaM2bMwMiRI7FkyRKsWLEC+/btQ1BQkN1FK6KuuPNburOdPcMHBtg8VhSBFQe+twke6V+cw/O/DO90J5CjwNPTpHdE9mzfvt277kNQKpVYu3atze9GjBhh/XnRokVYtGiRTXlISAj+8pe/SNEc6kPsTQ99MH+C3UtlBgX6dnjsu88+ZDd4qJTo9BIabiklb8eDaSQr9r6lr/rHf5A6a1yHNBQqJTo89ty1WruHzSaEBXd6CY2z/EdErpKYmIi33npLkroZEEhW7H1LL6uux9B7fHF4WTT+nvi/cHhZNJ6MHIwrNzo+dt+pH5Ey8xc2wSN11jhMCh/YaWoJZ/mPiFyluLgYpaWlktTNXEYkK47uHA4O0HbY2WPvsdfNjRgfdg8yEyfiyo0GDBngi8ghA+Dj0/l3J95mRt6OIwSSle58S7f32E1zH0bB5VrM2f41/rT7DOZs/xr/NlZ0eacQL7ohb8YRAslKd76l23usKALTNh7jTiHqkxgQSHa6c/Drzsfmna/iTiHyaDqdDjU1NZLUzYBA1I6jNQjuFCJPsWHDBt6HQNSeVLeOcacQ9WUcIZDXkfLWMXvrCqFB/h3yExG5y7x583Djxg0cPHjQ5XUzIJDXkTpFRPt1BUfBZ4RW1evXIeqJH3/80W6Ka1fglBF5HWcpItq4akrJUfCpU3KUQPLDEQJ5nc4Wfl05peQo+Fyvb+l9R4g8DAMCeZ3hAwOwae7D+P7HGxBEQKUAfnH/AOvcviunlBwFnyA/ThmR/HDKiLxSY7OI7bkl2PTFD9iWW4LG5ttTQl2ZUuoqR7uOAgTndyMQSeWxxx7DQw89JEndHCGQV2h/x4G/xgep2UaHIwBXniVwdPK5qOgnV3WNqFvWrVvnnvsQMjMzHZbNmTPH5Y0hssfemsCyJ0Zi19dluHKj9Vt/+9PEbd/q71xD6Ol2UV55SX2F04BQUlKCL7/8Ek899dTdag+RVduooLS6DkVXbyLIX2NNWZ156iLeiBuD4opaAK33HreNAJh1lORs1qxZqK2txb///W+X1+00ILzxxhsoKSlBTEwMxo0b5/IXJ3LE2agAAOZMCMVr+7+zlq2dMRahQf7W5/NbPclVdXW1ZOcQOl1DSE1NRX19vSQvTuSIvZ1CbXcbA0D6F+dsylb/Ix8TwoK6FADuvHOZoweiVp0GhODgYDzxxBNQKG7/DxMYGIh//OMfkjaM+jZHO4UUCkCpQI8zkkqZ9oLI23Vpl1F2djYAQBRF5Ofn4/PPP5e0UUSOdgpF/zwEvmoVtueW9GgXkdRpL4i8WZfOIWg0Gmg0Gmi1WowfPx4FBQVSt4v6OHv7/1NnjYPGRwGVEh3uPe7qLiJXnlEgcofY2FhMnDhRkrq7NEJIS0uzThlVVlZCqeR5NpJW+51CFTcb0NQiIv1IEZ4YPRjpX5xDkL8GiTHhiLivH8YM7o8HQrq2DsD7Dsjbvfnmm+45h9AmPDzc+vPo0aMRHR0tSWOI2mvbKQQAcenH8Pwvw62LyVduNCD9yA/wVStxeFl0l+f/XX1GgUhOuhQQZs6cKXU7iBxqm+ZR9GIxuQ3PKJC3mzp1KkwmE44dO+byupm6gjxe2zQPAJdM9/CMAnmz+vp6WCwWSermYgB5vLZpnoPflWPZEyN5vSWRRDhCII9nneYZ3A81dRZkJk6EubGFh8qIXIwBgbwCp3mIpMeAQC5jLyUEEbnW9OnTUVFRIUndkgQEQRCwZs0aFBUVQaPRQK/XIywsrMNjEhMTERsbi4SEBIiiiJiYGAwfPhwAoNPpkJSUJEXzSALtU0IE+Wvwuwn3I+Lefrg/cAAEQeS0DpGLvPrqq+49h9BdOTk5aGxsRGZmJgwGA9avX4+tW7faPGbDhg24efOm9b8vXryIyMhIvP/++1I0iSTWlhIiyF+D+RPDrOcFWhd+FcwVROQFJNlldPr0aevhNZ1Oh/z8fJvy7OxsKBQKmwNuBQUFqKiowPz58/HCCy+gpKREiqaRRNrOCjwTdX+HTKSv7DOgtJpXThK5wq9+9SssWLBAkrolGSGYTCYEBt5e/FOpVGhuboaPjw+Ki4tx6NAhpKenY/PmzdbHDBo0CImJiZg6dSpOnTqF5ORkHDhwwOnrWCwWyYZO7tDQ0OC1/fFV94evWunw8NiFqzWwVF1yU+tcz5vfK0fk2CdAfv0ym80QBEGSPkkSEAIDA1FXd/sboSAI8PFpfamsrCxUVFRgwYIFKC8vh1qtxtChQ/HII49ApVIBACZMmIBr165BFEWbtNt30mq1GDNmjBRdcAuj0ei1/REEEe/NVqLo6k27h8ceGByM8EGhbmyha3nze+WIHPsEyK9f/v7+MJvNPe6Ts0AiSUCIiorCl19+ibi4OBgMBkRERFjLXnvtNevPGzduREhICGJiYvDOO+/gnnvuwQsvvIDCwkIMGTLEaTAg9xMEEReq6lBWU4cAjQ8if9YPYwb3Q9jAAKz89D/MFUTkZSQJCFOmTMHx48cRHx8PURSRkpKCnTt3IjQ0FLGxsXafk5iYiOTkZBw9ehQqlQrr1q2TomnUTY5uFxMEEV8UVeD7H29AEAGVAhgYoMHPgvzw1LifQTfsHlyrbYBGsEAXzgVlIm8gSUBQKpVYu3atze9GjBjR4XEvvfSS9ecBAwZg+/btUjSHesjZ7WIXa+pwrsJkvajGV63Ey7EjUXLNhPCQQOshMqPRCKVyqLu7QiQbs2fPxtWrVyWpm7mMyCFHt4u1jhgs+PMR291Efz5yDkPu8edlM0QSevHFF5GQkCBJ3QwI5JCz28XqGpvtltU3NvOyGSIJmc1m1NfXS1I3U1eQQ85uFxNF+6mofxbkxwVkIgnFxcXBbDbjm2++cXndDAjkUGiQP7bPn4BTZTUQRODgd+VY8eQY6wf+nTeP/d9Z4zBx+EAuIBN5KQYEsksQRPzbWGHzgZ86axz+a8x91g983jxGJC8MCGSXvQXlFQe+xy+GDrCmoLaXkvrObarKW4cNicjzMSCQXc4WlB3dSWBvm+r6px/ESGY7JfIKDAhkl7MFZUfsjSpezzqLcaEDebENkYssXLgQly9flqRubjslu9ruMe7O/cXORhVE5BoLFy7EzJkzJambIwSyy3qPcTcWjXsyqiCi7qmqqsL169clqZsBgRzq7j3GbaOKO9cQeC6ByHWeffZZnkMgz2dvVNF0/TIXlIm8BAMCudSdowpjVYubW0REXcVFZSIiAsARgqw5usuAiMgeBgSZcnaXAYMCkfdasmQJysvLJambU0Yy5ewuAyLyXnPmzMHUqVMlqZsBQaYcHRKruMlDYkTe7NKlS7hy5YokdXPKSKb8NT52D4n5a5hsjsibzZ8/X7JzCBwhyFRjSwuWPTHSJvXEG0+ORmOzgLzzVSipNEEQRDe3kog8CUcIMjUwQIvMUxfx/C/DoVAAARoVLC0C5n/4DReZicgujhBkavjAAKx4cgz+8t8l2PTFD6hvasF7/6+Yi8xE5BBHCDJ1ZxoJc2NLt+83IKK+hQFBxtqnkSipNDETKZEMJCUl4dKlS5LUzSmjPqIn9xsQkef57W9/i1//+teS1M0RQh/Rk/sNiMjzFBUV4cKFCxgzZozL62ZA6EO6e78BEXmexYsXw2w2Iy4uzuV1c8qIiIgAMCAQEdEtDAhERASAawhejfcdEJErSRIQBEHAmjVrUFRUBI1GA71ej7CwsA6PSUxMRGxsLBISEtDQ0IDk5GRUV1cjICAAqampCA4OlqJ5ssD7Doj6plWrVqGsrEySuiWZMsrJyUFjYyMyMzORlJSE9evXd3jMhg0bcPPmTet/7927FxEREdizZw+efvppbNmyRYqmyQbvOyDqmyZPnoxJkyZJUrckAeH06dOIjo4GAOh0OuTn59uUZ2dnQ6FQWB9z53NiYmKQl5cnRdM8niCIKKk0dZqR1NF9B9dqed8BkZwZDAYYjUZJ6pZkyshkMiEw8PZed5VKhebmZvj4+KC4uBiHDh1Ceno6Nm/ebPOcfv36AQACAgJQW1vb6etYLBbJ/mHcoam5Gf/89iJezzprnQZa//SDGNPPAqGlxeaxvur+dlNRaATP+jdpaGjwqPa4ihz7Jcc+AfLrV2JiIgRB8J6DaYGBgairuz11IQgCfHxaXyorKwsVFRVYsGABysvLoVarMXToUJvn1NXVoX///p2+jlarleQfxV3O/FCO1/d9bzMN9HrWWRxeFt3hMJkgiHhvtrLDGoIufDCUyqHuaL5dRqNRVu9RGzn2S459AuTXL39/f5jN5h73yVlwlCQgREVF4csvv0RcXBwMBgMiIiKsZa+99pr1540bNyIkJAQxMTH44YcfcPToUYwbNw65ubkYP368FE3zaNfNXc9IylQURORqkgSEKVOm4Pjx44iPj4coikhJScHOnTsRGhqK2NhYu89JSEjAihUrkJCQALVajbS0NCma5tGC/VXdykjKVBRE5EqSBASlUom1a9fa/G7EiBEdHvfSSy9Zf/bz80N6eroUzfEa/kId3putQ2q2EdPHDYVKCTwSFozQIH93N42I+gAeTPMgQksL/mvMfWhqEbDiwPc8X0BEHaSkpKC0tFSSupm6wsNcvG62BgOA5wuIyNakSZPw8MMPS1I3A4KH4fkCInLmxIkT+PbbbyWpm1NGHua+/r686pKIHFq5ciXMZjPmzp3r8ro5QvAwvOqSiNyFIwQPIQgi6tX9cfJCNUbd1w/ZL0fj6k2eLyCiu4cBwQPczlzKnUVE5D6cMvIAzFxKRJ6AIwQP4GxnEU8hE1F7GzZsQElJiSR1MyB4AO4sIqKu0ul00Gq1ktTNKSMPwJ1FRNRVOTk5OHHihCR1c4TgAdoylw5eMA6NSi13FhGRQ3q9HmazGc8//7zL62ZA8BBKpQJ+TTcRJaO87UTkXThlREREABgQiIjoFgYEIiICwDUEIiKvsm3bNpw/f16SuhkQ7BAEEaXVdai42YD7+nPHDxF5jlGjRkEQhM4f2AMMCHe4nVfIwLxCRORxDh48iEuXLmGMBDsSuYZwB+YVIiJPlpaWhoyMDEnqZkC4A28sI6K+igHhDm15hdpjXiEi6gsYEO7AvEJE1FdxUfkObXmFRi+LxrVa3lhGRH0HA4IdSqUC4YMCeRcBEXmcXbt24dy5c5LUzYBARORFhg0bBpPJJEndDAhERF4kMzMT5eXlPIdARNTXbd26FX//+98lqZsBgYiIADAgEBHRLQwIREQEQKJFZUEQsGbNGhQVFUGj0UCv1yMsLMxa/re//Q2ffPIJFAoF/vjHPyIuLg6iKCImJgbDhw8HAOh0OiQlJUnRPCIiskOSgJCTk4PGxkZkZmbCYDBg/fr12Lp1KwCgpqYGe/fuxaeffgqLxYJp06Zh6tSpuHjxIiIjI/H+++9L0SQiIlnYv38/iouLJalbkoBw+vRpREdHA2j9pp+fn28tCw4ORlZWFnx8fFBeXg6tVguFQoGCggJUVFRg/vz58PX1xRtvvIHw8HCnr2OxWGA0GqXogls0NDTIqj+APPsEyLNfcuwTIM9++fn5SdInSQKCyWRCYODtU74qlQrNzc3w8Wl9OR8fH+zevRsbN27E/PnzAQCDBg1CYmIipk6dilOnTiE5ORkHDhxw+jparVaSvbjuYjQaZdUfQJ59AuTZLzn2CZBfvzIyMnD58mWsXLmyR893FkgkWVQODAxEXd3t+wMEQbAGgzbz5s3DsWPH8D//8z/4+uuvMXbsWMTGxgIAJkyYgGvXrkEURZe3TRBElFSakHe+CiWVJgiC61+DiEgqGRkZyMrKkqRuSQJCVFQUcnNzAQAGgwERERHWspKSEixduhSiKEKtVkOj0UCpVGLTpk346KOPAACFhYUYMmQIFArXJpRruw0tLv0YEj44ibj0Y8guuMqgQEQEiaaMpkyZguPHjyM+Ph6iKCIlJQU7d+5EaGgoYmNjMXr0aMyZMwcKhQLR0dF49NFHMWrUKCQnJ+Po0aNQqVRYt26dy9vl6Da00cuimciOiPo8SQKCUqnE2rVrbX43YsQI689Lly7F0qVLbcoHDBiA7du3S9EcK2e3oTEgEFFf16cOpvE2NCIix/pUQOBtaETk7Q4fPizZea0+lf6at6ERkbfz9/eHn5+fJHX3qRECcPs2tInhIQgfFMhgQEReZcuWLdi7d68kdfe5gEBE5M327duH7OxsSepmQCAiIgAMCEREdAsDAhERAWBAICKiWxSiFBnk7hKDwQCtVuvuZhAReQ2LxQKdTme3zKsDAhERuQ6njIiICAADAhER3cKAQEREABgQiIjoFgYEIiICwIBARES39Kn01+4iCALWrFmDoqIiaDQa6PV6hIWFWcszMjLwr3/9CwDw+OOPW++cjomJwfDhwwEAOp0OSUlJ7mi+Q531629/+xs++eQTKBQK/PGPf0RcXBwaGhqQnJyM6upqBAQEIDU1FcHBwW7sha2e9EkO71XbYxITExEbG4uEhASvf6/aHtO+T3J4r/R6Pc6cOYOAgNZ7XLZs2YKmpia8+uqraGhowL333ot169b1LEW2SJL7/PPPxRUrVoiiKIrffvut+Kc//cladvHiRXHmzJlic3OzKAiCOGfOHNFoNIqlpaXi4sWL3dXkLnHWr+rqanHatGliY2OjWFtbK8bExIiCIIgffvihmJ6eLoqiKB46dEh8++233dJ2R3rSJ29/r9qkpaWJv/vd78Q9e/aIoih69XvV5s4+yeG9io+PF6urq21+9/bbb4sHDhwQRVEUt23bJu7cubNHr80po7vg9OnTiI6OBtD6jSQ/P99aNnjwYOzYsQMqlQoKhQLNzc3QarUoKChARUUF5s+fjxdeeAElJSXuar5DzvoVHByMrKwsqNVqVFVVQavVQqFQ2DwnJiYGeXl5bmm7Iz3pk7e/VwCQnZ0NhUJhfcydz/G29wqw3ydvf68EQUBZWRlWr16N+Ph47N+/v8NzYmJicOLEiR69NgPCXWAymRAYGGj9b5VKhebmZgCAWq1GcHAwRFFEamoqHnzwQTzwwAMYNGgQEhMTsWvXLixevBjJycnuar5DzvoFAD4+Pti9ezfmzJmDp556yvqcfv36AQACAgJQW1t7dxvdiZ70ydvfq+LiYhw6dAgvv/xyh+d463vlqE/e/l6ZzWbMmzcP77zzDnbs2IE9e/agsLDQZe8V1xDugsDAQNTV1Vn/WxAE+Pjc/qe3WCxYuXIlAgIC8NZbbwEAxo4dC5VKBQCYMGECrl27BlEUoVB4zg1vnfULAObNm4fZs2fjhRdewNdff23znLq6OvTv3/+utrkzPenTQw895NXvVVZWFioqKrBgwQKUl5dDrVZj6NChXv1eOerTI4884tXvlZ+fH5577jnr+sDEiRNRWFhofY6vr2+v3iuOEO6CqKgo5ObmAmhNyBcREWEtE0URL774IkaNGoW1a9da/1g3bdqEjz76CABQWFiIIUOGeNQfLeC8XyUlJdbFcbVaDY1GA6VSiaioKBw9ehQAkJubi/Hjx7ul7Y70pE/e/l699tpr+Pjjj7Fr1y7MnDkTCxcuRExMjFe/V4765O3vVWlpKRISEtDS0oKmpiacOXMGkZGRLnuvmNzuLmjbNVBcXAxRFJGSkoLc3FyEhoZCEAS88sorNtkHX3nlFYSHhyM5ORlmsxkqlQqrV6/GiBEj3NcJO5z1KzY2Fps2bUJubq51Hnfp0qWor6/HihUrUFlZCbVajbS0NAwaNMjdXbHqSZ9u3Ljh9e9Vm40bNyIkJAQJCQle/161ad8nObxXO3bswGeffQa1Wo0ZM2YgISEBVVVVWLFiBerq6hAUFIS0tDT4+/t3+7UZEIiICACnjIiI6BYGBCIiAsCAQEREtzAgEBERAAYEIiK6hQfTiCRQVlaG119/HQqFAiNHjsRbb70FpZLfv8iz8S+USALr1q3D8uXLsWfPHoiiiCNHjri7SUSdYkAg6qVPPvkE7777LoDWNCRPPPEECgoK8OijjwLoXbIxoruJAYFIAu3z43hiYjgiexgQiFyo7eB/+/UCT0wMR2QPAwJRL2m1WlRWVgJozbcPAA8++CBOnjwJoDXZ2IQJE9zWPqKuYi4jol66efMmXnzxRbS0tCAyMhInT55Eeno63nzzTTQ1NSE8PBx6vd6ayZbIUzEgEBERAE4ZERHRLQwIREQEgAGBiIhuYUAgIiIADAhERHQLAwIREQFgQCAiolv+P7ASJt7LqS7gAAAAAElFTkSuQmCC\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "g = sns.scatterplot(data=aux, x='u0', y='u1')\n",
    "g.axvline(0.5, ls='--', color='black')\n",
    "g.axhline(0.5, ls='--', color='black');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In fact, if we fit the model so it instead has two intercepts, one for rural and one for urban, there is no strong correlation between those intercepts. Here's such a model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def model(did, urban, contra=None):\n",
    "    # fixed priors\n",
    "    a_bar = numpyro.sample('a_bar', dist.Normal(0, 1))\n",
    "    b_bar = numpyro.sample('b_bar', dist.Normal(0, 0.5))\n",
    "    sigma_ab = numpyro.sample('sigma_ab', dist.Exponential(1), sample_shape=(2,))\n",
    "    rho = numpyro.sample('rho', dist.LKJ(2,2))\n",
    "    # adaptive priors\n",
    "    mu = jnp.stack([a_bar,b_bar])\n",
    "    cov = jnp.outer(sigma_ab, sigma_ab)*rho\n",
    "    i = len(d.did.unique())\n",
    "    ab_did = numpyro.sample('ab_did', dist.MultivariateNormal(mu, cov), sample_shape=(i,))\n",
    "    # likelihood\n",
    "    theta = ab_did[did,urban]\n",
    "    contra_hat = numpyro.sample('contra_hat', dist.Bernoulli(logits=theta), obs=contra)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "sample: 100%|██████████| 1000/1000 [00:13<00:00, 71.67it/s, 15 steps of size 2.84e-01. acc. prob=0.86]\n",
      "sample: 100%|██████████| 1000/1000 [00:04<00:00, 228.54it/s, 15 steps of size 2.56e-01. acc. prob=0.78]\n",
      "sample: 100%|██████████| 1000/1000 [00:04<00:00, 236.13it/s, 15 steps of size 2.50e-01. acc. prob=0.88]\n",
      "sample: 100%|██████████| 1000/1000 [00:04<00:00, 239.82it/s, 15 steps of size 2.64e-01. acc. prob=0.85]\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[['did','urban','contra']].items()}\n",
    "mcmc.run(rng_key, **dat)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "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>mean</th>\n",
       "      <th>sd</th>\n",
       "      <th>hdi_5.5%</th>\n",
       "      <th>hdi_94.5%</th>\n",
       "      <th>mcse_mean</th>\n",
       "      <th>mcse_sd</th>\n",
       "      <th>ess_bulk</th>\n",
       "      <th>ess_tail</th>\n",
       "      <th>r_hat</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>a_bar</th>\n",
       "      <td>-0.72</td>\n",
       "      <td>0.10</td>\n",
       "      <td>-0.89</td>\n",
       "      <td>-0.56</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>1339.0</td>\n",
       "      <td>819.0</td>\n",
       "      <td>1.00</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>b_bar</th>\n",
       "      <td>0.02</td>\n",
       "      <td>0.12</td>\n",
       "      <td>-0.20</td>\n",
       "      <td>0.20</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>784.0</td>\n",
       "      <td>1178.0</td>\n",
       "      <td>1.00</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>sigma_ab[0]</th>\n",
       "      <td>0.62</td>\n",
       "      <td>0.10</td>\n",
       "      <td>0.46</td>\n",
       "      <td>0.78</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>549.0</td>\n",
       "      <td>920.0</td>\n",
       "      <td>1.01</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>sigma_ab[1]</th>\n",
       "      <td>0.48</td>\n",
       "      <td>0.15</td>\n",
       "      <td>0.25</td>\n",
       "      <td>0.72</td>\n",
       "      <td>0.01</td>\n",
       "      <td>0.01</td>\n",
       "      <td>101.0</td>\n",
       "      <td>92.0</td>\n",
       "      <td>1.02</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "             mean    sd  hdi_5.5%  hdi_94.5%  mcse_mean  mcse_sd  ess_bulk  \\\n",
       "a_bar       -0.72  0.10     -0.89      -0.56       0.00     0.00    1339.0   \n",
       "b_bar        0.02  0.12     -0.20       0.20       0.00     0.00     784.0   \n",
       "sigma_ab[0]  0.62  0.10      0.46       0.78       0.00     0.00     549.0   \n",
       "sigma_ab[1]  0.48  0.15      0.25       0.72       0.01     0.01     101.0   \n",
       "\n",
       "             ess_tail  r_hat  \n",
       "a_bar           819.0   1.00  \n",
       "b_bar          1178.0   1.00  \n",
       "sigma_ab[0]     920.0   1.01  \n",
       "sigma_ab[1]      92.0   1.02  "
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "samples_12 = az.from_numpyro(mcmc)\n",
    "az.summary(samples_12, var_names=['a_bar','b_bar', 'sigma_ab'], hdi_prob=.89).round(2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Correlation all gone."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> Now consider the predictor variables age.centered and living.children, also contained in data(`bangladesh`). Suppose that age influences contraceptive use (changing attitudes) and number of children (older people have had more time to have kids). Number of children may also directly influence contraceptive use. Draw a DAG that reflects these hypothetical relationships. Then build models needed to evaluate the DAG. You will need at least two models. Retain district and urban, as in Problem 1. What do you conclude about the causal influence of age and children?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's my DAG:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"./fig/w9_img1.png\" width=\"20%\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`A` is age, `K` is number of children, and `C` is contraception use. To study this DAG, we should estimate both the total causal influence of `A` and then condition also on `K` and see if the direct influence of `A` is smaller. Here's the model for the total influence of A:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "def model(did, urban, age, contra=None):\n",
    "    # fixed priors\n",
    "    a_bar = numpyro.sample('a_bar', dist.Normal(0, 1))\n",
    "    b_bar = numpyro.sample('b_bar', dist.Normal(0, 0.5))\n",
    "    sigma_ab = numpyro.sample('sigma_ab', dist.Exponential(1), sample_shape=(2,))\n",
    "    rho = numpyro.sample('rho', dist.LKJ(2,2))\n",
    "    bA = numpyro.sample('bA', dist.Normal(0,0.5))\n",
    "    # adaptive priors\n",
    "    mu = jnp.stack([a_bar,b_bar])\n",
    "    cov = jnp.outer(sigma_ab, sigma_ab)*rho\n",
    "    i, j = len(d.did.unique()), len(d.urban.unique())\n",
    "    ab_did = numpyro.sample('ab_did', dist.MultivariateNormal(mu, cov), sample_shape=(i,j))\n",
    "    # likelihood\n",
    "    a_did, b_did = ab_did[:,:,0], ab_did[:,:,1]\n",
    "    theta = a_did[did,urban] + b_did[did,urban]*urban + bA*age\n",
    "    contra_hat = numpyro.sample('contra_hat', dist.Bernoulli(logits=theta), obs=contra)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "sample: 100%|██████████| 1000/1000 [00:27<00:00, 36.38it/s, 31 steps of size 1.24e-01. acc. prob=0.82] \n",
      "sample: 100%|██████████| 1000/1000 [00:11<00:00, 85.01it/s, 63 steps of size 9.35e-02. acc. prob=0.88] \n",
      "sample: 100%|██████████| 1000/1000 [00:20<00:00, 49.23it/s, 63 steps of size 7.50e-02. acc. prob=0.77] \n",
      "sample: 100%|██████████| 1000/1000 [00:20<00:00, 49.44it/s, 63 steps of size 6.72e-02. acc. prob=0.95]\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[['did','urban','age','contra']].items()}\n",
    "mcmc.run(rng_key, **dat)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "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>mean</th>\n",
       "      <th>sd</th>\n",
       "      <th>hdi_5.5%</th>\n",
       "      <th>hdi_94.5%</th>\n",
       "      <th>mcse_mean</th>\n",
       "      <th>mcse_sd</th>\n",
       "      <th>ess_bulk</th>\n",
       "      <th>ess_tail</th>\n",
       "      <th>r_hat</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>a_bar</th>\n",
       "      <td>-0.67</td>\n",
       "      <td>0.10</td>\n",
       "      <td>-0.84</td>\n",
       "      <td>-0.52</td>\n",
       "      <td>0.01</td>\n",
       "      <td>0.01</td>\n",
       "      <td>76.0</td>\n",
       "      <td>311.0</td>\n",
       "      <td>1.05</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>b_bar</th>\n",
       "      <td>0.61</td>\n",
       "      <td>0.16</td>\n",
       "      <td>0.35</td>\n",
       "      <td>0.85</td>\n",
       "      <td>0.04</td>\n",
       "      <td>0.03</td>\n",
       "      <td>21.0</td>\n",
       "      <td>98.0</td>\n",
       "      <td>1.14</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>sigma_ab[0]</th>\n",
       "      <td>0.59</td>\n",
       "      <td>0.09</td>\n",
       "      <td>0.46</td>\n",
       "      <td>0.75</td>\n",
       "      <td>0.01</td>\n",
       "      <td>0.00</td>\n",
       "      <td>252.0</td>\n",
       "      <td>438.0</td>\n",
       "      <td>1.01</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>sigma_ab[1]</th>\n",
       "      <td>0.32</td>\n",
       "      <td>0.24</td>\n",
       "      <td>0.03</td>\n",
       "      <td>0.64</td>\n",
       "      <td>0.08</td>\n",
       "      <td>0.06</td>\n",
       "      <td>7.0</td>\n",
       "      <td>19.0</td>\n",
       "      <td>1.62</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "             mean    sd  hdi_5.5%  hdi_94.5%  mcse_mean  mcse_sd  ess_bulk  \\\n",
       "a_bar       -0.67  0.10     -0.84      -0.52       0.01     0.01      76.0   \n",
       "b_bar        0.61  0.16      0.35       0.85       0.04     0.03      21.0   \n",
       "sigma_ab[0]  0.59  0.09      0.46       0.75       0.01     0.00     252.0   \n",
       "sigma_ab[1]  0.32  0.24      0.03       0.64       0.08     0.06       7.0   \n",
       "\n",
       "             ess_tail  r_hat  \n",
       "a_bar           311.0   1.05  \n",
       "b_bar            98.0   1.14  \n",
       "sigma_ab[0]     438.0   1.01  \n",
       "sigma_ab[1]      19.0   1.62  "
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "samples_21 = az.from_numpyro(mcmc)\n",
    "az.summary(samples_21, var_names=['a_bar','b_bar', 'sigma_ab'], hdi_prob=.89).round(2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this model, the total causal effect of `age` is positive and very small. Older individuals use slightly more contraception. And now the model with both `K` and `A`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "def model(did, urban, age, children, contra=None):\n",
    "    # fixed priors\n",
    "    a_bar = numpyro.sample('a_bar', dist.Normal(0, 1))\n",
    "    b_bar = numpyro.sample('b_bar', dist.Normal(0, 0.5))\n",
    "    sigma_ab = numpyro.sample('sigma_ab', dist.Exponential(1), sample_shape=(2,))\n",
    "    rho = numpyro.sample('rho', dist.LKJ(2,2))\n",
    "    bA = numpyro.sample('bA', dist.Normal(0,0.5))\n",
    "    bK = numpyro.sample('bK', dist.Normal(0,0.5))\n",
    "    # adaptive priors\n",
    "    mu = jnp.stack([a_bar,b_bar])\n",
    "    cov = jnp.outer(sigma_ab, sigma_ab)*rho\n",
    "    i, j = len(d.did.unique()), len(d.urban.unique())\n",
    "    ab_did = numpyro.sample('ab_did', dist.MultivariateNormal(mu, cov), sample_shape=(i,j))\n",
    "    # likelihood\n",
    "    a_did, b_did = ab_did[:,:,0], ab_did[:,:,1]\n",
    "    theta = a_did[did,urban] + b_did[did,urban]*urban + bA*age + bK*children\n",
    "    contra_hat = numpyro.sample('contra_hat', dist.Bernoulli(logits=theta), obs=contra)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "sample: 100%|██████████| 1000/1000 [00:19<00:00, 51.75it/s, 31 steps of size 1.50e-01. acc. prob=0.87]\n",
      "sample: 100%|██████████| 1000/1000 [00:08<00:00, 123.12it/s, 31 steps of size 1.78e-01. acc. prob=0.52]\n",
      "sample: 100%|██████████| 1000/1000 [00:15<00:00, 64.93it/s, 319 steps of size 4.90e-02. acc. prob=0.96]\n",
      "sample: 100%|██████████| 1000/1000 [00:08<00:00, 113.71it/s, 31 steps of size 1.30e-01. acc. prob=0.81]\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[['did','urban','age', 'children', 'contra']].items()}\n",
    "mcmc.run(rng_key, **dat)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "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>mean</th>\n",
       "      <th>sd</th>\n",
       "      <th>hdi_5.5%</th>\n",
       "      <th>hdi_94.5%</th>\n",
       "      <th>mcse_mean</th>\n",
       "      <th>mcse_sd</th>\n",
       "      <th>ess_bulk</th>\n",
       "      <th>ess_tail</th>\n",
       "      <th>r_hat</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>a_bar</th>\n",
       "      <td>-0.72</td>\n",
       "      <td>0.10</td>\n",
       "      <td>-0.88</td>\n",
       "      <td>-0.56</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>423.0</td>\n",
       "      <td>795.0</td>\n",
       "      <td>1.01</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>b_bar</th>\n",
       "      <td>0.70</td>\n",
       "      <td>0.17</td>\n",
       "      <td>0.46</td>\n",
       "      <td>0.99</td>\n",
       "      <td>0.01</td>\n",
       "      <td>0.01</td>\n",
       "      <td>177.0</td>\n",
       "      <td>415.0</td>\n",
       "      <td>1.01</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>bA</th>\n",
       "      <td>-0.26</td>\n",
       "      <td>0.07</td>\n",
       "      <td>-0.36</td>\n",
       "      <td>-0.14</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>412.0</td>\n",
       "      <td>1208.0</td>\n",
       "      <td>1.01</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>bK</th>\n",
       "      <td>0.51</td>\n",
       "      <td>0.07</td>\n",
       "      <td>0.40</td>\n",
       "      <td>0.62</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>1547.0</td>\n",
       "      <td>1391.0</td>\n",
       "      <td>1.01</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "       mean    sd  hdi_5.5%  hdi_94.5%  mcse_mean  mcse_sd  ess_bulk  \\\n",
       "a_bar -0.72  0.10     -0.88      -0.56       0.00     0.00     423.0   \n",
       "b_bar  0.70  0.17      0.46       0.99       0.01     0.01     177.0   \n",
       "bA    -0.26  0.07     -0.36      -0.14       0.00     0.00     412.0   \n",
       "bK     0.51  0.07      0.40       0.62       0.00     0.00    1547.0   \n",
       "\n",
       "       ess_tail  r_hat  \n",
       "a_bar     795.0   1.01  \n",
       "b_bar     415.0   1.01  \n",
       "bA       1208.0   1.01  \n",
       "bK       1391.0   1.01  "
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "samples_22 = az.from_numpyro(mcmc)\n",
    "az.summary(samples_22, var_names=['a_bar','b_bar', 'bA', 'bK'], hdi_prob=.89).round(2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this model, the direct effect of `age` is negative, and much farther from zero than before. The effect of number of children is strong and positive. These results are consistent with the DAG, because they imply that the reason the total effect of age, from `model_21`, is positive is that older individuals also have more kids. Having more kids increases contraception. Being older, controlling for kids, actually makes contraception less likely."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise 3"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> Modify any models from Problem 2 that contained that `children` variable and model the variable now as a monotonic ordered category, like education from the week we did ordered categories. Education in that example had 8 categories. Children here will have fewer (no one in the sample had 8 children). So modify the code appropriately. What do you conclude about the causal influence of each additional child on use of contraception?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To build this model, you need the ordered categorical predictor machinery from the book example. The maximum observed number of kids in the sample is 4. So that means we need three parameters, for three transitions in number of kids. Remember for monotonic ordered variables we have to start the categories at zero. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "def model(did, urban, age, children, contra=None):\n",
    "    # fixed priors\n",
    "    a_bar = numpyro.sample('a_bar', dist.Normal(0, 1))\n",
    "    b_bar = numpyro.sample('b_bar', dist.Normal(0, 0.5))\n",
    "    sigma_ab = numpyro.sample('sigma_ab', dist.Exponential(1), sample_shape=(2,))\n",
    "    rho = numpyro.sample('rho', dist.LKJ(2,2))\n",
    "    bA = numpyro.sample('bA', dist.Normal(0,0.5))\n",
    "    bK = numpyro.sample('bK', dist.Normal(0,0.5))\n",
    "    children_coefs = numpyro.sample('children_coefs', dist.Dirichlet(np.ones(3)))\n",
    "    # adaptive priors\n",
    "    mu = jnp.stack([a_bar,b_bar])\n",
    "    cov = jnp.outer(sigma_ab, sigma_ab)*rho\n",
    "    i, j = len(d.did.unique()), len(d.urban.unique())\n",
    "    ab_did = numpyro.sample('ab_did', dist.MultivariateNormal(mu, cov).expand([i,j]))\n",
    "    # likelihood\n",
    "    delta_j = jnp.pad(children_coefs, (1, 0))\n",
    "    delta_K = jnp.sum(jnp.where(jnp.arange(4) <= children[..., None], delta_j, 0), -1)\n",
    "    a_did, b_did = ab_did[:,:,0], ab_did[:,:,1]\n",
    "    theta = a_did[did,urban] + b_did[did,urban]*urban + bA*age + bK*delta_K\n",
    "    contra_hat = numpyro.sample('contra_hat', dist.Bernoulli(logits=theta), obs=contra)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "sample: 100%|██████████| 1000/1000 [00:22<00:00, 43.51it/s, 31 steps of size 1.52e-01. acc. prob=0.88] \n",
      "sample: 100%|██████████| 1000/1000 [00:06<00:00, 143.78it/s, 31 steps of size 1.23e-01. acc. prob=0.85]\n",
      "sample: 100%|██████████| 1000/1000 [00:06<00:00, 148.84it/s, 3 steps of size 1.37e-01. acc. prob=0.43] \n",
      "sample: 100%|██████████| 1000/1000 [00:09<00:00, 107.73it/s, 31 steps of size 1.27e-01. acc. prob=0.79]\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[['did','urban','age', 'children', 'contra']].items()}\n",
    "mcmc.run(rng_key, **dat)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "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>mean</th>\n",
       "      <th>sd</th>\n",
       "      <th>hdi_5.5%</th>\n",
       "      <th>hdi_94.5%</th>\n",
       "      <th>mcse_mean</th>\n",
       "      <th>mcse_sd</th>\n",
       "      <th>ess_bulk</th>\n",
       "      <th>ess_tail</th>\n",
       "      <th>r_hat</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>a_bar</th>\n",
       "      <td>-0.76</td>\n",
       "      <td>0.11</td>\n",
       "      <td>-0.95</td>\n",
       "      <td>-0.59</td>\n",
       "      <td>0.01</td>\n",
       "      <td>0.00</td>\n",
       "      <td>343.0</td>\n",
       "      <td>626.0</td>\n",
       "      <td>1.03</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>b_bar</th>\n",
       "      <td>0.65</td>\n",
       "      <td>0.17</td>\n",
       "      <td>0.38</td>\n",
       "      <td>0.91</td>\n",
       "      <td>0.01</td>\n",
       "      <td>0.01</td>\n",
       "      <td>190.0</td>\n",
       "      <td>341.0</td>\n",
       "      <td>1.07</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>bA</th>\n",
       "      <td>0.03</td>\n",
       "      <td>0.06</td>\n",
       "      <td>-0.07</td>\n",
       "      <td>0.12</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>1046.0</td>\n",
       "      <td>1352.0</td>\n",
       "      <td>1.03</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>bK</th>\n",
       "      <td>0.48</td>\n",
       "      <td>0.30</td>\n",
       "      <td>0.04</td>\n",
       "      <td>0.97</td>\n",
       "      <td>0.01</td>\n",
       "      <td>0.01</td>\n",
       "      <td>968.0</td>\n",
       "      <td>991.0</td>\n",
       "      <td>1.06</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>children_coefs[0]</th>\n",
       "      <td>0.44</td>\n",
       "      <td>0.21</td>\n",
       "      <td>0.10</td>\n",
       "      <td>0.80</td>\n",
       "      <td>0.02</td>\n",
       "      <td>0.01</td>\n",
       "      <td>244.0</td>\n",
       "      <td>694.0</td>\n",
       "      <td>1.02</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>children_coefs[1]</th>\n",
       "      <td>0.28</td>\n",
       "      <td>0.19</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.54</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>1474.0</td>\n",
       "      <td>1101.0</td>\n",
       "      <td>1.06</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>children_coefs[2]</th>\n",
       "      <td>0.27</td>\n",
       "      <td>0.20</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.57</td>\n",
       "      <td>0.02</td>\n",
       "      <td>0.01</td>\n",
       "      <td>190.0</td>\n",
       "      <td>1143.0</td>\n",
       "      <td>1.03</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                   mean    sd  hdi_5.5%  hdi_94.5%  mcse_mean  mcse_sd  \\\n",
       "a_bar             -0.76  0.11     -0.95      -0.59       0.01     0.00   \n",
       "b_bar              0.65  0.17      0.38       0.91       0.01     0.01   \n",
       "bA                 0.03  0.06     -0.07       0.12       0.00     0.00   \n",
       "bK                 0.48  0.30      0.04       0.97       0.01     0.01   \n",
       "children_coefs[0]  0.44  0.21      0.10       0.80       0.02     0.01   \n",
       "children_coefs[1]  0.28  0.19      0.00       0.54       0.00     0.00   \n",
       "children_coefs[2]  0.27  0.20      0.00       0.57       0.02     0.01   \n",
       "\n",
       "                   ess_bulk  ess_tail  r_hat  \n",
       "a_bar                 343.0     626.0   1.03  \n",
       "b_bar                 190.0     341.0   1.07  \n",
       "bA                   1046.0    1352.0   1.03  \n",
       "bK                    968.0     991.0   1.06  \n",
       "children_coefs[0]     244.0     694.0   1.02  \n",
       "children_coefs[1]    1474.0    1101.0   1.06  \n",
       "children_coefs[2]     190.0    1143.0   1.03  "
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "samples_31 = az.from_numpyro(mcmc)\n",
    "az.summary(samples_31, var_names=['a_bar','b_bar', 'bA', 'bK', 'children_coefs'], hdi_prob=.89).round(2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So the general effects are the same: `age` reduces use and kids increase it (_parameter `bA`_). `children_coefs[0]` is the transition from 1 to 2 kids. It is much larger than the other two parameters. So most of the influence of kids on contraception comes from having a second child."
   ]
  }
 ],
 "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.9.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}