chapter10.Rmd

---
title: "chapter10"
author: "Scott Spencer"
date: "8/25/2018"
output: html_document
---

```{r setup, include=FALSE}
knitr::opts_chunk$set(echo = TRUE, 
                      warning = FALSE, message = FALSE, error = FALSE)
library(dplyr); library(tidyr); library(rstan); library(skimr); library(ggplot2); library(ggthemes)
theme_set(theme_tufte(base_family = 'sans'))
```

The code below is meant as a directly-in-Stan translation of the examples in Chapter 10 of McElreath's *Statistical Rethinking*.

## 10.1 binomial regression

### 10.1.1 Logistic regression: Prosocial chimpanzees

Load and review data.

```{r}
data('chimpanzees', package = 'rethinking')
d <- chimpanzees; rm(chimpanzees)
skim(d)
```

Fit first model.

```{stan output.var="m10_1"}
data {
  int N;
  int<lower=0, upper=1> L[N];  // pulled left lever
}
parameters {
  real a;
}
model {
  vector[N] p;
  target += normal_lpdf(a | 0, 10);
  for (i in 1:N) p[i] = a;
  target += binomial_logit_lpmf(L | 1, p);
}
generated quantities {
  vector[N] log_lik;
  {
  vector[N] p;
  for (n in 1:N) {
    p[n] = a;
    log_lik[n] = bernoulli_logit_lpmf(L[n] | p[n]);
  }
  }
}

```

Organize data and sample from the model.

```{r}
dat <- list(N = NROW(d), L = d$pulled_left)
fit10_1 <- sampling(m10_1, data = dat, iter = 1000, chains = 2, cores = 2)
print(fit10_1, include = F, pars = c("log_lik"), probs = c(.1, .5, .9))
```

Fit second model:

```{stan output.var="m10_2"}
data {
  int N;
  int<lower=0, upper=1> L[N]; // pulled left
  vector[N] P; // pro-social left
}
parameters {
  real a;
  real bp;
}
model {
  vector[N] p;
  target += normal_lpdf(a | 0, 10);
  target += normal_lpdf(bp | 0, 10);
  
  for (i in 1:N) p[i] = a + bp * P[i];
  
  target += binomial_logit_lpmf(L | 1, p);
}
generated quantities {
  vector[N] log_lik;
  {
  vector[N] p;
  for(n in 1:N) {
    p[n] = a + bp * P[n];
    log_lik[n] = bernoulli_logit_lpmf(L[n] | p[n]);
  }
  }
}

```

```{r}
dat <- list(N = NROW(d), L = d$pulled_left, P = d$prosoc_left)
fit10_2 <- sampling(m10_2, data = dat, iter = 1000, chains = 2, cores = 2)
print(fit10_2, include = F, pars = "log_lik", probs = c(.1, .5, .9))
```

Fit third model:

```{stan output.var="m10_3"}
data {
  int N;
  int<lower=0, upper=1> L[N];
  vector[N] P;
  vector[N] C;
}
parameters {
  real a;
  real bp;
  real bpc;
}
model {
  vector[N] p;
  target += normal_lpdf(a | 0, 10);
  target += normal_lpdf(bp | 0, 10);
  target += normal_lpdf(bpc | 0, 10);
  for (i in 1:N) p[i] = a + (bp + bpc * C[i]) * P[i];
  target += binomial_logit_lpmf(L | 1, p);
}
generated quantities {
  vector[N] log_lik;
  {
  vector[N] p;
  for(n in 1:N) {
    p[n] = a + (bp + bpc * C[n]) * P[n];
    log_lik[n] = binomial_logit_lpmf(L[n] | 1, p[n]);
  }
  }
}

```


```{r}
dat <- list(N = NROW(d), L = d$pulled_left, P = d$prosoc_left, C = d$condition)
fit10_3 <- sampling(m10_3, data = dat, iter = 1000, chains = 2, cores = 2)
print(fit10_3, include = F, pars = "log_lik", probs = c(.1, .5, .9))
```

Compare models

```{r}
# compare models
library(loo)

log_lik_10_1 <- extract_log_lik(fit10_1, merge_chains = FALSE)
r_eff        <- relative_eff(exp(log_lik_10_1))
loo_10_1     <- loo(log_lik_10_1, r_eff = r_eff, cores = 2)

log_lik_10_2 <- extract_log_lik(fit10_2, merge_chains = FALSE)
r_eff        <- relative_eff(exp(log_lik_10_2))
loo_10_2     <- loo(log_lik_10_2, r_eff = r_eff, cores = 2)

log_lik_10_3 <- extract_log_lik(fit10_3, merge_chains = FALSE)
r_eff        <- relative_eff(exp(log_lik_10_3))
loo_10_3     <- loo(log_lik_10_3, r_eff = r_eff, cores = 2)

loo::compare(loo_10_1, loo_10_2, loo_10_3)
```


Create model averaged predictions


