BRMS_Clarkia.Rmd

---
title: "Brms Clarkia"
author: "Julin Maloof"
date: "8/23/2020"
output: 
  html_document: 
    keep_md: yes
---

```{r setup, include=FALSE}
knitr::opts_chunk$set(echo = TRUE)
```

```{r}
library(rethinking)
library(tidyverse)
library(brms)
```


```{r}
data <- read_csv("clarkia_transplant_data.csv")
data %>% arrange(pop)
```

## Assignment

Redo Megan's Clarkia assignment using brms.  We did this for Feb 14.  I paste her email below:

I've attached some data from a common garden experiment, where plants from 15 different populations were planted out (locations are shown in Figure 1 here if you're curious). One goal for the experiment was to see if temperature of origin affected performance of plants in the common garden. Here are some practice questions, very similar to Julin's from last week. The data set is big-ish. I've already sub-sampled it, but there are still 3250 observations. The models are still running quickly on my computer, but if that's not the case for you, feel free to sub-sample it further. Please let me know if you have any questions.

## 1
_Fit a simple model with effects of temperature difference (temperature_diff_fall) on November germination (nov_germ). Temperature difference is already centered and scaled (i.e., negative values are the smallest temperature differences). Make sure to use the appropriate likelihood for the germination data (0 = no germ, 1  = germ). _


First let's take a look
```{r}
data %>%
  group_by(pop) %>%
  summarize(temp_diff_fall=unique(temperature_diff_fall), germ.prop=mean(nov_germ)) %>%
  ggplot(aes(x=temp_diff_fall, y=germ.prop, color=pop)) +
  geom_point()
```

### rethinking
fit a model
```{r}
datsmall <- data %>% select(nov_germ, temperature_diff_fall)
fm1 <- ulam(alist(nov_germ ~ dbinom(size=1, prob = p),
                  logit(p) <- a + b_temp*temperature_diff_fall,
                  a ~ dnorm(0, 2),
                  b_temp ~dnorm(0, 2)),
            data=datsmall,
            chains = 4,
            cores = 4,
            log_lik = TRUE,
            refresh = 0)
```

```{r}
precis(fm1)
```

### brms

```{r}
get_prior(nov_germ ~ temperature_diff_fall, data=data, family=bernoulli())
```

```{r}
fm1.brms <- brm(nov_germ ~ temperature_diff_fall,
                family = bernoulli(), 
                prior = c(set_prior("normal(0, 2)", class="Intercept"),
                          set_prior("normal(0, 2)", class = "b")),
                data=data,
                refresh=1000)

```

```{r}
summary(fm1.brms, prob = 0.89)
```

## 2
_2. Simulate from your priors to see if you've chosen reasonable priors, adjust them if necessary._

### rethinking


```{r}
prior <- extract.prior(fm1, n=100)
str(prior)
```


```{r}
pred.df <- data.frame(temperature_diff_fall=seq(-2,2,.1))
prior.pred <- link(fm1, data=pred.df, post=prior)
dim(prior.pred)
head(prior.pred[,1:10])
#reality check:
inv_logit(prior$a[1] + prior$b_temp[1]*-2)
```

```{r}
colnames(prior.pred) <- pred.df$temperature_diff_fall
prior.pred %>% as_tibble() %>%
  mutate(sample=1:nrow(.)) %>%
  gather(key="temp", value="germ", -sample) %>%
  mutate(temp=as.numeric(temp)) %>%
  ggplot(aes(x=temp, y=germ, group=sample)) +
  geom_line(alpha=.2)
```

### brms

To do this we can use sample_prior = "only" and create a data frame with some fake response data
```{r}
prior.data <- tibble(nov_germ=.1, temperature_diff_fall=seq(min(data$temperature_diff_fall), max(data$temperature_diff_fall), length.out = 50))

fm1.brms.priors <- brm(nov_germ ~ 0 + Intercept + temperature_diff_fall,
                       family = bernoulli(), 
                       prior = c(set_prior("normal(0, 2)", coef="Intercept"),
                                 set_prior("normal(0, 2)", class = "b")),
                       data=prior.data,
                       sample_prior = "only")
```

