<jupyter_start><jupyter_text>Rethinking Statistics course in NumPyro - Week 8 Lecture 15: Multilevel Models- [Video](https://www.youtube.com/watch?v=AALYPv5xSos)- [Slides](https://speakerdeck.com/rmcelreath/l15-statistical-rethinking-winter-2019)Lecture 16: Multilevel Models 2- [Video](https://www.youtube.com/watch?v=ZG3Oe35R5sY)- [Slides](https://speakerdeck.com/rmcelreath/l16-statistical-rethinking-winter-2019)[Proposed problems](https://github.com/gbosquechacon/statrethinking_winter2019/blob/master/homework/week08.pdf) and [solutions in R](https://github.com/gbosquechacon/statrethinking_winter2019/blob/master/homework/week08_solutions.pdf) for the exercises of the week.<jupyter_code>import pandas as pd

from jax import random
from jax.scipy.special import expit as logistic
import numpyro
from numpyro.infer import NUTS, MCMC, Predictive
import numpyro.distributions as dist
from numpyro.distributions.transforms import OrderedTransform

import seaborn as sns
import arviz as az
%load_ext watermark
%watermark -n -u -v -iv -w
sns.set_style('whitegrid')
rng_key = random.PRNGKey(0)<jupyter_output><empty_output><jupyter_text>Exercise 1 > Revisit the Reed frog survival data, `reedfrogs`, and add the predation and size treatment variables to the varying intercepts model. Consider models with either predictor alone, both predictors, as well as a model including their interaction. What do you infer about the causal influence of these predictor variables? Also focus on the inferred variation across tanks (the σ across tanks). Explain why it changes as it does across models with different predictors included. Let's get the data.<jupyter_code>d = pd.read_csv('./dat/reedfrogs.csv', header=0, sep=';')
d['tank'] = d.index
d['pred'] = pd.factorize(d['pred'])[0]
d['volume'] = d['size'].replace({'small':0, 'big':1})
d.head(3)<jupyter_output><empty_output><jupyter_text>Now to define a series of models. The first is just the varying intercepts model from the text. Model 1:<jupyter_code>def model(tank, density, surv=None):
    # priors
    a_bar = numpyro.sample('a_bar', dist.Normal(0, 1.5))
    sigma = numpyro.sample('sigma', dist.Exponential(1))
    a = numpyro.sample('a', dist.Normal(a_bar,sigma), sample_shape=tank.shape)
    # likelihood
    logits = a[tank]
    numpyro.sample('surv_hat', dist.Binomial(density, logits=logits), obs=surv)
kernel = NUTS(model)
mcmc = MCMC(kernel, num_warmup=500, num_samples=500, num_chains=4, chain_method='sequential')
dat = {k:v.to_numpy() for k,v in d[['tank', 'density', 'surv']].items()}
mcmc.run(rng_key, **dat)
samples_11 = az.from_numpyro(mcmc)<jupyter_output><empty_output><jupyter_text>Model 2 (pred):<jupyter_code>def model(tank, density, pred, surv=None):
    # priors
    a_bar = numpyro.sample('a_bar', dist.Normal(0, 1.5))
    sigma = numpyro.sample('sigma', dist.Exponential(1))
    a = numpyro.sample('a', dist.Normal(a_bar,sigma), sample_shape=tank.shape)
    bp = numpyro.sample('bp', dist.Normal(-0.5,1))
    # likelihood
    logits = a[tank] + bp*pred
    numpyro.sample('surv_hat', dist.Binomial(density, logits=logits), obs=surv)
kernel = NUTS(model)
mcmc = MCMC(kernel, num_warmup=500, num_samples=500, num_chains=4, chain_method='sequential')
dat = {k:v.to_numpy() for k,v in d[['tank', 'density', 'pred', 'surv']].items()}
mcmc.run(rng_key, **dat)
samples_12 = az.from_numpyro(mcmc)<jupyter_output><empty_output><jupyter_text>Model 3 (size):<jupyter_code>def model(tank, density, volume, surv=None):
    # priors
    a_bar = numpyro.sample('a_bar', dist.Normal(0, 1.5))
    sigma = numpyro.sample('sigma', dist.Exponential(1))
    a = numpyro.sample('a', dist.Normal(a_bar,sigma), sample_shape=tank.shape)
    s = numpyro.sample('s', dist.Normal(0,0.5), sample_shape=d.volume.unique().shape)
    # likelihood
    logits = a[tank] + s[volume]
    numpyro.sample('surv_hat', dist.Binomial(density, logits=logits), obs=surv)
kernel = NUTS(model)
mcmc = MCMC(kernel, num_warmup=500, num_samples=500, num_chains=4, chain_method='sequential')
dat = {k:v.to_numpy() for k,v in d[['tank', 'density', 'volume', 'surv']].items()}
mcmc.run(rng_key, **dat)
samples_13 = az.from_numpyro(mcmc)<jupyter_output><empty_output><jupyter_text>Model 4 (pred + size):<jupyter_code>def model(tank, density, pred, volume, surv=None):
    # priors
    a_bar = numpyro.sample('a_bar', dist.Normal(0, 1.5))
    sigma = numpyro.sample('sigma', dist.Exponential(1))
    a = numpyro.sample('a', dist.Normal(a_bar,sigma), sample_shape=tank.shape)
    bp = numpyro.sample('bp', dist.Normal(-0.5,1))
    s = numpyro.sample('s', dist.Normal(0,0.5), sample_shape=d.volume.unique().shape)
    # likelihood
    logits = a[tank] + bp*pred + s[volume]
    numpyro.sample('surv_hat', dist.Binomial(density, logits=logits), obs=surv)
kernel = NUTS(model)
mcmc = MCMC(kernel, num_warmup=500, num_samples=500, num_chains=4, chain_method='sequential')
dat = {k:v.to_numpy() for k,v in d[['tank', 'density', 'pred', 'volume', 'surv']].items()}
mcmc.run(rng_key, **dat)
samples_14 = az.from_numpyro(mcmc)<jupyter_output><empty_output><jupyter_text>Model 5 (pred + size + interaction):<jupyter_code>def model(tank, density, pred, volume, surv=None):
    # priors
    a_bar = numpyro.sample('a_bar', dist.Normal(0, 1.5))
    sigma = numpyro.sample('sigma', dist.Exponential(1))
    a = numpyro.sample('a', dist.Normal(a_bar,sigma), sample_shape=tank.shape)
    bp = numpyro.sample('bp', dist.Normal(-0.5,1), sample_shape=d.volume.unique().shape)
    s = numpyro.sample('s', dist.Normal(0,0.5), sample_shape=d.volume.unique().shape)
    # likelihood
    logits = a[tank] + bp[volume]*pred + s[volume]
    numpyro.sample('surv_hat', dist.Binomial(density, logits=logits), obs=surv)
kernel = NUTS(model)
mcmc = MCMC(kernel, num_warmup=500, num_samples=500, num_chains=4, chain_method='sequential')
dat = {k:v.to_numpy() for k,v in d[['tank', 'density', 'pred', 'volume', 'surv']].items()}
mcmc.run(rng_key, **dat)
samples_15 = az.from_numpyro(mcmc)<jupyter_output><empty_output><jupyter_text>I coded the interaction model (`samples_5`) using a non-centered parameterization. The interaction itself is done by creating a `bp` parameter for each size value. In this way, the effect of `pred` depends upon `volume`. First let's consider the WAIC scores:<jupyter_code>az.compare({'model_11': samples_11, 
            'model_12': samples_12, 
            'model_13': samples_13, 
            'model_14': samples_14,
            'model_15': samples_15}
           , scale='deviance')<jupyter_output>/Library/Frameworks/Python.framework/Versions/3.9/lib/python3.9/site-packages/arviz/stats/stats.py:145: UserWarning: The default method used to estimate the weights for each model,has changed from BB-pseudo-BMA to stacking
  warnings.warn(
/Library/Frameworks/Python.framework/Versions/3.9/lib/python3.9/site-packages/arviz/stats/stats.py:655: UserWarning: Estimated shape parameter of Pareto distribution is greater than 0.7 for one or more samples. You should consider using a more robust model, this is because importance sampling is less likely to work well if the marginal posterior and LOO posterior are very different. This is more likely to happen with a non-robust model and highly influential observations.
  warnings.warn(
/Library/Frameworks/Python.framework/Versions/3.9/lib/python3.9/site-packages/arviz/stats/stats.py:655: UserWarning: Estimated shape parameter of Pareto distribution is greater than 0.7 for one or more samples. You should consider using a more robust model, this is [...]<jupyter_text>These models are really very similar in expected out-of-sample accuracy. The tank variation is huge. But take a look at the posterior distributions for predation and size. You'll see that predation does seem to matter, as you'd expect. Size matters a lot less. So while predation doesn't explain much of the total variation, there is plenty of evidence that it is a real effect. Remember: We don't select a model using WAIC(or LOO). A predictor can make little difference in total accuracy but still be a real causal effect.Let's look at all the sigma posterior distributions:<jupyter_code>az.plot_forest({
    'model_11': samples_11.posterior['sigma'],
    'model_12': samples_12.posterior['sigma'],
    'model_13': samples_13.posterior['sigma'],
    'model_14': samples_14.posterior['sigma'],
    'model_15': samples_15.posterior['sigma']
}, combined=True, figsize=(9, 3));<jupyter_output><empty_output><jupyter_text>The two models that omit predation, `model_11` and `model_13`, have larger values of `sigma`. This is because predation explains some of the variation among tanks. So when you add it to the model, the variation in the tank intercepts gets smaller.  Exercise 2 > In 1980, a typical Bengali woman could have 5 or more children in her lifetime. By the year 2000, a typical Bengali woman had only 2 or 3. You're going to look at a historical set of data, when contraception was widely available but many families chose not to use it. These data reside in data `bangladesh` and come from the 1988 Bangladesh Fertility Survey. Each row is one of 1934 women. There are six variables, but you can focus on two of them for this practice problem:> 1. `district`: ID number of administrative district each woman resided in2. `use.contraception`: An indicator (0/1) of whether the woman was using contraception<jupyter_code>d = pd.read_csv('./dat/bangladesh.csv', header=0, sep=';')
d['did'] = pd.factorize(d['district'])[0]
d['conc'] = d['use.contraception']
d.head()<jupyter_output><empty_output><jupyter_text>Now for the ordinary fixed effect model:<jupyter_code>def model(did, conc=None):
    # priors
    a =  numpyro.sample('a', dist.Normal(0, 1.5), sample_shape=d.did.unique().shape)
    # likelihood
    logits = a[did]
    numpyro.sample('conc_hat', dist.Bernoulli(logits=logits), obs=conc)