```{r}
# setup new data
d_p <- data.frame(P = c(0, 1, 0, 1),
                  C = c(0, 0, 1, 1))

# get posterior of parameters
post10_1 <- as.data.frame(fit10_1) %>% select(a)
post10_2 <- as.data.frame(fit10_2) %>% select(a, bp)
post10_3 <- as.data.frame(fit10_3) %>% select(a, bp, bpc)

# posterior predictors for model 1
f_mu <- function(obs) plogis(post10_1$a)
mu <- mapply(f_mu, obs = 1:NROW(d_p))
mu.mean10_1 <- apply(mu, 2, mean)
mu.hpdi10_1 <- apply(mu, 2, HDInterval::hdi)

# posterior predictors for model 2
f_mu <- function(P) plogis(post10_2$a + post10_2$bp * P)
mu10_2 <- mapply(f_mu, P = d_p$P)
mu.mean10_2 <- apply(mu10_2, 2, mean)
mu.hpdi10_2 <- apply(mu10_2, 2, HDInterval::hdi)

# posterior predictors for model 3
f_mu <- function(C, P) plogis(post10_3$a + (post10_3$bp + post10_3$bpc * C) * P )
mu10_3 <- mapply(f_mu, C = d_p$C, P = d_p$P)
mu.mean10_3 <- apply(mu10_3, 2, mean)
mu.hpdi10_3 <- apply(mu10_3, 2, HDInterval::hdi)

# get model weights from loo above
m10_weights <- loo_model_weights(list(loo_10_1, loo_10_2, loo_10_3), method = "pseudobma", BB = F)

# create ensemble or model average for paramters stats
ensemble10_mean <- m10_weights[1] * mu.mean10_1 + m10_weights[2] * mu.mean10_2 + m10_weights[3] * mu.mean10_3
ensemble10_HPDI <- m10_weights[1] * mu.hpdi10_1 + m10_weights[2] * mu.hpdi10_2 + m10_weights[3] * mu.hpdi10_3

# organize data
d_p <- 
  d_p %>%
  mutate(action = row_number(),
         ens_mean = ensemble10_mean,
         ens_HPDI_l = ensemble10_HPDI[1,],
         ens_HPDI_h = ensemble10_HPDI[2,])
```


plot information

Figure 10.2

```{r}

d_ <- 
  d %>% 
  group_by(actor, prosoc_left, condition) %>% 
  summarise(prop_left = sum(pulled_left) / n()) %>%
  ungroup %>%
  mutate(action = rep(c(1, 3, 2, 4), times = 7))

ggplot() + 
  geom_line(data = d_,
            aes(x = action, y = prop_left, group = actor), color = 'dodgerblue') +
  geom_ribbon(data = d_p,
              aes(x = action,
                  ymin = ens_HPDI_l,
                  ymax = ens_HPDI_h),
              alpha = .1) +
  geom_line(data = d_p,
            aes(x = action, y = ens_mean)) +
  lims(y = c(0, 1)) +
  scale_x_continuous(labels = c("1" = "0/0", "2" = "1/0", "3" = "0/1", "4" = "1/1")) +
  labs(x = "Prosocial Left / Condition", y = "Proportion pulled left")
```

Include individuals in model.

```{stan output.var="m10_4"}
data {
  int N;
  int<lower=0, upper=1> L[N];
  vector[N] P;
  vector[N] C;
  
  int<lower=1, upper=N> N_chimps;
  int<lower=1, upper=N_chimps> chimp[N];
}
parameters {
  real a_chimp[N_chimps];
  real bp;
  real bpc;
}
model {
  vector[N] p;
  target += normal_lpdf(a_chimp | 0, 10);
  target += normal_lpdf(bp | 0, 10);
  target += normal_lpdf(bpc | 0, 10);
  for (i in 1:N) p[i] = a_chimp[chimp[i]] + (bp + bpc * C[i]) * P[i];
  target += binomial_logit_lpmf(L | 1, p);
}
generated quantities {
  vector[N] log_lik;
  {
  vector[N] p;
  for(n in 1:N) {
    p[n] = a_chimp[chimp[n]] + (bp + bpc * C[n]) * P[n];
    log_lik[n] = binomial_logit_lpmf(L[n] | 1, p[n]);
  }
  }
}

```

Organize data and sample from model.

```{r}
dat <- list(
  N = NROW(d),
  L = d$pulled_left,
  P = d$prosoc_left,
  C = d$condition,
  N_chimps = max(d$actor),
  chimp = d$actor
)

fit10_4 <- sampling(m10_4, data = dat, iter = 1000, chains = 2, cores = 2)
print(fit10_4, include = F, pars = "log_lik", probs = c(.1, .5, .9))
```

Review densities. Figure 10.3

```{r}
post10_4 <- as.data.frame(fit10_4)
ggplot() + 
  geom_density(aes(x = post10_4$`a_chimp[2]`), fill = 'skyblue') +
  labs(x = 'Chimp 2', y = 'density')

```

