Skip to content

Commit fc49318

Browse files
tansongchentansongchen
committed
Proof of concept on NASA paper
1 parent 9198662 commit fc49318

3 files changed

Lines changed: 196 additions & 35 deletions

File tree

‎examples/Project.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,11 +4,11 @@ GraphMakie = "1ecd5474-83a3-4783-bb4f-06765db800d2"
44
Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
55
LayeredLayouts = "f4a74d36-062a-4d48-97cd-1356bad1de4e"
66
NetworkLayout = "46757867-2c16-5918-afeb-47bfcb05e46a"
7+
NonlinearSolve = "8913a72c-1f9b-4ce2-8d82-65094dcecaec"
78
ODEProblemLibrary = "fdc4e326-1af4-4b90-96e7-779fcce2daa5"
89
OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
910
SymbolicUtils = "d1185830-fcd6-423d-90d6-eec64667417b"
1011
Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
1112
TaylorDiff = "b36ab563-344f-407b-a36a-4f200bebf99c"
1213
TaylorIntegration = "92b13dbe-c966-51a2-8445-caca9f8a7d42"
1314
TaylorSeries = "6aa5eb33-94cf-58f4-a9d0-e4b2c4fc25ea"
14-
TikzGraphs = "b4f28e30-c73f-5eaf-a395-8a9db949a742"

‎examples/implicit_integration.jl‎

