Skip to content

Commit e451f40

Browse files
tansongchentansongchen
committed
Adaptive order version of ExplicitTaylor
1 parent 2239fc8 commit e451f40

File tree

6 files changed

+225
-24
lines changed

6 files changed

+225
-24
lines changed

‎lib/OrdinaryDiffEqTaylorSeries/src/OrdinaryDiffEqTaylorSeries.jl

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,13 @@
11
module OrdinaryDiffEqTaylorSeries
22

3-
import OrdinaryDiffEqCore: alg_order, alg_stability_size, explicit_rk_docstring,
3+
import OrdinaryDiffEqCore: alg_order, get_current_adaptive_order, get_current_alg_order,
4+
alg_stability_size, explicit_rk_docstring,
45
OrdinaryDiffEqAdaptiveAlgorithm, OrdinaryDiffEqMutableCache,
56
alg_cache,
67
OrdinaryDiffEqConstantCache, @fold, trivial_limiter!,
8+
step_accept_controller!,
9+
stepsize_controller!,
10+
unwrap_alg,
711
constvalue, @unpack, perform_step!, calculate_residuals, @cache,
812
calculate_residuals!, _ode_interpolant, _ode_interpolant!,
913
CompiledFloats, @OnDemandTableauExtract, initialize!,
@@ -55,6 +59,6 @@ PrecompileTools.@compile_workload begin
5559
solver_list = nothing
5660
end
5761

58-
export ExplicitTaylor2, ExplicitTaylor, DAETS
62+
export ExplicitTaylor2, ExplicitTaylor, ExplicitTaylorAdaptiveOrder
5963

6064
end

‎lib/OrdinaryDiffEqTaylorSeries/src/TaylorSeries_caches.jl

Lines changed: 75 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -39,10 +39,10 @@ end
3939
get_fsalfirstlast(cache::ExplicitTaylor2Cache, u) = (cache.k1, cache.k1)
4040

4141
@cache struct ExplicitTaylorCache{
42-
P, jetType, uType, taylorType, uNoUnitsType, StageLimiter, StepLimiter,
42+
P, uType, taylorType, uNoUnitsType, StageLimiter, StepLimiter,
4343
Thread} <: OrdinaryDiffEqMutableCache
4444
order::Val{P}
45-
jet::jetType
45+
jet::Function
4646
u::uType
4747
uprev::uType
4848
utaylor::taylorType
@@ -54,23 +54,25 @@ get_fsalfirstlast(cache::ExplicitTaylor2Cache, u) = (cache.k1, cache.k1)
5454
thread::Thread
5555
end
5656

57-
function alg_cache(alg::ExplicitTaylor{P}, u, rate_prototype, ::Type{uEltypeNoUnits},
57+
function alg_cache(alg::ExplicitTaylor, u, rate_prototype, ::Type{uEltypeNoUnits},
5858
::Type{uBottomEltypeNoUnits}, ::Type{tTypeNoUnits}, uprev, uprev2, f, t,
5959
dt, reltol, p, calck,
60-
::Val{true}) where {P, uEltypeNoUnits, uBottomEltypeNoUnits, tTypeNoUnits}
61-
_, jet_iip = build_jet(f, p, Val(P), length(u))
62-
utaylor = TaylorDiff.make_seed(u, zero(u), Val(P))
60+
::Val{true}) where {uEltypeNoUnits, uBottomEltypeNoUnits, tTypeNoUnits}
61+
_, jet_iip = build_jet(f, p, alg.order, length(u))
62+
utaylor = TaylorDiff.make_seed(u, zero(u), alg.order)
6363
utilde = zero(u)
6464
atmp = similar(u, uEltypeNoUnits)
6565
recursivefill!(atmp, false)
6666
tmp = zero(u)
67-
ExplicitTaylorCache(Val(P), jet_iip, u, uprev, utaylor, utilde, tmp, atmp,
67+
ExplicitTaylorCache(alg.order, jet_iip, u, uprev, utaylor, utilde, tmp, atmp,
6868
alg.stage_limiter!, alg.step_limiter!, alg.thread)
6969
end
7070

71-
struct ExplicitTaylorConstantCache{P, jetType} <: OrdinaryDiffEqConstantCache
71+
get_fsalfirstlast(cache::ExplicitTaylorCache, u) = (cache.u, cache.u)
72+
73+
struct ExplicitTaylorConstantCache{P} <: OrdinaryDiffEqConstantCache
7274
order::Val{P}
73-
jet::jetType
75+
jet::Function
7476
end
7577
function alg_cache(::ExplicitTaylor{P}, u, rate_prototype, ::Type{uEltypeNoUnits},
7678
::Type{uBottomEltypeNoUnits}, ::Type{tTypeNoUnits}, uprev, uprev2, f, t,
@@ -84,5 +86,67 @@ function alg_cache(::ExplicitTaylor{P}, u, rate_prototype, ::Type{uEltypeNoUnits
8486
ExplicitTaylorConstantCache(Val(P), jet)
8587
end
8688

87-
# FSAL currently not used, providing dummy implementation to satisfy the interface
88-
get_fsalfirstlast(cache::ExplicitTaylorCache, u) = (cache.u, cache.u)
89+
@cache struct ExplicitTaylorAdaptiveOrderCache{
90+
uType, taylorType, uNoUnitsType, StageLimiter, StepLimiter,
91+
Thread} <: OrdinaryDiffEqMutableCache
92+
min_order::Int
93+
max_order::Int
94+
current_order::Ref{Int}
95+
jets::Vector{Function}
96+
u::uType
97+
uprev::uType
98+
utaylor::taylorType
99+
utilde::uType
100+
tmp::uType
101+
atmp::uNoUnitsType
102+
stage_limiter!::StageLimiter
103+
step_limiter!::StepLimiter
104+
thread::Thread
105+
end
106+
function alg_cache(
107+
alg::ExplicitTaylorAdaptiveOrder, u, rate_prototype, ::Type{uEltypeNoUnits},
108+
::Type{uBottomEltypeNoUnits}, ::Type{tTypeNoUnits}, uprev, uprev2, f, t,
109+
dt, reltol, p, calck,
110+
::Val{true}) where {uEltypeNoUnits, uBottomEltypeNoUnits, tTypeNoUnits}
111+
jets = Function[]
112+
for order in (alg.min_order):(alg.max_order)
113+
_, jet_iip = build_jet(f, p, Val(order), length(u))
114+
push!(jets, jet_iip)
115+
end
116+
utaylor = TaylorDiff.make_seed(u, zero(u), Val(alg.max_order))
117+
utilde = zero(u)
118+
atmp = similar(u, uEltypeNoUnits)
119+
recursivefill!(atmp, false)
120+
tmp = zero(u)
121+
current_order = Ref{Int}(alg.min_order)
122+
ExplicitTaylorAdaptiveOrderCache(alg.min_order, alg.max_order, current_order,
123+
jets, u, uprev, utaylor, utilde, tmp, atmp,
124+
alg.stage_limiter!, alg.step_limiter!, alg.thread)
125+
end
126+
127+
get_fsalfirstlast(cache::ExplicitTaylorAdaptiveOrderCache, u) = (cache.u, cache.u)
128+
129+
struct ExplicitTaylorAdaptiveOrderConstantCache <: OrdinaryDiffEqConstantCache
130+
min_order::Int
131+
max_order::Int
132+
current_order::Ref{Int}
133+
jets::Vector{Function}
134+
end
135+
function alg_cache(
136+
alg::ExplicitTaylorAdaptiveOrder, u, rate_prototype, ::Type{uEltypeNoUnits},
137+
::Type{uBottomEltypeNoUnits}, ::Type{tTypeNoUnits}, uprev, uprev2, f, t,
138+
dt, reltol, p, calck,
139+
::Val{false}) where {uEltypeNoUnits, uBottomEltypeNoUnits, tTypeNoUnits}
140+
jets = Function[]
141+
for order in (alg.min_order):(alg.max_order)
142+
if u isa AbstractArray
143+
jet, _ = build_jet(f, p, Val(order), length(u))
144+
else
145+
jet = build_jet(f, p, Val(order))
146+
end
147+
push!(jets, jet)
148+
end
149+
current_order = Ref{Int}(alg.min_order)
150+
ExplicitTaylorAdaptiveOrderConstantCache(
151+
alg.min_order, alg.max_order, current_order, jets)
152+
end

‎lib/OrdinaryDiffEqTaylorSeries/src/TaylorSeries_perform_step.jl

Lines changed: 111 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -60,7 +60,7 @@ end
6060
utaylor = jet(uprev, t)
6161
u = map(x -> evaluate_polynomial(x, dt), utaylor)
6262
if integrator.opts.adaptive
63-
utilde = TaylorDiff.get_coefficient(utaylor, P) * dt^(P + 1)
63+
utilde = TaylorDiff.get_coefficient(utaylor, P) * dt^P
6464
atmp = calculate_residuals(utilde, uprev, u, integrator.opts.abstol,
6565
integrator.opts.reltol, integrator.opts.internalnorm, t)
6666
integrator.EEst = integrator.opts.internalnorm(atmp, t)
@@ -90,11 +90,120 @@ end
9090
end
9191
if integrator.opts.adaptive
9292
@.. broadcast=false thread=thread utilde=TaylorDiff.get_coefficient(utaylor, P) *
93-
dt^(P + 1)
93+
dt^P
9494
calculate_residuals!(atmp, utilde, uprev, u, integrator.opts.abstol,
9595
integrator.opts.reltol, integrator.opts.internalnorm, t)
9696
integrator.EEst = integrator.opts.internalnorm(atmp, t)
9797
end
9898
OrdinaryDiffEqCore.increment_nf!(integrator.stats, P + 1)
9999
return nothing
100100
end
101+
102+
function initialize!(integrator, cache::ExplicitTaylorAdaptiveOrderCache)
103+
integrator.kshortsize = cache.max_order
104+
resize!(integrator.k, cache.max_order)
105+
# Setup k pointers
106+
for i in 1:(cache.max_order)
107+
integrator.k[i] = get_coefficient(cache.utaylor, i)
108+
end
109+
return nothing
110+
end
111+
112+
@muladd function perform_step!(
113+
integrator, cache::ExplicitTaylorAdaptiveOrderCache, repeat_step = false)
114+
@unpack t, dt, uprev, u, f, p = integrator
115+
alg = unwrap_alg(integrator, false)
116+
@unpack jets, current_order, min_order, max_order, utaylor, utilde, tmp, atmp, thread = cache
117+
118+
jet_index = current_order[] - min_order + 1
119+
# compute one additional order for adaptive order
120+
jet = jets[jet_index + 1]
121+
jet(utaylor, uprev, t)
122+
for i in eachindex(utaylor)
123+
u[i] = @inline evaluate_polynomial(utaylor[i], dt)
124+
end
125+
OrdinaryDiffEqCore.increment_nf!(integrator.stats, current_order[] + 1)
126+
if integrator.opts.adaptive
127+
min_work = Inf
128+
start_order = max(min_order, current_order[] - 1)
129+
end_order = min(max_order, current_order[] + 1)
130+
for i in start_order:end_order
131+
A = i * i
132+
@.. broadcast=false thread=thread utilde=TaylorDiff.get_coefficient(
133+
utaylor, i) * dt^i
134+
calculate_residuals!(atmp, utilde, uprev, u, integrator.opts.abstol,
135+
integrator.opts.reltol, integrator.opts.internalnorm, t)
136+
EEst = integrator.opts.internalnorm(atmp, t)
137+
138+
# backup
139+
e = integrator.EEst
140+
qold = integrator.qold
141+
# calculate dt
142+
integrator.EEst = EEst
143+
dtpropose = step_accept_controller!(integrator, alg,
144+
stepsize_controller!(integrator, alg))
145+
# restore
146+
integrator.EEst = e
147+
integrator.qold = qold
148+
149+
work = A / dtpropose
150+
if work < min_work
151+
cache.current_order[] = i
152+
min_work = work
153+
integrator.EEst = EEst
154+
end
155+
end
156+
end
157+
return nothing
158+
end
159+
160+
function initialize!(integrator, cache::ExplicitTaylorAdaptiveOrderConstantCache)
161+
integrator.kshortsize = cache.max_order
162+
integrator.k = typeof(integrator.k)(undef, cache.max_order)
163+
return nothing
164+
end
165+
166+
@muladd function perform_step!(
167+
integrator, cache::ExplicitTaylorAdaptiveOrderConstantCache, repeat_step = false)
168+
@unpack t, dt, uprev, u, f, p = integrator
169+
alg = unwrap_alg(integrator, false)
170+
@unpack jets, current_order, min_order, max_order = cache
171+
172+
jet_index = current_order[] - min_order + 1
173+
# compute one additional order for adaptive order
174+
jet = jets[jet_index + 1]
175+
utaylor = jet(uprev, t)
176+
u = map(x -> evaluate_polynomial(x, dt), utaylor)
177+
OrdinaryDiffEqCore.increment_nf!(integrator.stats, current_order[] + 1)
178+
if integrator.opts.adaptive
179+
min_work = Inf
180+
start_order = max(min_order, current_order[] - 1)
181+
end_order = min(max_order, current_order[] + 1)
182+
for i in start_order:end_order
183+
A = i * i
184+
utilde = TaylorDiff.get_coefficient(utaylor, i) * dt^i
185+
atmp = calculate_residuals(utilde, uprev, u, integrator.opts.abstol,
186+
integrator.opts.reltol, integrator.opts.internalnorm, t)
187+
EEst = integrator.opts.internalnorm(atmp, t)
188+
189+
# backup
190+
e = integrator.EEst
191+
qold = integrator.qold
192+
# calculate dt
193+
integrator.EEst = EEst
194+
dtpropose = step_accept_controller!(integrator, alg,
195+
stepsize_controller!(integrator, alg))
196+
# restore
197+
integrator.EEst = e
198+
integrator.qold = qold
199+
200+
work = A / dtpropose
201+
if work < min_work
202+
cache.current_order[] = i
203+
min_work = work
204+
integrator.EEst = EEst
205+
end
206+
end
207+
end
208+
return nothing
209+
end

‎lib/OrdinaryDiffEqTaylorSeries/src/alg_utils.jl

Lines changed: 12 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -4,13 +4,19 @@ alg_stability_size(alg::ExplicitTaylor2) = 1
44
alg_order(::ExplicitTaylor{P}) where {P} = P
55
alg_stability_size(alg::ExplicitTaylor) = 1
66

7-
JET_CACHE = IdDict()
7+
alg_order(alg::ExplicitTaylorAdaptiveOrder) = alg.min_order
8+
get_current_adaptive_order(::ExplicitTaylorAdaptiveOrder, cache) = cache.current_order[]
9+
get_current_alg_order(::ExplicitTaylorAdaptiveOrder, cache) = cache.current_order[]
810

9-
function build_jet(f::ODEFunction{iip}, p, order, length = nothing) where {iip}
10-
build_jet(f, Val{iip}(), p, order, length)
11+
TaylorScalar{T, P}(x::TaylorScalar{T, Q}) where {T, P, Q} = TaylorScalar{P}(x)
12+
13+
const JET_CACHE = IdDict()
14+
15+
function make_term(a)
16+
term(TaylorScalar, Symbolics.unwrap(a.value), map(Symbolics.unwrap, a.partials))
1117
end
1218

13-
function build_jet(f, ::Val{iip}, p, order::Val{P}, length = nothing) where {P, iip}
19+
function build_jet(f::ODEFunction{iip}, p, order::Val{P}, length = nothing) where {P, iip}
1420
if haskey(JET_CACHE, f)
1521
list = JET_CACHE[f]
1622
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,
3743
d = get_coefficient(fu, index - 1) / index
3844
u = append_coefficient(u, d)
3945
end
40-
jet = build_function(u, u0, t0; expression = Val(false), cse = true)
46+
u_term = make_term.(u)
47+
jet = build_function(u_term, u0, t0; expression = Val(false), cse = true)
4148
if !haskey(JET_CACHE, f)
4249
JET_CACHE[f] = []
4350
end

‎lib/OrdinaryDiffEqTaylorSeries/src/algorithms.jl

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,11 +15,23 @@ end
1515

1616
@doc explicit_rk_docstring(
1717
"An arbitrary-order explicit Taylor series method.",
18-
"ExplicitTaylor2")
18+
"ExplicitTaylor")
1919
Base.@kwdef struct ExplicitTaylor{P, StageLimiter, StepLimiter, Thread} <:
2020
OrdinaryDiffEqAdaptiveAlgorithm
2121
order::Val{P} = Val{1}()
2222
stage_limiter!::StageLimiter = trivial_limiter!
2323
step_limiter!::StepLimiter = trivial_limiter!
2424
thread::Thread = False()
2525
end
26+
27+
@doc explicit_rk_docstring(
28+
"An adaptive order explicit Taylor series method.",
29+
"ExplicitTaylorAdaptiveOrder")
30+
Base.@kwdef struct ExplicitTaylorAdaptiveOrder{StageLimiter, StepLimiter, Thread} <:
31+
OrdinaryDiffEqAdaptiveAlgorithm
32+
min_order::Int = 1
33+
max_order::Int = 10
34+
stage_limiter!::StageLimiter = trivial_limiter!
35+
step_limiter!::StepLimiter = trivial_limiter!
36+
thread::Thread = False()
37+
end

‎lib/OrdinaryDiffEqTaylorSeries/test/runtests.jl

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ using Test
1111
@test sim.𝒪est[:final]≈2 atol=testTol
1212
end
1313

14-
@testset "TaylorN Convergence Tests" begin
14+
@testset "Taylor Convergence Tests" begin
1515
# Test convergence
1616
dts = 2. .^ (-8:-4)
1717
testTol = 0.2
@@ -24,7 +24,12 @@ end
2424
end
2525
end
2626

27-
@testset "TaylorN Adaptive Tests" begin
28-
sol = solve(prob_ode_linear, ExplicitTaylor(order=Val(2)))
27+
@testset "Taylor Adaptive time-step Tests" begin
28+
sol = solve(prob_ode_linear, ExplicitTaylor(order=Val(4)))
29+
@test length(sol) < 20
30+
end
31+
32+
@testset "Taylor Adaptive time-step Adaptive order Tests" begin
33+
sol = solve(prob_ode_linear, ExplicitTaylorAdaptiveOrder())
2934
@test length(sol) < 20
3035
end

0 commit comments

Comments
 (0)