```{r}
# setup new data
d_p <- data.frame(P = c(0, 1, 0, 1),
                  C = c(0, 0, 1, 1),
                  chimp = rep(1:7, each = 4),
                  action = rep(c(1, 2, 3, 4), times = 7))

# posterior predictors for model 3
f_mu <- function(chimp, C, P) plogis(post10_4[, chimp] + (post10_4$bp + post10_4$bpc * C) * P )
mu10_4 <- mapply(f_mu, chimp = d_p$chimp, C = d_p$C, P = d_p$P)
mu.mean10_4 <- apply(mu10_4, 2, mean)
mu.hpdi10_4 <- apply(mu10_4, 2, HDInterval::hdi)

# organize data
d_p <- d_p %>%
  mutate(prop_left_pred = mu.mean10_4,
         prop_left_hpdi_l = mu.hpdi10_4[1,],
         prop_left_hpdi_h = mu.hpdi10_4[2,])

d_p <- d_p %>% left_join(d_[,c("actor", "action", "prop_left")], 
                         by = c("chimp" = "actor", "action" = "action"))
```

Plot each chimp actual and average prediction info.

Modified Figure 10.4

```{r}
ggplot(d_p) + 
  geom_line(aes(x = action, y = prop_left), color = 'dodgerblue') + 
  geom_ribbon(aes(x = action, ymin = prop_left_hpdi_l, ymax = prop_left_hpdi_h), alpha = .1) + 
  geom_line(aes(x = action, y = prop_left_pred), color = 'black') + 
  facet_wrap(~chimp, nrow = 2) +   
  lims(y = c(0, 1)) +
  scale_x_continuous(labels = c("1" = "0/0", "2" = "1/0", "3" = "0/1", "4" = "1/1")) +
  labs(x = "Prosocial Left / Condition", y = "Proportion pulled left") +
  theme(panel.border = element_rect(colour = "gray90", fill=NA, size=1),
        panel.spacing.x = unit(0, "lines"),
        panel.spacing.y = unit(2, "lines"))
```

### 10.1.2. Aggregated binomial: Chimpanzees again, condensed

Re-code model 10.3 using aggregated data.

```{stan output.var="m10_5"}
data {
  int N;
  int<lower=0> L[N]; // changed to indicate number of pulls
  int<lower=0> O[N]; // changed to indicate opportunities to pull
  vector[N] P;
  vector[N] C;
}
parameters {
  real a;
  real bp;
  real bpc;
}
model {
  vector[N] p;
  target += normal_lpdf(a | 0, 10);
  target += normal_lpdf(bp | 0, 10);
  target += normal_lpdf(bpc | 0, 10);
  for (i in 1:N) p[i] = a + (bp + bpc * C[i]) * P[i];
  target += binomial_logit_lpmf(L | O, p); // changed to use opportunities to pull
}
generated quantities {
  vector[N] log_lik;
  {
  vector[N] p;
  for(n in 1:N) {
    p[n] = a + (bp + bpc * C[n]) * P[n];
    log_lik[n] = binomial_logit_lpmf(L[n] | O[n], p[n]);
  }
  }
}

```

Organize data and sample from model.

```{r}
d1 <- d %>% 
  group_by(actor, prosoc_left, condition) %>% 
  summarise(pulled_left = sum(pulled_left), 
            Opportunities = n()) %>%
  ungroup

dat <- list(
  N = NROW(d1),
  L = d1$pulled_left,
  O = d1$Opportunities,
  P = d1$prosoc_left,
  C = d1$condition
)

fit10_5 <- sampling(m10_5, data = dat, iter = 1000, chains = 2, cores = 2)
print(fit10_5, include = F, pars = "log_lik", probs = c(.1, .5, .9))
```

### 10.1.3. Aggregatedbinomial:Graduateschooladmissions

Load data

```{r}
data('UCBadmit', package = 'rethinking')
d <- UCBadmit; rm(UCBadmit)
```

Two models in Stan, one with male, one without.

```{stan output.var="m10_6"}
data {
  int N;
  int<lower=0, upper=1> male[N];
  int N_admit[N];
  int N_applications[N];
}
parameters {
  real a;
  real bm;
}
model {
  vector[N] p_admit;
  for (i in 1:N) p_admit[i] = a + bm * male[i];
  target += binomial_logit_lpmf(N_admit | N_applications, p_admit);
  target += cauchy_lpdf(a | 0, 5);
  target += normal_lpdf(bm | 0, 5);
}
generated quantities {
  vector[N] log_lik;
  {
  vector[N] p_admit;
  for(n in 1:N) {
    p_admit[n] = a + bm * male[n];
    log_lik[n] = binomial_logit_lpmf(N_admit[n] | N_applications[n], p_admit[n]);
  }
  }
}

```

Organize data and sample from model

```{r}
d <- d %>% mutate(male = as.integer(applicant.gender == 'male'))
dat <- list(
  N = NROW(d),
  N_applications = d$applications,
  N_admit = d$admit,
  male = d$male
)

fit10_6 <- sampling(m10_6, data = dat, iter = 1000, chains = 2, cores = 2)
```