Lines changed: 158 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,158 @@
1+
include("integration.jl")
2+
3+
struct ImplicitFirstOrder{T}
4+
μ::T
5+
end
6+
7+
function (solver::ImplicitFirstOrder)(prob; h)
8+
μ = solver.μ
9+
fast_oop, _ = build_jetcoeffs(prob.f, prob.p, Val(1), length(prob.u0))
10+
t0, t_end = prob.tspan
11+
t = t0
12+
u = copy(prob.u0)
13+
while t < t_end
14+
nlf = (u_next, u_curr) -> begin
15+
du = get_coefficient(fast_oop(u_next, t + μ * h), 1)
16+
u_next .- u_curr .- μ * h .* du
17+
end
18+
guess = copy(u)
19+
nlprob = NonlinearProblem(nlf, guess, u)
20+
sol = solve(nlprob)
21+
f = get_coefficient(fast_oop(sol.u, t + μ * h), 1)
22+
u .= sol.u + (1 - μ) * h * f
23+
t += h
24+
end
25+
return u
26+
end
27+
28+
struct ImplicitSecondOrder{T}
29+
μ::T
30+
end
31+
32+
function (solver::ImplicitSecondOrder)(prob; h)
33+
μ = solver.μ
34+
fast_oop, _ = build_jetcoeffs(prob.f, prob.p, Val(2), length(prob.u0))
35+
t0, t_end = prob.tspan
36+
t = t0
37+
u = copy(prob.u0)
38+
while t < t_end
39+
nlf = (u_next, u_curr) -> begin
40+
poly = fast_oop(u_next, t + μ * h)
41+
u1 = get_coefficient(poly, 1)
42+
u2 = get_coefficient(poly, 2)
43+
u_curr + μ * h * u1 - (μ * h)^2 * u2 - u_next
44+
end
45+
guess = copy(u)
46+
nlprob = NonlinearProblem(nlf, guess, u)
47+
sol = solve(nlprob)
48+
poly = fast_oop(sol.u, t + μ * h)
49+
u1 = get_coefficient(poly, 1)
50+
u2 = get_coefficient(poly, 2)
51+
u .= sol.u + (1 - μ) * h * u1 + ((1 - μ) * h)^2 * u2
52+
t += h
53+
end
54+
return u
55+
end
56+
57+
struct ImplicitThirdOrder{T}
58+
μ::T
59+
end
60+
61+
function (solver::ImplicitThirdOrder)(prob; h)
62+
μ = solver.μ
63+
uh = μ * h
64+
cuh = (1 - μ) * h
65+
fast_oop, _ = build_jetcoeffs(prob.f, prob.p, Val(3), length(prob.u0))
66+
t0, t_end = prob.tspan
67+
t = t0
68+
u = copy(prob.u0)
69+
while t < t_end
70+
nlf = (u_next, u_curr) -> begin
71+
poly = fast_oop(u_next, t + uh)
72+
u1 = get_coefficient(poly, 1)
73+
u2 = get_coefficient(poly, 2)
74+
u3 = get_coefficient(poly, 3)
75+
u_curr + uh * u1 - uh^2 * u2 + uh^3 * u3 - u_next
76+
end
77+
guess = copy(u)
78+
nlprob = NonlinearProblem(nlf, guess, u)
79+
sol = solve(nlprob)
80+
poly = fast_oop(sol.u, t + uh)
81+
u1 = get_coefficient(poly, 1)
82+
u2 = get_coefficient(poly, 2)
83+
u3 = get_coefficient(poly, 3)
84+
u .= sol.u + cuh * u1 + cuh^2 * u2 + cuh^3 * u3
85+
t += h
86+
end
87+
return u
88+
end
89+
90+
struct ImplicitFourthOrder{T}
91+
μ::T
92+
end
93+
94+
function (solver::ImplicitFourthOrder)(prob; h)
95+
μ = solver.μ
96+
uh = μ * h
97+
cuh = (1 - μ) * h
98+
fast_oop, _ = build_jetcoeffs(prob.f, prob.p, Val(4), length(prob.u0))
99+
t0, t_end = prob.tspan
100+
t = t0
101+
u = copy(prob.u0)
102+
while t < t_end
103+
nlf = (u_next, u_curr) -> begin
104+
poly = fast_oop(u_next, t + uh)
105+
u1 = get_coefficient(poly, 1)
106+
u2 = get_coefficient(poly, 2)
107+
u3 = get_coefficient(poly, 3)
108+
u4 = get_coefficient(poly, 4)
109+
u_curr + uh * u1 - uh^2 * u2 + uh^3 * u3 - uh^4 * u4 - u_next
110+
end
111+
guess = copy(u)
112+
nlprob = NonlinearProblem(nlf, guess, u)
113+
sol = solve(nlprob)
114+
poly = fast_oop(sol.u, t + uh)
115+
u1 = get_coefficient(poly, 1)
116+
u2 = get_coefficient(poly, 2)
117+
u3 = get_coefficient(poly, 3)
118+
u4 = get_coefficient(poly, 4)
119+
u .= sol.u + cuh * u1 + cuh^2 * u2 + cuh^3 * u3 + cuh^4 * u4
120+
t += h
121+
end
122+
return u
123+
end
124+
125+
function estimate_order(solver, prob; dt1 = 0.01, dt2 = 0.005)
126+
sol = solve(prob, Vern7(); abstol = 1e-12, reltol = 1e-12)
127+
ref = sol.u[end]
128+
println("Reference: ", ref)
129+
u1 = solver(prob; h = dt1)
130+
println("u1: ", u1)
131+
u2 = solver(prob; h = dt2)
132+
println("u2: ", u2)
133+
err1 = norm(u1 - ref)
134+
err2 = norm(u2 - ref)
135+
order = log(err1 / err2) / log(dt1 / dt2)
136+
return order
137+
end
138+
139+
prob = prob_ode_lotkavolterra
140+
prob_complex = ODEProblem(ODEProblemLibrary.lotka, complex([1.0, 1.0]), (0.0, 1.0), [1.5, 1.0, 3.0, 1.0])
141+
142+
ImplicitEuler = ImplicitFirstOrder(1.0)
143+
ImplicitMidpoint = ImplicitFirstOrder(0.5)
144+
estimate_order(ImplicitEuler, prob)
145+
estimate_order(ImplicitMidpoint, prob)
146+
147+
ImplicitTaylor2 = ImplicitSecondOrder(1.0)
148+
ImplicitTaylor2Midpoint = ImplicitSecondOrder(0.5)
149+
ImplicitTaylor2Complex = ImplicitSecondOrder(0.5 + √3im/6)
150+
estimate_order(ImplicitTaylor2, prob)
151+
estimate_order(ImplicitTaylor2Midpoint, prob)
152+
estimate_order(ImplicitTaylor2Complex, prob_complex)
153+
154+
ImplicitTaylor3Midpoint = ImplicitThirdOrder(0.5)
155+
estimate_order(ImplicitTaylor3Midpoint, prob)
156+
157+
ImplicitTaylor4Midpoint = ImplicitFourthOrder(0.5)
158+
estimate_order(ImplicitTaylor4Midpoint, prob)

