Adaptive order version of ExplicitTaylor

Author: tansongchen

lib/OrdinaryDiffEqTaylorSeries/src/OrdinaryDiffEqTaylorSeries.jl

@@ -1,9 +1,13 @@
 module OrdinaryDiffEqTaylorSeries
 
-import OrdinaryDiffEqCore: alg_order, alg_stability_size, explicit_rk_docstring,
+import OrdinaryDiffEqCore: alg_order, get_current_adaptive_order, get_current_alg_order,
+                           alg_stability_size, explicit_rk_docstring,
                            OrdinaryDiffEqAdaptiveAlgorithm, OrdinaryDiffEqMutableCache,
                            alg_cache,
                            OrdinaryDiffEqConstantCache, @fold, trivial_limiter!,
+                           step_accept_controller!,
+                           stepsize_controller!,
+                           unwrap_alg,
                            constvalue, @unpack, perform_step!, calculate_residuals, @cache,
                            calculate_residuals!, _ode_interpolant, _ode_interpolant!,
                            CompiledFloats, @OnDemandTableauExtract, initialize!,
@@ -55,6 +59,6 @@ PrecompileTools.@compile_workload begin
     solver_list = nothing
 end
 
-export ExplicitTaylor2, ExplicitTaylor, DAETS
+export ExplicitTaylor2, ExplicitTaylor, ExplicitTaylorAdaptiveOrder
 
 end

lib/OrdinaryDiffEqTaylorSeries/src/TaylorSeries_caches.jl

@@ -39,10 +39,10 @@ end
 get_fsalfirstlast(cache::ExplicitTaylor2Cache, u) = (cache.k1, cache.k1)
 
 @cache struct ExplicitTaylorCache{
-    P, jetType, uType, taylorType, uNoUnitsType, StageLimiter, StepLimiter,
+    P, uType, taylorType, uNoUnitsType, StageLimiter, StepLimiter,
     Thread} <: OrdinaryDiffEqMutableCache
     order::Val{P}
-    jet::jetType
+    jet::Function
     u::uType
     uprev::uType
     utaylor::taylorType
@@ -54,17 +54,17 @@ get_fsalfirstlast(cache::ExplicitTaylor2Cache, u) = (cache.k1, cache.k1)
     thread::Thread
 end
 
-function alg_cache(alg::ExplicitTaylor{P}, u, rate_prototype, ::Type{uEltypeNoUnits},
+function alg_cache(alg::ExplicitTaylor, u, rate_prototype, ::Type{uEltypeNoUnits},
         ::Type{uBottomEltypeNoUnits}, ::Type{tTypeNoUnits}, uprev, uprev2, f, t,
         dt, reltol, p, calck,
-        ::Val{true}) where {P, uEltypeNoUnits, uBottomEltypeNoUnits, tTypeNoUnits}
-    _, jet_iip = build_jet(f, p, Val(P), length(u))
-    utaylor = TaylorDiff.make_seed(u, zero(u), Val(P))
+        ::Val{true}) where {uEltypeNoUnits, uBottomEltypeNoUnits, tTypeNoUnits}
+    _, jet_iip = build_jet(f, p, alg.order, length(u))
+    utaylor = TaylorDiff.make_seed(u, zero(u), alg.order)
     utilde = zero(u)
     atmp = similar(u, uEltypeNoUnits)
     recursivefill!(atmp, false)
     tmp = zero(u)