```{stan output.var="m10_7"}
data {
  int N;
  int N_admit[N];
  int N_applications[N];
}
parameters {
  real a;
}
model {
  vector[N] p_admit;
  for (i in 1:N) p_admit[i] = a;
  target += binomial_logit_lpmf(N_admit | N_applications, p_admit);
  target += cauchy_lpdf(a | 0, 5);
}
generated quantities {
  vector[N] log_lik;
  {
  vector[N] p_admit;
  for(n in 1:N) {
    p_admit[n] = a;
    log_lik[n] = binomial_logit_lpmf(N_admit[n] | N_applications[n], p_admit[n]);
  }
  }
}

```

```{r}
dat <- list(
  N = NROW(d),
  N_applications = d$applications,
  N_admit = d$admit
)

fit10_7 <- sampling(m10_7, data = dat, iter = 1000, chains = 2, cores = 2)
```

Compare the two models.

```{r}
# compare models
library(loo)

log_lik_10_6 <- extract_log_lik(fit10_6, merge_chains = FALSE)
r_eff        <- relative_eff(exp(log_lik_10_6))
loo_10_6     <- loo(log_lik_10_6, r_eff = r_eff, cores = 2)
waic10_6 <- waic(log_lik_10_6)

log_lik_10_7 <- extract_log_lik(fit10_7, merge_chains = FALSE)
r_eff        <- relative_eff(exp(log_lik_10_7))
loo_10_7     <- loo(log_lik_10_7, r_eff = r_eff, cores = 2)
waic10_7 <- waic(log_lik_10_7)

loo::compare(waic10_6, waic10_7)
```

Review model.

```{r}
print(fit10_6, pars = c('a', 'bm', 'lp__'), probs = c(.1, .5, .9))
```

Two new models, this time adding department.

```{stan output.var="m10_8"}
data {
  int N;
  int N_admit[N];
  int N_applications[N];
  int N_dept;
  int<lower=1,upper=N_dept> dept[N];
}
parameters {
  real a[N_dept];
}
model {
  vector[N] p_admit;
  for (i in 1:N) p_admit[i] = a[dept[i]];
  target += binomial_logit_lpmf(N_admit | N_applications, p_admit);
  target += normal_lpdf(a | 0, 10);
}
generated quantities {
  vector[N] log_lik;
  {
  vector[N] p_admit;
  for(n in 1:N) {
    p_admit[n] = a[dept[n]];
    log_lik[n] = binomial_logit_lpmf(N_admit[n] | N_applications[n], p_admit[n]);
  }
  }
}

```

Organize data and sample from model

```{r}
dat <- list(
  N = NROW(d),
  N_applications = d$applications,
  N_admit = d$admit,
  N_dept = length(unique(d$dept)),
  dept = as.integer(d$dept)
)

fit10_8 <- sampling(m10_8, data = dat, iter = 1000, chains = 2, cores = 2)
```


```{stan output.var="m10_9"}
data {
  int N;
  int<lower=0,upper=1> male[N];
  int N_admit[N];
  int N_applications[N];
  int N_dept;
  int<lower=1,upper=N_dept> dept[N];
}
parameters {
  real a[N_dept];
  real bm;
}
model {
  vector[N] p_admit;
  for (i in 1:N) p_admit[i] = a[dept[i]] + bm * male[i];
  target += binomial_logit_lpmf(N_admit | N_applications, p_admit);
  target += normal_lpdf(a | 0, 10);
  target += normal_lpdf(bm | 0, 10);
}
generated quantities {
  vector[N] log_lik;
  {
  vector[N] p_admit;
  for(n in 1:N) {
    p_admit[n] = a[dept[n]] + bm * male[n];
    log_lik[n] = binomial_logit_lpmf(N_admit[n] | N_applications[n], p_admit[n]);
  }
  }
}

```

```{r}
dat <- list(
  N = NROW(d),
  N_applications = d$applications,
  N_admit = d$admit,
  male = d$male,
  N_dept = length(unique(d$dept)),
  dept = as.integer(d$dept)
)

fit10_9 <- sampling(m10_9, data = dat, iter = 1000, chains = 2, cores = 2)
```


```{r}
# compare models
library(loo)

log_lik_10_8 <- extract_log_lik(fit10_8, merge_chains = FALSE)
r_eff        <- relative_eff(exp(log_lik_10_8))
loo_10_8     <- loo(log_lik_10_8, r_eff = r_eff, cores = 2)
waic10_8 <- waic(log_lik_10_8)

log_lik_10_9 <- extract_log_lik(fit10_9, merge_chains = FALSE)
r_eff        <- relative_eff(exp(log_lik_10_9))
loo_10_9     <- loo(log_lik_10_9, r_eff = r_eff, cores = 2)
waic10_9 <- waic(log_lik_10_9)

loo::compare(waic10_6, waic10_7, waic10_8, waic10_9)
```

