[WIP] Adaptive order version of ExplicitTaylor

lib/OrdinaryDiffEqTaylorSeries/Project.toml

@@ -6,6 +6,7 @@ version = "1.1.0"
 [deps]
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
 FastBroadcast = "7034ab61-46d4-4ed7-9d0f-46aef9175898"
+FunctionWrappers = "069b7b12-0de2-55c6-9aab-29f3d0a68a2e"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MuladdMacro = "46d2c3a1-f734-5fdb-9937-b9b9aeba4221"
 OrdinaryDiffEqCore = "bbf590c4-e513-4bbe-9b18-05decba2e5d8"
@@ -23,6 +24,7 @@ TruncatedStacktraces = "781d530d-4396-4725-bb49-402e4bee1e77"
 DiffEqBase = "6.152.2"
 DiffEqDevTools = "2.44.4"
 FastBroadcast = "0.3.5"
+FunctionWrappers = "1.1.3"
 LinearAlgebra = "<0.0.1, 1"
 MuladdMacro = "0.2.4"
 OrdinaryDiffEqCore = "1.1"

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!,
@@ -21,6 +25,8 @@ using TaylorDiff, Symbolics
 using TaylorDiff: make_seed, get_coefficient, append_coefficient, flatten
 import DiffEqBase: @def
 import OrdinaryDiffEqCore
+using FunctionWrappers
+import FunctionWrappers: FunctionWrapper
 
 using Reexport
 @reexport using DiffEqBase
@@ -55,6 +61,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, tType, uType, taylorType, uNoUnitsType, StageLimiter, StepLimiter,
     Thread} <: OrdinaryDiffEqMutableCache
     order::Val{P}
-    jet::jetType
+    jet::FunctionWrapper{Nothing, Tuple{taylorType, uType, tType}}
     u::uType
     uprev::uType
     utaylor::taylorType
@@ -58,31 +58,105 @@ function alg_cache(alg::ExplicitTaylor{P}, u, rate_prototype, ::Type{uEltypeNoUn
         ::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))
+    _, jet_iip = build_jet(f, p, P, length(u))
+    utaylor = TaylorDiff.make_seed(u, zero(u), alg.order)
+    jet_wrapped = FunctionWrapper{Nothing, Tuple{typeof(utaylor), typeof(u), typeof(t)}}(jet_iip)
     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_wrapped, u, uprev, utaylor, utilde, tmp, atmp,
         alg.stage_limiter!, alg.step_limiter!, alg.thread)
 end
 
-struct ExplicitTaylorConstantCache{P, jetType} <: OrdinaryDiffEqConstantCache
+get_fsalfirstlast(cache::ExplicitTaylorCache, u) = (cache.u, cache.u)
+
+struct ExplicitTaylorConstantCache{P, taylorType, uType, tType} <:
+       OrdinaryDiffEqConstantCache
     order::Val{P}
-    jet::jetType
+    jet::FunctionWrapper{taylorType, Tuple{uType, tType}}
 end
-function alg_cache(::ExplicitTaylor{P}, u, rate_prototype, ::Type{uEltypeNoUnits},
+function alg_cache(alg::ExplicitTaylor{P}, u, rate_prototype, ::Type{uEltypeNoUnits},
         ::Type{uBottomEltypeNoUnits}, ::Type{tTypeNoUnits}, uprev, uprev2, f, t,
         dt, reltol, p, calck,
         ::Val{false}) where {P, uEltypeNoUnits, uBottomEltypeNoUnits, tTypeNoUnits}
     if u isa AbstractArray
-        jet, _ = build_jet(f, p, Val(P), length(u))
+        jet, _ = build_jet(f, p, P, length(u))
     else
-        jet = build_jet(f, p, Val(P))
+        jet = build_jet(f, p, P)
     end
-    ExplicitTaylorConstantCache(Val(P), jet)
+    utaylor = TaylorDiff.make_seed(u, zero(u), alg.order) # not used, but needed for type
+    jet_wrapped = FunctionWrapper{typeof(utaylor), Tuple{typeof(u), typeof(t)}}(jet)
+    ExplicitTaylorConstantCache(alg.order, jet_wrapped)
 end
 
-# FSAL currently not used, providing dummy implementation to satisfy the interface
-get_fsalfirstlast(cache::ExplicitTaylorCache, u) = (cache.u, cache.u)
+@cache struct ExplicitTaylorAdaptiveOrderCache{P, Q,
+    tType, uType, taylorType, uNoUnitsType, StageLimiter, StepLimiter,
+    Thread} <: OrdinaryDiffEqMutableCache
+    min_order::Val{P}
+    max_order::Val{Q}
+    current_order::Base.RefValue{Int}
+    jets::Vector{FunctionWrapper{Nothing, Tuple{taylorType, uType, tType}}}
+    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}
+    utaylor = TaylorDiff.make_seed(u, zero(u), alg.max_order)
+    jets = FunctionWrapper{Nothing, Tuple{typeof(utaylor), typeof(u), typeof(t)}}[]
+    min_order_value = get_value(alg.min_order)
+    max_order_value = get_value(alg.max_order)
+    for order in min_order_value:max_order_value
+        jet_iip = build_jet(f, p, order, length(u))[2]
+        push!(jets, jet_iip)
+    end
+    utilde = zero(u)
+    atmp = similar(u, uEltypeNoUnits)
+    recursivefill!(atmp, false)
+    tmp = zero(u)
+    current_order = Ref(min_order_value)
+    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)
+
+struct ExplicitTaylorAdaptiveOrderConstantCache{P, Q, taylorType, uType, tType} <:
+       OrdinaryDiffEqConstantCache
+    min_order::Val{P}
+    max_order::Val{Q}
+    current_order::Base.RefValue{Int}
+    jets::Vector{FunctionWrapper{taylorType, Tuple{uType, tType}}}
+end
+function alg_cache(
+        alg::ExplicitTaylorAdaptiveOrder, u, rate_prototype, ::Type{uEltypeNoUnits},
+        ::Type{uBottomEltypeNoUnits}, ::Type{tTypeNoUnits}, uprev, uprev2, f, t,
+        dt, reltol, p, calck,
+        ::Val{false}) where {uEltypeNoUnits, uBottomEltypeNoUnits, tTypeNoUnits}
+    utaylor = TaylorDiff.make_seed(u, zero(u), alg.max_order) # not used, but needed for type
+    jets = FunctionWrapper{typeof(utaylor), Tuple{typeof(u), typeof(t)}}[]
+    min_order_value = get_value(alg.min_order)
+    max_order_value = get_value(alg.max_order)
+    for order in min_order_value:max_order_value
+        if u isa AbstractArray
+            jet, _ = build_jet(f, p, order, length(u))
+        else
+            jet = build_jet(f, p, order)
+        end
+        push!(jets, jet)
+    end
+    current_order = Ref(min_order_value)
+    ExplicitTaylorAdaptiveOrderConstantCache(
+        alg.min_order, alg.max_order, current_order, jets)
+end