```{r}
prior_predict_samples <- posterior_epred(fm1.brms.priors, newdata = prior.data, nsamples = 100)
str(prior_predict_samples)
```
```{r}
prior_predict_samples %>% as_tibble() %>%
  magrittr::set_colnames(prior.data$temperature_diff_fall) %>% 
  mutate(sample=1:n()) %>%
  pivot_longer(cols=-sample, names_to="temperature_diff_fall", values_to = "nov_germ", names_transform=list(temperature_diff_fall=as.numeric)) %>%
  ggplot(aes(x=temperature_diff_fall, y = nov_germ, group=sample)) +
  geom_line(alpha=.2)
```
Most of these have the temperature response going all the way from 0 to full germination.  Would like more in the middle, so try narrowing the priors on slope to .5 and alpha to 1.5.


```{r}
fm2.brms.priors <- brm(nov_germ ~ 0 + Intercept + temperature_diff_fall,
                       family = bernoulli(), 
                       prior = c(set_prior("normal(0, 1.5)", coef="Intercept"),
                                 set_prior("normal(0, .5)", class = "b")),
                       data=prior.data,
                       sample_prior = "only")
```

```{r}
priorpred <- posterior_epred(fm2.brms.priors, newdata = prior.data, nsamples = 100) %>% 
  as_tibble()
```


```{r}
priorpred %>%
  magrittr::set_colnames(prior.data$temperature_diff_fall) %>% 
  mutate(sample=1:n()) %>%
  pivot_longer(cols=-sample, names_to="temperature_diff_fall", values_to = "nov_germ", names_transform=list(temperature_diff_fall=as.numeric)) %>%
  ggplot(aes(x=temperature_diff_fall, y = nov_germ, group=sample)) +
  geom_line(alpha=.2)
```
better...

## 3.
_These blocks were set up in the field, and had differences in soil depth, slope, and competitive environment. So maybe a model that includes block will describe the data better. _

_Fit a model that includes an effect of block (blk), with no pooling._

### rethinking

```{r}
datsmall <- data %>% select(nov_germ, temperature_diff_fall, blk)
sort(unique(datsmall$blk))

fm3 <- ulam(alist(nov_germ ~ dbinom(size=1, prob = p),
                  logit(p) <- a[blk] + b_temp*temperature_diff_fall,
                  a[blk] ~ dnorm(0, 1.5),
                  b_temp ~ dnorm(0, .5)),
            data=datsmall,
            chains = 4,
            cores = 4,
            log_lik = TRUE,
            refresh=0)
```

```{r}
precis(fm3, depth=2)
```

### brms

```{r}
data$blk2 <- as.character(data$blk)
fm3.brms <- brm(nov_germ ~ 0 + blk2 + temperature_diff_fall,
                family = bernoulli(), 
                prior = c(set_prior("normal(0, 1.5)"),
                          set_prior("normal(0, .5)", coef = "temperature_diff_fall")),
                data=data,
                refresh=0)
```

```{r}
summary(fm3.brms)
```

```{r}
prior_summary(fm3.brms)
```

## 4. 
_Fit a model that includes block, and allows partial pooling._

### rethinking

```{r}
datsmall <- data %>% select(nov_germ, temperature_diff_fall, blk)
sort(unique(datsmall$blk))

fm4 <- ulam(alist(nov_germ ~ dbinom(size=1, prob = p),
                  logit(p) <- a[blk] + b_temp*temperature_diff_fall,
                  a[blk] ~ dnorm(ablk_bar, sigma),
                  b_temp ~ dnorm(0, .5),
                  ablk_bar ~ dnorm(0, 1.5),
                  sigma ~ dexp(1)),
            data=datsmall,
            chains = 4,
            cores = 4,
            log_lik = TRUE,
            refresh=0)
```


```{r}
precis(fm4, depth=2)
```

### brms