Modified Figure 10.6


```{r}
post10_9 <- as.data.frame(fit10_9)
f_mu <- function(dept, male) plogis(post10_9[,dept] + post10_9$bm * male)
p_admit_hat <- mapply(f_mu, dept = as.integer(d$dept), male = (d$male == "male") )

# get expectation and 89% intervals of the expectation
d <- 
  d %>%
  mutate(p_admit = admit / applications,
         p_hat_mean = colMeans(p_admit_hat),
         p_hat_hpdi_l = apply(p_admit_hat, 2, rethinking::HPDI)[1,],
         p_hat_hpdi_h = apply(p_admit_hat, 2, rethinking::HPDI)[2,])

# get 89% intervals of simulated samples
f_mu <- function(dept, male) rbinom(n = 1e3, size = 18, prob = plogis(post10_9[,dept] + post10_9$bm * male) )
n_admit_sim <- mapply(f_mu, dept = as.integer(d$dept), male = (d$male == 1) )
p_admit_sim <- n_admit_sim / 18

d <- d %>%
  mutate(p_sim_mean = colMeans(p_admit_sim),
         p_sim_hpdi_l = apply(p_admit_sim, 2, rethinking::HPDI)[1,],
         p_sim_hpdi_h = apply(p_admit_sim, 2, rethinking::HPDI)[2,])
```

```{r}
ggplot(d) + 
  facet_wrap(~paste0('Dept ', dept)) +
  geom_point(aes(x = male, y = p_admit), color = 'dodgerblue') + 
  geom_line(aes(x = male, y = p_admit), color = 'dodgerblue') + 
  geom_segment(aes(x = male, xend = male, y = p_hat_hpdi_l, yend = p_hat_hpdi_h)) + 
  geom_point(aes(x = male, y = p_hat_mean), shape = 21, fill = 'white') + 
  geom_point(aes(x = male, y = p_sim_hpdi_l), shape = 3) +
  geom_point(aes(x = male, y = p_sim_hpdi_h), shape = 3) +
  scale_y_continuous(limits = c(0, 1)) +
  scale_x_continuous(limits = c(-0.5, 1.5), breaks = c(0, 1), labels = c("Female", "Male")) +
  theme(panel.border = element_rect(colour = "gray90", fill=NA, size=1),
        panel.spacing.x = unit(-0.5, "mm"),
        panel.spacing.y = unit(2, "lines")) + 
  labs(x = '', y = 'Probability of admission')

```

## 10.2 Poisson regression

### 10.2.1. Example: Oceanic tool complexity

load data.

```{r}
data("Kline", package = "rethinking")
d <- Kline; rm(Kline)
d <- d %>% 
  mutate(log_pop = log(population),
         contact_high = as.integer(contact == "high"))
```

Code model in Stan

```{stan output.var="m10_10"}
data {
  int N;
  int T[N];
  int P[N];
  int C[N];
}
parameters {
  real a;
  real bp;
  real bc;
  real bpc;
}
model {
  vector[N] lambda;
  for(i in 1:N) lambda[i] = a + bp * log(P[i]) + bc * C[i] + bpc * C[i] * log(P[i]);
  target += poisson_log_lpmf(T | lambda);
  target += normal_lpdf(a | 0, 100);
  target += normal_lpdf(bp | 0, 1);
  target += normal_lpdf(bc | 0, 1);
  target += normal_lpdf(bpc | 0, 1);
}
generated quantities {
  vector[N] log_lik;
  {
  vector[N] lambda;
  for(i in 1:N) {
    lambda[i] = a + bp * log(P[i]) + bc * C[i] + bpc * C[i] * log(P[i]);
    log_lik[i] = poisson_log_lpmf(T[i] | lambda[i]);
  }
  }
}

```

Organize data and sample from model.


```{r}
dat = list(N = NROW(d), T = d$total_tools, P = d$population, C = d$contact_high)
fit10_10 <- sampling(m10_10, data = dat, iter = 1000, chains = 2, cores = 2)
```

```{r}
print(fit10_10, include = F, pars = 'log_lik', probs = c(.1, .5, .9))
```

```{r}
post10_10 <- as.data.frame(fit10_10)
post10_10 %>% select(a, bp, bc, bpc) %>% cor() %>% print(digits = 2)
```

Check whether the interaction has an effect.

Figure 10.8

```{r}
lambda_high <- exp(post10_10$a + post10_10$bc + (post10_10$bp + post10_10$bpc) * 8 )
lambda_low <- exp(post10_10$a + post10_10$bp * 8 )

p <- ggplot()
p1 <- p +
  geom_density(aes(lambda_high - lambda_low), fill = 'dodgerblue') +
  geom_vline(xintercept = 0, linetype = 'dashed')

p2 <- p +
  geom_point(aes(post10_10$bc, post10_10$bpc), color = 'dodgerblue', alpha = .1) +
  geom_segment(aes(x = rethinking::HPDI(post10_10$bc)[1],
                   xend = rethinking::HPDI(post10_10$bc)[2],
                   y = -.2, yend = -.2)) +
  geom_segment(aes(x = -2.5, xend = -2.5, 
                   y = rethinking::HPDI(post10_10$bpc)[1],
                   yend = rethinking::HPDI(post10_10$bpc)[2])) + labs(x = 'bc', y = 'bpc')

library(gridExtra)
grid.arrange(p1, p2, nrow = 1)
```

