chapter13.Rmd

---
title: "Chapter 13"
author: "Scott Spencer"
date: "9/1/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 13 of McElreath's *Statistical Rethinking*.

## 13.1 Varying slopes by construction

### 13.1.1 simulate the population

Simulate data

```{r}
a <- 3.5
b <- (-1)
sigma_a <- 1
sigma_b <- 0.5
rho <- (-0.7)
Mu <- c(a, b)
cov_ab <- sigma_a * sigma_b * rho
```

Setup the covariance matrix

```{r}
# approach 1
Sigma <- matrix( c(sigma_a^2, cov_ab, cov_ab, sigma_b^2), ncol = 2 )

# approach 2
sigmas <- c(sigma_a, sigma_b)
Rho <- matrix(c(1, rho, rho, 1), nrow = 2)
Sigma <- diag(sigmas) %*% Rho %*% diag(sigmas)

```

simulate cafes

```{r}
N_cafes <- 20
library(MASS)
set.seed(5)
vary_effects <- mvrnorm(N_cafes, Mu, Sigma)
```

separate the intercepts and slopes

```{r}
a_cafe <- vary_effects[,1]
b_cafe <- vary_effects[,2]
```

Figure 13.2

```{r}
ggplot() + 
  geom_point(aes(a_cafe, b_cafe), shape = 21, color = 'dodgerblue') +
  stat_ellipse(aes(a_cafe, b_cafe), level = .50, alpha = .1) +
  stat_ellipse(aes(a_cafe, b_cafe), level = .89, alpha = .1) +
  stat_ellipse(aes(a_cafe, b_cafe), level = .97, alpha = .1) +
  labs(x = 'intercepts (a_cafe)', y = 'slopes (b_cafe)')
```

### 13.1.2 simulate observations

```{r}
N_visits <- 10
afternoon <- rep(0:1, N_visits * N_cafes / 2)
cafe_id <- rep(1:N_cafes, each = N_visits)
mu <- a_cafe[cafe_id] + b_cafe[cafe_id] * afternoon
sigma <- 0.5
wait <- rnorm(N_visits * N_cafes, mu, sigma)
d <- data.frame(cafe = cafe_id, afternoon = afternoon, wait = wait)
```

### 13.1.3 the varying slopes model

Code a model in Stan.

```{stan output.var="m13_1"}
data {
  int N;
  int N_cafe;
  int cafe_id[N];
  real W[N];
  int A[N];
}
parameters {
  real a;
  real b;
  vector[N_cafe] a_cafe;
  vector[N_cafe] b_cafe;
  real<lower=0> sigma;
  vector<lower=0>[2] sigma_cafe;
  corr_matrix[2] Rho;
}
transformed parameters {
  vector[2] Mu_ab = [a, b]';
  vector[2] v_a_cafeb_cafe[N_cafe];
  cov_matrix[2] SRS_sigma_cafeRho;
  for ( j in 1:N_cafe ) v_a_cafeb_cafe[j,1:2] = [a_cafe[j], b_cafe[j]]';
  SRS_sigma_cafeRho = quad_form_diag(Rho,sigma_cafe);
}
model {
  vector[N] mu;
  
  //priors
  target += normal_lpdf(a | 0, 10);
  target += normal_lpdf(b | 0, 10);
  target += cauchy_lpdf(sigma | 0 , 2 );
  target += cauchy_lpdf(sigma_cafe | 0 , 2 );
  target += lkj_corr_lpdf(Rho | 2 );
  
  // linear model  
  for(i in 1:N) mu[i] = a_cafe[cafe_id[i]] + b_cafe[cafe_id[i]] * A[i];
  
  target += normal_lpdf(W | mu, sigma);
  target += multi_normal_lpdf(v_a_cafeb_cafe | Mu_ab , SRS_sigma_cafeRho );
}
generated quantities {
  vector[N] log_lik;
  {
  vector[N] mu;
  for(i in 1:N) {
    mu[i] = a_cafe[cafe_id[i]] + b_cafe[cafe_id[i]] * A[i];
    log_lik[i] = normal_lpdf(W[i] | mu[i], sigma);
  }
  }
}

```

