#' Create an integrator.
#'
#' Create either a Leapfrog integrator or an Implicit Midpoint integrator with a particular nonlinear solve function.
#'
#' @param is_implicit Whether to use Implicit Midpoint or Leapfrog
#' @param do_nonlin_solve a function that solves the nonlin system g(x) = 0
#'
#' @return A function that integrates a single step.
#' @export
#'
#' @examples
create_integrator <- function(is_implicit, midpoint_formulation = TRUE, do_nonlin_solve = NULL) {
# set integrator to either explicit or implicit
# if implicit then the function we return
integrate_one_step <- take_one_step_lf
if (is_implicit) {
# use either Backward Euler formulation or Midpoint Formulation
if(midpoint_formulation) {
integrate_one_step <- function(z0, z_1, z_2, direction, ham_system, H0) {
take_one_step_im_midpoint(z0, z_1, z_2, direction, ham_system, do_nonlin_solve)
}
} else {
integrate_one_step <- function(z0, z_1, z_2, direction, ham_system, H0) {
take_one_step_im(z0, z_1, z_2, direction, ham_system, do_nonlin_solve)
}
}
}
# define and return stepper function
function(z0, z_1, z_2, direction, ham_system, H0) {
result <- integrate_one_step(z0, z_1, z_2, direction, ham_system, H0)
list(z1 = result$z1,
integrator_error = result$error,
num_grad = result$num_grad,
num_hess = ifelse(is_implicit, result$num_hess, as.integer(NA)),
num_hess_vec = ifelse(is_implicit, result$num_hess_vec_prod, as.integer(NA)),
num_newton = ifelse(is_implicit, result$num_newton, as.integer(NA)))
}
}
#' Integrate multiple steps.
#'
#' Uses a take step function to integrate multiple steps then format the result as a tibble.
#'
#' @param q0 initial position
#' @param p0 initial momentum
#' @param h0 initial stepsize
#' @param ham_system the Hamiltonian system to integrate
#' @param take_step the take single step function
#' @param num_steps the number of steps to integrate
#'
#' @return a formatted tibble containing information of the steps
#' @export
#'
#' @examples
integrate_multiple_steps <- function(q0, p0, h0, ham_system, take_step, num_steps) {
D <- length(q0)
h <- rep(NA, num_steps+1)
H <- rep(NA, num_steps+1)
integrator_error <- rep(NA, num_steps+1)
num_grad <- rep(NA, num_steps+1)
num_hess <- rep(NA, num_steps+1)
num_hess_vec <- rep(NA, num_steps+1)
num_newton <- rep(NA, num_steps+1)
q <- matrix(NA, nrow = num_steps+1, ncol = D)
p <- matrix(NA, nrow = num_steps+1, ncol = D)
h[1] <- h0
H[1] <- ham_system$compute_H(list(q = q0, p = p0, h = h0))
q[1,] <- q0
p[1,] <- p0
# start solving
for(n in 1:num_steps) {
# pack variables
z0 <- list(q = q[n,], p = p[n,], h = h[n])
z_1 <- NULL
z_2 <- NULL
if(n > 2) {
z_1 <- list(q = q[n-1,], p = p[n-1,], h = h[n-1])
z_2 <- list(q = q[n-2,], p = p[n-2,], h = h[n-2])
}
# take step and catch any integrator errors
one_step_result <- take_step(z0, z_1, z_2, direction = 1, ham_system, H[1])
# unpack result and save
h[n+1] <- one_step_result$z1$h
H[n+1] <- ham_system$compute_H(list(q = one_step_result$z1$q, p = one_step_result$z1$p, h = one_step_result$z1$h))
integrator_error[n+1] <- one_step_result$integrator_error
num_grad[n+1] <- one_step_result$num_grad
num_hess[n+1] <- one_step_result$num_hess
num_hess_vec[n+1] <- one_step_result$num_hess_vec
num_newton[n+1] <- one_step_result$num_newton
q[n+1,] <- one_step_result$z1$q
p[n+1,] <- one_step_result$z1$p
if(!is.na(one_step_result$integrator_error)) break
}
# return as a tibble
soln <- tibble(t = cumsum(c(0,tail(h, num_steps)))) %>%
mutate(h = h,
H0 = H[1],
H = H,
integrator_error = integrator_error,
num_grad = num_grad,
num_hess = num_hess,
num_hess_vec = num_hess_vec,
num_newton = num_newton) %>%
bind_cols(as_tibble(q) %>% set_names(paste0("q", 1:D))) %>%
bind_cols(as_tibble(p) %>% set_names(paste0("p", 1:D)))
return(soln)
}
Add the following code to your website.
For more information on customizing the embed code, read Embedding Snippets.