Reside-IC

We are working on automatic differentiation of odin models, which requires support for differentiating expressions symbolically, in order to write new equations that can be used to numerically propagate derivatives of a model.

R already has support for doing this via the D function:

D(quote(2 * x^2 * log(sqrt(x))), "x")

## 2 * (2 * x) * log(sqrt(x)) + 2 * x^2 * (0.5 * x^-0.5/sqrt(x))

and the Deriv package provides an extensible interface. However, odin has peculiar syntax with arrays and we’re interested in trying to differentiate through stochastic functions, so a bespoke solution felt useful.

Symbolic differentiation turns out to be surprisingly easy, and quite elegant, to implement; this post shows the general idea.

To start, consider differentiating the expression x^2 + x^3 with respect to x. Recall the mechanical rules of differentiation from school that we can write this as d/dx x^2 + d/dx x^3 and then that we differentiate functions of the form x^n as n x^(n - 1) – this is the primary insight we need: that the process is recursive as we break down every operation into smaller chunks and keep on differentiating interior expressions with respect to x until there’s nothing left.

The simplest possible differentiation rules concern numbers; d/dx n for any number n is zero (that is, the gradient of n with respect to x is zero). Similarly, for any symbol (say a but not a + b) except x the derivative is also zero. And the derivative of x with respect to x is one. With this, we have the edge case for a recursive function:

differentiate <- function(expr, name) {
  if (!is.recursive(expr)) {
    if (identical(expr, as.symbol(name))) 1 else 0
  } else {
    stop("not yet implemented")
  }
}

which we can apply like so:

differentiate(quote(x), "x")

## [1] 1

differentiate(quote(a), "x")

## [1] 0

differentiate(quote(1), "x")

## [1] 0

Interesting expressions are not supported yet:

differentiate(quote(x + x * x), "x")

## Error in differentiate(quote(x + x * x), "x"): not yet implemented

To implement the case where we have compound expressions (is.recursive(expr) returning TRUE), consider the way we can represent these expressions:

expr <- quote(x + x * x)
as.list(expr)

## [[1]]
## `+`
## 
## [[2]]
## x
## 
## [[3]]
## x * x

Every call can be represented this way - the first element is the function being called and the remaining elements are its arguments. This structure is recursive:

as.list(expr[[3]])

## [[1]]
## `*`
## 
## [[2]]
## x
## 
## [[3]]
## x

To apply our differentiation rules we need to describe how to handle each function (here, + and *) and put together the results, descending into the subexpressions with differentiate() again until we get our edge cases.

The rule for differentiating sums is very straightforward, as noted above: we take the sum of the derivatives!

d_plus <- function(expr, name) {
  call("+", differentiate(expr[[2]], name), differentiate(expr[[3]], name))
}

The call() function constructs expression arguments (so call("+", quote(x), 1) returns x + 1) and here we are descending into each expression with differentiate(). We then rewrite differentiate() to call d_plus() when required:

differentiate <- function(expr, name) {
  if (!is.recursive(expr)) {
    if (identical(expr, as.symbol(name))) 1 else 0
  } else {
    fn <- as.character(expr[[1]])
    switch(fn,
           "+" = d_plus(expr, name),
           stop("not yet implemented"))
  }
}

With this we can differentiate a sum of any depth:

differentiate(quote(x + 5), "x")

## 1 + 0

differentiate(quote(x + y + x), "x")

## 1 + 0 + 1

We can then proceed, writing out rules for different functions as we need them. For example, the product rule:

d_product <- function(expr, name) {
  a <- expr[[2]]
  b <- expr[[3]]
  da <- differentiate(a, name)
  db <- differentiate(b, name)
  call("+", call("*", da, b), call("*", a, db))
}

or the quotient rule

d_quotient <- function(expr, name) {
  a <- expr[[2]]
  b <- expr[[3]]
  da <- differentiate(a, name)
  db <- differentiate(b, name)
  ## da / b - a * db / (b * b)
  call("-", call("/", da, b), call("/", call("*", a, db), call("*", b, b)))
}

For subtraction, we need to distinguish between unary minus (e.g., -a) and subtraction (e.g., a - b)

d_minus <- function(expr, name) {
  if (length(expr) == 2) {
    call("-", differentiate(expr[[2]], name))
  } else {
    call("-", differentiate(expr[[2]], name), differentiate(expr[[3]], name))
  }
}