Organize data and sample from model.

```{r}
dat <- list(
  N = NROW(d),
  N_cafe = max(d$cafe),
  cafe_id = d$cafe,
  W = d$wait,
  A = d$afternoon
)

fit13_1 <- sampling(m13_1, data = dat, iter = 5000, chains = 2, cores = 2)
```

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

Examine correlation.

Figure 13.4

```{r}
rho_prior <- rethinking::rlkjcorr(5000, 2, 2)

post13_1 <- as.data.frame(fit13_1)

ggplot(post13_1) + 
  geom_density(aes(x = `Rho[1,2]`), color = 'dodgerblue') +
  geom_density(aes(x = rho_prior[, 1, 2]), linetype = 'dashed') +
  annotate('text', x=0.2, y = .75, label = 'prior') +
  annotate('text', x=0, y = 1.2, label = 'posterior', color = 'dodgerblue') +
  labs(x = 'correlation')
```

Compare slope and intercept for regularized varying effects to unpooled estimates.

Figure 13.5 left side

```{r}
# calculate unpooled estimates from the data
a1 <- sapply(1:N_cafes, function(i) mean(wait[cafe_id == i & afternoon == 0]))
b1 <- sapply(1:N_cafes, function(i) mean(wait[cafe_id == i & afternoon == 1])) - a1

# extract posterior means from partially pooled estimates
a2 <- apply(post13_1[,grep('^a_cafe', colnames(post13_1))], 2, mean)
b2 <- apply(post13_1[,grep('^b_cafe', colnames(post13_1))], 2, mean)

df <- data.frame(a1, b1, a2, b2)

# plot both
ggplot(df) + theme_tufte(base_family = 'sans') +
  stat_ellipse(aes(x = a1, y = b1), level = .1, alpha = .2) +
  stat_ellipse(aes(x = a1, y = b1), level = .3, alpha = .2) +
  stat_ellipse(aes(x = a1, y = b1), level = .5, alpha = .2) +
  stat_ellipse(aes(x = a1, y = b1), level = .8, alpha = .2) +
  stat_ellipse(aes(x = a1, y = b1), level = .99, alpha = .2) +
  geom_point(aes(x = a1, y = b1), color = 'dodgerblue') +
  geom_point(aes(x = a2, y = b2), shape = 21) +
  geom_segment(aes(x = a1, xend = a2, y = b1, yend = b2)) +
  labs(x = 'intercept', y = 'slope')
```

Compare wait times for regularized varying effects to unpooled estimates.

Figure 13.5 Right side

```{r}
df <- df %>% 
  mutate(w_morning1 = a1,
         w_afternoon1 = a1 + b1,
         w_morning2 = a2, 
         w_afternoon2 = a2 + b2)

ggplot(df) + theme_tufte(base_family = 'sans') +
  stat_ellipse(aes(x = w_morning1, y = w_afternoon1), level = .1, alpha = .2) +
  stat_ellipse(aes(x = w_morning1, y = w_afternoon1), level = .3, alpha = .2) +
  stat_ellipse(aes(x = w_morning1, y = w_afternoon1), level = .5, alpha = .2) +
  stat_ellipse(aes(x = w_morning1, y = w_afternoon1), level = .8, alpha = .2) +
  stat_ellipse(aes(x = w_morning1, y = w_afternoon1), level = .99, alpha = .2) +
  geom_point(aes(x = w_morning1, y = w_afternoon1), color = 'dodgerblue') +
  geom_point(aes(x = a2, y = w_afternoon2), shape = 21) +
  geom_segment(aes(x = w_morning1, xend = w_morning2, y = w_afternoon1, yend = w_afternoon2)) +
  geom_line(data = data.frame(x = seq(6)), aes(x = x, y = x), linetype = 'dashed') +
  labs(x = 'morning wait (min)', y = 'afternoon wait (min)')
```

## 13.2 example: admissions decision and gender

Load data.

```{r}
data('UCBadmit', package = 'rethinking')
d <- UCBadmit; rm(UCBadmit)
d <- d %>% mutate(male = applicant.gender == 'male',
                  dept_id = as.integer(dept))
```

