From e451f4036f5bb7a1020299aa491ec05ca05f95ec Mon Sep 17 00:00:00 2001 From: Songchen Tan Date: Wed, 9 Apr 2025 14:49:01 -0400 Subject: [PATCH 1/2] Adaptive order version of ExplicitTaylor --- .../src/OrdinaryDiffEqTaylorSeries.jl | 8 +- .../src/TaylorSeries_caches.jl | 86 +++++++++++-- .../src/TaylorSeries_perform_step.jl | 113 +++++++++++++++++- .../src/alg_utils.jl | 17 ++- .../src/algorithms.jl | 14 ++- .../test/runtests.jl | 11 +- 6 files changed, 225 insertions(+), 24 deletions(-) diff --git a/lib/OrdinaryDiffEqTaylorSeries/src/OrdinaryDiffEqTaylorSeries.jl b/lib/OrdinaryDiffEqTaylorSeries/src/OrdinaryDiffEqTaylorSeries.jl index 08b724af5f..e48d82f605 100644 --- a/lib/OrdinaryDiffEqTaylorSeries/src/OrdinaryDiffEqTaylorSeries.jl +++ b/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 diff --git a/lib/OrdinaryDiffEqTaylorSeries/src/TaylorSeries_caches.jl b/lib/OrdinaryDiffEqTaylorSeries/src/TaylorSeries_caches.jl index b07e9e912d..d1921b362e 100644 --- a/lib/OrdinaryDiffEqTaylorSeries/src/TaylorSeries_caches.jl +++ b/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,23 +54,25 @@ 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 -struct ExplicitTaylorConstantCache{P, jetType} <: OrdinaryDiffEqConstantCache +get_fsalfirstlast(cache::ExplicitTaylorCache, u) = (cache.u, cache.u) + +struct ExplicitTaylorConstantCache{P} <: OrdinaryDiffEqConstantCache order::Val{P} - jet::jetType + jet::Function end function alg_cache(::ExplicitTaylor{P}, u, rate_prototype, ::Type{uEltypeNoUnits}, ::Type{uBottomEltypeNoUnits}, ::Type{tTypeNoUnits}, uprev, uprev2, f, t, @@ -84,5 +86,67 @@ function alg_cache(::ExplicitTaylor{P}, u, rate_prototype, ::Type{uEltypeNoUnits ExplicitTaylorConstantCache(Val(P), jet) 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) + +struct ExplicitTaylorAdaptiveOrderConstantCache <: OrdinaryDiffEqConstantCache + min_order::Int + max_order::Int + current_order::Ref{Int} + jets::Vector{Function} +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) + if u isa AbstractArray + jet, _ = build_jet(f, p, Val(order), length(u)) + else + jet = build_jet(f, p, Val(order)) + end + push!(jets, jet) + end + current_order = Ref{Int}(alg.min_order) + ExplicitTaylorAdaptiveOrderConstantCache( + alg.min_order, alg.max_order, current_order, jets) +end diff --git a/lib/OrdinaryDiffEqTaylorSeries/src/TaylorSeries_perform_step.jl b/lib/OrdinaryDiffEqTaylorSeries/src/TaylorSeries_perform_step.jl index 3b48979993..7f247337a3 100644 --- a/lib/OrdinaryDiffEqTaylorSeries/src/TaylorSeries_perform_step.jl +++ b/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,7 +90,7 @@ 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) @@ -98,3 +98,112 @@ 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 + +function initialize!(integrator, cache::ExplicitTaylorAdaptiveOrderConstantCache) + integrator.kshortsize = cache.max_order + integrator.k = typeof(integrator.k)(undef, cache.max_order) + 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 + + jet_index = current_order[] - min_order + 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) + 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 diff --git a/lib/OrdinaryDiffEqTaylorSeries/src/alg_utils.jl b/lib/OrdinaryDiffEqTaylorSeries/src/alg_utils.jl index 06907a8888..04ddfce10b 100644 --- a/lib/OrdinaryDiffEqTaylorSeries/src/alg_utils.jl +++ b/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 diff --git a/lib/OrdinaryDiffEqTaylorSeries/src/algorithms.jl b/lib/OrdinaryDiffEqTaylorSeries/src/algorithms.jl index 68694319d5..189d11ec35 100644 --- a/lib/OrdinaryDiffEqTaylorSeries/src/algorithms.jl +++ b/lib/OrdinaryDiffEqTaylorSeries/src/algorithms.jl @@ -15,7 +15,7 @@ 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}() @@ -23,3 +23,15 @@ Base.@kwdef struct ExplicitTaylor{P, StageLimiter, StepLimiter, Thread} <: 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 diff --git a/lib/OrdinaryDiffEqTaylorSeries/test/runtests.jl b/lib/OrdinaryDiffEqTaylorSeries/test/runtests.jl index ee42ac54d2..8022d9aedf 100644 --- a/lib/OrdinaryDiffEqTaylorSeries/test/runtests.jl +++ b/lib/OrdinaryDiffEqTaylorSeries/test/runtests.jl @@ -11,7 +11,7 @@ using Test @test sim.𝒪est[:final]≈2 atol=testTol end -@testset "TaylorN Convergence Tests" begin +@testset "Taylor Convergence Tests" begin # Test convergence dts = 2. .^ (-8:-4) testTol = 0.2 @@ -24,7 +24,12 @@ end end end -@testset "TaylorN Adaptive Tests" begin - sol = solve(prob_ode_linear, ExplicitTaylor(order=Val(2))) +@testset "Taylor Adaptive time-step Tests" begin + sol = solve(prob_ode_linear, ExplicitTaylor(order=Val(4))) + @test length(sol) < 20 +end + +@testset "Taylor Adaptive time-step Adaptive order Tests" begin + sol = solve(prob_ode_linear, ExplicitTaylorAdaptiveOrder()) @test length(sol) < 20 end From bf876245482105ec40a0558e28bcf70707375745 Mon Sep 17 00:00:00 2001 From: Songchen Tan Date: Mon, 5 May 2025 16:59:39 -0400 Subject: [PATCH 2/2] Use FunctionWrapper to improve type stability --- lib/OrdinaryDiffEqTaylorSeries/Project.toml | 2 + .../src/OrdinaryDiffEqTaylorSeries.jl | 2 + .../src/TaylorSeries_caches.jl | 76 +++++++++++-------- .../src/TaylorSeries_perform_step.jl | 28 ++++--- .../src/alg_utils.jl | 22 +++--- .../src/algorithms.jl | 6 +- .../test/runtests.jl | 8 +- 7 files changed, 79 insertions(+), 65 deletions(-) diff --git a/lib/OrdinaryDiffEqTaylorSeries/Project.toml b/lib/OrdinaryDiffEqTaylorSeries/Project.toml index 416f8311b3..339eee49a8 100644 --- a/lib/OrdinaryDiffEqTaylorSeries/Project.toml +++ b/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" diff --git a/lib/OrdinaryDiffEqTaylorSeries/src/OrdinaryDiffEqTaylorSeries.jl b/lib/OrdinaryDiffEqTaylorSeries/src/OrdinaryDiffEqTaylorSeries.jl index e48d82f605..6f82c3616d 100644 --- a/lib/OrdinaryDiffEqTaylorSeries/src/OrdinaryDiffEqTaylorSeries.jl +++ b/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 diff --git a/lib/OrdinaryDiffEqTaylorSeries/src/TaylorSeries_caches.jl b/lib/OrdinaryDiffEqTaylorSeries/src/TaylorSeries_caches.jl index d1921b362e..35d2468f3b 100644 --- a/lib/OrdinaryDiffEqTaylorSeries/src/TaylorSeries_caches.jl +++ b/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,17 +112,19 @@ 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) @@ -126,27 +132,31 @@ 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 diff --git a/lib/OrdinaryDiffEqTaylorSeries/src/TaylorSeries_perform_step.jl b/lib/OrdinaryDiffEqTaylorSeries/src/TaylorSeries_perform_step.jl index 7f247337a3..9fd745d978 100644 --- a/lib/OrdinaryDiffEqTaylorSeries/src/TaylorSeries_perform_step.jl +++ b/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,7 +165,9 @@ 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) @@ -177,8 +175,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, current_order[] + 1) for i in start_order:end_order A = i * i utilde = TaylorDiff.get_coefficient(utaylor, i) * dt^i diff --git a/lib/OrdinaryDiffEqTaylorSeries/src/alg_utils.jl b/lib/OrdinaryDiffEqTaylorSeries/src/alg_utils.jl index 04ddfce10b..640886adcf 100644 --- a/lib/OrdinaryDiffEqTaylorSeries/src/alg_utils.jl +++ b/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 diff --git a/lib/OrdinaryDiffEqTaylorSeries/src/algorithms.jl b/lib/OrdinaryDiffEqTaylorSeries/src/algorithms.jl index 189d11ec35..c6dc0aa7ae 100644 --- a/lib/OrdinaryDiffEqTaylorSeries/src/algorithms.jl +++ b/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() diff --git a/lib/OrdinaryDiffEqTaylorSeries/test/runtests.jl b/lib/OrdinaryDiffEqTaylorSeries/test/runtests.jl index 8022d9aedf..529585a4bf 100644 --- a/lib/OrdinaryDiffEqTaylorSeries/test/runtests.jl +++ b/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