```{r}
get_prior(nov_germ ~ (1|blk2) + temperature_diff_fall,
          family = bernoulli(),
          data=data)
```


```{r}
fm4.brms <-  brm(nov_germ ~ (1|blk2) + temperature_diff_fall,
                 family = bernoulli(), 
                 prior = c(set_prior("normal(0, 1.5)", class="Intercept"),
                           set_prior("normal(0, .5)", coef = "temperature_diff_fall"),
                           set_prior("exponential(1)", class="sd")),
                 control = list(adapt_delta=0.9),
                 iter = 4000,
                 data=data,
                 refresh=0)
```

```{r}
summary(fm4.brms)
coef(fm4.brms)
```


The experiment included many individuals from each of the 15 populations. So, each individual is not an independent representative of a given temperature, but might be similar to other plants from that population for reasons besides temperature. 

5. Build a model that accounts for this by including population (pop) and allowing partial pooling between populations A) without block, and B) with block included as in the model above. How does including population affect the temperature estimate?

### rethinking

```{r}
datsmall <- data %>% 
  mutate(pop_i = as.numeric(as.factor(pop))) %>%
  select(nov_germ, temperature_diff_fall, pop_i)

fm5a <- ulam(alist(nov_germ ~ dbinom(size=1, prob = p),
                   logit(p) <- a[pop_i] + b_temp*temperature_diff_fall,
                   a[pop_i] ~ dnorm(apop_bar, sigma_pop),
                   apop_bar ~ dnorm(0,1.5),
                   sigma_pop ~ dexp(1),
                   b_temp ~ dnorm(0, .5)),
             data=datsmall,
             chains = 4,
             cores = 4,
             iter=2000,
             log_lik = TRUE,
             refresh = 0)
```

```{r}
precis(fm5a, depth=2)
```
temp estimate the same, but confidence interval much wider

with block.  had to adjust sigma_pop and sigma_blk for narrower priors, and increase iter
```{r}
datsmall <- data %>% 
  mutate(pop_i = as.numeric(as.factor(pop))) %>%
  select(nov_germ, temperature_diff_fall, pop_i, blk)

fm5b <- ulam(alist(nov_germ ~ dbinom(size=1, prob = p),
                   logit(p) <- a[pop_i] + b_temp*temperature_diff_fall + b_blk[blk],
                   a[pop_i] ~ dnorm(apop_bar, sigma_pop),
                   b_blk[blk] ~ dnorm(0, sigma_blk),
                   apop_bar ~ dnorm(0, .5),
                   sigma_pop ~ dcauchy(0, .25),
                   sigma_blk ~ dcauchy(0, .25),
                   b_temp ~ dnorm(0, .5)),
             data=datsmall,
             chains = 4,
             cores = 4,
             iter=4000,
             log_lik = TRUE,
             refresh = 0)
```

```{r}
precis(fm5b, depth=2)
extract.samples(fm5b) %>%
  as.data.frame() %>%
  cor() %>%
  round(2)
```

### brms

```{r}
fm5a.brms <-  brm(nov_germ ~ (1|pop) + temperature_diff_fall,
                  family = bernoulli(), 
                  prior = c(set_prior("normal(0, 1.5)", class="Intercept"),
                            set_prior("normal(0, .5)", coef = "temperature_diff_fall"),
                            set_prior("exponential(1)", class="sd")),
                  control = list(adapt_delta=0.9),
                  iter = 4000,
                  data=data,
                  refresh=0)
```
```{r}
summary(fm5a.brms, prob=.89)
coef(fm5a.brms, probs=c(0.045, 0.955))$pop[,,"Intercept"] %>% round(3)
```

```{r}
precis(fm5a, depth = 2)
```


```{r}
fm5b.brms <-  brm(nov_germ ~ (1|pop) + (1|blk2) + temperature_diff_fall,
                  family = bernoulli(), 
                  prior = c(set_prior("normal(0, .5)", class="Intercept"),
                            set_prior("normal(0, .5)", coef = "temperature_diff_fall"),
                            set_prior("cauchy(0, 0.25)", class="sd")),
                  control = list(adapt_delta=0.9),
                  iter = 4000,
                  data=data,
                  refresh=0)
```