kernel = NUTS(model)
mcmc = MCMC(kernel, num_warmup=500, num_samples=500, num_chains=4, chain_method='sequential')
mcmc.run(random.PRNGKey(0), d.did.values, d.conc.values)

samples_21 = az.from_numpyro(mcmc)
def model(did, conc=None):
    # priors
    a_bar = numpyro.sample('a_bar', dist.Normal(0,1.5))
    sigma = numpyro.sample('sigma', dist.Exponential(1))
    a = numpyro.sample('a', dist.Normal(a_bar,sigma), sample_shape=d.did.unique().shape)
    # likelihood
    logits = a[did]
    numpyro.sample('conc_hat', dist.Bernoulli(logits=logits), obs=conc)

kernel = NUTS(model)
mcmc = MCMC(kernel, num_warmup=500, num_samples=500, num_chains=4, chain_method='sequential')
mcmc.run(rng_key, d.did.values, d.conc.values)

samples_22 = az.from_numpyro(mcmc)<jupyter_output>sample: 100%|██████████| 1000/1000 [00:06<00:00, 151.97it/s, 7 steps of size 4.27e-01. acc. prob=0.88]
sample: 100%|██████████| 1000/1000 [00:02<00:00, 422.27it/s, 31 steps of size 4.29e-01. acc. prob=0.87]
sample: 100%|██████████| 1000/1000 [00:02<00:00, 350.04it/s, 15 steps of size 3.71e-01. acc. prob=0.90]
sample: 100%|██████████| 1000/1000 [00:02<00:00, 413.46it/s, 7 steps of size 4.34e-01. acc. prob=0.86]<jupyter_text>Now let's extract the samples, compute posterior mean probabilities in each district, and plot it all:<jupyter_code>model_21 = pd.DataFrame(logistic(samples_21.posterior['a'].mean(axis=0).mean(axis=0)))
model_22 = pd.DataFrame(logistic(samples_22.posterior['a'].mean(axis=0).mean(axis=0)))
aux = pd.concat([model_21.assign(model='model_21'),
                 model_22.assign(model='model_22')]).rename(columns={0:'prob'}).reset_index()