‎examples/integration.jl‎

Lines changed: 37 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,10 @@
11
using TaylorDiff: TaylorDiff, TaylorScalar, make_seed, flatten, get_coefficient,
22
set_coefficient, append_coefficient
33
using TaylorSeries, TaylorIntegration
4+
using LinearAlgebra
45
using ODEProblemLibrary, OrdinaryDiffEq, BenchmarkTools, Symbolics
6+
using CairoMakie
7+
58
SymbolicUtils.ENABLE_HASHCONSING[] = true
69

710
# There are two ways to compute the Taylor coefficients of a ODE solution
@@ -138,39 +141,39 @@ end
138141
return :($(Expr(:meta, :inline)); v = flatten(t); $ex)
139142
end
140143

141-
prob = prob_ode_lotkavolterra
142-
t0 = prob.tspan[1]
143-
raw = Float64[]
144-
optimized = Float64[]
145-
orders = [2, 4, 6, 8, 10, 12]
146-
for P in orders
147-
raw_time = @belapsed jetcoeffs($prob.f, $prob.u0, $prob.p, $t0, Val($P))
148-
fast_oop, fast_iip = build_jetcoeffs(prob.f, prob.p, Val(P), length(prob.u0))
149-
optimized_time = @belapsed $fast_oop($prob.u0, $t0)
150-
push!(raw, raw_time)
151-
push!(optimized, optimized_time)
152-
end
153-
154-
using CairoMakie
144+
function plot_simplification_effect()
145+
prob = prob_ode_lotkavolterra
146+
t0 = prob.tspan[1]
147+
raw = Float64[]
148+
optimized = Float64[]
149+
orders = [2, 4, 6, 8, 10, 12]
150+
for P in orders
151+
raw_time = @belapsed jetcoeffs($prob.f, $prob.u0, $prob.p, $t0, Val($P))
152+
fast_oop, fast_iip = build_jetcoeffs(prob.f, prob.p, Val(P), length(prob.u0))
153+
optimized_time = @belapsed $fast_oop($prob.u0, $t0)
154+
push!(raw, raw_time)
155+
push!(optimized, optimized_time)
156+
end
155157

156-
colors = Makie.wong_colors()
157-
f = begin
158-
f = Figure(resolution = (700, 400))
159-
ax = Axis(f[1, 1],
160-
xlabel = "Order",
161-
ylabel = "Time for computing truncated Taylor series (s)",
162-
title = "Effect of Symbolic Simplification on Computing Truncated Taylor Series",
163-
xticks = orders,
164-
yscale = log10
165-
)
166-
167-
group = [fill(1, length(orders));
168-
fill(2, length(orders))]
169-
barplot!(ax, [orders; orders], [raw; optimized],
170-
dodge = group,
171-
color = colors[group])
172-
elements = [PolyElement(polycolor = colors[i]) for i in 1:2]
173-
Legend(f[1, 2], elements, ["Raw", "Optimized"], "Groups")
174-
f
158+
colors = Makie.wong_colors()
159+
f = begin
160+
f = Figure(resolution = (700, 400))
161+
ax = Axis(f[1, 1],
162+
xlabel = "Order",
163+
ylabel = "Time for computing truncated Taylor series (s)",
164+
title = "Effect of Symbolic Simplification on Computing Truncated Taylor Series",
165+
xticks = orders,
166+
yscale = log10
167+
)
168+
169+
group = [fill(1, length(orders));
170+
fill(2, length(orders))]
171+
barplot!(ax, [orders; orders], [raw; optimized],
172+
dodge = group,
173+
color = colors[group])
174+
elements = [PolyElement(polycolor = colors[i]) for i in 1:2]
175+
Legend(f[1, 2], elements, ["Raw", "Optimized"], "Groups")
176+
f
177+
end
178+
save("simplification_effect.png", f)
175179
end
176-
f

0 commit comments

Comments
 (0)