```{r}
summary(fm5b.brms, prob=.89)
coef(fm5b.brms, probs=c(0.045, 0.955))$pop[,,"Intercept"]
coef(fm5b.brms, probs=c(0.045, 0.955))$blk2[,,"Intercept"]
```

```{r}
precis(fm5b, depth = 2)
```
6. Compare the five models you built using WAIC. Which fits best?

```{r}
compare(fm1, fm3, fm4, fm5a, fm5b)
```


```{r}
for (b in ls(pattern="brms$")) {
  assign(b, add_criterion(get(b), "waic"))
}
```

```{r}
print(loo_compare(fm1.brms, fm3.brms, fm4.brms, fm5a.brms, fm5b.brms, criterion = "waic"), simplify = FALSE)
```

7. Plot effects of temperature difference for the average block, and also make a plot that includes the variability across blocks.

### rethinking


Should I do this for the different populations?

first, average block, all pops:

```{r}
post <- extract.samples(fm5b)
str(post)
```

```{r}
link_avg_blk <- function(post, temp) {
  m <- with(post, inv_logit(a + as.vector(b_temp*temp)))
  m
}
```

```{r}
pred.df <- tibble(temp = seq(-2,2,.1))

pred.df <- pred.df %>%
  mutate(posterior.pred=map(temp, ~ link_avg_blk(post, temp=.)))

pred.df # a tibble of tables
```

now compute mean and hpdi for each temperature
```{r}
pred.df <- pred.df %>%
  mutate(mean=map(posterior.pred, ~ apply(., 2, mean)),
         low89=map(posterior.pred, ~ apply(., 2, HPDI)[1,]),
         high89=map(posterior.pred, ~ apply(., 2, HPDI)[2,]),
  )
pred.df
```

```{r}
plot.df <- pred.df %>% select(-posterior.pred) %>% 
  unnest(mean, low89, high89) %>%
  mutate(pop=rep_along(mean, levels(factor(data$pop)) )) # an unsatisfying way to handle this
plot.df
```

```{r}
plot.df %>%
  ggplot(aes(x=temp, y=mean, ymin=low89, ymax=high89, color=pop, group=pop)) +
  geom_line() 
```

alternative:

```{r}
plot.df %>%
  ggplot(aes(x=temp, y=mean, ymin=low89, ymax=high89)) +
  geom_ribbon(alpha=.2) +
  geom_line() +
  facet_wrap(~pop, nrow=3)
```


### average block, average pop:

```{r}
link_avg_blk_avg_pop <- function(post, temp) {
  m <- with(post, inv_logit(rowMeans(a) + as.vector(b_temp*temp)))
  m
}
```

```{r}
pred.df <- tibble(temp = seq(-2,2,.1))

pred.df <- pred.df %>%
  mutate(posterior.pred=map(temp, ~ link_avg_blk_avg_pop(post, temp=.)))

pred.df # a tibble of tables
```

now compute mean and hpdi for each temperature
```{r}
pred.df <- pred.df %>%
  mutate(mean=map_dbl(posterior.pred, ~ mean(.)),
         low89=map_dbl(posterior.pred, ~ HPDI(.)[1]),
         high89=map_dbl(posterior.pred, ~ HPDI(.)[2]),
  )
pred.df
```


```{r}
pred.df %>%
  select(-posterior.pred) %>%
  ggplot(aes(x=temp, y=mean, ymin=low89, ymax=high89)) +
  geom_ribbon(alpha=.2) +
  geom_line() 
```

### include uncertainity about block

```{r}
pickOnePerRow <- function(m) { # pick one value from each row of a matrix and return as vector
                               # is there a better way to do this?
  result <- vector(mode=mode(m), length=nrow(m))
  for(i in 1:nrow(m)) result[i] <- m[i, sample(ncol(m), size=1)]
  return(result)
}

link_blk <- function(post, temp) {
  m <- with(post, 
            inv_logit(
              a + 
                as.vector(b_temp*temp) + 
                pickOnePerRow(b_blk) # pick a flat at random
            )
  )
  return(m)
}
```

