In mikldk/taldi: Teaching Algorithmic Differentiation

knitr::opts_chunk$set(
  collapse = TRUE,
  comment = "#>",
  message = FALSE,
  fig.width = 7, 
  fig.height = 7
)

library(taldi)
library(caracas)

Summary

1. infer_ast(): infer abstract syntex tree
2. make_dag(): construct directed acyclic graph
3. (collect_leaves(): Collect symbol leaves)
4. Compute
5. forward_computation()
6. reverse_ad()

Infering abstract syntex tree

e <- quote(cos(2*x1 + x2) + x2^2)
symlist <- infer_ast(e)
str(symlist, 2)

Construct directed acyclic graph (DAG):

dag <- make_dag(symlist)

Plot DAG:

ggdag(dag)

Collect symbol leaves (such that each symbol only occurs once):

dag2 <- collect_leaves(dag)
ggdag(dag2)

Evaluating function in point

Forward computation:

e <- quote(cos(2*x1 + x2) + x2^2)
vals <- list(x1 = 1, x2 = 2)
ast <- infer_ast(e)
dag <- make_dag(ast)
dag <- collect_leaves(dag)
dag <- forward_computation(dag, values = vals)

The value in the node is in the second line (prefixed by d =, which can be changed by the value_prefix argument) - note that add_value must be set to TRUE:

ggdag(dag, add_value = TRUE)

The numeric value can be extracted as follows:

root <- get_root(dag)
get_values(dag, root)
eval(e, vals)

Finding derivatives

Symbolically

e <- quote(cos(2*x1 + x2) + x2^2)
vals <- list(x1 = 1, x2 = 2)

Derivatives can be found symbolically using caracas (R code hidden here, please refer to the vignette source for details):

f <- as_sym(deparse1(e))
#f
dfdx1 <- der(f, "x1")
dfdx2 <- der(f, "x2")

\begin{align} f(x_1, x_2) &= r tex(f) \ \frac{\partial f}{\partial x_1} &= r tex(dfdx1) \ \frac{\partial f}{\partial x_2} &= r tex(dfdx2) \end{align}

And evaluated in the point $(x_1, x_2) = (r vals["x1"], r vals["x2"])$:

#val <- eval_to_symbol("UnevaluatedExpr(pi/2)")
#val1 <- eval_to_symbol("UnevaluatedExpr(1)")
#val2 <- eval_to_symbol("UnevaluatedExpr(2)")
val1 <- eval_to_symbol(paste0("UnevaluatedExpr(", vals[1L], ")"))
val2 <- eval_to_symbol(paste0("UnevaluatedExpr(", vals[2L], ")"))

f_val_uneval <- subs_lst(f, list("x1" = val1, "x2" = val2))
#f_val_uneval
f_val <- doit(f_val_uneval)
#f_val

dfdx1_val_uneval <- subs_lst(dfdx1, list("x1" = val1, "x2" = val2))
#dfdx1_val_uneval
dfdx1_val <- doit(dfdx1_val_uneval)
#dfdx1_val

dfdx2_val_uneval <- subs_lst(dfdx2, list("x1" = val1, "x2" = val2))
#dfdx2_val_uneval
dfdx2_val <- doit(dfdx2_val_uneval)
#dfdx2_val

\begin{align} f\left (r tex(val1), r tex(val2) \right ) &= r tex(f_val_uneval) = r tex(f_val) \approx r N(f_val, 5)\ \frac{\partial f}{\partial x_1} \Bigg \vert_{x_1 = x_2 = r tex(val1)} &= r tex(dfdx1_val_uneval) = r tex(dfdx1_val) \approx r N(dfdx1_val, 5) \ \frac{\partial f}{\partial x_2} \Bigg \vert_{x_1 = x_2 = r tex(val2)} &= r tex(dfdx2_val_uneval) = r tex(dfdx2_val) \approx r N(dfdx2_val, 5) \end{align}

Reverse algorithmic differentiation (AD)

Using reverse algorithmic differentiation (AD):

ast <- infer_ast(e)
dag <- make_dag(ast)
dag <- collect_leaves(dag)
dag <- reverse_ad(dag, values = vals)
ggdag(dag, add_value = TRUE, add_deriv = TRUE)

Note that the value is the second line and the derivative the third line.

Advanced

Find leaves

is_leaf <- sapply(V(dag), function(x) {
  length(neighbors(dag, x, mode = "out")) == 0L
})
V(dag)[is_leaf]
leaves_labels <- vertex_attr(dag, "label", V(dag)[is_leaf])
leaves_types <- vertex_attr(dag, "type", V(dag)[is_leaf])
cbind(leaves_labels, leaves_types)

Preparing

• init_graph()
• bind_literals()
• bind_symbols()

ast <- infer_ast("cos(2*x1 + x2) + x2^2")
dag <- make_dag(ast)
dag <- collect_leaves(dag)

Init all node values to NA:

dag2 <- init_graph(dag)
ggdag(dag2, add_value = TRUE)

Bind literals:

dag3 <- bind_literals(dag2)
ggdag(dag3, add_value = TRUE)

Bind symbols:

dag4 <- bind_symbols(dag2, values = list(x1 = 1, x2 = 2))
ggdag(dag4, add_value = TRUE)

Bind both literals and symbols:

dag5 <- bind_literals(dag2)
dag5 <- bind_symbols(dag5, values = list(x1 = 1, x2 = 2))
ggdag(dag5, add_value = TRUE)

Plot using ggraph and tidygraph

library(tidygraph)
library(ggraph)

ast <- infer_ast("cos(2*x1 + x2) + x2^2")
dag <- make_dag(ast)
dag <- collect_leaves(dag)

tg <- as_tbl_graph(dag)
l <- create_layout(tg, layout = "sugiyama")

p <- ggraph(l) + 
  geom_edge_link(arrow = arrow(length = unit(0.25, 'cm')), end_cap = circle(0.85, 'cm')) + 
  geom_node_circle(aes(r = 0.3, fill = type)) + 
  geom_node_text(aes(label = label)) + 
  labs(fill = NULL) + 
  theme_bw() + 
  theme_graph()

print(p)  

This is essentially how ggdag() is implemented.

mikldk/taldi documentation built on March 26, 2022, 1:47 a.m.