foward mode autodiff in R, completing the unfinished example in Simon Wood's "Core Statistics" 5.5.3

foward_ad.R

rm(list=ls());gc()

ad <- function(x,diff = c(1,1)) {
  ## create class "ad" object. diff[1] is length of grad
  ## diff[2] is element of grad to set to 1.
  grad <- rep(0,diff[1])
  if (diff[2]>0 && diff[2]<=diff[1]) grad[diff[2]] <- 1
  attr(x,"grad") <- grad
  class(x) <- "ad"
  x
}

# chain rule is for taking d/dx of functions like f(x) = p(q(x))
# f'(x) = p'(q(x)) * q'(x)
sin.ad <- function(a) {
  grad.a <- attr(a,"grad")
  a <- as.numeric(a) # the value of input
  d <- sin(a) # the value of output
  attr(d,"grad") <- cos(a) * grad.a # chain rule; p' = cos, q' = grad.a
  class(d) <- "ad"
  d
}

"*.ad" <- function(a,b) { ## ad multiplication
  grad.a <- attr(a,"grad")
  grad.b <- attr(b,"grad")
  a <- as.numeric(a)
  b <- as.numeric(b)
  d <- a*b ## evaluation
  attr(d,"grad") <- a * grad.b + b * grad.a ## chain rule
  class(d) <- "ad"
  d
}

# missing stuff
# "exercise to the reader" =]
"/.ad" <- function(a,b) { ## 
  grad.a <- attr(a,"grad")
  grad.b <- attr(b,"grad")
  a <- as.numeric(a)
  b <- as.numeric(b)
  d <- a/b ## evaluation
  attr(d,"grad") <- (b * grad.a - a * grad.b)/(b^2) ## chain rule
  class(d) <- "ad"
  d
}

"+.ad" <- function(a,b) {
  grad.a <- attr(a,"grad")
  grad.b <- attr(b,"grad")
  a <- as.numeric(a)
  b <- as.numeric(b)
  d <- a + b ## evaluation
  attr(d,"grad") <- grad.a + grad.b
  class(d) <- "ad"
  d
}

"-.ad" <- function(a,b) {
  grad.a <- attr(a,"grad")
  grad.b <- attr(b,"grad")
  a <- as.numeric(a)
  b <- as.numeric(b)
  d <- a - b ## evaluation
  attr(d,"grad") <- grad.a - grad.b
  class(d) <- "ad"
  d
}

exp.ad <- function(a) {
  grad.a <- attr(a,"grad")
  a <- as.numeric(a) # the value of input
  d <- exp(a) # the value of output
  attr(d,"grad") <- exp(a) * grad.a # chain rule; p' = exp, q' = grad.a
  class(d) <- "ad"
  d
}

# 1st element of diff is length of gradient vector, 2nd element is which element
# of the gradient vector does this variable correspond to
x1 <- ad(1,c(3,1))
x2 <- ad(2,c(3,2))
x3 <- ad(pi/2,c(3,3))
(x1*x2*sin(x3)+ exp(x1*x2))/x3

fn <- function(x) {
  x1 <- x[1]; x2 <- x[2]; x3 <- x[3]
  (x1*x2*sin(x3)+ exp(x1*x2))/x3
}

library(numDeriv)
numDeriv::grad(func = fn, x = c(x1, x2, x3))