```{r}
pred.df <- tibble(temp = seq(-2,2,.1)) 

pred.df <- pred.df %>%
  mutate(posterior.pred=map(temp, ~ link_blk(post, temp=.)))

pred.df # a tibble of tables
```

now compute mean and hpdi for each temperature
```{r}
pred.df <- pred.df %>%
  mutate(mean=map(posterior.pred, ~ apply(., 2, mean)), 
         low89=map(posterior.pred, ~ apply(., 2, HPDI)[1,]),
         high89=map(posterior.pred, ~ apply(., 2, HPDI)[2,]),
  )
pred.df
```


```{r}
plot.df <- pred.df %>% select(-posterior.pred) %>%  
  unnest(c(mean, low89, high89)) %>%
  mutate(pop=rep_along(mean, levels(factor(data$pop)) )) # an unsatisfying way to handle this
plot.df
```


```{r}
plot.df %>% 
  ggplot(aes(x=temp, y=mean, ymin=low89, ymax=high89)) +
  geom_ribbon(alpha=.2) +
  geom_line() +
  facet_wrap(~pop, nrow=3)
```

### brms

```{r}
newdata <- tibble(temperature_diff_fall=seq(min(data$temperature_diff_fall), 
                                            max(data$temperature_diff_fall), length.out=25))

avg.pred <- posterior_epred(fm5b.brms, newdata = newdata, re_formula=NA)

dim(avg.pred)
```

```{r}
avg.pred.df <- tibble(newdata,
                      nov_germ=apply(avg.pred,2,mean),
                      lower=apply(avg.pred,2,HPDI)[1,],
                      upper=apply(avg.pred,2,HPDI)[2,])
avg.pred.df
avg.pred.df %>% 
  ggplot(aes(x=temperature_diff_fall, y=nov_germ, ymin=lower, ymax=upper)) +
  geom_ribbon(fill="gray70")+
  geom_line()
```

For variability in our blks

```{r}
newdata <- expand_grid(
  temperature_diff_fall=seq(
    min(data$temperature_diff_fall), 
    max(data$temperature_diff_fall), length.out=25),
  blk2=unique(data$blk2))

avg.pred <- posterior_epred(fm5b.brms, newdata = newdata, re_formula=~(1|blk2))

dim(avg.pred)
```

```{r}
avg.pred.df <- tibble(newdata,
                      nov_germ=apply(avg.pred,2,mean),
                      lower=apply(avg.pred,2,HPDI)[1,],
                      upper=apply(avg.pred,2,HPDI)[2,])
avg.pred.df
avg.pred.df %>% 
  ggplot(aes(x=temperature_diff_fall, y=nov_germ, ymin=lower, ymax=upper, group=blk2, fill=blk2)) +
  geom_ribbon(alpha=.3)+
  geom_line()
```

For variability in new blocks

```{r}
newdata <- expand_grid(
  temperature_diff_fall=seq(
    min(data$temperature_diff_fall), 
    max(data$temperature_diff_fall), length.out=25),
  blk2=as.character(1:10))

avg.pred <- posterior_epred(fm5b.brms, newdata = newdata, re_formula=~(1|blk2),
                            allow_new_levels=TRUE)

dim(avg.pred)
```

```{r}
avg.pred.df <- tibble(newdata,
                      nov_germ=apply(avg.pred,2,mean),
                      lower=apply(avg.pred,2,HPDI)[1,],
                      upper=apply(avg.pred,2,HPDI)[2,])
avg.pred.df
avg.pred.df %>% 
  ggplot(aes(x=temperature_diff_fall, y=nov_germ, ymin=lower, ymax=upper, group=blk2, fill=blk2)) +
  geom_ribbon(alpha=.3)+
  geom_line()
```
Hmmm, not what I expected for the new blocks.  Something isn't working, or I don't understand...