logistic(samples_22.posterior['a_bar'].mean()) # use it for the red line in the next plot<jupyter_output><empty_output><jupyter_text>The blue points are the fixed estimations. The _orange ones_ points are the varying effects. As you'd expect, they are shrunk towards the mean (the _red_ line). Some are shrunk more than others. The third district from the left shrunk a lot. Let's look at the sample size in each district:<jupyter_code>g = sns.scatterplot(data=aux, x='index', y='prob', hue='model', s=50)
g.set(ylim=(0,1), xlabel='district', ylabel='prob. use contraception')
g.figure.set_size_inches(12,2.5)
g.axhline(0.3679, ls='--', c='r');<jupyter_output><empty_output><jupyter_text>District 3 has only 2 women sampled. So it shrinks a lot. There are couple of other districts, like 49 and 54, that also have very few women sampled. But their fixed estimates aren't as extreme, so they don't shrink as much as district 3 does.All of this is explained by partial pooling, of course.<jupyter_code>pd.DataFrame(d.district.value_counts().sort_index()).T<jupyter_output><empty_output><jupyter_text>Exercise 3 Return to the Trolley data, (`Trolley`), from Chapter 12. Define and fit a varying intercepts model for these data. By this I mean to add an intercept parameter for the individual to the linear model. Cluster the varying intercepts on individual participants, as indicated by the unique values in the id variable. Include action, intention, and contact as before. Compare the varying intercepts model and a model that ignores individuals, using both WAIC/LOO and posterior predictions. What is the impact of individual variation in these data?<jupyter_code>d = pd.read_csv('./dat/Trolley.csv', header=0, sep=';')
elvl = d['edu'].unique()
idx = [7 , 0 , 6 , 4 , 2 , 1, 3, 5]
cat = pd.Categorical(d.edu, categories=list(elvl[idx]), ordered=True)
d['edu_new'] = pd.factorize(cat, sort=True)[0].astype('int')
d['age_std'] = (d['age'] - d['age'].mean())/d['age'].std()
d['id'] = pd.factorize(d['id'])[0]
d.head(3)<jupyter_output><empty_output><jupyter_text>First, let's load the data and re-run the old model from Chapter 12:<jupyter_code>def model(action, contact, intention, response=None):
    # priors
    cutpoints = numpyro.sample('cutpoints', 
                               dist.TransformedDistribution( 
                                   dist.Normal(0, 1.5).expand([6]), 
                                   OrderedTransform()
                                   ),
                               )
    norm = lambda label: numpyro.sample(label, dist.Normal(0,0.5))
    bA = norm('bA')
    bC = norm('bC')
    bI = norm('bI')
    bIA = norm('bIA')
    bIC = norm('bIC')
    # likelihood
    BI = bI + bIA*action + bIC*contact
    phi = bA*action + bC*contact + BI*intention
    response_hat = numpyro.sample('response_hat', dist.OrderedLogistic(phi, cutpoints), obs=response-1)

