Use FunctionWrapper to improve type stability

Author: tansongchen

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

@@ -25,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

lib/OrdinaryDiffEqTaylorSeries/src/TaylorSeries_caches.jl

@@ -39,10 +39,10 @@ end
 get_fsalfirstlast(cache::ExplicitTaylor2Cache, u) = (cache.k1, cache.k1)
 
 @cache struct ExplicitTaylorCache{
-    P, uType, taylorType, uNoUnitsType, StageLimiter, StepLimiter,
+    P, tType, uType, taylorType, uNoUnitsType, StageLimiter, StepLimiter,
     Thread} <: OrdinaryDiffEqMutableCache
     order::Val{P}
-    jet::Function
+    jet::FunctionWrapper{Nothing, Tuple{taylorType, uType, tType}}
     u::uType
     uprev::uType
     utaylor::taylorType
@@ -54,45 +54,49 @@ get_fsalfirstlast(cache::ExplicitTaylor2Cache, u) = (cache.k1, cache.k1)
     thread::Thread
 end
 
-function alg_cache(alg::ExplicitTaylor, 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{true}) where {uEltypeNoUnits, uBottomEltypeNoUnits, tTypeNoUnits}
-    _, jet_iip = build_jet(f, p, alg.order, length(u))
+        ::Val{true}) where {P, uEltypeNoUnits, uBottomEltypeNoUnits, tTypeNoUnits}
+    _, 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(alg.order, 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
 
 get_fsalfirstlast(cache::ExplicitTaylorCache, u) = (cache.u, cache.u)
 
-struct ExplicitTaylorConstantCache{P} <: OrdinaryDiffEqConstantCache
+struct ExplicitTaylorConstantCache{P, taylorType, uType, tType} <:
+       OrdinaryDiffEqConstantCache
     order::Val{P}
-    jet::Function
+    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
 
-@cache struct ExplicitTaylorAdaptiveOrderCache{
-    uType, taylorType, uNoUnitsType, StageLimiter, StepLimiter,
+@cache struct ExplicitTaylorAdaptiveOrderCache{P, Q,
+    tType, uType, taylorType, uNoUnitsType, StageLimiter, StepLimiter,
     Thread} <: OrdinaryDiffEqMutableCache
-    min_order::Int
-    max_order::Int
-    current_order::Ref{Int}
-    jets::Vector{Function}
+    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
@@ -108,45 +112,51 @@ function alg_cache(
         ::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))
+    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
-    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)
+    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 <: OrdinaryDiffEqConstantCache
-    min_order::Int
-    max_order::Int
-    current_order::Ref{Int}
-    jets::Vector{Function}
+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}
-    jets = Function[]
-    for order in (alg.min_order):(alg.max_order)
+    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, Val(order), length(u))
+            jet, _ = build_jet(f, p, order, length(u))
         else
-            jet = build_jet(f, p, Val(order))
+            jet = build_jet(f, p, order)
         end
         push!(jets, jet)
     end
-    current_order = Ref{Int}(alg.min_order)
+    current_order = Ref(min_order_value)
     ExplicitTaylorAdaptiveOrderConstantCache(
         alg.min_order, alg.max_order, current_order, jets)
 end

lib/OrdinaryDiffEqTaylorSeries/src/TaylorSeries_perform_step.jl

@@ -100,13 +100,6 @@ end
 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!(
@@ -115,7 +108,9 @@ end
     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
+    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)
@@ -125,8 +120,8 @@ 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)
+        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(
@@ -158,8 +153,9 @@ end
 end
 
 function initialize!(integrator, cache::ExplicitTaylorAdaptiveOrderConstantCache)
-    integrator.kshortsize = cache.max_order
-    integrator.k = typeof(integrator.k)(undef, cache.max_order)
+    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
 
@@ -169,16 +165,18 @@ end
     alg = unwrap_alg(integrator, false)
     @unpack jets, current_order, min_order, max_order = cache
 
-    jet_index = current_order[] - min_order + 1
+    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, current_order[] - 1)
-        end_order = min(max_order, current_order[] + 1)
+        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

lib/OrdinaryDiffEqTaylorSeries/src/alg_utils.jl

@@ -4,7 +4,7 @@ alg_stability_size(alg::ExplicitTaylor2) = 1
 alg_order(::ExplicitTaylor{P}) where {P} = P
 alg_stability_size(alg::ExplicitTaylor) = 1
 
-alg_order(alg::ExplicitTaylorAdaptiveOrder) = alg.min_order
+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[]
 
@@ -16,11 +16,14 @@ function make_term(a)
     term(TaylorScalar, Symbolics.unwrap(a.value), map(Symbolics.unwrap, a.partials))
 end
 
-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)
-        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)
@@ -32,7 +35,7 @@ function build_jet(f::ODEFunction{iip}, p, order::Val{P}, length = nothing) wher
         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)
@@ -45,10 +48,9 @@ function build_jet(f::ODEFunction{iip}, p, order::Val{P}, length = nothing) wher
     end
     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] = []
+    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

@@ -27,10 +27,10 @@ end
 @doc explicit_rk_docstring(
     "An adaptive order explicit Taylor series method.",
     "ExplicitTaylorAdaptiveOrder")
-Base.@kwdef struct ExplicitTaylorAdaptiveOrder{StageLimiter, StepLimiter, Thread} <:
+Base.@kwdef struct ExplicitTaylorAdaptiveOrder{P, Q, StageLimiter, StepLimiter, Thread} <:
                    OrdinaryDiffEqAdaptiveAlgorithm
-    min_order::Int = 1
-    max_order::Int = 10
+    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()

lib/OrdinaryDiffEqTaylorSeries/test/runtests.jl

@@ -3,7 +3,7 @@ using Test
 
 @testset "Taylor2 Convergence Tests" begin
     # Test convergence
-    dts = 2. .^ (-8:-4)
+    dts = 2.0 .^ (-8:-4)
     testTol = 0.2
     sim = test_convergence(dts, prob_ode_linear, ExplicitTaylor2())
     @test sim.𝒪est[:final]≈2 atol=testTol
@@ -13,10 +13,10 @@ end
 
 @testset "Taylor Convergence Tests" begin
     # Test convergence
-    dts = 2. .^ (-8:-4)
+    dts = 2.0 .^ (-8:-4)
     testTol = 0.2
     for N in 3:4
-        alg = ExplicitTaylor(order=Val(N))
+        alg = ExplicitTaylor(order = Val(N))
         sim = test_convergence(dts, prob_ode_linear, alg)
         @test sim.𝒪est[:final]≈N atol=testTol
         sim = test_convergence(dts, prob_ode_2Dlinear, alg)
@@ -25,7 +25,7 @@ end
 end
 
 @testset "Taylor Adaptive time-step Tests" begin
-    sol = solve(prob_ode_linear, ExplicitTaylor(order=Val(4)))
+    sol = solve(prob_ode_linear, ExplicitTaylor(order = Val(4)))
     @test length(sol) < 20
 end