Skip to content

Commit 125184e

Browse files
ChrisRackauckasChrisRackauckas
authored
Merge pull request #2620 from tansongchen/master
Explicit Taylor solvers
2 parents f137dfa + d40514a commit 125184e

File tree

8 files changed

+433
-0
lines changed

8 files changed

+433
-0
lines changed

‎lib/OrdinaryDiffEqTaylorSeries/LICENSE.md

+24
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
The OrdinaryDiffEq.jl package is licensed under the MIT "Expat" License:
2+
3+
> Copyright (c) 2016-2020: ChrisRackauckas, Yingbo Ma, Julia Computing Inc, and
4+
> other contributors:
5+
>
6+
> https://github.com/SciML/OrdinaryDiffEq.jl/graphs/contributors
7+
>
8+
> Permission is hereby granted, free of charge, to any person obtaining a copy
9+
> of this software and associated documentation files (the "Software"), to deal
10+
> in the Software without restriction, including without limitation the rights
11+
> to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
12+
> copies of the Software, and to permit persons to whom the Software is
13+
> furnished to do so, subject to the following conditions:
14+
>
15+
> The above copyright notice and this permission notice shall be included in all
16+
> copies or substantial portions of the Software.
17+
>
18+
> THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
19+
> IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
20+
> FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
21+
> AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
22+
> LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
23+
> OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
24+
> SOFTWARE.

‎lib/OrdinaryDiffEqTaylorSeries/Project.toml

+51
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
name = "OrdinaryDiffEqTaylorSeries"
2+
uuid = "9c7f1690-dd92-42a3-8318-297ee24d8d39"
3+
authors = ["ParamThakkar123 <[email protected]>"]
4+
version = "1.1.0"
5+
6+
[deps]
7+
DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
8+
FastBroadcast = "7034ab61-46d4-4ed7-9d0f-46aef9175898"
9+
LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
10+
MuladdMacro = "46d2c3a1-f734-5fdb-9937-b9b9aeba4221"
11+
OrdinaryDiffEqCore = "bbf590c4-e513-4bbe-9b18-05decba2e5d8"
12+
PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
13+
Preferences = "21216c6a-2e73-6563-6e65-726566657250"
14+
RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd"
15+
Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
16+
SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
17+
Static = "aedffcd0-7271-4cad-89d0-dc628f76c6d3"
18+
Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
19+
TaylorDiff = "b36ab563-344f-407b-a36a-4f200bebf99c"
20+
TruncatedStacktraces = "781d530d-4396-4725-bb49-402e4bee1e77"
21+
22+
[compat]
23+
DiffEqBase = "6.152.2"
24+
DiffEqDevTools = "2.44.4"
25+
FastBroadcast = "0.3.5"
26+
LinearAlgebra = "<0.0.1, 1"
27+
MuladdMacro = "0.2.4"
28+
OrdinaryDiffEqCore = "1.1"
29+
PrecompileTools = "1.2.1"
30+
Preferences = "1.4.3"
31+
Random = "<0.0.1, 1"
32+
RecursiveArrayTools = "3.27.0"
33+
Reexport = "1.2.2"
34+
SafeTestsets = "0.1.0"
35+
SciMLBase = "2.72.2"
36+
Static = "1.1.1"
37+
Symbolics = "6.28.0"
38+
TaylorDiff = "0.3.1"
39+
Test = "<0.0.1, 1"
40+
TruncatedStacktraces = "1.4.0"
41+
julia = "1.10"
42+
43+
[extras]
44+
DiffEqDevTools = "f3b72e0c-5b89-59e1-b016-84e28bfd966d"
45+
ODEProblemLibrary = "fdc4e326-1af4-4b90-96e7-779fcce2daa5"
46+
Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
47+
SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
48+
Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
49+
50+
[targets]
51+
test = ["DiffEqDevTools", "Random", "SafeTestsets", "Test", "ODEProblemLibrary"]

‎lib/OrdinaryDiffEqTaylorSeries/src/OrdinaryDiffEqTaylorSeries.jl

