Runge-Kutta methods

Consider an ordinary differential equation (ODE) of the form

\[u'(t) = f(t, u(t)).\]

A Runge-Kutta (RK) method with $s$ stages is given by its Butcher coefficients

\[A = (a_{i,j})_{i,j} \in \mathbb{R}^{s \times s}, \quad b = (b_i)_i \in \mathbb{R}^{s}, \quad c = (c_i)_i \in \mathbb{R}^{s},\]

which are often written in form of the Butcher tableau

\[\begin{array}{c|cc} c & A \\ \hline & b^T \\ \end{array}\]

Usually, the consistency condition

\[\forall i\colon \quad c_i = \sum_j a_{i,j}\]

is assumed, which reduces all analysis to autonomous problems.

The step from $u^{n}$ to $u^{n+1}$ is given by

\[\begin{aligned} y^i &= u^n + \Delta t \sum_j a_{i,j} f(t^n + c_i \Delta t, y^i), \\ u^{n+1} &= u^n + \Delta t \sum_i b_{i} f(t^n + c_i \Delta t, y^i), \end{aligned}\]

where $y^i$ are the stage values.

In RootedTrees.jl, RK methods are represented as RungeKuttaMethods.

Order conditions

The order conditions of RK methods can be derived using rooted trees. In RootedTrees.jl, this functionality is implemented in residual_order_condition. Thus, a RungeKuttaMethod is of order $p$ if the residual_order_condition vanishes for all rooted trees with order up to $p$.

For example, the classical fourth-order RK method can be written as follows.

using RootedTrees

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]
rk = RungeKuttaMethod(A, b)

RungeKuttaMethod{Rational{Int64}} with
A: 4×4 Matrix{Rational{Int64}}:
  0     0    0  0
 1//2   0    0  0
  0    1//2  0  0
  0     0    1  0
b: 4-element Vector{Rational{Int64}}:
 1//6
 1//3
 1//3
 1//6
c: 4-element Vector{Rational{Int64}}:
  0
 1//2
 1//2
  1

To verify that this method is at least fourth-order accurate, we can check the residual_order_conditions up to this order.

using Test

@testset "RK4, order 4" begin
  for o in 1:4
    for t in RootedTreeIterator(o)
      @test iszero(residual_order_condition(t, rk))
    end
  end
end

Test Summary: | Pass  Total  Time
RK4, order 4  |    8      8  0.1s

To verify that this method does not satisfy any of the order conditions for an order of accuracy of five, we can use the following code.

using Test

@testset "RK4, not order 5" begin
  for t in RootedTreeIterator(5)
    @test !iszero(residual_order_condition(t, rk))
  end
end

Test Summary:    | Pass  Total  Time
RK4, not order 5 |    9      9  0.0s

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.

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

Order 1
t = RootedTree{Int64}: [1]
b_1 + b_2 + b_3 - 1

Order 2
t = RootedTree{Int64}: [1, 2]
b_1*(a_11 + a_12 + a_13) + b_2*(a_21 + a_22 + a_23) + b_3*(a_31 + a_32 + a_33) - 1/2

Order 3
t = RootedTree{Int64}: [1, 2, 3]
b_1*(a_11*(a_11 + a_12 + a_13) + a_12*(a_21 + a_22 + a_23) + a_13*(a_31 + a_32 + a_33)) + b_2*(a_21*(a_11 + a_12 + a_13) + a_22*(a_21 + a_22 + a_23) + a_23*(a_31 + a_32 + a_33)) + b_3*(a_31*(a_11 + a_12 + a_13) + a_32*(a_21 + a_22 + a_23) + a_33*(a_31 + a_32 + a_33)) - 1/6
t = RootedTree{Int64}: [1, 2, 2]
b_1*(a_11 + a_12 + a_13)^2/2 + b_2*(a_21 + a_22 + a_23)^2/2 + b_3*(a_31 + a_32 + a_33)^2/2 - 1/6

Automatic differentiation

The order conditions can be differentiated with respect to the Runge-Kutta coefficients. For example, we can use ForwardDiff.jl and Zygote.jl as follows.

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)

8×20 Matrix{Rational{Int64}}:
  0      0      0      0      0     …   0      0     1   1      1      1
 1//6   1//3   1//3   1//6   1//6      1//3   1//6   0  1//2   1//2    1
 1//6   1//6   1//6    0     1//4      1//2   1//6   0   0     1//4   1//2
  0     1//6   1//6   1//6    0        1//6   1//6   0  1//8   1//8   1//2
 1//12  1//12   0      0     1//6      1//3   1//12  0   0      0     1//4
  0     1//12  1//12   0     1//48  …  1//4   1//12  0   0     1//16  1//8
 1//12  1//12  1//4   1//12  1//12     5//12  1//4   0   0     1//8   1//2
  0     1//24  1//24  1//12   0        1//24  1//12  0  1//48  1//48  1//6

Zygote.jacobian(coeffs -> all_order_conditions(trees, coeffs), coeffs)

(Rational{Int64}[0 0 … 1 1; 1//6 1//3 … 1//2 1; … ; 1//12 1//12 … 1//8 1//2; 0 1//24 … 1//48 1//6],)

ForwardDiff.jacobian(coeffs -> all_order_conditions(trees, coeffs), coeffs) == first(Zygote.jacobian(coeffs -> all_order_conditions(trees, coeffs), coeffs))

true