### 13.2.1. Varying intercepts

Code model in Stan.

```{stan output.var="m13_2"}
data {
  int N;
  int n[N];
  int m[N];
  int A[N];
  int N_depts;
  int dept[N];
}
parameters {
  real alpha;
  real beta;
  real<lower=0> sigma;
  vector[N_depts] a_dept;
}
model {
  vector[N] p;
  target += normal_lpdf(a_dept | alpha, sigma);
  target += normal_lpdf(alpha | 0, 10);
  target += normal_lpdf(beta | 0, 1);
  target += cauchy_lpdf(sigma | 0, 2);
  for(i in 1:N) p[i] = a_dept[dept[i]] + beta * m[i];
  target += binomial_logit_lpmf(A | n, p);
}
generated quantities {
  vector[N] log_lik;
  {
  vector[N] p;
  for (i in 1:N) {
    p[i] = a_dept[dept[i]] + beta * m[i];
    log_lik[i] = binomial_logit_lpmf(A[i] | n[i], p[i]);
  }
  }
}

```

Organize data and sample from model.

```{r}
dat <- list(
  N = NROW(d),
  n = d$applications,
  m = d$male,
  A = d$admit,
  N_depts = max(d$dept_id),
  dept = d$dept_id
)

fit13_2 <- sampling(m13_2, data = dat, iter = 1000, chains = 2, cores = 2)
```

Summarize the model

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

### 13.2.2 varying effects of being male

Code model in Stan.

```{stan output.var="m13_3"}
data {
  int N;
  int n[N];
  real m[N];
  int A[N];
  int N_depts;
  int dept[N];
}
parameters {
  real a;
  real b;
  vector[N_depts] a_dept;
  vector[N_depts] b_dept;
  vector<lower=0>[2] sigma_dept;
  corr_matrix[2] rho;
}
transformed parameters {
  vector[2] mu_ab;
  cov_matrix[2] Sigma;
  vector[2] v_a_deptb_dept[N_depts];
  
  for(j in 1:N_depts) v_a_deptb_dept[j, 1:2] = [a_dept[j], b_dept[j]]';
  
  mu_ab = [a, b]';
  Sigma = quad_form_diag(rho, sigma_dept);
}
model {
  vector[N] p;
  
  // priors
  target += normal_lpdf(a | 0, 10);
  target += normal_lpdf(b | 0, 1);
  target += cauchy_lpdf(sigma_dept | 0, 2);
  target += lkj_corr_lpdf(rho | 2);
  
  for (i in 1:N) p[i] = a_dept[dept[i]] + b_dept[dept[i]] * m[i];

  target += binomial_logit_lpmf(A | n, p);
  target += multi_normal_lpdf(v_a_deptb_dept | mu_ab, Sigma);
}
generated quantities {
  vector[N] log_lik;
  {
  vector[N] p;
  for(i in 1:N) {
    p[i] = a_dept[dept[i]] + b_dept[dept[i]] * m[i];
    log_lik[i] = binomial_logit_lpmf(A[i] | n[i], p[i]);
  }
  }
}

```

Sample from model.

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

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

13.2.3. Shrinkage

Figure 13.6 left side

```{r}
post13_3 <- as.data.frame(fit13_3)
ggplot() + 
  geom_density(aes(post13_3$`Sigma[1,2]`)) +
  labs(x = 'correlation')
```

Figure 13.6 Right side