# posterior sampling
mcmc = MCMC(NUTS(model), num_warmup=500, num_samples=500, num_chains=4, chain_method='sequential')
mcmc.run(rng_key, d.action.values, d.contact.values, d.intention.values, d.response.values)
mcmc.print_summary()
samples_31 = az.from_numpyro(mcmc)<jupyter_output>sample: 100%|██████████| 1000/1000 [00:29<00:00, 34.45it/s, 31 steps of size 2.09e-01. acc. prob=0.95]
sample: 100%|██████████| 1000/1000 [17:08<00:00,  1.03s/it, 1023 steps of size 5.34e-06. acc. prob=0.83]
sample: 100%|██████████| 1000/1000 [00:17<00:00, 56.43it/s, 15 steps of size 2.82e-01. acc. prob=0.93]
sample: 100%|██████████| 1000/1000 [00:18<00:00, 53.22it/s, 23 steps of size 2.20e-01. acc. prob=0.93]<jupyter_text>Now to run the varying intercept model, we need to build a valid individual ID variable. The IDs in the data are long tags, so we can coerce them to integers in many ways. What is important is that the index values go from 1 (_zero in python_) to the number of individuals.<jupyter_code>def model(action, contact, intention, id, response=None):
    # priors
    cutpoints = numpyro.sample('cutpoints', 
                               dist.TransformedDistribution( 
                                   dist.Normal(0, 1.5).expand([6]), 
                                   OrderedTransform()
                                   ),
                               )
    norm = lambda label: numpyro.sample(label, dist.Normal(0,0.5))
    bA = norm('bA')
    bC = norm('bC')
    bI = norm('bI')
    bIA = norm('bIA')
    bIC = norm('bIC')
    sigma = numpyro.sample('sigma', dist.Exponential(1))
    a = numpyro.sample('a', dist.Normal(0, sigma), sample_shape=d.id.unique().shape)
    # likelihood
    BI = bI + bIA*action + bIC*contact
    phi = a[id] + bA*action + bC*contact + BI*intention
    response_hat = numpyro.sample('response_hat', dist.OrderedLogistic(phi, cutpoints), obs=response-1)