+60
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
module OrdinaryDiffEqTaylorSeries
2+
3+
import OrdinaryDiffEqCore: alg_order, alg_stability_size, explicit_rk_docstring,
4+
OrdinaryDiffEqAdaptiveAlgorithm, OrdinaryDiffEqMutableCache,
5+
alg_cache,
6+
OrdinaryDiffEqConstantCache, @fold, trivial_limiter!,
7+
constvalue, @unpack, perform_step!, calculate_residuals, @cache,
8+
calculate_residuals!, _ode_interpolant, _ode_interpolant!,
9+
CompiledFloats, @OnDemandTableauExtract, initialize!,
10+
perform_step!, OrdinaryDiffEqAlgorithm,
11+
CompositeAlgorithm, _ode_addsteps!, copyat_or_push!,
12+
AutoAlgSwitch, get_fsalfirstlast,
13+
full_cache, DerivativeOrderNotPossibleError
14+
import Static: False
15+
import MuladdMacro: @muladd
16+
import FastBroadcast: @..
17+
import RecursiveArrayTools: recursivefill!, recursive_unitless_bottom_eltype
18+
import LinearAlgebra: norm
19+
using TruncatedStacktraces
20+
using TaylorDiff, Symbolics
21+
using TaylorDiff: make_seed, get_coefficient, append_coefficient, flatten
22+
import DiffEqBase: @def
23+
import OrdinaryDiffEqCore
24+
25+
using Reexport
26+
@reexport using DiffEqBase
27+
28+
include("algorithms.jl")
29+
include("alg_utils.jl")
30+
include("TaylorSeries_caches.jl")
31+
include("TaylorSeries_perform_step.jl")
32+
33+
import PrecompileTools
34+
import Preferences
35+
36+
PrecompileTools.@compile_workload begin
37+
lorenz = OrdinaryDiffEqCore.lorenz
38+
lorenz_oop = OrdinaryDiffEqCore.lorenz_oop
39+
solver_list = [ExplicitTaylor2()]
40+
prob_list = []
41+
42+
if Preferences.@load_preference("PrecompileNoSpecialize", false)
43+
push!(prob_list,
44+
ODEProblem{true, SciMLBase.NoSpecialize}(lorenz, [1.0; 0.0; 0.0], (0.0, 1.0)))
45+
push!(prob_list,
46+
ODEProblem{true, SciMLBase.NoSpecialize}(lorenz, [1.0; 0.0; 0.0], (0.0, 1.0),
47+
Float64[]))
48+
end
49+
50+
for prob in prob_list, solver in solver_list
51+
solve(prob, solver)(5.0)
52+
end
53+
54+
prob_list = nothing
55+
solver_list = nothing
56+
end
57+
58+
export ExplicitTaylor2, ExplicitTaylor, DAETS
59+
60+
end

‎lib/OrdinaryDiffEqTaylorSeries/src/TaylorSeries_caches.jl

+88
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,88 @@
1+
@cache struct ExplicitTaylor2Cache{
2+
uType, rateType, uNoUnitsType, StageLimiter, StepLimiter,
3+
Thread} <: OrdinaryDiffEqMutableCache
4+
u::uType
5+
uprev::uType
6+
k1::rateType
7+
k2::rateType
8+
k3::rateType
9+
utilde::uType
10+
tmp::uType
11+
atmp::uNoUnitsType
12+
stage_limiter!::StageLimiter
13+
step_limiter!::StepLimiter
14+
thread::Thread
15+
end
16+
17+
function alg_cache(alg::ExplicitTaylor2, u, rate_prototype, ::Type{uEltypeNoUnits},
18+
::Type{uBottomEltypeNoUnits}, ::Type{tTypeNoUnits}, uprev, uprev2, f, t,
19+
dt, reltol, p, calck,
20+
::Val{true}) where {uEltypeNoUnits, uBottomEltypeNoUnits, tTypeNoUnits}
21+
k1 = zero(rate_prototype)
22+
k2 = zero(rate_prototype)
23+
k3 = zero(rate_prototype)
24+
utilde = zero(u)
25+
atmp = similar(u, uEltypeNoUnits)
26+
recursivefill!(atmp, false)
27+
tmp = zero(u)
28+
ExplicitTaylor2Cache(u, uprev, k1, k2, k3, utilde, tmp, atmp,
29+
alg.stage_limiter!, alg.step_limiter!, alg.thread)
30+
end
31+
struct ExplicitTaylor2ConstantCache <: OrdinaryDiffEqConstantCache end
32+
function alg_cache(alg::ExplicitTaylor2, u, rate_prototype, ::Type{uEltypeNoUnits},
33+
::Type{uBottomEltypeNoUnits}, ::Type{tTypeNoUnits}, uprev, uprev2, f, t,
34+
dt, reltol, p, calck,
35+
::Val{false}) where {uEltypeNoUnits, uBottomEltypeNoUnits, tTypeNoUnits}
36+
ExplicitTaylor2ConstantCache()
37+
end
38+
# FSAL currently not used, providing dummy implementation to satisfy the interface
39+
get_fsalfirstlast(cache::ExplicitTaylor2Cache, u) = (cache.k1, cache.k1)
40+
41+
@cache struct ExplicitTaylorCache{
42+
P, jetType, uType, taylorType, uNoUnitsType, StageLimiter, StepLimiter,
43+
Thread} <: OrdinaryDiffEqMutableCache
44+
order::Val{P}
45+
jet::jetType
46+
u::uType
47+
uprev::uType
48+
utaylor::taylorType
49+
utilde::uType
50+
tmp::uType
51+
atmp::uNoUnitsType
52+
stage_limiter!::StageLimiter
53+
step_limiter!::StepLimiter
54+
thread::Thread
55+
end
56+
57+
function alg_cache(alg::ExplicitTaylor{P}, u, rate_prototype, ::Type{uEltypeNoUnits},
58+
::Type{uBottomEltypeNoUnits}, ::Type{tTypeNoUnits}, uprev, uprev2, f, t,
59+
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))
63+
utilde = zero(u)
64+
atmp = similar(u, uEltypeNoUnits)
65+
recursivefill!(atmp, false)
66+
tmp = zero(u)
67+
ExplicitTaylorCache(Val(P), jet_iip, u, uprev, utaylor, utilde, tmp, atmp,
68+
alg.stage_limiter!, alg.step_limiter!, alg.thread)
69+
end
70+
71+
struct ExplicitTaylorConstantCache{P, jetType} <: OrdinaryDiffEqConstantCache
72+
order::Val{P}
73+
jet::jetType
74+
end
75+
function alg_cache(::ExplicitTaylor{P}, u, rate_prototype, ::Type{uEltypeNoUnits},
76+
::Type{uBottomEltypeNoUnits}, ::Type{tTypeNoUnits}, uprev, uprev2, f, t,
77+
dt, reltol, p, calck,
78+
::Val{false}) where {P, uEltypeNoUnits, uBottomEltypeNoUnits, tTypeNoUnits}
79+
if u isa AbstractArray
80+
jet, _ = build_jet(f, p, Val(P), length(u))
81+
else
82+
jet = build_jet(f, p, Val(P))
83+
end
84+
ExplicitTaylorConstantCache(Val(P), jet)
85+
end
86+
87+
# FSAL currently not used, providing dummy implementation to satisfy the interface
88+
get_fsalfirstlast(cache::ExplicitTaylorCache, u) = (cache.u, cache.u)

