animate_diffeq_in_rgl.R

# animate a saddle differential equation system 
# 
#     dX = -x - y
#     dY = x - y
#     dZ = .25 * z
# 
# 25 trajectories are simulated and animated -- rendered using rgl in R
#
# Read more about rgl: https://dmurdoch.github.io/rgl/

library(deSolve)
library(rgl)
library(tidyverse)
library(magrittr)
library(purrr)

# minimum example : a single simulated trajectory, plotted statically

#' @param State initial state of the model
#' @param Params parameters for the model
saddle_model <- function(Time, State, Params) {
  with(as.list(c(State, Params)), {
    dX <- -x - y
    dY <- x - y
    dZ <- .25 * z

    return(list(c(dX, dY, dZ)))
  })
}

pars <- c()
yini <- c(x = 1, y = 1, z = 1)
times <- seq(0, 30, length.out = 800)

out <- ode(yini, times, saddle_model, pars)



# data wrangling  - create line segment pairs
# this is so that when we go to plot with segment3d()
# we can pass it pairs of (x_prior, x), (y_prior, y),
# and (z_prior, z) which make up each of the tiny little
# segments that are part of the simulated trajectory.

out_mat <- as.matrix(out)
out_mat <- as.data.frame(out_mat)

out_mat_lagged <- out_mat %>% mutate(
  time = lag(time)) %>% 
  rename(
    x_prior = x,
    y_prior = y,
    z_prior = z)

out_mat %<>% left_join(
  out_mat_lagged)

out_mat %<>% na.omit() # drop the first observation which has no prior (x,y,z) position



# plotting in 3d!  the fun part! 

open3d(scale = c(1,1,3))

pairify <- function(x) { as.vector(t(x)) }

get_coordinate_steps <- function(out_mat, var) {
  pairify(out_mat[, c(str_c(var, '_prior'), var)])
}


segments3d(x = get_coordinate_steps(out_mat, 'x'),
  y = get_coordinate_steps(out_mat, 'y'),
  z = get_coordinate_steps(out_mat, 'z'))

last_position <- function(out_mat, var) { out_mat[nrow(out_mat), var] }

points3d(
  x = last_position(out_mat, 'x'),
  y = last_position(out_mat, 'y'),
  z = last_position(out_mat, 'z')
  )


# now we want to add more simulated trajectories --- N=25 of them
# we use purrr to do this in a functional programming style

N <- 25
yinits <- map(1:N, ~ c(x = rnorm(n=1, sd = 10), y = rnorm(n=1, sd = 10), z = rnorm(n=1))) 

outs <- map(1:N, ~ ode(yinits[[.]], times, saddle_model, pars))

outs <- map(outs, as.data.frame)

outs %<>% map(~ mutate(.,
    x_prior = lag(x),
    y_prior = lag(y),
    z_prior = lag(z)) %>% 
    na.omit())

blues <- RColorBrewer::brewer.pal(9, 'Blues')[3:9]

outs %<>% map(~ mutate(.,
  color = blues[[sample.int(n = 7, size=1)]]
))


# get the scale dimensions from manually configuring them and 
# saving them -- use dput() to hard-code them after extracting them:

# zoom<-par3d()$zoom
# userMatrix<-par3d()$userMatrix
# windowRect<-par3d()$windowRect

zoom <- .325
userMatrix <- 
	structure(c(0.883055090904236, -0.273544698953629, 0.381296038627625,
	0, 0.423031091690063, 0.112328171730042, -0.899125576019287,
	0, 0.203120812773705, 0.955277800559998, 0.214909911155701, 0,
	0, 0, 0, 1), .Dim = c(4L, 4L))
windowRect <-  c(0L, 45L, 972L, 948L)


# create a directory to store animation frames in
dir.create("frames/")

# for each of the 800 iterations in our simulations, plot in 3d the ode solutions up 
# to that iteration

for (frame_i in 1:800) {

	theta <- seq(0, 2*pi, len = 400)[(frame_i %% 400) + 1]


	open3d(zoom = zoom, userMatrix = 
		rotate3d(userMatrix, theta, 0, 0, 1), windowRect = windowRect)

	map(outs, ~ 
		segments3d(
			x = get_coordinate_steps(head(., frame_i), 'x'),
			y = get_coordinate_steps(head(., frame_i), 'y'),
			z = get_coordinate_steps(head(., frame_i), 'z'),
			col = head(., frame_i)[['color']]))

	map(outs, ~ 
		points3d(
			x = last_position(head(., frame_i), 'x'),
			y = last_position(head(., frame_i), 'y'),
			z = last_position(head(., frame_i), 'z'),
			col = head(., frame_i)[['color']]))

  Sys.sleep(0.05) # this sleep is so we don't accidentally capture a blank frame before 
  # rgl finishes plotting

  rgl.snapshot(
    str_c('frames/', 
			stringr::str_pad(frame_i, width = 3, side = 'left', pad = '0'), 
			'.png'),
    fmt = 'png')

  # it's important to have rgl.close() so you don't open ~800 rgl windows on 
  # your computer and have to close them one by one...
  rgl.close()
}

# convert is from the ImageMagick tool which we use to render the animation
system("convert -delay 10 frames/*.png animation.gif")

Click to view the animation on YouTube: