Add support for Two-Derivative RK Methods

src/BSeries.jl

@@ -19,7 +19,7 @@ end
 
 using Latexify: Latexify, LaTeXString
 using Combinatorics: Combinatorics, permutations
-using LinearAlgebra: LinearAlgebra, rank
+using LinearAlgebra: LinearAlgebra, rank, dot
 using SparseArrays: SparseArrays, sparse
 
 @reexport using Polynomials: Polynomials, Polynomial
@@ -46,6 +46,8 @@ export is_energy_preserving, energy_preserving_order
 
 export order_of_symplecticity, is_symplectic
 
+export TwoDerivativeRungeKuttaMethod
+
 # Types used for traits
 # These traits may decide between different algorithms based on the
 # corresponding complexity etc.
@@ -1253,6 +1255,226 @@ end
 # TODO: bseries(mis::MultirateInfinitesimalSplitMethod)
 # should create a lazy version, optionally a memoized one
 
+"""
+    TwoDerivativeRungeKuttaMethod(A1, b1, A2, b2, c = vec(sum(A1, dims=2)))
+
+Represent a two-derivative Runge-Kutta method with Butcher coefficients
+`A1`, `b1`, and `c` for the first derivative and `A2`, `b2` for the second
+derivative.
+If `c` is not provided, the usual "row sum" requirement of consistency with
+autonomous problems is applied.
+
+Given an ODE ``u'(t) = f(t, u(t))`` with ``u''(t) = g(t, u(t))``,
+one step from ``u^{n}`` to ``u^{n+1}`` is given by
+
+```math
+\\begin{aligned}
+  y^i &= u^n + \\Delta t \\sum_j a^{1}_{i,j} f(t^n + c_j \\Delta t, y^j) + \\Delta t^2 \\sum_j a^{2}_{i,j} g(t^n + c_j \\Delta t, y^j), \\\\
+  u^{n+1} &= u^n + \\Delta t \\sum_i b^1_{i} f(t^n + c_i \\Delta t, y^i) + \\Delta t^2 \\sum_i b^2_{i} g(t^n + c_i \\Delta t, y^i).
+\\end{aligned}
+```
+
+# References
+
+- Chan, R.P.K., Tsai, A.Y.J.
+  "On explicit two-derivative Runge-Kutta methods."
+  Numer Algor 53, 171–194 (2010):
+  [DOI: 10.1007/s11075-009-9349-1](https://doi.org/10.1007/s11075-009-9349-1)
+
+"""
+struct TwoDerivativeRungeKuttaMethod{T,
+                                     MatT <: AbstractMatrix{T},
+                                     VecT <: AbstractVector{T}} <:RootedTrees.AbstractTimeIntegrationMethod
+    A1::MatT
+    b1::VecT
+    A2::MatT
+    b2::VecT
+    c::VecT
+end
+
+function TwoDerivativeRungeKuttaMethod(A1, b1, A2, b2, c = vec(sum(A1, dims=2)))
+    # promote all numeric types together
+    T = promote_type(eltype(A1), eltype(b1), eltype(A2), eltype(b2), eltype(c))
+
+    A1T = T.(A1)
+    b1T = T.(b1)
+    A2T = T.(A2)
+    b2T = T.(b2)
+    cT = T.(c)
+
+    return TwoDerivativeRungeKuttaMethod{T, typeof(A1T), typeof(b1T)}(
+        A1T, b1T, A2T, b2T, cT
+    )
+end
+
+
+Base.eltype(tdrk::TwoDerivativeRungeKuttaMethod{T}) where {T} = T
+
+"""
+    bseries(tdrk::TwoDerivativeRungeKuttaMethod, order) -> TruncatedBSeries
+
+Construct the truncated B-series of a two-derivative Runge–Kutta method `tdrk`
+up to the specified integer `order`.
+
+See also [`TwoDerivativeRungeKuttaMethod`](@ref).
+
+!!! note "Normalization by elementary differentials"
+    The coefficients of the B-series returned by this method need to be
+    multiplied by a power of the time step divided by the `symmetry` of the
+    rooted tree and multiplied by the corresponding elementary differential
+    of the input vector field ``f``.
+    See also [`evaluate`](@ref).
+"""
+function bseries(tdrk::TwoDerivativeRungeKuttaMethod, order)
+    # determine coefficient type
+    V_tmp = eltype(tdrk)
+    if V_tmp <: Integer
+        # If people use integer coefficients, they will likely want to have results
+        # as exact as possible. However, general terms are not integers. Thus, we
+        # use rationals instead.
+        V = Rational{V_tmp}
+    else
+        V = V_tmp
+    end
+    series = TruncatedBSeries{RootedTree{Int, Vector{Int}}, V}()
+
+    series[rootedtree(Int[])] = one(V)
+    for o in 1:order
+        for t in RootedTreeIterator(o)
+            series[copy(t)] = elementary_weight(t, tdrk)
+        end
+    end
+
+    return series
+end
+
+"""
+    elementary_weight(t::RootedTree, tdrk::TwoDerivativeRungeKuttaMethod)
+
+Compute the elementary weight \$\Phi(t)\$ for a TDRK method.
+
+Following Chan & Tsai (2010, Eq. 16), the weight is:
+\$\$ \Phi(t) = b_1 \cdot \eta(t) + b_2 \cdot \bar{\eta}(t) \$\$
+"""
+function RootedTrees.elementary_weight(t::RootedTree, tdrk::TwoDerivativeRungeKuttaMethod)
+    b1 = tdrk.b1
+    b2 = tdrk.b2
+    # alpha(t) = b1*eta(subtrees(t)) +b2*eta(collapse_trees(nu))
+    dot(b1, derivative_weight(t, tdrk)) + dot(b2, collapsed_derivative_weight(t, tdrk))
+end
+
+"""
+    derivative_weight(t::RootedTree, tdrk::TwoDerivativeRungeKuttaMethod)
+
+Compute the derivative weight \$\eta(t)\$ for a TDRK method.
+
+Following Chan & Tsai (2010, Eq. 15):
+\$\$ \eta(t) = A_1 \cdot \prod \eta(t_i) + A_2 \cdot \prod \bar{\eta}(t_j) \$\$
+"""
+function RootedTrees.derivative_weight(t::RootedTree, tdrk::TwoDerivativeRungeKuttaMethod)
+    A1 = tdrk.A1
+    A2 = tdrk.A2
+    c = tdrk.c
+
+    result1 = zero(c) .+ one(eltype(c))
+
+    if t == rootedtree(Int64[]) || t == rootedtree([1])
+        return zero(c) .+ one(eltype(c))
+    else
+        subtrees_arr = subtrees(t)
+        l = 1
+        for n in SubtreeIterator(t)
+            tmp = A1 * derivative_weight(subtrees_arr[l], tdrk) .+
+                  A2 * collapsed_derivative_weight(subtrees_arr[l], tdrk)
+            result1 = result1 .* tmp
+            l += 1
+        end
+        return result1
+    end
+end
+"""
+    collapsed_derivative_weight(t::RootedTree, tdrk::TwoDerivativeRungeKuttaMethod)
+
+
+Compute the derivative weight for the standered trees `t` in a two-derivative Runge–Kutta (TDRK) method.
+
+this corresponds to formula 15 in Chan, R.P.K., Tsai, A.Y.J. On explicit two-derivative Runge-Kutta methods.
+
+eta(t) = A1*eta(t) + A2*eta(t\\[1,2])
+
+where we are evaluating eta(t\\[1/2]) part for the elementary weight of the tree t
+"""
+function collapsed_derivative_weight(t::RootedTree, tdrk::TwoDerivativeRungeKuttaMethod)
+    A1 = tdrk.A1
+    A2 = tdrk.A2
+    c = tdrk.c
+
+    result = zero(c)
+
+    if t == rootedtree(Int64[])
+        return zero(c) .+ one(eltype(c))
+    else
+        collapsed_trees = collapse_tree(t)
+        number_of_trees = length(collapsed_trees)
+
+        for k in 1:number_of_trees
+            treecombinations = collapsed_trees[k]
+            number2 = length(treecombinations)
+            sum = zero(c) .+ one(eltype(c))
+
+            for m in 1:number2
+                step = A1 * derivative_weight(treecombinations[m], tdrk) .+
+                       A2 * collapsed_derivative_weight(treecombinations[m], tdrk)
+                sum = sum .* step
+            end
+
+            result = result .+ sum
+        end
+
+        return result
+    end
+end
+
+# Multi-Derivative Features
+
+"""
+    collapse_tree(t::RootedTree)
+
+recursively collapse a rooted tree `t` by removing [1,2] type branches.
+
+A collapse groups the children of `t` into all possible merged subsets,
+corresponding to the combinatorial partitions needed for the collapsed
+derivative weights 'eta(t\\[1,2])` in two-derivative B-series.
+
+Each element of the returned array is one valid list of collapsed
+subtrees of `t`. No modification of `t` is performed.
+"""
+function collapse_tree(t::RootedTree)
+    CollapsedArray = []
+
+    subtrees_arr = subtrees(t)
+    subtrees_multiplicity = length(subtrees_arr)
+
+    # Recustive approach to create all the possibilities of subtrees
+    for i in 1:subtrees_multiplicity
+        subsubtrees = subtrees(subtrees_arr[i])
+        numberofsubsubtrees = length(subsubtrees)
+
+        for j in 1:subtrees_multiplicity
+            if j == i
+            elseif j > i
+                push!(subsubtrees, subtrees_arr[j])
+            else j < i
+                pushfirst!(subsubtrees, subtrees_arr[i-j])
+            end
+        end
+        push!(CollapsedArray, subsubtrees)
+    end
+
+    return CollapsedArray
+end
+
+
 """
     substitute(b, a, t::AbstractRootedTree)
 

test/runtests.jl

@@ -3,7 +3,7 @@ using BSeries
 
 using BSeries.Latexify: latexify
 
-using LinearAlgebra: I
+using LinearAlgebra: I, dot
 using StaticArrays: @SArray, @SMatrix, @SVector
 
 using Symbolics: Symbolics
@@ -3245,4 +3245,71 @@ using Aqua: Aqua
             @test substituted == s
         end
     end
+
+    @testset "TwoDerivativeRungeKuttaMethod" begin
+
+        # Example two-derivative RK method (order 3) from "On explicit two-derivative Runge-Kutta methods" from R.P.K. Chan and A.Y.J. Tsai (2010)
+        A1 = [0 0;
+            1//2 0]
+
+        b1 = [1, 0]
+
+        A2 = [0 0;
+            1//8 0]
+
+        b2 = [1//6, 1//3]
+
+        #constructor
+
+        tdrk = @inferred TwoDerivativeRungeKuttaMethod(A1, b1, A2, b2)
+
+        @test tdrk isa TwoDerivativeRungeKuttaMethod
+        @test size(tdrk.A1) == (2, 2)
+        @test size(tdrk.A2) == (2, 2)
+        @test length(tdrk.b1) == 2
+        @test length(tdrk.b2) == 2
+
+        # row-sum default for c1
+        @test tdrk.c1 == vec(sum(A1, dims = 2))
+
+        # bseries
+        tdrk_series = @inferred bseries(tdrk, 5)
+
+        #should be 4th order
+        @test @inferred(order_of_accuracy(tdrk_series)) == 4
+
+        #now test with 5 stage 7th order method from Chan and Tsai (2010)
+
+        A_1 = [ 0      0       0       0       0;
+                2//7   0       0       0       0;
+                2//5   0       0       0       0;
+                4//7   0       0       0       0;
+                1       0       0       0       0
+        ]
+
+        b_1 = [1,0,0,0,0]
+
+        A_2 = [
+            0       0       0       0       0;
+            2//49   0       0       0       0;
+            2//25   0       0       0       0;
+            4//49   4//49   0       0       0;
+            -159//832  1715//832  -1875//832  735//832   0
+        ]
+
+        b_2 = [
+            71//960,
+            2401//4800,
+            -625//1728,
+            2401//8640,
+            13//1350
+        ]
+
+        tdrk_2 = @inferred TwoDerivativeRungeKuttaMethod(A_1, b_1, A_2, b_2)
+        # bseries
+        tdrk_series_2 = @inferred bseries(tdrk_2, 8)
+        #should be 7th order
+        @test @inferred(order_of_accuracy(tdrk_series_2)) == 7
+    end
+
 end # @testset "BSeries"