sleepstudy_borrowingstrength_RStudio.Rmd

---
title: "Borrowing Strength with Linear Mixed Models: R-based lme4"
output: html_document
editor_options: 
  chunk_output_type: console
---

## Setup 

We produce the subject facets for the `sleepstudy` data including

+ observations
+ conditional means
+ within-subject slopes
+ estimate for pooled data


```{r}
library(lme4)
library(tidyverse)

data(sleepstudy)
sleepstudy <- mutate(sleepstudy, Subject = fct_reorder(Subject, Reaction, mean)) 
```

## Basics

### Fit the model

```{r}
fm <- lmer(Reaction ~ 1 + Days + (Days | Subject), sleepstudy, REML=FALSE,
           control=lmerControl(calc.derivs=FALSE))
print(summary(fm))
```

## Extract conditional means

```{r}
cms <- coef(fm)[["Subject"]] %>% 
  tibble::rownames_to_column("Subject") %>% 
  rename(Day_0 = `(Intercept)`) %>% 
  tibble::add_column(Estimate = "Conditional mean")
cms
```

## Facet plot

```{r}
plot1 <- 
  sleepstudy %>%
  mutate(Subject = forcats::fct_reorder(Subject, Reaction, mean)) %>%
  ggplot(aes(x = Days, y = Reaction)) + 
  geom_point() +
  geom_abline(data=cms, aes(intercept = Day_0, slope = Days)) +
  scale_x_continuous("Day", breaks=seq(0,9)) +
  scale_y_continuous("Reaction time [ms]", limits=c(100,500)) +
  facet_wrap(~ fct_rev(fct_rev(Subject))) + 
  theme_bw()
plot1
```

## Demonstration of borrowing strength

### Estimates of within-subject slopes

```{r}
wss <- lmList(Reaction ~ 1 + Days | Subject, sleepstudy) %>% 
  coef() %>% 
  tibble::rownames_to_column("Subject") %>% 
  rename(Day_0 = `(Intercept)`) %>% 
  tibble::add_column(Estimate = "Within-subject")
```

### Estimate for pooled data

```{r}
pld <- tibble(
  Subject = factor(levels(sleepstudy$Subject)),
  Day_0 = coef(lm(Reaction ~ 1 + Days, sleepstudy))[1],
  Days = coef(lm(Reaction ~ 1 + Days, sleepstudy))[2],
  Estimate = "Pooled"
)
```

### Combine the estimates

```{r}
cms_wss_pld <- bind_rows(cms, wss, pld) 
cms_wss_pld$Estimate <- factor(cms_wss_pld$Estimate, levels=c("Conditional mean", "Within-subject", "Pooled"))
cms_wss_pld
```

### Combined facet plot

```{r}
plot2 <- 
  sleepstudy %>%
  ggplot(aes(x = Days, y = Reaction)) + 
  geom_point() +
  geom_abline(data=cms_wss_pld, aes(intercept = Day_0, slope = Days, color = Estimate)) +
  facet_wrap(~ fct_rev(fct_rev(Subject))) +
  scale_x_continuous("Day", breaks=seq(0,9)) +
  scale_y_continuous("Reaction time [ms]", limits=c(100,500)) +
  theme_bw() +
  theme(legend.position = c(.9, .1))

# save the plot
ggsave("plot2.svg", plot2, width = 8, height = 7)

# show the plot
plot2
```

### Shrinkage plot

```{r}
bind_rows(wss, cms) %>% 
  ggplot(aes(x = Day_0, y = Days, color = Estimate)) + 
  geom_point() + 
  geom_point(data = pld) + 
  geom_path(aes(group = Subject), 
            arrow = arrow(length = unit(.02, "npc"),end="last")) + 
  scale_x_continuous("Reaction time @ day 0 [ms]", limits=c(200,300)) +
  scale_y_continuous("Change per day [ms]", limits=c(-5,25)) +
  theme_bw() + theme(legend.position = "top") +
  ggtitle("'Borrowing Strength' (Shrinkage)")
```

## Appendix

```{r}
sessionInfo()
```