‎lib/OrdinaryDiffEqTaylorSeries/src/TaylorSeries_perform_step.jl

+100
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,100 @@
1+
using TaylorDiff: TaylorDiff, extract_derivative, extract_derivative!
2+
3+
@inline make_taylor(all::Vararg{X, P}) where {P, X <: AbstractArray} = TaylorArray(
4+
Base.first(all), Base.tail(all))
5+
@inline make_taylor(all::Vararg{X, P}) where {P, X} = TaylorScalar(all)
6+
7+
function initialize!(integrator, cache::ExplicitTaylor2ConstantCache)
8+
integrator.kshortsize = 3
9+
integrator.k = typeof(integrator.k)(undef, integrator.kshortsize)
10+
end
11+
12+
@muladd function perform_step!(
13+
integrator, cache::ExplicitTaylor2ConstantCache, repeat_step = false)
14+
@unpack t, dt, uprev, u, f, p = integrator
15+
k1 = f(uprev, p, t)
16+
u1 = make_taylor(uprev, k1)
17+
t1 = TaylorScalar{1}(t, one(t))
18+
k2 = f(u1, p, t1).partials[1]
19+
u = @.. uprev + dt * k1 + dt^2 / 2 * k2
20+
OrdinaryDiffEqCore.increment_nf!(integrator.stats, 3)
21+
integrator.k[1] = k1
22+
integrator.k[2] = k2
23+
integrator.u = u
24+
end
25+
26+
function initialize!(integrator, cache::ExplicitTaylor2Cache)
27+
integrator.kshortsize = 3
28+
resize!(integrator.k, integrator.kshortsize)
29+
# Setup k pointers
30+
integrator.k[1] = cache.k1
31+
integrator.k[2] = cache.k2
32+
integrator.k[3] = cache.k3
33+
return nothing
34+
end
35+
36+
@muladd function perform_step!(integrator, cache::ExplicitTaylor2Cache, repeat_step = false)
37+
@unpack t, dt, uprev, u, f, p = integrator
38+
@unpack k1, k2, k3, utilde, tmp = cache
39+
40+
# The following code is written to be fully non-allocating
41+
f(k1, uprev, p, t)
42+
u1 = make_taylor(uprev, k1)
43+
t1 = TaylorScalar{1}(t, one(t))
44+
out1 = make_taylor(k1, k2)
45+
f(out1, u1, p, t1)
46+
@.. u = uprev + dt * k1 + dt^2 / 2 * k2
47+
OrdinaryDiffEqCore.increment_nf!(integrator.stats, 3)
48+
return nothing
49+
end
50+
51+
function initialize!(integrator, cache::ExplicitTaylorConstantCache{P}) where {P}
52+
integrator.kshortsize = P
53+
integrator.k = typeof(integrator.k)(undef, P)
54+
end
55+
56+
@muladd function perform_step!(
57+
integrator, cache::ExplicitTaylorConstantCache{P}, repeat_step = false) where {P}
58+
@unpack t, dt, uprev, u, f, p = integrator
59+
@unpack jet = cache
60+
utaylor = jet(uprev, t)
61+
u = map(x -> evaluate_polynomial(x, dt), utaylor)
62+
if integrator.opts.adaptive
63+
utilde = TaylorDiff.get_coefficient(utaylor, P) * dt^(P + 1)
64+
atmp = calculate_residuals(utilde, uprev, u, integrator.opts.abstol,
65+
integrator.opts.reltol, integrator.opts.internalnorm, t)
66+
integrator.EEst = integrator.opts.internalnorm(atmp, t)
67+
end
68+
OrdinaryDiffEqCore.increment_nf!(integrator.stats, P + 1)
69+
integrator.u = u
70+
end
71+
72+
function initialize!(integrator, cache::ExplicitTaylorCache{P}) where {P}
73+
integrator.kshortsize = P
74+
resize!(integrator.k, P)
75+
# Setup k pointers
76+
for i in 1:P
77+
integrator.k[i] = get_coefficient(cache.utaylor, i)
78+
end
79+
return nothing
80+
end
81+
82+
@muladd function perform_step!(
83+
integrator, cache::ExplicitTaylorCache{P}, repeat_step = false) where {P}
84+
@unpack t, dt, uprev, u, f, p = integrator
85+
@unpack jet, utaylor, utilde, tmp, atmp, thread = cache
86+
87+
jet(utaylor, uprev, t)
88+
for i in eachindex(utaylor)
89+
u[i] = @inline evaluate_polynomial(utaylor[i], dt)
90+
end
91+
if integrator.opts.adaptive
92+
@.. broadcast=false thread=thread utilde=TaylorDiff.get_coefficient(utaylor, P) *
93+
dt^(P + 1)
94+
calculate_residuals!(atmp, utilde, uprev, u, integrator.opts.abstol,
95+
integrator.opts.reltol, integrator.opts.internalnorm, t)
96+
integrator.EEst = integrator.opts.internalnorm(atmp, t)
97+
end
98+
OrdinaryDiffEqCore.increment_nf!(integrator.stats, P + 1)
99+
return nothing
100+
end