lib/OrdinaryDiffEqTaylorSeries/src/TaylorSeries_perform_step.jl

@@ -60,7 +60,7 @@ end
     utaylor = jet(uprev, t)
     u = map(x -> evaluate_polynomial(x, dt), utaylor)
     if integrator.opts.adaptive
-        utilde = TaylorDiff.get_coefficient(utaylor, P) * dt^(P + 1)
+        utilde = TaylorDiff.get_coefficient(utaylor, P) * dt^P
         atmp = calculate_residuals(utilde, uprev, u, integrator.opts.abstol,
             integrator.opts.reltol, integrator.opts.internalnorm, t)
         integrator.EEst = integrator.opts.internalnorm(atmp, t)
@@ -90,11 +90,118 @@ 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)
+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
+
+    min_order_value = get_value(min_order)
+    max_order_value = get_value(max_order)
+    jet_index = current_order[] - min_order_value + 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_value, current_order[] - 1)
+        end_order = min(max_order_value - 1, 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
+
+function initialize!(integrator, cache::ExplicitTaylorAdaptiveOrderConstantCache)
+    max_order_value = get_value(cache.max_order)
+    integrator.kshortsize = max_order_value
+    integrator.k = typeof(integrator.k)(undef, max_order_value)
+    return nothing
+end
+
+@muladd function perform_step!(
+        integrator, cache::ExplicitTaylorAdaptiveOrderConstantCache, repeat_step = false)
+    @unpack t, dt, uprev, u, f, p = integrator
+    alg = unwrap_alg(integrator, false)
+    @unpack jets, current_order, min_order, max_order = cache
+
+    min_order_value = get_value(min_order)
+    max_order_value = get_value(max_order)
+    jet_index = current_order[] - min_order_value + 1
+    # compute one additional order for adaptive order
+    jet = jets[jet_index + 1]
+    utaylor = jet(uprev, t)
+    u = map(x -> evaluate_polynomial(x, dt), utaylor)
+    OrdinaryDiffEqCore.increment_nf!(integrator.stats, current_order[] + 1)
+    if integrator.opts.adaptive
+        min_work = Inf
+        start_order = max(min_order_value, current_order[] - 1)
+        end_order = min(max_order_value, current_order[] + 1)
+        for i in start_order:end_order
+            A = i * i
+            utilde = TaylorDiff.get_coefficient(utaylor, i) * dt^i
+            atmp = calculate_residuals(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,17 +4,26 @@ 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) = get_value(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}
-    if haskey(JET_CACHE, f)
-        list = JET_CACHE[f]
-        index = findfirst(x -> x[1] == order && x[2] == p, list)
-        index !== nothing && return list[index][3]
+function get_value(::Val{P}) where {P}
+    return P
+end
+
+function build_jet(f::ODEFunction{iip}, p, order, length = nothing) where {iip}
+    key = (f, order, p)
+    if haskey(JET_CACHE, key)
+        return JET_CACHE[key]
     end
     @variables t0::Real
     u0 = isnothing(length) ? Symbolics.variable(:u0) : Symbolics.variables(:u0, 1:length)
@@ -26,7 +35,7 @@ function build_jet(f, ::Val{iip}, p, order::Val{P}, length = nothing) where {P,
         f0 = f(u0, p, t0)
     end
     u = TaylorDiff.make_seed(u0, f0, Val(1))
-    for index in 2:P
+    for index in 2:order
         t = TaylorScalar{index - 1}(t0, one(t0))
         if iip
             fu = similar(u)
@@ -37,11 +46,11 @@ 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)
-    if !haskey(JET_CACHE, f)
-        JET_CACHE[f] = []
+    u_term = make_term.(u)
+    jet = build_function(u_term, u0, t0; expression = Val(false), cse = true)
+    if !haskey(JET_CACHE, key)
+        JET_CACHE[key] = jet
     end
-    push!(JET_CACHE[f], (order, p, jet))
     return jet
 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{P, Q, StageLimiter, StepLimiter, Thread} <:
+                   OrdinaryDiffEqAdaptiveAlgorithm
+    min_order::Val{P} = Val{1}()
+    max_order::Val{Q} = Val{10}()
+    stage_limiter!::StageLimiter = trivial_limiter!
+    step_limiter!::StepLimiter = trivial_limiter!
+    thread::Thread = False()
+end