```{r}
a1 <- d %>% filter(male == 0) %>% group_by(dept_id) %>% summarise(avg = admit/applications) %>% .$avg
b1 <- d %>% filter(male == 1) %>% group_by(dept_id) %>% summarise(avg = admit/applications) %>% .$avg - a1

a2 <- apply(post13_3[, grep('^a_dept', colnames(post13_3))], 2, mean) 
b2 <- apply(post13_3[, grep('^b_dept', colnames(post13_3))], 2, mean) 
a2 <- a2 %>% plogis
df <- data.frame(a1, b1, a2, b2)

ggplot(df) + theme_tufte(base_family = 'sans') +
  stat_ellipse(aes(x = a1, y = b1), level = .1, alpha = .2) +
  stat_ellipse(aes(x = a1, y = b1), level = .3, alpha = .2) +
  stat_ellipse(aes(x = a1, y = b1), level = .5, alpha = .2) +
  stat_ellipse(aes(x = a1, y = b1), level = .7, alpha = .2) +
  stat_ellipse(aes(x = a1, y = b1), level = .97, alpha = .2) +
  geom_point(aes(a1, b1), color = 'dodgerblue') +
  geom_point(aes(a2, b2), shape = 21) +
  labs(x = 'intercept (a_dept)', y = 'slope (b_dept)')
```

### 13.2.4 Model comparisons

Code model without gender.

```{stan output.var="m13_4"}
data {
  int N;
  int N_depts;
  int A[N];
  int n[N];
  int dept[N];
}
parameters {
  vector[N_depts] a_dept;
  real a;
  real<lower=0> sigma_dept;
}
model {
  vector[N] p;
  target += normal_lpdf(a_dept[dept] | a, sigma_dept);
  target += normal_lpdf(a | 0, 10);
  target += cauchy_lpdf(sigma_dept | 0, 2);
  p = a_dept[dept];
  target += binomial_logit_lpmf(A | n, p);
}
generated quantities {
  vector[N] log_lik;
  {
  vector[N] p;
  for(i in 1:N) {
    p = a_dept[dept];
    log_lik[i] = binomial_logit_lpmf(A[i] | n[i], p[i]);
  }
  }
}

```

Sample from model.

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

Compare models.

```{r}
# extract log likelihoods
library(loo)
ll13_2 <- extract_log_lik(fit13_2)
ll13_3 <- extract_log_lik(fit13_3)
ll13_4 <- extract_log_lik(fit13_4)

# calculate reff
neff13_2 <- relative_eff(ll13_2, chain_id = c(rep(1, 500), rep(2, 500)), cores =2)
neff13_3 <- relative_eff(ll13_3, chain_id = c(rep(1, 500), rep(2, 500)), cores =2)
neff13_4 <- relative_eff(ll13_4, chain_id = c(rep(1, 500), rep(2, 500)), cores =2)

# calculate waics
waic13_2 <- waic(ll13_2, r_eff = reff13_2, cores = 2)
waic13_3 <- waic(ll13_3, r_eff = reff13_3, cores = 2)
waic13_4 <- waic(ll13_4, r_eff = reff13_4, cores = 2)

# compare
compare(waic13_2, waic13_3, waic13_4)
```

## 13.3 Example: cross-classified chimpanzees with varying slopes

load the data.

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

Code the model in Stan.

```{stan output.var="m13_6"}
data {
  int N;
  int N_actors;
  int N_blocks;
  int L[N]; // pulled left
  int C[N]; // condition
  int P[N]; // prosocial left
  int actor_id[N];
  int block_id[N];
}
parameters {
  real a;
  real bp;
  real bpc;
  vector[N_actors] a_actor;
  vector[N_actors] bp_actor;
  vector[N_actors] bpc_actor;
  
  vector[N_blocks] a_block;
  vector[N_blocks] bp_block;
  vector[N_blocks] bpc_block;
  
  vector<lower=0>[3] sigma_block;
  vector<lower=0>[3] sigma_actor;
  corr_matrix[3] Rho_actor;
  corr_matrix[3] Rho_block;
}
transformed parameters {
  vector[3] v_block[N_blocks];
  cov_matrix[3] SRS_block;
  vector[3] v_actor[N_actors];
  cov_matrix[3] SRS_actor;
  
  for(j in 1:N_blocks) v_block[j] = [a_block[j], bp_block[j], bpc_block[j]]';
  SRS_block = quad_form_diag(Rho_block, sigma_block);
  
  for(j in 1:N_actors) v_actor[j] = [a_actor[j], bp_actor[j], bpc_actor[j]]';
  SRS_actor = quad_form_diag(Rho_actor, sigma_actor);
}
model {
  vector[N] p;
  vector[N] A;
  vector[N] Bp;
  vector[N] Bpc;
  
  target += multi_normal_lpdf(v_block | rep_vector(0, 3), SRS_block);
  target += multi_normal_lpdf(v_actor | rep_vector(0, 3), SRS_actor);
  
  // linear models
  A = a + a_actor[actor_id] + a_block[block_id];
  Bp = bp + bp_actor[actor_id] + bp_block[block_id];
  Bpc = bpc + bpc_actor[actor_id] + bpc_block[block_id];
  for (i in 1:N) p[i] = A[i] + (Bp[i] + Bpc[i] * C[i]) * P[i];
  
  // likelihood
  target += binomial_logit_lpmf(L | 1, p);
}
generated quantities {
  vector[N] log_lik;
  {
  vector[N] p;
  vector[N] A;
  vector[N] Bp;
  vector[N] Bpc;
  
  A = a + a_actor[actor_id] + a_block[block_id];
  Bp = bp + bp_actor[actor_id] + bp_block[block_id];
  Bpc = bpc + bpc_actor[actor_id] + bpc_block[block_id];
  for (i in 1:N) {
    p[i] = A[i] + (Bp[i] + Bpc[i] * C[i]) * P[i];
    log_lik[i] = binomial_logit_lpmf(L[i] | 1, p[i]);
  }
  }
}

```