-    ExplicitTaylorCache(Val(P), jet_iip, u, uprev, utaylor, utilde, tmp, atmp,
+    ExplicitTaylorCache(alg.order, jet_iip, u, uprev, utaylor, utilde, tmp, atmp,
         alg.stage_limiter!, alg.step_limiter!, alg.thread)
 end
 
@@ -86,3 +86,42 @@ end
 
 # FSAL currently not used, providing dummy implementation to satisfy the interface
 get_fsalfirstlast(cache::ExplicitTaylorCache, u) = (cache.u, cache.u)
+
+@cache struct ExplicitTaylorAdaptiveOrderCache{
+    uType, taylorType, uNoUnitsType, StageLimiter, StepLimiter,
+    Thread} <: OrdinaryDiffEqMutableCache
+    min_order::Int
+    max_order::Int
+    current_order::Ref{Int}
+    jets::Vector{Function}
+    u::uType
+    uprev::uType
+    utaylor::taylorType
+    utilde::uType
+    tmp::uType
+    atmp::uNoUnitsType
+    stage_limiter!::StageLimiter
+    step_limiter!::StepLimiter
+    thread::Thread
+end
+function alg_cache(
+        alg::ExplicitTaylorAdaptiveOrder, u, rate_prototype, ::Type{uEltypeNoUnits},
+        ::Type{uBottomEltypeNoUnits}, ::Type{tTypeNoUnits}, uprev, uprev2, f, t,
+        dt, reltol, p, calck,
+        ::Val{true}) where {uEltypeNoUnits, uBottomEltypeNoUnits, tTypeNoUnits}
+    jets = Function[]
+    for order in alg.min_order:alg.max_order
+        _, jet_iip = build_jet(f, p, Val(order), length(u))
+        push!(jets, jet_iip)
+    end
+    utaylor = TaylorDiff.make_seed(u, zero(u), Val(alg.max_order))
+    utilde = zero(u)
+    atmp = similar(u, uEltypeNoUnits)
+    recursivefill!(atmp, false)
+    tmp = zero(u)
+    current_order = Ref{Int}(alg.min_order)
+    ExplicitTaylorAdaptiveOrderCache(alg.min_order, alg.max_order, current_order, jets, u, uprev, utaylor, utilde, tmp, atmp,
+        alg.stage_limiter!, alg.step_limiter!, alg.thread)
+end
+
+get_fsalfirstlast(cache::ExplicitTaylorAdaptiveOrderCache, u) = (cache.u, cache.u)

lib/OrdinaryDiffEqTaylorSeries/src/TaylorSeries_perform_step.jl

@@ -90,11 +90,69 @@ end
     end
     if integrator.opts.adaptive
         @.. broadcast=false thread=thread utilde=TaylorDiff.get_coefficient(utaylor, P) *
-                                                 dt^(P + 1)
+                                                 dt^P
         calculate_residuals!(atmp, utilde, uprev, u, integrator.opts.abstol,
             integrator.opts.reltol, integrator.opts.internalnorm, t)
         integrator.EEst = integrator.opts.internalnorm(atmp, t)
     end
     OrdinaryDiffEqCore.increment_nf!(integrator.stats, P + 1)
     return nothing
 end
+
+function initialize!(integrator, cache::ExplicitTaylorAdaptiveOrderCache)
+    integrator.kshortsize = cache.max_order
+    resize!(integrator.k, cache.max_order)
+    # Setup k pointers
+    for i in 1:(cache.max_order)
+        integrator.k[i] = get_coefficient(cache.utaylor, i)
+    end
+    return nothing
+end
+
+@muladd function perform_step!(
+        integrator, cache::ExplicitTaylorAdaptiveOrderCache, repeat_step = false)
+    @unpack t, dt, uprev, u, f, p = integrator
+    alg = unwrap_alg(integrator, false)
+    @unpack jets, current_order, min_order, max_order, utaylor, utilde, tmp, atmp, thread = cache
+
+    jet_index = current_order[] - min_order + 1
+    # compute one additional order for adaptive order
+    jet = jets[jet_index + 1]
+    jet(utaylor, uprev, t)
+    for i in eachindex(utaylor)
+        u[i] = @inline evaluate_polynomial(utaylor[i], dt)
+    end
+    OrdinaryDiffEqCore.increment_nf!(integrator.stats, current_order[] + 1)
+    if integrator.opts.adaptive
+        min_work = Inf
+        start_order = max(min_order, current_order[] - 1)
+        end_order = min(max_order, current_order[] + 1)
+        for i in start_order:end_order
+            A = i * i
+            @.. broadcast=false thread=thread utilde=TaylorDiff.get_coefficient(
+                utaylor, i) * dt^i
+            calculate_residuals!(atmp, utilde, uprev, u, integrator.opts.abstol,
+                integrator.opts.reltol, integrator.opts.internalnorm, t)
+            EEst = integrator.opts.internalnorm(atmp, t)
+
+            # backup
+            e = integrator.EEst
+            qold = integrator.qold
+            # calculate dt
+            integrator.EEst = EEst
+            dtpropose = step_accept_controller!(integrator, alg,
+                stepsize_controller!(integrator, alg))
+            # restore
+            integrator.EEst = e
+            integrator.qold = qold
+
+            work = A / dtpropose
+            if work < min_work
+                cache.current_order[] = i
+                min_work = work
+                integrator.EEst = EEst
+            end
+        end
+    end
+    return nothing
+end

lib/OrdinaryDiffEqTaylorSeries/src/alg_utils.jl

@@ -4,13 +4,19 @@ alg_stability_size(alg::ExplicitTaylor2) = 1
 alg_order(::ExplicitTaylor{P}) where {P} = P
 alg_stability_size(alg::ExplicitTaylor) = 1
 
-JET_CACHE = IdDict()
+alg_order(alg::ExplicitTaylorAdaptiveOrder) = alg.min_order
+get_current_adaptive_order(::ExplicitTaylorAdaptiveOrder, cache) = cache.current_order[]
+get_current_alg_order(::ExplicitTaylorAdaptiveOrder, cache) = cache.current_order[]
 
-function build_jet(f::ODEFunction{iip}, p, order, length = nothing) where {iip}
-    build_jet(f, Val{iip}(), p, order, length)
+TaylorScalar{T, P}(x::TaylorScalar{T, Q}) where {T, P, Q} = TaylorScalar{P}(x)
+
+const JET_CACHE = IdDict()
+
+function make_term(a)
+    term(TaylorScalar, Symbolics.unwrap(a.value), map(Symbolics.unwrap, a.partials))
 end
 
-function build_jet(f, ::Val{iip}, p, order::Val{P}, length = nothing) where {P, iip}
+function build_jet(f::ODEFunction{iip}, p, order::Val{P}, length = nothing) where {P, iip}
     if haskey(JET_CACHE, f)
         list = JET_CACHE[f]
         index = findfirst(x -> x[1] == order && x[2] == p, list)
@@ -37,7 +43,8 @@ function build_jet(f, ::Val{iip}, p, order::Val{P}, length = nothing) where {P,
         d = get_coefficient(fu, index - 1) / index
         u = append_coefficient(u, d)
     end
-    jet = build_function(u, u0, t0; expression = Val(false), cse = true)
+    u_term = make_term.(u)
+    jet = build_function(u_term, u0, t0; expression = Val(false), cse = true)
     if !haskey(JET_CACHE, f)
         JET_CACHE[f] = []
     end

lib/OrdinaryDiffEqTaylorSeries/src/algorithms.jl

@@ -15,11 +15,23 @@ end
 
 @doc explicit_rk_docstring(
     "An arbitrary-order explicit Taylor series method.",
-    "ExplicitTaylor2")
+    "ExplicitTaylor")
 Base.@kwdef struct ExplicitTaylor{P, StageLimiter, StepLimiter, Thread} <:
                    OrdinaryDiffEqAdaptiveAlgorithm
     order::Val{P} = Val{1}()
     stage_limiter!::StageLimiter = trivial_limiter!
     step_limiter!::StepLimiter = trivial_limiter!
     thread::Thread = False()
 end
+
+@doc explicit_rk_docstring(
+    "An adaptive order explicit Taylor series method.",
+    "ExplicitTaylorAdaptiveOrder")
+Base.@kwdef struct ExplicitTaylorAdaptiveOrder{StageLimiter, StepLimiter, Thread} <:
+                   OrdinaryDiffEqAdaptiveAlgorithm
+    min_order::Int = 1
+    max_order::Int = 10
+    stage_limiter!::StageLimiter = trivial_limiter!
+    step_limiter!::StepLimiter = trivial_limiter!
+    thread::Thread = False()
+end