```{r}
sum(lambda_high - lambda_low > 0) / length(lambda_high)
```

Remove interaction

```{stan output.var="m10_11"}
data {
  int N;
  int T[N];
  int P[N];
  int C[N];
}
parameters {
  real a;
  real bp;
  real bc;
}
model {
  vector[N] lambda;
  for(i in 1:N) lambda[i] = a + bp * log(P[i]) + bc * C[i];
  target += poisson_log_lpmf(T | lambda);
  target += normal_lpdf(a | 0, 100);
  target += normal_lpdf(bp | 0, 1);
  target += normal_lpdf(bc | 0, 1);
}
generated quantities {
  vector[N] log_lik;
  {
  vector[N] lambda;
  for(i in 1:N) {
    lambda[i] = a + bp * log(P[i]) + bc * C[i];
    log_lik[i] = poisson_log_lpmf(T[i] | lambda[i]);
  }
  }
}

```

Remove contact

```{stan output.var="m10_12"}
data {
  int N;
  int T[N];
  int P[N];
}
parameters {
  real a;
  real bp;
}
model {
  vector[N] lambda;
  for(i in 1:N) lambda[i] = a + bp * log(P[i]);
  target += poisson_log_lpmf(T | lambda);
  target += normal_lpdf(a | 0, 100);
  target += normal_lpdf(bp | 0, 1);
}
generated quantities {
  vector[N] log_lik;
  {
  vector[N] lambda;
  for(i in 1:N) {
    lambda[i] = a + bp * log(P[i]);
    log_lik[i] = poisson_log_lpmf(T[i] | lambda[i]);
  }
  }
}

```

Remove log population

```{stan output.var="m10_13"}
data {
  int N;
  int T[N];
  int C[N];
}
parameters {
  real a;
  real bc;
}
model {
  vector[N] lambda;
  for(i in 1:N) lambda[i] = a + bc * C[i];
  target += poisson_log_lpmf(T | lambda);
  target += normal_lpdf(a | 0, 100);
  target += normal_lpdf(bc | 0, 1);
}
generated quantities {
  vector[N] log_lik;
  {
  vector[N] lambda;
  for(i in 1:N) {
    lambda[i] = a + bc * C[i];
    log_lik[i] = poisson_log_lpmf(T[i] | lambda[i]);
  }
  }
}

```

Intercept only

```{stan output.var="m10_14"}
data {
  int N;
  int T[N];
}
parameters {
  real a;
}
model {
  vector[N] lambda;
  for(i in 1:N) lambda[i] = a;
  target += poisson_log_lpmf(T | lambda);
  target += normal_lpdf(a | 0, 100);
}
generated quantities {
  vector[N] log_lik;
  {
  vector[N] lambda;
  for(i in 1:N) {
    lambda[i] = a;
    log_lik[i] = poisson_log_lpmf(T[i] | lambda[i]);
  }
  }
}

```

```{r}
fit10_11 <- sampling(m10_11, data = dat, iter = 1000, chains = 2, cores = 2)
```

```{r}
dat = list(N = NROW(d), T = d$total_tools, P = d$population)
fit10_12 <- sampling(m10_12, data = dat, iter = 1000, chains = 2, cores = 2)
```

```{r}
dat = list(N = NROW(d), T = d$total_tools, C = d$contact_high)
fit10_13 <- sampling(m10_13, data = dat, iter = 1000, chains = 2, cores = 2)
```

```{r}
dat = list(N = NROW(d), T = d$total_tools)
fit10_14 <- sampling(m10_14, data = dat, iter = 1000, chains = 2, cores = 2)
```

Compare with waic

```{r}
library(loo)

fit_list <- list(fit10_10, fit10_11, fit10_12, fit10_13, fit10_14)

# extract log likelihoods
ll_list <- lapply(fit_list, extract_log_lik) 

# exponentiate
exp_ll_list <- lapply(ll_list, exp)

# get relative neff
rel_n_eff_list <- lapply(exp_ll_list, relative_eff, chain_id = c(rep(1, 500), rep(2, 500)))

# loo
waic_list <- list()
for(i in 1:5) {
  waic_list[[i]] <- waic(ll_list[[i]], r_eff = rel_n_eff_list[[i]], cores = 4)
}
names(waic_list) <- c('fit10_10', 'fit10_11', 'fit10_12', 'fit10_13', 'fit10_14')

loo::compare(x = waic_list)
```

Get model weights

