R/ode.R

In targeted: Targeted Inference

#' Define compiled code for ordinary differential equation.
#'
#' The model (\code{code}) should be specified as the body of of C++ function.
#' The following variables are defined bye default
#' (see the argument \code{pname})
#' \describe{
#' \item{dy}{Vector with derivatives, i.e. the rhs of the ODE (the result).}
#' \item{x}{Vector with the first element being the time, and the following
#' elements additional exogenous input variables,}
#' \item{y}{Vector with the dependent variable}
#' \item{p}{Parameter vector}
#' }
#' \eqn{y'(t) = f_{p}(x(t), y(t))}
#' All variables are treated as Armadillo (http://arma.sourceforge.net/)
#' vectors/matrices.
#'
#' As an example consider the *Lorenz Equations*
#' \eqn{\frac{dx_{t}}{dt} = \sigma(y_{t}-x_{t})}
#' \eqn{\frac{dy_{t}}{dt} = x_{t}(\rho-z_{t})-y_{t}}
#' \eqn{\frac{dz_{t}}{dt} = x_{t}y_{t}-\beta z_{t}}
#'
#' We can specify this model as
#' \code{ode <- 'dy(0) = p(0)*(y(1)-y(0));
#'       dy(1) = y(0)*(p(1)-y(2));
#'       dy(2) = y(0)*y(1)-p(2)*y(2);'}
#' \code{dy <- specify_ode(ode)}
#'
#' As an example of model with exogenous inputs consider the following ODE:
#' \eqn{y'(t) = \beta_{0} + \beta_{1}y(t) +
#' \beta_{2}y(t)x(t) + \beta_{3}x(t)\cdot t}
#' This could be specified as
#' \code{mod <- 'double t = x(0);
#'               dy = p(0) + p(1)*y + p(2)*x(1)*y + p(3)*x(1)*t;'}
#' \code{dy <- specify_ode(mod)}
#' @title Specify Ordinary Differential Equation (ODE)
#' @param code string with the body of the function definition (see details)
#' @param fname Optional name of the exported C++ function
#' @param pname Vector of variable names (results, inputs, states, parameters)
#' @return pointer (externalptr) to C++ function
#' @author Klaus Kähler Holst
#' @seealso solve_ode
#' @examples
#' ode <- paste0(
#'   "dy(0) = p(0)*(y(1)-y(0));",
#'   "dy(1) = y(0)*(p(1)-y(2));",
#'   "dy(2) = y(0)*y(1)-p(2)*y(2);", collapse="\n"
#' )
#' \donttest{ # Reduce test time
#' dy <- specify_ode(ode)
#' tt <- seq(0, 100, length.out=2e4)
#' yy <- solve_ode(dy, input=tt, init=c(1, 1, 1), par=c(10, 28, 8/3))
#' }
#' @export
specify_ode <- function(code, fname=NULL, pname=c("dy", "x", "y", "p")) {
  if (is.null(fname))
    fname <- paste0("dy_", rlang::hash(code))
  hd <- "// [[Rcpp::depends(RcppArmadillo)]]
// [[Rcpp::plugins(cpp11)]]
// [[Rcpp::depends(targeted)]]
#include <RcppArmadillo.h>
#include <targeted.h>
using namespace arma;
"
  fun <- paste0("arma::rowvec ", fname,
                "(const arma::rowvec &", pname[2], ",
const arma::rowvec &", pname[3], ",
const arma::rowvec &", pname[4], ") {
  arma::rowvec ", pname[1], "(", pname[3], ".n_elem, arma::fill::zeros);
")
  fptr <- paste0("
  return ", pname[1], ";
}
// [[Rcpp::export]]
odeptr_t make_dy() {
  return odeptr_t(new odefunc_t(")
  tl <- "));
}
"
  make_dy <- NULL # Avoid warning about missing global variable
  rcpp_code <- paste(hd, fun, code, fptr, fname, tl)
  res <- Rcpp::sourceCpp(code=rcpp_code) # nolint
  f <- make_dy()
  return(f)
}


#' Solve ODE with Runge-Kutta method (RK4)
#'
#' The external point should be created with the function
#' \code{targeted::specify_ode}.
#' @title Solve ODE
#' @param ode_ptr pointer (externalptr) to C++ function or an R function
#' @param input Input matrix. 1st column specifies the time points
#' @param init Initial conditions
#' @param par Parameters defining the ODE (parsed to ode_ptr)
#' @seealso specify_ode
#' @return Matrix with solution
#' @author Klaus Kähler Holst
#' @export
#' @examples
#' example(specify_ode)
solve_ode <- function(ode_ptr, input, init, par=0) {
  if (is.function(ode_ptr)) {
    return(.ode_solve2(ode_ptr, cbind(input), rbind(init), rbind(par)))
  }
  return(.ode_solve(ode_ptr, cbind(input), rbind(init), rbind(par)))
}

## Dormand-Prince method. Adaptive step-size, see E. Hairer, S. P. Norsett G.
## Wanner, "Solving Ordinary Differential Equations I: Nonstiff Problems", Sec.
## II.
solve_ode_adaptive <- function(f, y0, t0, t1, # nolint
                               h0=0.1,
                               control=list()) {

  control0 <- list(hmax=h0, hmin=1e-20,
                   fac=0.9, fac_max=5, fac_min=1/5,
                   atol=1e-6, rtol=1e-6)
  control0[names(control)] <- control
  control <- control0

  # Adaptive step size
  cs <- c(0, 1/5, 3/10, 4/5, 8/9, 1, 1)
  as <- matrix(0, 7, 7)
  as[, 1] <- c(0, 1/5, 3/40, 44/45, 19327/6561, 9017/3168, 35/384)
  as[, 2] <- c(0, 0, 9/40, -56/15, -25360/2187, -355/33, 0)
  as[, 3] <- c(0, 0, 0, 32/9, 64448/6561, 46732/5247, 500/1113)
  as[, 4] <- c(0, 0, 0, 0, -212/729, 49/176, 125/192)
  as[, 5] <- c(0, 0, 0, 0, 0, -5103/18656, -2187/6784)
  as[, 6] <- c(0, 0, 0, 0, 0, 0, 11/84)
  bs <- rbind(
    as[7, ],
    c(
      5179 / 57600, 0, 7571 / 16695, 393 / 640,
      -92097 / 339200, 187 / 2100, 1 / 40
    )
  )
  q <- nrow(as) - 1

  intstep <- function(t, h, y) {
    k <- matrix(0, nrow = q + 1, ncol = length(y0))
    for (i in seq_len(nrow(k))) {
      yval <- y
      for (j in seq_along(i-1))
        yval <- yval + as[i, j]*k[j, ]
      k[i, ] <- h*f(t+cs[i]*h, yval)
    }
    # Estimate y(x+h) as weighted average of the q increments
    if (identical(as[, q+1], bs[1, ])) {
      y1 <- yval
    } else {
      y1 <- y
      for (j in seq_len(q+1))
        y1 <- y1 + bs[1, j]*k[j, ]
    }
    ## Estimate y(x+h) as weighted average of the q+1 increments
    y2 <- y
    for (j in seq_len(q+1))
        y2 <- y2 + bs[2, j]*k[j, ]

    return(list(y1, y2))
  }

  step <- function(t, h, y, fac_max) {
    accept <- FALSE
    while (!accept) {
      teval <- t + h
      step <- intstep(t, h, y)
      y1 <- step[[1]]
      y2 <- step[[2]]
      d <- abs(y2-y1)
      sc <- control$atol + control$rtol*pmax(abs(y1), abs(y2))
      err <- mean((d/sc)^2)^.5 ## err ~ C*h^(q+1)
      if (err<=1) { # accepts
        if (err<1e-30) {
          fac <- fac_max
        }
        accept <- TRUE
      }
      fac <- min(fac_max, max(control$fac_min, control$fac*(1/err)^(1/(q+1))))
      h <- min(control$hmax, h*fac)
      if (accept) {
        fac_max <- control$fac_max
      } else {
        fac_max <- 1
      }
    }
    return(structure(c(teval, y2), h=h, fac_max=fac_max))
  }

  t <- t0
  h <- max(h0, 1e-30)
  y <- y0
  r0 <- c(t, y)
  fac <- control$fac_max
  while (t<t1) {
    r1 <- step(t, h, y, fac)
    r0 <- rbind(r0, r1)
    t <- r1[1]
    h <- min(t1-t, attr(r1, "h"))
    fac <- attr(r1, "fac_max")
    y <- r1[-1]
  }
  return(r0)
}