spline_overfitting_animation.R

# Load packages
library(tidyverse)
library(gganimate)
library(tidymodels)

set.seed(2021)

# Generate true relationship between x and y as piecewise
piecewise <- function(x) {
  case_when(
    between(x, 0, 0.2) ~ x^2,
    between(x, 0.2, 0.4) ~ 0.2^2 + 3*(x-0.2),
    between(x, 0.4, 0.6) ~ 0.2^2 + 3*(x-0.2) + 20*(x-0.4)^3,
    between(x,0.6, 1) ~ 0.2^2 + 3*(x-0.2) + 20*(x-0.4)^3 - 10*(x-0.6)
  )
}

# Generate a data.frame of x and y data
gen_data <- function(){
  x <- runif(300)
  y <- piecewise(x) + rnorm(300, sd=0.5)
  return(tibble(x,y))
}

get_splines_predictions <- function(deg_free, train_dataset, test_dataset){
  
  mod <- lm(y ~ splines::ns(x, df = deg_free),
            data = train_dataset)
  
  grid <- tibble(x=seq(0,1, length.out=1e4))
  
  grid$ypred <- predict(mod, grid)
  grid$deg_free <- deg_free
  
  train_rmse <- sum((train_dataset$y - predict(mod, train_dataset))^2)
  test_rmse <- sum((test_dataset$y - predict(mod, test_dataset))^2)
  grid$train_rmse <- train_rmse
  grid$test_rmse <- test_rmse

  return(grid)
}

train <- gen_data()
test <- gen_data()

degs <- c(1,2,3,4,5,10,15,20,25,30,50,75,100,125,150)

df <-
  purrr::map(degs, ~get_splines_predictions(.x,
                                            train_dataset = train,
                                            test_dataset = test), .progress = TRUE) |>
  purrr::list_rbind()

df$deg_free <- factor(df$deg_free, levels=degs)

# Animate splines on top of training data
anim <-
  train |>
  ggplot() +
  aes(x=x, y=y) +
  geom_point(size=0.5, alpha=0.8) +
  geom_function(fun=piecewise, linetype="longdash", linewidth=0.7, color="gray") +
  geom_line(mapping=aes(x=x, y=ypred, group=1), data=df, color="#ae2012", linewidth=0.7, alpha=0.9) +
  cowplot::theme_cowplot(font_size = 10, font_family = "Helvetica") +
  labs(title = "Training data",
       subtitle = "Degrees of freedom = {closest_state}",
       x="",
       y="") +
  transition_states(deg_free, wrap = FALSE) +
  ease_aes('quartic-in-out')

gif <- animate(anim,
               duration=15,
               fps=60,
               height = 3.5,
               width = 3.5,
               units = "in",
               res = 300,
               renderer = ffmpeg_renderer())

# Save to mp4
anim_save(animation = gif, filename = "spline1.mp4")

# Animate splines on top of test data
anim <-
  test |>
  ggplot() +
  aes(x=x, y=y) +
  geom_point(size=0.5, alpha=0.8) +
  geom_function(fun=piecewise, linetype="longdash", linewidth=0.7, color="gray") +
  geom_line(mapping=aes(x=x, y=ypred, group=1), data=df, color="#ae2012", linewidth=0.7, alpha=0.9) +
  cowplot::theme_cowplot(font_size = 10, font_family = "Helvetica") +
  labs(title = "Training data",
       subtitle = "Degrees of freedom = {closest_state}",
       x="",
       y="") +
  transition_states(deg_free, wrap = FALSE) +
  ease_aes('quartic-in-out')

gif <- animate(anim,
               duration=15,
               fps=60,
               height = 3.5,
               width = 3.5,
               units = "in",
               res = 300,
               renderer = ffmpeg_renderer())

# Save to mp4
anim_save(animation = gif, filename = "spline2.mp4")

## Animate RMSE

rmse <-
  df |>
  distinct(deg_free, train_rmse, test_rmse)

rmse <-
  rmse |>
  pivot_longer(cols=`train_rmse`:`test_rmse`) |>
  mutate(
    name = case_match(
      name,
      "train_rmse" ~ "Training Error",
      "test_rmse"  ~ "Test Error",
    )
  )

rmse$name <- factor(rmse$name,
                    levels=c("Training Error",
                             "Test Error"))

rmse$deg_free <-
  rmse$deg_free |>
  as.character() |>
  as.integer()

anim <-
  rmse |>
  ggplot() +
  aes(x=deg_free,
      y=value,
      color=name,
      group=1) +
  geom_point(size=2) +
  cowplot::theme_minimal_hgrid(font_size = 10,
                               font_family = "Helvetica") +
  labs(x="Degrees of Freedom", 
       y="Mean squared error") +
  scale_x_continuous(breaks=c(1,20,50,75,100,125,150)) +
  facet_wrap(~name, ncol = 2) +
  theme(legend.position = "none",
        legend.title = element_blank()) +
  transition_states(deg_free, wrap = FALSE) +
  shadow_mark() +
  ease_aes('quartic-in-out')

gif <- animate(anim,
               duration=15,
               fps=60,
               height = 3,
               width = 7,
               units = "in",
               res = 300,
               renderer = ffmpeg_renderer())

# Save to mp4
anim_save(animation = gif, filename = "spline3.mp4")

####
## Run ffmpeg terminal commands to merge the three videos into one panel
### Stack horizontal
system("ffmpeg -i spline1.mp4 -i spline2.mp4 -filter_complex hstack=inputs=2 temp.mp4")
### Stack vertical
system("ffmpeg -i temp.mp4 -i spline3.mp4 -filter_complex vstack=inputs=2 spline_animation.mp4")
## Clean up intermediate files
fs::file_delete(c("spline1.mp4", "spline2.mp4", "spline3.mp4", "temp.mp4"))
####