```{r}
# get model weights from loo above

loo_list <- list()
for(i in 1:5) {
  loo_list[[i]] <- loo(ll_list[[i]], r_eff = rel_n_eff_list[[i]], cores = 4)
}
names(loo_list) <- c('fit10_10', 'fit10_11', 'fit10_12', 'fit10_13', 'fit10_14')

weights <- loo_model_weights(loo_list, method = "pseudobma", BB = F)
weights
```

Plot counterfactual predictions with an ensemble

```{r}
# setup new data from which to predict total tools
nd <- 
  expand.grid(log_pop = seq(6, 13, length.out = 30),
              contact_high = 0:1)

# get posterior parameters of each
post10_11 <- as.data.frame(fit10_11)
post10_12 <- as.data.frame(fit10_12)

# get mu, mu interval, prediction interval for three top models

# model 10_10
f_mu_10_10 <- function(P, C) with(post10_10, 
              a + bp * P + bc * C + bpc * C * P )
mu_10_10 <- mapply(f_mu_10_10, P = nd$log_pop, C = nd$contact_high)

# model 10_11

f_mu_10_11 <- function(P, C) with(post10_11, 
              a + bp * P + bc * C )
mu_10_11 <- mapply(f_mu_10_11, P = nd$log_pop, C = nd$contact_high)

# model 10_12

f_mu_10_12 <- function(P) with(post10_12, 
              a + bp * P )
mu_10_12 <- mapply(f_mu_10_12, P = nd$log_pop)

# create ensemble: weight mu, interval, and pred interval together
mu_ensemble <- weights[1] * exp(mu_10_10) + weights[2] * exp(mu_10_11) + weights[3] * exp(mu_10_12) 
mu_ensemble_mean <- colMeans(mu_ensemble) 
mu_ens_hpdi <- apply(mu_ensemble, 2, rethinking::HPDI)

nd <- nd %>%
  mutate(mu = mu_ensemble_mean,
         mu_hpdi_l = mu_ens_hpdi[1,],
         mu_hpdi_h = mu_ens_hpdi[2,],
         contact_high = as.factor(contact_high))

```

```{r}
ggplot() + 
  geom_point(data = d, aes(x = log_pop, y = total_tools, 
                 color = as.factor(contact_high))) +
  geom_ribbon(data = nd, aes(x = log_pop, 
                  ymin = mu_hpdi_l, ymax = mu_hpdi_h, 
                  group = contact_high, color = NA, fill = contact_high), alpha = 0.1) +
  geom_line(data = nd, aes(x = log_pop, y = mu, 
                group = contact_high, color = contact_high)) +
  scale_color_manual(values = c("black", "dodgerblue"), aesthetics = c("colour", "fill")) +
  theme(legend.position = "") + 
  coord_fixed(ratio = .1, xlim = c(6, 13), ylim = c(0, 70)) +
  scale_x_continuous(breaks = seq(7, 12, by = 1)) + 
  scale_y_continuous(breaks = seq(0, 70, by = 10))
```

### 10.2.2. MCMC islands

Review correlations among posterior of parameters

Figure 10.10 left side

```{r}
library(GGally)
post10_10 %>% select(a, bp, bc, bpc) %>% ggpairs() + theme_tufte(base_family = 'sans') 
```

Two predictors are highly correlated. Centering data may help this.

```{stan output.var="m10_10c"}
data {
  int N;
  int T[N];
  int P[N];
  int C[N];
}
transformed data {
  vector[N] log_P_c;
  for(i in 1:N) log_P_c[i] = log(P[i]) - mean(log(P));
}
parameters {
  real a;
  real bp;
  real bc;
  real bpc;
}
model {
  vector[N] lambda;
  for (i in 1:N) lambda[i] = a + bp * log_P_c[i] + bc * C[i] + bpc * log_P_c[i] * C[i];
  target += poisson_log_lpmf(T | lambda);
  target += normal_lpdf(a| 0, 100);
  target += normal_lpdf(bp | 0, 10);
  target += normal_lpdf(bc | 0, 10);
  target += normal_lpdf(bpc | 0, 10);
}
```

```{r}
dat = list(N = NROW(d), T = d$total_tools, P = d$population, C = d$contact_high)
fit10_10c <- sampling(m10_10c, data = dat, iter = 1000, chains = 2, cores = 2)
```

```{r}
print(fit10_10c, include = F, pars = 'log_lik', probs = c(.1, .5, .9))
```

Centering removed the correlations between parameters:

Figure 10.10 Right side

```{r}
post10_10c <- as.data.frame(fit10_10c)
post10_10c %>% select(a, bp, bc, bpc) %>% ggpairs() + theme_tufte(base_family = 'sans') 
```

### 10.2.3. Example: Exposure and the offset.

simulate dummy data

```{r}
num_days <- 30
y <- rpois(num_days, 1.5)
num_weeks <- 4
y_new <- rpois(num_weeks, .5 * 7)
y_all <- c(y, y_new)
exposure <- c(rep(1, 30), rep(7, 4))
monastery <- c(rep(0, 30), rep(1, 4))
d <- data.frame(y = y_all, days = exposure, monastery = monastery)
d <- d %>% mutate(log_days = log(days))
```