Organize data and sample from model.

```{r}
dat <- list(
  N = NROW(d),
  N_actors = max(d$actor),
  N_blocks = max(d$block),
  L = d$pulled_left,
  C = d$condition,
  P = d$prosoc_left,
  actor_id = d$actor,
  block_id = d$block
)

fit13_6 <- sampling(m13_6, data = dat, iter = 1000, chains = 2, cores = 2)
```

```{r, message = TRUE}
check_hmc_diagnostics(fit13_6)
```

Code model using non-centered parameterization.

```{stan output.var="m13_6NC1"}
data {
  int N;
  int N_actors;
  int N_blocks;
  int L[N];
  int C[N];
  int P[N];
  int actor_id[N];
  int block_id[N];
}
parameters {
  real a;
  real bp;
  real bpc;
  vector[N_actors] za_actor;
  vector[N_actors] zbp_actor;
  vector[N_actors] zbpc_actor;
  
  vector[N_blocks] za_block;
  vector[N_blocks] zbp_block;
  vector[N_blocks] zbpc_block;
  
  vector<lower=0>[3] sigma_block;
  vector<lower=0>[3] sigma_actor;
  corr_matrix[3] Rho_actor;
  corr_matrix[3] Rho_block;
}
transformed parameters {
  vector[3] v_block[N_blocks];
  vector[3] v_actor[N_actors];

  for(j in 1:N_blocks) v_block[j] = [za_block[j], zbp_block[j], zbpc_block[j]]';
  for(j in 1:N_actors) v_actor[j] = [za_actor[j], zbp_actor[j], zbpc_actor[j]]';
}
model {
  vector[N] p;
  vector[N] A;
  vector[N] Bp;
  vector[N] Bpc;
  
  // priors
  target += lkj_corr_lpdf(Rho_block | 4);
  target += lkj_corr_lpdf(Rho_actor | 4);
  target += cauchy_lpdf(sigma_block | 0, 2);
  target += cauchy_lpdf(sigma_actor | 0, 2);
  target += normal_lpdf(a | 0, 1);
  target += normal_lpdf(bp | 0, 1);
  target += normal_lpdf(bpc | 0, 1);
  
  target += multi_normal_lpdf(v_block | rep_vector(0, 3), Rho_block);
  target += multi_normal_lpdf(v_actor | rep_vector(0, 3), Rho_actor);
  
  // linear models
  A = a + za_actor[actor_id] * sigma_actor[1] + za_block[block_id] * sigma_block[1];
  Bp = bp + zbp_actor[actor_id] * sigma_actor[2] + zbp_block[block_id] * sigma_block[2];
  Bpc = bpc + zbpc_actor[actor_id] * sigma_actor[3] + zbpc_block[block_id] * sigma_block[3];
  for (i in 1:N) p[i] = A[i] + (Bp[i] + Bpc[i] * C[i]) * P[i];
  
  // likelihood
  target += binomial_logit_lpmf(L | 1, p);
}
generated quantities {
  vector[N] log_lik;
  {
  vector[N] p;
  vector[N] A;
  vector[N] Bp;
  vector[N] Bpc;
  
  A = a + za_actor[actor_id] * sigma_actor[1] + za_block[block_id] * sigma_block[1];
  Bp = bp + zbp_actor[actor_id] * sigma_actor[2] + zbp_block[block_id] * sigma_block[2];
  Bpc = bpc + zbpc_actor[actor_id] * sigma_actor[3] + zbpc_block[block_id] * sigma_block[3];
  for (i in 1:N) {
    p[i] = A[i] + (Bp[i] + Bpc[i] * C[i]) * P[i];
    log_lik[i] = binomial_logit_lpmf(L[i] | 1, p[i]);
  }
  }
}

```

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


