multilevel-predictions.R

library(tidyverse)
library(gapminder)
library(brms)
library(tidybayes)

# Bayes stuff
options(mc.cores = 4,
        brms.backend = "cmdstanr")

# Countries to focus on (two per continent)
countries <- tribble(
  ~country, ~continent,
  "Egypt",  "Africa",
  "Sierra Leone", "Africa",
  "Pakistan", "Asia",
  "Yemen, Rep.", "Asia",
  "Bolivia", "Americas",
  "Canada", "Americas",
  "Italy", "Europe",
  "Portugal", "Europe"
)

# Clean the data
gapminder <- gapminder %>%
  filter(continent != "Oceania") %>%
  mutate(log_gdpPercap = log(gdpPercap)) %>% 
  # Make 1952 be year = 0 so intercept is interpretable and so the model fits faster
  mutate(year_orig = year,
         year = year - 1952) %>% 
  mutate(highlight = country %in% countries$country)

# Extract just the highlighted countries
original_points <- gapminder %>% 
  filter(highlight) %>% 
  mutate(year = year_orig)

# Each country gets its own slope and intercept for the year trend
model <- brm(
  bf(lifeExp ~ year + (1 + year | country)),
  data = gapminder,
  chains = 4
)

# Generate predictions
pred_model <- model %>%
  epred_draws(newdata = expand_grid(country = countries$country,
                                    year = unique(gapminder$year)),
              re_formula = NULL) %>% 
  mutate(year = year + 1952)

# Plot
ggplot(pred_model, aes(x = year, y = .epred)) +
  geom_point(data = original_points, aes(y = lifeExp)) +
  stat_lineribbon(alpha = 0.5) +
  labs(title = "Intercepts and slopes for year trend vary by country",
       subtitle = "lifeExp ~ year + (1 + year | country)",
       caption = "Source: Gapminder dataset",
       x = NULL, y = "Predicted life expectancy") +
  guides(fill = "none") +
  facet_wrap(vars(country), nrow = 2) +
  theme_bw() +
  theme(legend.position = "bottom",
        plot.subtitle = element_text(family = "Consolas"),
        axis.text.x = element_text(angle = 45, hjust = 0.5, vjust = 0.5))