Code the model

```{stan output.var="m10_15"}
data {
  int N;
  int y[N];
  vector[N] log_days;
  int monastery[N];
}
parameters {
  real a;
  real b;
}
model {
  vector[N] lambda;
  for(i in 1:N) lambda[i] = log_days[i] + a + b * monastery[i];
  target += poisson_log_lpmf(y | lambda);
  target += normal_lpdf(a | 0, 100);
  target += normal_lpdf(b | 0, 1);
}
```

Organize data and sample from model.

```{r}
dat <- list(N = NROW(d), y = d$y, log_days = d$log_days, monastery = d$monastery)
fit10_15 <- sampling(m10_15, data = dat, iter = 1000, chains = 2, cores = 2)
```

Summarise the model

```{r}
print(fit10_15, probs = c(.1, .5, .9))
```


```{r}
post10_15 <- as.data.frame(fit10_15)
lambda_old <- with(post10_15, exp(a))
lambda_new <- with(post10_15, exp(a + b))
skim(data.frame(lambda_old, lambda_new))
```

## 10.3 other count regressions

### 10.3.1. Multinomial

#### 10.3.1.1. Explicit multinomial models

Simulate career choices

```{r}
N <- 500
income <- 1:3
score <- 0.5 * income
softmax <- function(x) exp(x) / sum(exp(x))
p <- softmax(score)

career <- rep(NA, N)
for(i in 1:N) career[i] <- sample(1:3, size = 1, prob = p)
```

Code a model in Stan

```{stan output.var="m10_16"}
data{
    int<lower=1> N;
    int career[N];
}
parameters{
    real b;
}
model{
    vector[N] s2;
    vector[N] s3;
    target += normal_lpdf(b | 0 , 5 );

    for(i in 1:N) s2[i] = b * 2;
    for(i in 1:N) s3[i] = b * 3;

    for ( i in 1:N )
    {
        vector[3] theta;
        theta[1] = 0; theta[2] = s2[i]; theta[3] = s3[i]; 
        target += categorical_lpmf(career[i] | softmax(theta) );
    }
}
```

Organized data and sample from model.

```{r}
dat <- list(N = N, career = career)
fit10_16 <- sampling(m10_16, data = dat, iter = 1000, chains = 2, cores = 2)
```

Summarise model

```{r}
print(fit10_16, probs = c(.1, .5, .9))
```

second example

```{r}
N <- 100
family_income <- runif(N)
b <- (1:-1)
career <- rep(NA, N)
for(i in 1:N) {
  score <- 0.5 * (1:3) + b * family_income[i]
  p <- softmax(score)
  career[i] <- sample(1:3, size = 1, prob = p)
}
```

```{stan output.var="m10_17"}
data {
  int N;
  int career[N];
  vector[N] family_income;
}
parameters {
  real a2;
  real a3;
  real b2;
  real b3;
}
model {
    vector[N] s2;
    vector[N] s3;
    target += normal_lpdf(a2 | 0 , 5 );
    target += normal_lpdf(a3 | 0 , 5 );
    target += normal_lpdf(b2 | 0 , 5 );
    target += normal_lpdf(b3 | 0 , 5 );
    
    for(i in 1:N) s2[i] = a2 + b2 * family_income[i];
    for(i in 1:N) s3[i] = a3 + b3 * family_income[i];
    
    for(i in 1:N) {
        vector[3] theta;
        theta[1] = 0; theta[2] = s2[i]; theta[3] = s3[i]; 
        target += categorical_lpmf(career[i] | softmax(theta) );
    }
}

```

Organize data and sample from model.

```{r}
dat <- list(N = N, career = career, family_income = family_income)
fit10_17 <- sampling(m10_17, data = dat, iter = 1000, chains = 2, cores = 2)
```

Summarise the model

```{r}
print(fit10_17, probs = c(.1, .5, .9))
```


### 10.3.2. Geometric

Simulate data

```{r}
N <- 100
x <- runif(N)
y <- stats::rgeom( N , prob = rethinking::logistic( -1 + 2 * x ) )
```

Code the model

```{stan output.var="m10_18"}
data{
  int N;
  vector[N] x;
  vector[N] y;
}
parameters{
  real a;
  real b;
}
model{
  vector[N] p;
  target += normal_lpdf(a | 0 , 10 );
  target += normal_lpdf(b | 0 , 1 );
  p = inv_logit(a + b * x);
  target += bernoulli_lpmf(1 | p) + y * bernoulli_lpmf(0 | p);
}

```

Organize data and sample from model

```{r}
dat <- list(N = N, x = x, y = y)
fit10_18 <- sampling(m10_18, data = dat, iter = 1000, chains = 2, cores = 2)
```

Summarise model

```{r}
print(fit10_18, probs = c(.1, .5, .9))
```

Hmmm, this doesn't match the book.