It turns out that ( is a function too, and also needs a rule, but it is very simple:

d_parenthesis <- function(expr, name) {
  call("(", differentiate(expr[[2]], name))
}

We can put all these rules into a list:

rules <- list(
  "+" = d_plus,
  "-" = d_minus,
  "*" = d_product,
  "/" = d_quotient,
  "(" = d_parenthesis)

and rewrite our differentiate() implementation again:

differentiate <- function(expr, name) {
  if (!is.recursive(expr)) {
    if (identical(expr, as.symbol(name))) 1 else 0
  } else {
    fn <- as.character(expr[[1]])
    if (!(fn %in% names(rules))) {
      stop(sprintf("Differentiation of '%s' not yet implemented", fn))
    }
    rules[[fn]](expr, name)
  }
}

and with this we can differentiate all sorts of things:

differentiate(quote(-2 * x / (x * x - 3 * x)), "x")

## (-0 * x + -2 * 1)/(x * x - 3 * x) - -2 * x * (1 * x + x * 1 - 
##     (0 * x + 3 * 1))/((x * x - 3 * x) * (x * x - 3 * x))

This is fine, except that the generated expressions are fairly ugly, with lots of obviously redundant expressions (e.g., -0 * x which is obviously 0 and + -2 * 1 which is just - 2). However, the expressions agree with those from D once evaluated:

eval(differentiate(quote(-2 * x / (x * x - 3 * x)), "x"), list(x = pi))

## [1] 99.75819

eval(D(quote(-2 * x / (x * x - 3 * x)), "x"), list(x = pi))

## [1] 99.75819

Extending the implementation by adding more rules

We could extend this easily now by adding more rules, and the implementation will even tell us what we need to add. So if we try and evaluate

differentiate(quote(exp(2 * x)), "x")

## Error in differentiate(quote(exp(2 * x)), "x"): Differentiation of 'exp' not yet implemented

we get told to we need to implement the rule for exp which is simply:

d_exp <- function(expr, name) {
  call("*", differentiate(expr[[2]], name), expr)
}

(that is, d/dx exp(f(x)) is f'(x) exp(f(x))). We add this to our set of rules:

rules$exp <- d_exp

and now we can differentiate this new expression:

differentiate(quote(exp(2 * x)), "x")

## (0 * x + 2 * 1) * exp(2 * x)

Improving the implementation by writing sensible expressions

Simplifying the expressions turns out to be much more work than the differentiation. The trick that we use is to avoid using call() directly to build expressions and create a simplifying expression builder that applies some simple rules to avoid building overly complicated expressions. This does not simplify everything, but cuts out the most egregious bits of noise.

This is probably not important for the efficiency of generated code (we’re going to send this to an optimising compiler via some C++ code eventually) but it does make the resulting expressions easier to think about.

Consider replacing call("+", a, b) with something that will avoid creating silly expressions. If given numeric arguments for a and b we should sum them, and if either of a or b is zero we should return the other argument:

m_plus <- function(a, b) {
  if (is.numeric(a) && is.numeric(b)) {
    a + b
  } else if (is.numeric(b)) {
    m_plus(b, a)
  } else if (is.numeric(a) && a == 0) {
    b
  } else {
    call("+", a, b)
  }
}

The pattern here is that only the last branch gives up and actually builds an expression with call()

m_plus(3, 4)

## [1] 7

m_plus(quote(a), 0)

## a

m_plus(quote(a), 1)

## 1 + a

m_plus(0, quote(b))

## b

m_plus(1, quote(b))

## 1 + b

m_plus(quote(a), quote(b))

## a + b

We can do similar things with multiplication:

m_product <- function(a, b) {
  if (is.numeric(a) && is.numeric(b)) {
    a * b
  } else if (is.numeric(b)) {
    m_product(b, a)
  } else if (is.numeric(a) && a == 0) {
    0
  } else if (is.numeric(a) && a == 1) {
    b
  } else {
    call("*", a, b)
  }
}
m_product(3, 4)

## [1] 12

m_product(quote(a), 0)

## [1] 0

m_product(quote(a), 1)

## a

m_product(quote(a), 2)

## 2 * a

m_product(quote(a), quote(b))

## a * b

unary minus

is_call <- function(x, name) {
  is.recursive(x) && as.character(x[[1]]) == name
}
m_uminus <- function(a) {
  if (is.numeric(a)) {
    -a
  } else if (length(a) == 2 && identical(a[[1]], quote(`-`))) {
    a[[2]]
  } else if (is_call(a, "(")) {
    m_uminus(a[[2]])
  } else {
    call("-", a)
  }
}

subtraction

m_minus <- function(a, b) {
  if (is.numeric(a) && is.numeric(b)) {
    a - b
  } else if (is.numeric(a) && a == 0) {
    m_uminus(b)
  } else if (is.numeric(b) && b == 0) {
    a
  } else {
    call("-", a, b)
  }
}

and division

m_quotient <- function(a, b) {
  if (is.numeric(a) && is.numeric(b)) {
    a / b
  } else if (is.numeric(a) && a == 0) {
    0
  } else if (is.numeric(b) && b == 0) {
    Inf
  } else if (is.numeric(b) && b == 1) {
    a
  } else {
    call("/", a, b)
  }
}

Finally, parentheses (this is done differently in our implementation but this is a little simpler):

m_parenthesis <- function(a) {
  if (is.symbol(a) || is.numeric(a)) {
    a
  } else {
    call("(", a)
  }
}

We can then rewrite all our rules to use these functions instead of call() directly:

d_plus <- function(expr, name) {
  m_plus(differentiate(expr[[2]], name), differentiate(expr[[3]], name))
}

d_minus <- function(expr, name) {
  if (length(expr) == 2) {
    m_uminus(differentiate(expr[[2]], name))
  } else {
    m_minus(differentiate(expr[[2]], name), differentiate(expr[[3]], name))
  }
}

d_product <- function(expr, name) {
  a <- expr[[2]]
  b <- expr[[3]]
  da <- differentiate(a, name)
  db <- differentiate(b, name)
  m_plus(m_product(da, b), m_product(a, db))
}

d_quotient <- function(expr, name) {
  a <- expr[[2]]
  b <- expr[[3]]
  da <- differentiate(a, name)
  db <- differentiate(b, name)
  ## da / b - a * db / (b * b)
  m_minus(m_quotient(da, b), m_quotient(m_product(a, db), m_product(b, b)))
}

d_parenthesis <- function(expr, name) {
  m_parenthesis(differentiate(expr[[2]], name))
}

rules <- list(
  "+" = d_plus,
  "-" = d_minus,
  "*" = d_product,
  "/" = d_quotient,
  "(" = d_parenthesis)

Now, when we call differentiate(), things look much nicer:

differentiate(quote(-2 * x / (x * x - 3 * x)), "x")

## -2/(x * x - 3 * x) - -2 * x * (x + x - 3)/((x * x - 3 * x) * 
##     (x * x - 3 * x))

There are still some weirdnesses here (e.g., a - -2 * x * (x + x - 3) which are surprisingly hard to undo; the rhs of this expression is a tree with structure:

lobstr::ast(-2 * x * (x + x - 3))

## █─`*` 
## ├─█─`*` 
## │ ├─█─`-` 
## │ │ └─2 
## │ └─x 
## └─█─`(` 
##   └─█─`-` 
##     ├─█─`+` 
##     │ ├─x 
##     │ └─x 
##     └─3

so to simplify subtraction we have to extract the - from within two layers of multiplication, which is yet more recursion.

Our full implementation can be seen in this pull request, which follows the presentation here fairly closely.

Why go to this effort?

Given that the D function (and the Deriv package) can do all this (and more) already, it’s not very obvious why you might want to this. We don’t really want to differentiate R, but instead the domain specific language that supports odin. This is a small subset of R but there are two things that have quite different semantics that we need considerable control over (they’re not yet supported in the PR linked above).

Firstly, odin (via odin.dust) supports converting stochastic models into deterministic ones by taking expectations of the stochastic components; the underlying stochastic support already looks different to the call expected from R. So for example, we might write:

m <- rbinom(n, p)

to represent a binomial random draw with size n and probability p; R’s first argument (the number of draws to take) does not appear here. The expectation of this draw is simply n * p and we can now easily add rules to our differentiation support to differentiate this part of a model with respect to any other model quantity.

The other tricky part we have is the way that odin interprets array expressions, especially those that contain sums, as these have specific semantics. For example, the valid odin expression

lambda[] <- beta * sum(s_ij[, i])

could be rewritten as in something more R-ish

for (i in seq_along(lambda)) {
  lambda[i] <- beta * sum(s_ij[, i])
}

and differentiating the odin DSL version requires knowledge of these semantics. The solution to this is left as an exercise to the reader, and possibly a future blog post…