```{r, message = TRUE}
check_hmc_diagnostics(fit13_6NC1)
```


```{r}
print(fit13_6NC1, include = T, pars = c('sigma_actor', 'sigma_block'))
```

Cholesky decomposition version.

```{stan output.var="m13_6NC"}
data{
    int<lower=1> N;

    int L[N];
    int C[N];
    int P[N];
    
    int<lower=1> N_actors;
    int actor_id[N];
    int<lower=1> N_blocks;
    int block_id[N];
}
parameters{
    real a;
    real bp;
    real bpc;
    
    vector<lower=0>[3] sigma_actor;
    vector<lower=0>[3] sigma_block;

    matrix[3,N_actors] z_actor;
    cholesky_factor_corr[3] L_Rho_actor;

    matrix[3,N_blocks] z_block;
    cholesky_factor_corr[3] L_Rho_block;
}
transformed parameters{
    matrix[N_actors,3] v_actor;
    vector[N_actors] a_actor;
    vector[N_actors] bp_actor;
    vector[N_actors] bpc_actor;
    matrix[3,3] Rho_actor;

    matrix[N_blocks,3] v_block;
    vector[N_blocks] a_block;
    vector[N_blocks] bp_block;
    vector[N_blocks] bpc_block;
    matrix[3,3] Rho_block;
    
    v_actor = (diag_pre_multiply(sigma_actor,L_Rho_actor) * z_actor)';  
    a_actor   = v_actor[,1];
    bp_actor  = v_actor[,2];
    bpc_actor = v_actor[,3];
    Rho_actor = L_Rho_actor * L_Rho_actor';

    v_block = (diag_pre_multiply(sigma_block,L_Rho_block) * z_block)';
    a_block   = v_block[, 1];
    bp_block  = v_block[, 2];
    bpc_block = v_block[, 3];
    Rho_block = L_Rho_block * L_Rho_block';
}
model{
    vector[N] A;
    vector[N] BP;
    vector[N] BPC;
    vector[N] p;

    // priors
    target += lkj_corr_cholesky_lpdf(L_Rho_actor | 4 );
    target += lkj_corr_cholesky_lpdf(L_Rho_block | 4 );
    target += cauchy_lpdf(sigma_actor | 0 , 2 );
    target += cauchy_lpdf(sigma_block | 0 , 2 );
    target += normal_lpdf(a | 0 , 1 );
    target += normal_lpdf(bp |0 , 1 );
    target += normal_lpdf(bpc | 0 , 1 );
    target += normal_lpdf(to_vector(z_block) | 0 , 1 ); 
    target += normal_lpdf(to_vector(z_actor) | 0 , 1 );
    
    // linear models
    for ( i in 1:N ) {
      A[i] = a + a_actor[actor_id[i]] + a_block[block_id[i]];
      BP[i] = bp + bp_actor[actor_id[i]] + bp_block[block_id[i]];
      BPC[i] = bpc + bpc_actor[actor_id[i]] + bpc_block[block_id[i]];
      p[i] = A[i] + (BP[i] + BPC[i] * C[i]) * P[i];    
    }
    
    // likelihood
    target += binomial_logit_lpmf(L | 1 , p );

}

generated quantities{
    vector[N] log_lik;
  {
    vector[N] A;
    vector[N] BP;
    vector[N] BPC;
    vector[N] p;
    for ( i in 1:N ) {
        A[i] = a + a_actor[actor_id[i]] + a_block[block_id[i]];
        BP[i] = bp + bp_actor[actor_id[i]] + bp_block[block_id[i]];
        BPC[i] = bpc + bpc_actor[actor_id[i]] + bpc_block[block_id[i]];
        p[i] = A[i] + (BP[i] + BPC[i] * C[i]) * P[i];
        log_lik[i] = binomial_logit_lpmf(L[i] | 1, p[i] );
    }
  }
}

```

