AD examples (#64)

Author: ranocha

Project.toml

@@ -1,7 +1,7 @@
 name = "RootedTrees"
 uuid = "47965b36-3f3e-11e9-0dcf-4570dfd42a8c"
 authors = ["Hendrik Ranocha <mail@ranocha.de> and contributors"]
-version = "2.8.2"
+version = "2.8.3"
 
 [deps]
 Latexify = "23fbe1c1-3f47-55db-b15f-69d7ec21a316"

docs/Project.toml

@@ -1,5 +1,11 @@
 [deps]
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+SymPy = "24249f21-da20-56a4-8eb1-6a02cf4ae2e6"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
 Documenter = "0.27"
+ForwardDiff = "0.10"
+SymPy = "1"
+Zygote = "0.6"

docs/src/tutorials/RK_order_conditions.md

@@ -89,3 +89,73 @@ using Test
 end
 nothing # hide
 ```
+
+
+## Symbolic computation and automatic differentiation
+
+The implementation is fully generic using plain Julia code. In particular,
+this enables automatic differentiation (AD) and symbolic computations.
+
+### Symbolic computations
+
+For example, you can determine the order conditions symbolically as follows.
+```@example
+using RootedTrees, SymPy
+
+s = 3 # number of stages
+A = [symbols("a_$(i)$(j)", real=true) for i in 1:s, j in 1:s]
+b = [symbols("b_$(i)", real=true) for i in 1:s]
+rk = RungeKuttaMethod(A, b)
+
+for o in 1:3
+  println("Order ", o)
+  for t in RootedTreeIterator(o)
+    println("t = ", t)
+    println(residual_order_condition(t, rk))
+  end
+  println()
+end
+
+nothing # hide
+```
+
+### Automatic differentiation
+
+The order conditions can be differentiated with respect to the Runge-Kutta
+coefficients. For example, we can use
+[ForwardDiff.jl](https://github.com/JuliaDiff/ForwardDiff.jl) and
+[Zygote.jl](https://github.com/FluxML/Zygote.jl)
+as follows.
+
+```@example AD-Jacobian
+using RootedTrees, ForwardDiff, Zygote
+
+# collect all rooted trees up to order 4
+trees = [copy(t) for o in 1:4 for t in RootedTreeIterator(o)]
+
+# classical RK4 method
+A = [0 0 0 0; 1//2 0 0 0; 0 1//2 0 0; 0 0 1 0]
+b = [1//6, 1//3, 1//3, 1//6]
+coeffs = vcat(vec(A), vec(b)) # one vector of all parameters
+
+function all_order_conditions(trees, coeffs)
+  # extract Butcher coefficients from the vector of all coefficients
+  # for an RK method with four stages
+  A = reshape(view(coeffs, 1:16), 4, 4)
+  b = view(coeffs, 17:20)
+  rk = RungeKuttaMethod(A, b)
+  map(t -> residual_order_condition(t, rk), trees)
+end
+
+@assert iszero(all_order_conditions(trees, coeffs)) # fourth-order accurate
+
+ForwardDiff.jacobian(coeffs -> all_order_conditions(trees, coeffs), coeffs)
+```
+
+```@example AD-Jacobian
+Zygote.jacobian(coeffs -> all_order_conditions(trees, coeffs), coeffs)
+```
+
+```@example AD-Jacobian
+ForwardDiff.jacobian(coeffs -> all_order_conditions(trees, coeffs), coeffs) == first(Zygote.jacobian(coeffs -> all_order_conditions(trees, coeffs), coeffs))
+```