Visualize effect of letting slopes and intercepts vary by groups in linear regression

regression_simulation.R

library(tidyverse)
set.seed(1)

make_regression_equation <- function(df) {
  df %>% 
    summarise(eq = paste0(round(estimate, 1), term, collapse = " + ")) %>% 
    mutate(eq = paste0("y = ", eq),
           eq = str_remove(eq, "\\(Intercept\\)")) %>% 
    pull(eq)
}

make_simulated_data <- function(true_equation, n, eps_dist = rnorm, ...) {
  data_frame(
    x   = rep(seq_len(n), 2),
    grp = rep(0:1, n),
    eps = eps_dist(n*2, ...),
    y   = eval(true_equation)
  )
}

true_equation <- quote(200 + 100*x - 3000*grp + 3*(grp*x)^2 + eps)
data <- make_simulated_data(true_equation, 100, mean = 0, sd = 1000)

models <- tribble(
  ~formula,                       ~desc,
  "y ~ x",                        "1: same slope, same intercept",
  "y ~ x + grp",                  "2: same slope, different intercept",
  "y ~ x + x:grp",                "3: different slope, same intercept",
  "y ~ x + grp + x:grp",          "4: different slope, different intercept",
  "y ~ 1 + x + grp + I(grp*x^2)", "5: different slope, different intercept, quadratic"
) %>% 
  mutate(formula = map(formula, as.formula),
         model   = map(formula, lm, data = data),
         preds   = map(model, modelr::add_predictions, data = data),
         coeffs  = map(model, broom::tidy),
         eq      = map_chr(coeffs, make_regression_equation))

ggplot(unnest(models, preds), aes(x = x, y = y)) +
  geom_point(aes(color = factor(grp)), alpha = 0.25) +
  geom_line(aes(y = pred, color = factor(grp)), size = 1) +
  geom_text(aes(x = -Inf, y = Inf, label = eq), models, 
            size = 3, hjust = 0, vjust = 1.5) +
  facet_wrap(~desc, 3, scales = "free") +
  scale_y_continuous(labels = scales::comma_format()) +
  theme_minimal() +
  theme(legend.position = c(0.55, 0.25)) +
  labs(title = "Regression simulation", 
       caption = paste0("True equation: y = ", deparse(true_equation)),
       color = "Group")