Sample from model.

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

```{r}
print(fit13_6NC, include = T, pars = c('sigma_actor', 'sigma_block'))
```

## 13.4 Continuous categories and the gaussian process

### 13.4.1 example: spatial autocorrelation in Oceanic tools

Load the data.

```{r}
# Oceanic data now with coordinates
data('Kline2', package = 'rethinking')
d <- Kline2; rm(Kline2)
d$society <- 1:10

# Distance matrix
data('islandsDistMatrix', package = 'rethinking')
Dmat <- islandsDistMatrix; rm(islandsDistMatrix)
colnames(Dmat) <- c("Ml","Ti","SC","Ya","Fi","Tr","Ch","Mn","To","Ha")
```

Review shapes of covariance using linear and squared representations.

Figure 13.8

```{r}
ex <- data.frame(dist = seq(0, 5, by = .01)) %>% 
     mutate(shape_linear = exp(-1 * dist), shape_squared = exp(-1 * dist^2))

ggplot(ex) + theme_tufte(base_family = 'sans') +
  geom_line(aes(dist, shape_linear), linetype = 'dashed') +
  geom_line(aes(dist, shape_squared)) + labs(x = 'Distance', y = 'Covariance shape')
```

Code model in Stan.

```{stan output.var="m13_7"}
data {
  int N;            // observations
  int T[N];         // tools
  vector[N] logpop; // log population
  
  int N_societies;  // number of societies
  int society[N];   // index of societies
  matrix[N_societies,N_societies] Dmat; // distance matrix
}
parameters {
  real alpha;
  real beta_lpop;
  vector[N_societies] gamma;
  real<lower=0> etasq;
  real<lower=0> rhosq;
}
model {
  vector[N] lambda;
  matrix[N_societies,N_societies] K;
  
  // create Sigma (K) matrix
  for ( i in 1:(N_societies - 1) )
    for ( j in (i+1):N_societies ) {
      K[i,j] = etasq * exp( -rhosq * pow(Dmat[i,j],2) );
      K[j,i] = K[i,j];
      }
  for ( k in 1:N_societies )
    K[k,k] = etasq + 0.01;
        
  // priors
  target += multi_normal_lpdf(gamma | rep_vector(0,N_societies) , K );
  target += normal_lpdf(alpha | 0, 10);
  target += normal_lpdf(beta_lpop | 0, 1);
  target += cauchy_lpdf(etasq | 0, 1);
  target += cauchy_lpdf(rhosq | 0, 1);
  
  // linear model
  lambda = alpha + gamma[society] + beta_lpop * logpop;
  
  // likelihood
  target += poisson_log_lpmf(T | lambda);
}

```

Organize data and sample from model.

```{r}
dat <- list(
  N = NROW(d),
  T = d$total_tools,
  logpop = d$logpop,
  N_societies = max(d$society),
  society = d$society,
  Dmat = Dmat
)

fit13_7 <- sampling(m13_7, data = dat, iter = 1000, chains = 2, cores = 2)
```

Summary of model.

```{r}
print(fit13_7, probs = c(0.1, 0.5, 0.9))
```

Posterior distribution of the spatial covariance.

Figure 13.9

