ranocha · Sondar74 · May 17, 2023 · May 21, 2023 · Jun 26, 2023 · Jun 26, 2023
diff --git a/src/BSeries.jl b/src/BSeries.jl
@@ -17,7 +17,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
@@ -40,6 +40,8 @@ export renormalize!
 
 export is_energy_preserving, energy_preserving_order
 
+export CSRK
+
 # Types used for traits
 # These traits may decide between different algorithms based on the
 # corresponding complexity etc.
@@ -74,6 +76,7 @@ end
 
 TruncatedBSeries{T, V}() where {T, V} = TruncatedBSeries{T, V}(OrderedDict{T, V}())
 
+
 # general interface methods of `AbstractDict` for `TruncatedBSeries`
 @inline Base.iterate(series::TruncatedBSeries) = iterate(series.coef)
 @inline Base.iterate(series::TruncatedBSeries, state) = iterate(series.coef, state)
@@ -612,6 +615,74 @@ function bseries(ros::RosenbrockMethod, order)
     return series
 end
 
+"""
+    ContinuousStageRungeKuttaMethod
+
+A struct that describes a CSRK method. This kind of 'struct' should be constructed 
+via [`CSRK`] 
+        'csrk = CSRK(M)' 
+in order to later call the 'bseries' function. 
+"""
+struct ContinuousStageRungeKuttaMethod{MatT <: AbstractMatrix}
+    matrix::MatT
+end
+
+function CSRK(matrix::AbstractMatrix)
+    T = promote_type(eltype(matrix))
+    _M = T.(matrix)
+    return ContinuousStageRungeKuttaMethod(_M)
+end
+
+"""
+        bseries(csrk::ContinuousStageRungeKuttaMethod, order)
+
+Return a truncated B-series up to the specified `order` with coefficients
+determined by the square matrix located in `csrk`. The coefficients for every 
+RootedTree are obtained via the 'elementary_differentials_csrk' function,
+which is calculated according to Miyatake & Butcher (2015). [See @ref csrk]
+
+# Example:
+The energy-preserving 4x4 matrix given by Miyatake & Butcher (2015) is
+```
+M = [-6//5   72//5  -36//1  24//1;
+        72//5  -144//5 -48//1  72//1;
+        -36//1 -48//1   720//1 -720//1;
+        24//1   72//1  -720//1  720//1]
+```
+
+Then, we calculate the bseries with the following code:
+
+```
+csrk = CSRK(M)
+series = bseries(csrk, 4)
+TruncatedBSeries{RootedTree{Int64, Vector{Int64}}, Rational{Int64}} with 9 entries:
+  RootedTree{Int64}: Int64[]      => 1//1
+  RootedTree{Int64}: [1]          => 1//1
+  RootedTree{Int64}: [1, 2]       => 1//2
+  RootedTree{Int64}: [1, 2, 3]    => 6004799503160661//36028797018963968
+  RootedTree{Int64}: [1, 2, 2]    => 6004799503160661//18014398509481984
+  RootedTree{Int64}: [1, 2, 3, 4] => 6004799503160661//144115188075855872
+  RootedTree{Int64}: [1, 2, 3, 3] => 6004799503160661//72057594037927936
+  RootedTree{Int64}: [1, 2, 3, 2] => 1//8
+  RootedTree{Int64}: [1, 2, 2, 2] => 1//4
+
+```
+# References
+Butcher, John & Miyatake, Yuto. (2015). A Characterization of Energy-Preserving
+Methods and the Construction of Parallel Integrators for Hamiltonian Systems. 
+SIAM Journal on Numerical Analysis. 54. 10.1137/15M1020861. 
+"""
+function bseries(csrk::ContinuousStageRungeKuttaMethod, order)
+    csrk = csrk.matrix
+    bseries(o) do t, series
+        if order(t) in (0, 1)
+            return one(csrk)
+        else
+            return elementary_differentials_csrk(csrk, t)
+        end
+    end
+end
+
 # TODO: bseries(ros::RosenbrockMethod)
 # should create a lazy version, optionally a memoized one
 
