Animating the line fitting process for a simple linear model

animate-lm-optimization.R

library(tidyverse)
library(gganimate)
library(modelr)

options(na.action = na.warn)

slopes <- 0
intercepts <- 0
fit <- 0


model1 <- function(a, data) {
  a[1] + data$x * a[2]
}

measure_distance <- function(mod, data) {
  slopes <<- c(slopes, mod[2])
  intercepts <<- c(intercepts, mod[1])
  diff <- data$y - model1(mod, data)
  fit <<- c(fit, sqrt(mean(diff ^ 2)))
  sqrt(mean(diff ^ 2))
}

sim1_dist <- function(a1, a2) {
  dist <- measure_distance(c(a1, a2), sim1)
  dist
}


best <- optim(c(0, 0), measure_distance, data = sim1)

DF <- 
data.frame(
  iteration = 0:(length(slopes) - 1)
  ,slope = slopes
  ,intercept = intercepts
  ,rmse = fit
) %>%
  filter(iteration > 0)


ggplot(DF, aes(x = iteration, y = rmse)) +
  geom_line()


DF$plot_text = paste0(DF$iteration, "\nRMSE: ", DF$rmse)
DF$plot_text <- factor(DF$plot_text, levels = DF$plot_text, labels = DF$plot_text)

g <- 
ggplot(DF) + 
  geom_abline(aes(intercept = intercept, slope = slope)) +
  geom_point(data = sim1, aes(x, y)) +
  transition_states(plot_text, transition_length = 100, state_length = 10) +
  ease_aes("elastic-in") +
  labs(title = "Regression Fit Lines During Optimization"
       ,subtitle = "Iteration: {closest_state}")


p <- animate(g, nframes = 350, end_pause = 10)
anim_save(filename = "Regression Fit Line.gif", animation = p)