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 08b724af5f..6f82c3616d 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!, @@ -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 diff --git a/lib/OrdinaryDiffEqTaylorSeries/src/TaylorSeries_caches.jl b/lib/OrdinaryDiffEqTaylorSeries/src/TaylorSeries_caches.jl index b07e9e912d..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, 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 diff --git a/lib/OrdinaryDiffEqTaylorSeries/src/TaylorSeries_perform_step.jl b/lib/OrdinaryDiffEqTaylorSeries/src/TaylorSeries_perform_step.jl index 3b48979993..9fd745d978 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,110 @@ 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 diff --git a/lib/OrdinaryDiffEqTaylorSeries/src/alg_utils.jl b/lib/OrdinaryDiffEqTaylorSeries/src/alg_utils.jl index 06907a8888..640886adcf 100644 --- a/lib/OrdinaryDiffEqTaylorSeries/src/alg_utils.jl +++ b/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 diff --git a/lib/OrdinaryDiffEqTaylorSeries/src/algorithms.jl b/lib/OrdinaryDiffEqTaylorSeries/src/algorithms.jl index 68694319d5..c6dc0aa7ae 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{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 diff --git a/lib/OrdinaryDiffEqTaylorSeries/test/runtests.jl b/lib/OrdinaryDiffEqTaylorSeries/test/runtests.jl index ee42ac54d2..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 @@ -11,12 +11,12 @@ 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) + 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) @@ -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