Code to generate example animations with leapfrog integrator

divergence_example.Rmd

```{r}
library(tidyverse)
library(gganimate)
theme_set(cowplot::theme_cowplot())
```


```{r}
quad <- function(x) {
  -x^2
}

quad_grad <- function(x) {
  -2*x
}

wonky <- function(x) {
  -x^2 + dnorm((x-1) * 20) * 30
}

wonky_grad <- function(x) {
  x_shift <- (x-1) * 20
  -2*x + 30 * exp(-x_shift^2 / 2) * x / sqrt(2 * pi)
}


leapfrog_my <- function(start, start_v, step_size, steps, my_func, my_func_grad) {
  x <- array(NA_real_, steps + 1)
  v <- array(NA_real_, steps + 1)
  energy <- array(NA_real_, steps + 1)
  expected_y <- array(NA_real_, steps + 1)
  x[1] <- start
  v[1] <- start_v
  expected_y[1] <- my_func(start)
  
  start_energy <- 0.5 *start_v^2 - my_func(start)
  energy[1] <- start_energy
  for(s in 2:(steps + 1)) {
    a <- my_func_grad(x[s - 1])
    v_half <- v[s - 1] + my_func_grad(x[s - 1]) * step_size * 0.5
    x[s] <- x[s - 1] + v_half * step_size
    v[s] <- v_half + my_func_grad(x[s]) * step_size * 0.5
    
    energy[s] <- 0.5 * v[s] ^ 2 - my_func(x[s])
    expected_y[s] <-  0.5 * v[s] ^ 2 - start_energy
  }
  
  data.frame(x = x, y = my_func(x), expected_y = expected_y,
             energy = energy,
             v = v, step = 1:(steps + 1)) %>%
    mutate(v_end_x = x + v * step_size, v_end_y = y )
}


#leapfrog_iter <- function(
```

```{r}
my_plot <- function(start_x, start_v, step_size, steps, func, func_grad, show_expected = TRUE, animate = TRUE,
                    range = 3) {
  a_width <- 600
  a_height <- 350
  a_fps <- 3
  a_end_pause <- 5

  
  x <- seq(-range, range, length.out = 200)
  #p1 <- leapfrog_my(2, -1, step_size = 0.2, steps = 5) 
  #p1 <- leapfrog_my(2, -0.8, step_size = 0.1, steps = 20) 
  p1 <- leapfrog_my(start_x, start_v, step_size = step_size, steps = steps, func, func_grad) 
  
  if(show_expected) {
    expected_geom <- geom_point(aes(y = expected_y), color = "#00bf7d", shape = 4, size = 5)
  } else {
    expected_geom <- NULL
  }
  
  p <- ggplot(p1, aes(x = x, y = y, group = seq_along(step))) + geom_line(data = data.frame(x, y = func(x))) +
    expected_geom +
    geom_segment(aes(xend =v_end_x, yend = v_end_y), color = "#00b0f6") +
    geom_point(color = "#f8766d", size = 3) +
    scale_y_continuous("Log density")
   
  if(animate) {
    p <- p + transition_reveal(step)
    p <- animate(p, width = a_width, height = a_height, fps = a_fps, end_pause = a_end_pause, nframes = steps + a_end_pause + 1)
  }
  
  p
}
```



```{r}


p <- my_plot(start_x = 2, start_v = -1.6, step_size = 0.25, steps = 10, func = quad, func_grad = quad_grad, show_expected = FALSE, animate = TRUE)

anim_save("trajectory.gif", p)
p
```


```{r}
p <- my_plot(start_x = 2, start_v = -1.6, step_size = 0.25, steps = 10, func = quad, func_grad = quad_grad, show_expected = TRUE, animate = TRUE)
anim_save("trajectory_expected.gif", p)
p
```


```{r}
steep <- 50
edge <- function(x) {
  -x^2 - log1p(exp( (-2.5 -x) * steep)) 
}

edge_grad <- function(x) {
  -2*x + steep /(1 + exp(steep * (x + 2.5))) 
}

p <- my_plot(start_x = -1, start_v = -4, step_size = 0.25, steps = 4, func = edge, func_grad = edge_grad, show_expected = TRUE, animate = TRUE)

anim_save("trajectory_edge.gif", p)

p
```


```{r}
p <- my_plot(start_x = 2, start_v = -1.5, step_size = 0.25, steps = 10, func = wonky, func_grad = wonky_grad, show_expected = TRUE, animate = TRUE)

anim_save("trajectory_wonky_missed.gif", p)
p

```


```{r}
p <- my_plot(start_x = 2, start_v = -1.2, step_size = 0.25, steps = 7, func = wonky, func_grad = wonky_grad, show_expected = TRUE, animate = TRUE)

anim_save("trajectory_wonky_hit.gif", p)
p

```
```{r}
p <- my_plot(start_x = 2, start_v = -1.6, step_size = 0.4, steps = 20, func = wonky, func_grad = wonky_grad, show_expected = TRUE, animate = TRUE, range = 5)

anim_save("trajectory_wonky_diverge.gif", p)
p

```