‎lib/OrdinaryDiffEqTaylorSeries/src/alg_utils.jl

+55
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
alg_order(::ExplicitTaylor2) = 2
2+
alg_stability_size(alg::ExplicitTaylor2) = 1
3+
4+
alg_order(::ExplicitTaylor{P}) where {P} = P
5+
alg_stability_size(alg::ExplicitTaylor) = 1
6+
7+
JET_CACHE = IdDict()
8+
9+
function build_jet(f::ODEFunction{iip}, p, order, length = nothing) where {iip}
10+
build_jet(f, Val{iip}(), p, order, length)
11+
end
12+
13+
function build_jet(f, ::Val{iip}, p, order::Val{P}, length = nothing) where {P, iip}
14+
if haskey(JET_CACHE, f)
15+
list = JET_CACHE[f]
16+
index = findfirst(x -> x[1] == order && x[2] == p, list)
17+
index !== nothing && return list[index][3]
18+
end
19+
@variables t0::Real
20+
u0 = isnothing(length) ? Symbolics.variable(:u0) : Symbolics.variables(:u0, 1:length)
21+
if iip
22+
@assert length isa Integer
23+
f0 = similar(u0)
24+
f(f0, u0, p, t0)
25+
else
26+
f0 = f(u0, p, t0)
27+
end
28+
u = TaylorDiff.make_seed(u0, f0, Val(1))
29+
for index in 2:P
30+
t = TaylorScalar{index - 1}(t0, one(t0))
31+
if iip
32+
fu = similar(u)
33+
f(fu, u, p, t)
34+
else
35+
fu = f(u, p, t)
36+
end
37+
d = get_coefficient(fu, index - 1) / index
38+
u = append_coefficient(u, d)
39+
end
40+
jet = build_function(u, u0, t0; expression = Val(false), cse = true)
41+
if !haskey(JET_CACHE, f)
42+
JET_CACHE[f] = []
43+
end
44+
push!(JET_CACHE[f], (order, p, jet))
45+
return jet
46+
end
47+
48+
# evaluate using Qin Jiushao's algorithm
49+
@generated function evaluate_polynomial(t::TaylorScalar{T, P}, z) where {T, P}
50+
ex = :(v[$(P + 1)])
51+
for i in P:-1:1
52+
ex = :(v[$i] + z * $ex)
53+
end
54+
return :($(Expr(:meta, :inline)); v = flatten(t); $ex)
55+
end

0 commit comments

Comments
 (0)