forked from JuliaLang/julia
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathode.jl
432 lines (368 loc) · 15.7 KB
/
ode.jl
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
#ODE23 Solve non-stiff differential equations.
#
# ODE23(F,TSPAN,Y0) with TSPAN = [T0 TFINAL] integrates the system
# of differential equations dy/dt = f(t,y) from t = T0 to t = TFINAL.
# The initial condition is y(T0) = Y0.
#
# The first argument, F, is a function handle or an anonymous function
# that defines f(t,y). This function must have two input arguments,
# t and y, and must return a column vector of the derivatives, dy/dt.
#
# With two output arguments, [T,Y] = ODE23(...) returns a column
# vector T and an array Y where Y(:,k) is the solution at T(k).
#
# More than four input arguments, ODE23(F,TSPAN,Y0,RTOL,P1,P2,...),
# are passed on to F, F(T,Y,P1,P2,...).
#
# ODE23 uses the Runge-Kutta (2,3) method of Bogacki and Shampine (BS23).
#
# Example
# tspan = [0 2*pi]
# y_0 = [1 0]'
# F = (t, y) -> [0 1; -1 0]*y
# ode23(F, tspan, y_0)
#
# See also ODE23.
# Initialize variables.
# Adapted from Cleve Moler's textbook
# http://www.mathworks.com/moler/ncm/ode23tx.m
load("poly.jl")
function ode23(F::Function, tspan::AbstractVector, y_0::AbstractVector)
rtol = 1.e-5
atol = 1.e-8
t0 = tspan[1]
tfinal = tspan[end]
tdir = sign(tfinal - t0)
threshold = atol / rtol
hmax = abs(0.1*(tfinal-t0))
t = t0
y = y_0[:]
tout = t
yout = y.'
tlen = length(t)
# Compute initial step size.
s1 = F(t, y)
r = norm(s1./max(abs(y), threshold), Inf) + realmin() # TODO: fix type bug in max()
h = tdir*0.8*rtol^(1/3)/r
# The main loop.
while t != tfinal
hmin = 16*eps()*abs(t)
if abs(h) > hmax; h = tdir*hmax; end
if abs(h) < hmin; h = tdir*hmin; end
# Stretch the step if t is close to tfinal.
if 1.1*abs(h) >= abs(tfinal - t)
h = tfinal - t;
end
# Attempt a step.
s2 = F(t+h/2, y+h/2*s1)
s3 = F(t+3*h/4, y+3*h/4*s2)
tnew = t + h
ynew = y + h*(2*s1 + 3*s2 + 4*s3)/9
s4 = F(tnew, ynew)
# Estimate the error.
e = h*(-5*s1 + 6*s2 + 8*s3 - 9*s4)/72
err = norm(e./max(max(abs(y), abs(ynew)), threshold), Inf) + realmin()
# Accept the solution if the estimated error is less than the tolerance.
if err <= rtol
t = tnew
y = ynew
tout = [tout; t]
yout = [yout; y.']
s1 = s4 # Reuse final function value to start new step
end
# Compute a new step size.
h = h*min(5, 0.8*(rtol/err)^(1/3))
# Exit early if step size is too small.
if abs(h) <= hmin
println("Step size ", h, " too small at t = ", t)
t = tfinal
end
end # while (t != tfinal)
return (tout, yout)
end # ode23
# ode45 adapted from http://users.powernet.co.uk/kienzle/octave/matcompat/scripts/ode_v1.11/ode45.m
# ode45 (v1.11) integrates a system of ordinary differential equations using
# 4th & 5th order embedded formulas from Dormand & Prince or Fehlberg.
#
# The Fehlberg 4(5) pair is established and works well, however, the
# Dormand-Prince 4(5) pair minimizes the local truncation error in the
# 5th-order estimate which is what is used to step forward (local extrapolation.)
# Generally it produces more accurate results and costs roughly the same
# computationally. The Dormand-Prince pair is the default.
#
# This is a 4th-order accurate integrator therefore the local error normally
# expected would be O(h^5). However, because this particular implementation
# uses the 5th-order estimate for xout (i.e. local extrapolation) moving
# forward with the 5th-order estimate should yield errors on the order of O(h^6).
#
# The order of the RK method is the order of the local *truncation* error, d.
# The local truncation error is defined as the principle error term in the
# portion of the Taylor series expansion that gets dropped. This portion of
# the Taylor series exapansion is within the group of terms that gets multipled
# by h in the solution definition of the general RK method. Therefore, the
# order-p solution created by the RK method will be roughly accurate to
# within O(h^(p+1)). The difference between two different-order solutions is
# the definition of the "local error," l. This makes the local error, l, as
# large as the error in the lower order method, which is the truncation
# error, d, times h, resulting in O(h^(p+1)).
# Summary: For an RK method of order-p,
# - the local truncation error is O(h^p)
# - the local error (used for stepsize adjustment) is O(h^(p+1))
#
# This requires 6 function evaluations per integration step.
#
# The error estimate formula and slopes are from
# Numerical Methods for Engineers, 2nd Ed., Chappra & Cannle, McGraw-Hill, 1985
#
# Usage:
# (tout, xout) = ode45(F, tspan, x0)
#
# INPUT:
# F - User-defined function
# Call: xprime = F(t,x)
# t - Time (scalar).
# x - Solution column-vector.
# xprime - Returned derivative COLUMN-vector; xprime(i) = dx(i)/dt.
# tspan - [ tstart, tfinal ]
# x0 - Initial value COLUMN-vector.
#
# OUTPUT:
# tout - Returned integration time points (column-vector).
# xout - Returned solution, one solution column-vector per tout-value.
#
# Original Octave implementation:
# Marc Compere
# created : 06 October 1999
# modified: 17 January 2001
function oderkf{T}(F::Function, tspan::AbstractVector, x0::AbstractVector{T}, a, b4, b5)
tol = 1.0e-5
# see p.91 in the Ascher & Petzold reference for more infomation.
pow = 1/6;
c = sum(a, 2)
# Initialization
t = tspan[1]
tfinal = tspan[end]
hmax = (tfinal - t)/2.5
hmin = (tfinal - t)/1e9
h = (tfinal - t)/100 # initial guess at a step size
x = x0
tout = [t] # first output time
xout = x.' # first output solution
k = zeros(eltype(x), (length(c),length(x)))
while (t < tfinal) & (h >= hmin)
if t + h > tfinal
h = tfinal - t
end
# Compute the slopes by computing the k[:,j+1]'th column based on the previous k[:,1:j] columns
# notes: k needs to end up as an Nxs, a is 7x6, which is s by (s-1),
# s is the number of intermediate RK stages on [t (t+h)] (Dormand-Prince has s=7 stages)
if -eps() < c[end]-1 < eps() && t != tspan[1]
# Assign the last stage for x(k) as the first stage for computing x[k+1].
# This is part of the Dormand-Prince pair caveat.
# k[:,7] has already been computed, so use it instead of recomputing it
# again as k[:,1] during the next step.
k[1,:] = k[end,:]
else
k[1,:] = F(t,x) # first stage
end
for j = 2:length(c)
k[j,:] = F(t + h.*c[j], x + h.*(a[j,1:j-1]*k[1:j-1,:]).')
end
# compute the 4th order estimate
x4 = x + h.*(b4*k).'
# compute the 5th order estimate
x5 = x + h.*(b5*k).'
# estimate the local truncation error
gamma1 = x5 - x4
# Estimate the error and the acceptable error
delta = norm(gamma1, Inf) # actual error
tau = tol*max(norm(x,Inf), 1.0) # allowable error
# Update the solution only if the error is acceptable
if (delta <= tau)
t = t + h
x = x5 # <-- using the higher order estimate is called 'local extrapolation'
tout = [tout; t]
xout = [xout; x.']
end
# Update the step size
if delta == 0.0
delta = 1e-16
end
h = min(hmax, 0.8*h*(tau/delta)^pow)
end # while (t < tfinal) & (h >= hmin)
if (t < tfinal)
println("Step size grew too small. t=", t, ", h=", h, ", x=", x)
end
return (tout, xout)
end
# Both the Dormand-Prince and Fehlberg 4(5) coefficients are from a tableau in
# U.M. Ascher, L.R. Petzold, Computer Methods for Ordinary Differential Equations
# and Differential-Agebraic Equations, Society for Industrial and Applied Mathematics
# (SIAM), Philadelphia, 1998
#
# Dormand-Prince coefficients
dp_coefficients = ([ 0 0 0 0 0 0
1/5 0 0 0 0 0
3/40 9/40 0 0 0 0
44/45 -56/15 32/9 0 0 0
19372/6561 -25360/2187 64448/6561 -212/729 0 0
9017/3168 -355/33 46732/5247 49/176 -5103/18656 0
35/384 0 500/1113 125/192 -2187/6784 11/84],
# 4th order b-coefficients
[5179/57600 0 7571/16695 393/640 -92097/339200 187/2100 1/40],
# 5th order b-coefficients
[35/384 0 500/1113 125/192 -2187/6784 11/84 0],
)
ode45_dp(F, tspan, x0) = oderkf(F, tspan, x0, dp_coefficients...)
# Fehlberg coefficients
fb_coefficients = ([ 0 0 0 0 0
1/4 0 0 0 0
3/32 9/32 0 0 0
1932/2197 -7200/2197 7296/2197 0 0
439/216 -8 3680/513 -845/4104 0
-8/27 2 -3544/2565 1859/4104 -11/40],
# 4th order b-coefficients
[25/216 0 1408/2565 2197/4104 -1/5 0],
# 5th order b-coefficients
[16/135 0 6656/12825 28561/56430 -9/50 2/55],
)
ode45_fb(F, tspan, x0) = oderkf(F, tspan, x0, fb_coefficients...)
# Cash-Karp coefficients
# Numerical Recipes in Fortran 77
ck_coefficients = ([ 0 0 0 0 0
1/5 0 0 0 0
3/40 9/40 0 0 0
3/10 -9/10 6/5 0 0
-11/54 5/2 -70/27 35/27 0
1631/55296 175/512 575/13824 44275/110592 253/4096],
# 4th order b-coefficients
[37/378 0 250/621 125/594 0 512/1771],
# 5th order b-coefficients
[2825/27648 0 18575/48384 13525/55296 277/14336 1/4],
)
ode45_ck(F, tspan, x0) = oderkf(F, tspan, x0, ck_coefficients...)
# Use Dormand Prince version of ode45 by default
ode45 = ode45_dp
#ODE4 Solve non-stiff differential equations, fourth order
# fixed-step Runge-Kutta method.
#
# [T,X] = ODE4(ODEFUN, TSPAN, X0) with TSPAN = [T0:H:TFINAL]
# integrates the system of differential equations x' = f(t,x) from time
# T0 to TFINAL in steps of H with initial conditions X0. Function
# ODEFUN(T,X) must return a column vector corresponding to f(t,x). Each
# row in the solution array X corresponds to a time returned in the
# column vector T.
function ode4{T}(F::Function, tspan::AbstractVector, x0::AbstractVector{T})
h = diff(tspan)
x = Array(T, (length(tspan), length(x0)))
x[1,:] = x0'
midxdot = Array(T, (4, length(x0)))
for i = 1:(length(tspan)-1)
# Compute midstep derivatives
midxdot[1,:] = F(tspan[i], x[i,:]')
midxdot[2,:] = F(tspan[i]+h[i]./2, x[i,:]' + midxdot[1,:]'.*h[i]./2)
midxdot[3,:] = F(tspan[i]+h[i]./2, x[i,:]' + midxdot[2,:]'.*h[i]./2)
midxdot[4,:] = F(tspan[i]+h[i], x[i,:]' + midxdot[3,:]'.*h[i])
# Integrate
x[i+1,:] = x[i,:] + 1./6.*h[i].*[1 2 2 1]*midxdot
end
return (tspan, x)
end
#ODEROSENBROCK Solve stiff differential equations, Rosenbrock method
# with provided coefficients.
function oderosenbrock{T}(F::Function, G::Function, tspan::AbstractVector, x0::AbstractVector{T}, gamma, a, b, c)
h = diff(tspan)
x = Array(T, length(tspan), length(x0))
x[1,:] = x0'
solstep = 1
while tspan[solstep] < max(tspan)
ts = tspan[solstep]
hs = h[solstep]
xs = reshape(x[solstep,:], size(x0))
dFdx = G(ts, xs)
jac = eye(size(dFdx)[1])./gamma./hs-dFdx
g = zeros(size(a)[1], length(x0))
g[1,:] = jac \ F(ts + b[1].*hs, xs)
for i = 2:size(a)[1]
g[i,:] = jac \ (F(ts + b[i].*hs, xs + (a[i,1:i-1]*g[1:i-1,:]).') + (c[i,1:i-1]*g[1:i-1,:]).'./hs)
end
x[solstep+1,:] = x[solstep,:] + b*g
solstep += 1
end
return (tspan, x)
end
function oderosenbrock{T}(F::Function, tspan::AbstractVector, x0::AbstractVector{T}, gamma, a, b, c)
# Crude forward finite differences estimator as fallback
function jacobian(F::Function, t::Number, x::AbstractVector)
ftx = F(t, x)
dFdx = zeros(length(x), length(x))
for j = 1:length(x)
dx = zeros(size(x))
# The 100 below is heuristic
dx[j] = (x[j]+(x[j]==0))./100
dFdx[:,j] = (F(t,x+dx)-ftx)./dx[j]
end
return dFdx
end
oderosenbrock(F, (t, x)->jacobian(F, t, x), tspan, x0, gamma, a, b, c)
end
# Kaps-Rentrop coefficients
kr4_coefficients = (0.231,
[0 0 0 0
2 0 0 0
4.4524708 4.163528 0 0
4.4524708 4.163528 0 0],
[3.957037 4.624892 0.617477 1.282613],
[ 0 0 0 0
-5.071675 0 0 0
6.020153 0.159750 0 0
-1.856344 -8.505381 -2.084075 0],)
ode4s_kr(F, tspan, x0) = oderosenbrock(F, tspan, x0, kr4_coefficients...)
ode4s_kr(F, G, tspan, x0) = oderosenbrock(F, G, tspan, x0, kr4_coefficients...)
# Shampine coefficients
s4_coefficients = (0.5,
[ 0 0 0 0
2 0 0 0
48/25 6/25 0 0
48/25 6/25 0 0],
[19/9 1/2 25/108 125/108],
[ 0 0 0 0
-8 0 0 0
372/25 12/5 0 0
-112/125 -54/125 -2/5 0],)
ode4s_s(F, tspan, x0) = oderosenbrock(F, tspan, x0, s4_coefficients...)
ode4s_s(F, G, tspan, x0) = oderosenbrock(F, G, tspan, x0, s4_coefficients...)
# Use Shampine coefficients by default (matching Numerical Recipes)
const ode4s = ode4s_s
# ODE_MS Fixed-step, fixed-order multi-step numerical method with Adams-Bashforth-Moulton coefficients
function ode_ms{T}(F::Function, tspan::AbstractVector, x0::AbstractVector{T}, order::Integer)
h = diff(tspan)
x = zeros(T,(length(tspan), length(x0)))
x[1,:] = x0
if 1 <= order <= 4
b = [ 1 0 0 0
-1/2 3/2 0 0
5/12 -16/12 23/12 0
-9/24 37/24 -59/24 55/24];
else
for steporder = size(b,1):order
s = steporder - 1;
for j = 0:s
# Assign in correct order for multiplication below
# (a factor depending on j and s) .* (an integral of a polynomial with -(0:s), except -j, as roots)
b[steporder,s-j+1] = (-1).^j./factorial(j)./factorial(s-j).*diff(polyval(polyint(poly(diagm(-[0:j-1; j+1:s]))),0:1));
end
end
end
# TODO: use a better data structure here (should be an order-element circ buffer)
xdot = similar(x)
for i = 1:length(tspan)-1
# Need to run the first several steps at reduced order
steporder = min(i, order)
xdot[i,:] = F(tspan[i], x[i,:]')
x[i+1,:] = x[i,:] + b[steporder,1:steporder]*xdot[i-(steporder-1):i,:].*h[i]
end
return (tspan, x)
end
ode4ms(F, tspan, x0) = ode_ms(F, tspan, x0, 4)