```{r}
post13_7 <- as.data.frame(fit13_7, pars = c('etasq', 'rhosq'))

dist     <- seq(0, 10, by = .1) 

cov_dist <- diag(post13_7$etasq) %*% exp(-post13_7$rhosq %*% t(dist^2) )
colnames(cov_dist) <- dist

cov_dist <- 
  as.data.frame(cov_dist) %>%
  mutate(Iter = row_number()) %>%
  tidyr::gather(key = 'dist', value = 'covariance', -Iter) %>%
  mutate(dist = as.numeric(as.character(dist)))

cov_dist_med <- 
  cov_dist %>% group_by(dist) %>% 
  summarise(covariance = median(covariance))

ggplot() + theme_tufte(base_family = 'sans') +
  geom_line(data = group_by(cov_dist, dist) %>% filter(row_number() <= 100), 
            aes(x = dist, y = covariance, group = Iter), alpha = .1) + 
  geom_line(data = cov_dist_med, aes(x = dist, y = covariance), lwd = 1) + 
  scale_y_continuous(limits = c(0, 1)) +
  labs(x = 'distance (thousand km)', y = 'covariance') 
```

Review correlation among the 10 societies.

```{r}
K <- median(post13_7$etasq) * exp(-median(post13_7$rhosq) * Dmat^2)
diag(K) <-  median(post13_7$etasq) + 0.01
Rho <- cov2cor(K)
round(Rho, 2) 
```

Map of societies.

Figure 13.10 left side

```{r}
colnames(Rho) <- rownames(Rho)
Rho[upper.tri(Rho, diag = T)] <- NA
Rho <- reshape2::melt(Rho, na.rm = TRUE)
colnames(Rho) <- c('culture', 'adjacent', 'correlation')
Rho <- Rho %>% left_join(d[,c('culture', 'lon2', 'lat', 'logpop', 'total_tools')], by = c('culture' = 'culture'))
Rho <- Rho %>% left_join(d[,c('culture', 'lon2', 'lat', 'logpop', 'total_tools')], by = c('adjacent' = 'culture'), suffix = c('', '.adj'))

ggplot() + theme_tufte(base_family = 'sans') +
  geom_segment(data = Rho, aes(x = lon2, xend = lon2.adj, y = lat, yend = lat.adj, alpha = correlation)) +
  geom_point(data = d, aes(lon2, lat, size = logpop), color = 'skyblue') + 
  geom_text(data = d, aes(lon2+1.8, lat, label = culture), hjust = 0, size = 3) +
  scale_x_continuous(limits = c(-50, 30)) +
  scale_alpha(range = c(0, 1)) +
  theme(legend.position = '') +
  labs(x = 'longitude', y = 'latitude')
```

Plot tools versus log population.

Figure 13.10 right side

```{r}
post <- as.data.frame(fit13_7, pars = c('alpha', 'beta_lpop'))
tot_tools_mu <- post$alpha + post$beta_lpop %*% t(seq(6, 13))
tt_mean <- apply(tot_tools_mu, 2, mean) %>% exp()
tt_pi <- apply(tot_tools_mu, 2, rethinking::PI) %>% exp()


ggplot() + theme_tufte(base_family = 'sans') +
  geom_line(aes(x = seq(6, 13), y = tt_pi[2,]), linetype = 'dashed', alpha = .3) +
  geom_line(aes(x = seq(6, 13), y = tt_mean), linetype = 'dashed', alpha = .3) +
  geom_line(aes(x = seq(6, 13), y = tt_pi[1,]), linetype = 'dashed', alpha = .3) +
  geom_segment(data = Rho, aes(x = logpop, xend = logpop.adj, 
                               y = total_tools, yend = total_tools.adj, 
                               alpha = correlation)) +
  geom_point(data = d, aes(logpop, total_tools, size = logpop), color = 'skyblue') + 
  geom_text(data = d, aes(logpop + .1, total_tools, label = culture), hjust = 0, size = 3) +
  scale_x_continuous(limits = c(7, 13), breaks = seq(7, 13)) +
  coord_cartesian(ylim = c(9, 80)) +
  scale_alpha(range = c(0, 1)) +
  theme(legend.position = '') +
  labs(x = 'log population', y = 'total tools')

```