@@ -2016,4 +2087,86 @@ function equivalent_trees(tree)
     return equivalent_trees_set
 end
 
+
+# This function generates a polynomial 
+#       A_{t,z} = [t,t^2/2,..., t^s/s]*M*[1, z, ..., z^(s-1)]^T
+# for a given square matrix M of dimension s and chars 't' and 'z'.  
+function PolynomialA(M,t,z)
+    # get the dimension of the matrix
+    s = size(M,1)
+    # we need symbolic variables to work with
+    variable1 = Sym(t)
+    variable2 = Sym(z)
+    # conjugate the variable 1, since  this will be the variable of the left polynomial
+    # and the function 'dot' assumes it to be conjugated
+    variable1 = conjugate(variable1)
+    # generate the components of the polynomial with powers of t
+    poli_z = Array{SymPy.Sym}(undef, s)
+    for i in 1:s
+        poli_z[i] = variable2^(i-1)
+    end
+    # generate the components of the polynomial with powers of z
+    poli_t = Array{SymPy.Sym}(undef, s)
+    for i in 1:s
+        poli_t[i] = (1 // i)*(variable1^i)
+    end
+    # multiply matrix times vector
+    result = M * poli_z
+    # use dot product for the two vectors
+    return dot(poli_t,result)
+end
+
+
+"""
+    elementary_differentials_csrk(M,tree)
+
+This function calculates the CSRK elementary differential for a given 
+square matrix 'M' and a given RootedTree according to [@ref].
+
+# References
+Butcher, John & Miyatake, Yuto. (2015). A Characterization of Energy-Preserving
+Methods and the Construction of Parallel Integrators for Hamiltonian Systems. 
+SIAM Journal on Numerical Analysis. 54. 10.1137/15M1020861. 
+(https://www.researchgate.net/publication/276211444_A_Characterization_of_Energy-Preserving_Methods_and_the_Construction_of_Parallel_Integrators_for_Hamiltonian_Systems)
+
+"""
+function elementary_differentials_csrk(M,rootedtree)
+    # we extract the level_sequence of 'rootedtree'
+    tree = rootedtree.level_sequence
+    m = maximum(tree)
+    l = length(tree)
+    # Since we will integrate with symbolic variables for 
+    # every node in the tree, we create the variables called 
+    # 'xi' for 1 <= i <= m, because m is the most distant node from the root. 
+    variables = []
+    for i in 1:m
+        var_name = "x$i"
+        var = Sym(var_name)
+        push!(variables, var)
+    end
+    # we will calculate an integral for every node in the level_sequence from right to left
+    inverse_counter = l-1
+    # stablish initial integrand, which is the rightmost leaf (last node of the level sequence)
+    if l > 1
+        integrand = integrate(PolynomialA(M,variables[tree[end]-1],variables[tree[end]]),(variables[tree[end]],0,1))
+    else
+        # if the RootedTree is [1] or [], the elementary differential will be 1.
+        return 1
+    end
+    # Start a cycle for integrating 
+    while inverse_counter > 1
+        # we define the pseudo_integrand as the product between the last integral and the new polynomial (since the polynomials are
+        # multiplying each others inside the biggest integral). For a node 'i', this new Polynomial is computed for the variables
+        # 'xi' and 'x(i-1)'. 
+        pseudo_integrand = PolynomialA(M,variables[tree[inverse_counter]-1],variables[tree[inverse_counter]])*integrand
+        # integrate this new pseudo_integrand with respect to the variable 'xi'
+        integrand = integrate(pseudo_integrand,(variables[tree[inverse_counter]],0,1))
+        inverse_counter -= 1
+    end
+    # Once we have covered every node except for the base, multiply for the Basis_Polynomial, i.e. the Polynomial B
+    # defined by B_{x1} = A_{1, x1}.
+    # return the integral with respect to x1.
+    return integrate(PolynomialA(M,1,variables[1])*integrand,(variables[1],0,1))
+end
+
 end # module