Skip to content

Commit c0981ef

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

3 files changed

Lines changed: 316 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: 278 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,278 @@
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 .= real(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 ImplicitThirdOrderExtrapolated{T1, T2}
91+
μ1::T1
92+
c1::T1
93+
μ2::T2
94+
c2::T1
95+
end
96+
97+
function (solver::ImplicitThirdOrderExtrapolated)(prob; h)
98+
uh1 = solver.μ1 * h
99+
cuh1 = (1 - solver.μ1) * h
100+
uh2 = solver.μ2 * h
101+
cuh2 = (1 - solver.μ2) * h
102+
fast_oop, _ = build_jetcoeffs(prob.f, prob.p, Val(3), length(prob.u0))
103+
t0, t_end = prob.tspan
104+
t = t0
105+
u = copy(prob.u0)
106+
while t < t_end
107+
nlf = (u_next, u_curr) -> begin
108+
poly = fast_oop(u_next, t + uh1)
109+
u1 = get_coefficient(poly, 1)
110+
u2 = get_coefficient(poly, 2)
111+
u3 = get_coefficient(poly, 3)
112+
u_curr + uh1 * u1 - uh1^2 * u2 + uh1^3 * u3 - u_next
113+
end
114+
guess = copy(u)
115+
nlprob = NonlinearProblem(nlf, guess, u)
116+
sol = solve(nlprob)
117+
poly = fast_oop(sol.u, t + uh1)
118+
u1 = get_coefficient(poly, 1)
119+
u2 = get_coefficient(poly, 2)
120+
u3 = get_coefficient(poly, 3)
121+
u_trial1 = real(sol.u + cuh1 * u1 + cuh1^2 * u2 + cuh1^3 * u3)
122+
nlf2 = (u_next, u_curr) -> begin
123+
poly = fast_oop(u_next, t + uh2)
124+
u1 = get_coefficient(poly, 1)
125+
u2 = get_coefficient(poly, 2)
126+
u3 = get_coefficient(poly, 3)
127+
u_curr + uh2 * u1 - uh2^2 * u2 + uh2^3 * u3 - u_next
128+
end
129+
guess2 = copy(u)
130+
nlprob2 = NonlinearProblem(nlf2, guess2, u)
131+
sol2 = solve(nlprob2)
132+
poly2 = fast_oop(sol2.u, t + uh2)
133+
u1_2 = get_coefficient(poly2, 1)
134+
u2_2 = get_coefficient(poly2, 2)
135+
u3_2 = get_coefficient(poly2, 3)
136+
u_trial2 = real(sol2.u + cuh2 * u1_2 + cuh2^2 * u2_2 + cuh2^3 * u3_2)
137+
u .= (0.8 * u_trial1 + 0.2 * u_trial2)
138+
t += h
139+
end
140+
return u
141+
end
142+
143+
struct ImplicitFourthOrder{T}
144+
μ::T
145+
end
146+
147+
function (solver::ImplicitFourthOrder)(prob; h)
148+
μ = solver.μ
149+
uh = μ * h
150+
cuh = (1 - μ) * h
151+
fast_oop, _ = build_jetcoeffs(prob.f, prob.p, Val(4), length(prob.u0))
152+
t0, t_end = prob.tspan
153+
t = t0
154+
u = copy(prob.u0)
155+
while t < t_end
156+
nlf = (u_next, u_curr) -> begin
157+
poly = fast_oop(u_next, t + uh)
158+
u1 = get_coefficient(poly, 1)
159+
u2 = get_coefficient(poly, 2)
160+
u3 = get_coefficient(poly, 3)
161+
u4 = get_coefficient(poly, 4)
162+
u_curr + uh * u1 - uh^2 * u2 + uh^3 * u3 - uh^4 * u4 - u_next
163+
end
164+
guess = copy(u)
165+
nlprob = NonlinearProblem(nlf, guess, u)
166+
sol = solve(nlprob)
167+
poly = fast_oop(sol.u, t + uh)
168+
u1 = get_coefficient(poly, 1)
169+
u2 = get_coefficient(poly, 2)
170+
u3 = get_coefficient(poly, 3)
171+
u4 = get_coefficient(poly, 4)
172+
u .= sol.u + cuh * u1 + cuh^2 * u2 + cuh^3 * u3 + cuh^4 * u4
173+
t += h
174+
end
175+
return u
176+
end
177+
178+
struct ImplicitFourthOrderExtrapolated{T1, T2}
179+
μ1::T2
180+
c1::T1
181+
μ2::T2
182+
c2::T1
183+
end
184+
185+
function (solver::ImplicitFourthOrderExtrapolated)(prob; h)
186+
uh1 = solver.μ1 * h
187+
cuh1 = (1 - solver.μ1) * h
188+
uh2 = solver.μ2 * h
189+
cuh2 = (1 - solver.μ2) * h
190+
fast_oop, _ = build_jetcoeffs(prob.f, prob.p, Val(4), length(prob.u0))
191+
t0, t_end = prob.tspan
192+
t = t0
193+
u = copy(prob.u0)
194+
while t < t_end
195+
nlf = (u_next, u_curr) -> begin
196+
poly = fast_oop(u_next, t + uh1)
197+
u1 = get_coefficient(poly, 1)
198+
u2 = get_coefficient(poly, 2)
199+
u3 = get_coefficient(poly, 3)
200+
u4 = get_coefficient(poly, 4)
201+
u_curr + uh1 * u1 - uh1^2 * u2 + uh1^3 * u3 - uh1^4 * u4 - u_next
202+
end
203+
guess = copy(u)
204+
nlprob = NonlinearProblem(nlf, guess, u)
205+
sol = solve(nlprob)
206+
poly = fast_oop(sol.u, t + uh1)
207+
u1 = get_coefficient(poly, 1)
208+
u2 = get_coefficient(poly, 2)
209+
u3 = get_coefficient(poly, 3)
210+
u4 = get_coefficient(poly, 4)
211+
u_trial1 = real(sol.u + cuh1 * u1 + cuh1^2 * u2 + cuh1^3 * u3 + cuh1^4 * u4)
212+
nlf2 = (u_next, u_curr) -> begin
213+
poly = fast_oop(u_next, t + uh2)
214+
u1 = get_coefficient(poly, 1)
215+
u2 = get_coefficient(poly, 2)
216+
u3 = get_coefficient(poly, 3)
217+
u4 = get_coefficient(poly, 4)
218+
u_curr + uh2 * u1 - uh2^2 * u2 + uh2^3 * u3 - uh2^4 * u4 - u_next
219+
end
220+
guess2 = copy(u)
221+
nlprob2 = NonlinearProblem(nlf2, guess2, u)
222+
sol2 = solve(nlprob2)
223+
poly2 = fast_oop(sol2.u, t + uh2)
224+
u1_2 = get_coefficient(poly2, 1)
225+
u2_2 = get_coefficient(poly2, 2)
226+
u3_2 = get_coefficient(poly2, 3)
227+
u4_2 = get_coefficient(poly2, 4)
228+
u_trial2 = real(sol2.u + cuh2 * u1_2 + cuh2^2 * u2_2 + cuh2^3 * u3_2 +
229+
cuh2^4 * u4_2)
230+
u .= (solver.c1 * u_trial1 + solver.c2 * u_trial2) / (solver.c1 + solver.c2)
231+
t += h
232+
end
233+
return u
234+
end
235+
236+
function estimate_order(solver, prob; dt1 = 0.01, dt2 = 0.005)
237+
ref = prob.f.analytic(prob.u0, prob.p, prob.tspan[2])
238+
println("Reference: ", ref)
239+
u1 = solver(prob; h = dt1)
240+
println("u1: ", u1)
241+
u2 = solver(prob; h = dt2)
242+
println("u2: ", u2)
243+
err1 = norm(u1 - ref)
244+
err2 = norm(u2 - ref)
245+
order = log(err1 / err2) / log(dt1 / dt2)
246+
return order
247+
end
248+
249+
init = rand(2)
250+
linear_f = ODEFunction(
251+
ODEProblemLibrary.f_2dlinear, analytic = ODEProblemLibrary.f_2dlinear_analytic)
252+
prob = ODEProblem(linear_f, init, (0.0, 1.0), 1.01)
253+
prob_complex = ODEProblem(
254+
linear_f, complex(init), (0.0, 1.0), 1.01)
255+
256+
ImplicitEuler = ImplicitFirstOrder(1.0)
257+
ImplicitMidpoint = ImplicitFirstOrder(0.5)
258+
estimate_order(ImplicitEuler, prob)
259+
estimate_order(ImplicitMidpoint, prob)
260+
261+
ImplicitTaylor2 = ImplicitSecondOrder(1.0)
262+
ImplicitTaylor2Midpoint = ImplicitSecondOrder(0.5)
263+
ImplicitTaylor2Complex = ImplicitSecondOrder(0.5 + sqrt(3) * im / 6)
264+
estimate_order(ImplicitTaylor2, prob)
265+
estimate_order(ImplicitTaylor2Midpoint, prob)
266+
estimate_order(ImplicitTaylor2Complex, prob_complex)
267+
268+
ImplicitTaylor3Midpoint = ImplicitThirdOrder(0.5)
269+
ImplicitTaylor3Extrapolated = ImplicitThirdOrderExtrapolated(0.5, 0.8, 0.5 + 0.5im, 0.2)
270+
estimate_order(ImplicitTaylor3Midpoint, prob)
271+
estimate_order(ImplicitTaylor3Extrapolated, prob_complex; dt1 = 1 / 12, dt2 = 1 / 16)
272+
273+
ImplicitTaylor4Midpoint = ImplicitFourthOrder(0.5)
274+
ImplicitTaylor4Extrapolated = ImplicitFourthOrderExtrapolated(
275+
0.5 + 0.5im * tan(π / 10), tan(π / 10) * sec(π / 10)^5,
276+
0.5 + 0.5im * tan(3π / 10), tan(3π / 10) * sec(3π / 10)^5)
277+
estimate_order(ImplicitTaylor4Midpoint, prob)
278+
estimate_order(ImplicitTaylor4Extrapolated, prob_complex; dt1 = 1 / 9, dt2 = 1 / 12)

‎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)