# posterior sampling
mcmc = MCMC(NUTS(model), num_warmup=500, num_samples=500, num_chains=4, chain_method='sequential')
mcmc.run(rng_key, d.action.values, d.contact.values, d.intention.values, d.id.values, d.response.values)
samples_32 = az.from_numpyro(mcmc)<jupyter_output>sample: 100%|██████████| 1000/1000 [01:02<00:00, 16.01it/s, 31 steps of size 1.60e-01. acc. prob=0.93]
sample: 100%|██████████| 1000/1000 [00:49<00:00, 20.18it/s, 31 steps of size 1.54e-01. acc. prob=0.92] 
sample: 100%|██████████| 1000/1000 [01:00<00:00, 16.62it/s, 31 steps of size 1.55e-01. acc. prob=0.93]
sample: 100%|██████████| 1000/1000 [00:46<00:00, 21.42it/s, 31 steps of size 1.49e-01. acc. prob=0.91]<jupyter_text>We can begin by comparing the posterior distributions. The original coefficients are:<jupyter_code>az.summary(samples_31, var_names=['bA', 'bC', 'bI', 'bIA', 'bIC'], hdi_prob=.89).round(2)<jupyter_output><empty_output><jupyter_text>And the new ones, having added the individual IDs, are:<jupyter_code>az.summary(samples_32, var_names=['bA', 'bC', 'bI', 'bIA', 'bIC', 'sigma'], hdi_prob=.89).round(2)<jupyter_output><empty_output><jupyter_text>Everything has gotten more negative. This is because there is a lot of individual variation in average rating, look at the distribution for sigma. That is on the $logit$ scale, so that's a lot of variation on the probability scale. That variation in average rating was hiding some of the effect of the treatments. We get more precision by conditioning on individual.The WAIC comparison can also help show how much variation comes from individual differences in average rating:<jupyter_code>az.compare({'model_31': samples_31, 'model_32': samples_32}, ic='waic', scale='deviance')<jupyter_output>/Library/Frameworks/Python.framework/Versions/3.9/lib/python3.9/site-packages/arviz/stats/stats.py:145: UserWarning: The default method used to estimate the weights for each model,has changed from BB-pseudo-BMA to stacking
  warnings.warn(
/Library/Frameworks/Python.framework/Versions/3.9/lib/python3.9/site-packages/arviz/stats/stats.py:1405: UserWarning: For one or more samples the posterior variance of the log predictive densities exceeds 0.4. This could be indication of WAIC starting to fail. 
See http://arxiv.org/abs/1507.04544 for details
  warnings.warn(
/Library/Frameworks/Python.framework/Versions/3.9/lib/python3.9/site-packages/arviz/stats/stats.py:1405: UserWarning: For one or more samples the posterior variance of the log predictive densities exceeds 0.4. This could be indication of WAIC starting to fail. 
See http://arxiv.org/abs/1507.04544 for details
  warnings.warn(
{'fraction_non_alphanumeric': 0.13988048019461632, 'fraction_numerical': 0.042995399016341425, 'mean_word_length': 4.765243902439025, 'pattern_counts': {'":': 0, '<': 88, '<?xml version=': 0, '>': 79, 'https://': 6, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 8, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '20', 'dir': 'jupyter-structured-clean-dedup', 'id': '514354', 'license': 'MIT', 'n_tokens_mistral': 7153, 'n_tokens_neox': 6445, 'n_words': 1723, 'path': '_statrethink_numpyro_w08.ipynb', 'repo_name': 'robom/statrethink-course-numpyro-2019', 'size': '18909'}