Skip to content
This repository was archived by the owner on Sep 9, 2024. It is now read-only.

RFC: Duck typing and consistent return types for all solvers. #14

Merged
merged 2 commits into from
Mar 12, 2014
Merged

RFC: Duck typing and consistent return types for all solvers. #14

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
9ca6b34
Duck typing and consistent return types (Vector{typeof(y0)}) for all …
acroy Feb 26, 2014
5686490
Fixed deprecation warning for + in fdjacobian and ode4.
acroy Mar 12, 2014
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
158 changes: 88 additions & 70 deletions src/ODE.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,8 +43,7 @@ export ode23, ode4, ode45, ode4s, ode4ms
# Initialize variables.
# Adapted from Cleve Moler's textbook
# http://www.mathworks.com/moler/ncm/ode23tx.m
function ode23{T}(F::Function, tspan::AbstractVector, y0::AbstractVector{T}; reltol = 1.e-5, abstol = 1.e-8)

function ode23(F, y0, tspan; reltol = 1.e-5, abstol = 1.e-8)
if reltol == 0
warn("setting reltol = 0 gives a step size of zero")
end
Expand All @@ -58,7 +57,8 @@ function ode23{T}(F::Function, tspan::AbstractVector, y0::AbstractVector{T}; rel
y = y0

tout = t
yout = y.'
yout = Array(typeof(y0),1)
yout[1] = y

tlen = length(t)

Expand Down Expand Up @@ -101,7 +101,7 @@ function ode23{T}(F::Function, tspan::AbstractVector, y0::AbstractVector{T}; rel
t = tnew
y = ynew
tout = [tout; t]
yout = [yout; y.']
push!(yout, y)
s1 = s4 # Reuse final function value to start new step
end

Expand All @@ -118,11 +118,12 @@ function ode23{T}(F::Function, tspan::AbstractVector, y0::AbstractVector{T}; rel

end # while (t != tfinal)

return (tout, yout)
return tout, yout

end # ode23



# ode45 adapted from http://users.powernet.co.uk/kienzle/octave/matcompat/scripts/ode_v1.11/ode45.m
# (a newer version (v1.15) can be found here https://sites.google.com/site/comperem/home/ode_solvers)
#
Expand Down Expand Up @@ -180,9 +181,7 @@ end # ode23
# [email protected]
# created : 06 October 1999
# modified: 17 January 2001

function oderkf{T}(F::Function, tspan::AbstractVector, x0::AbstractVector{T}, a, b4, b5; reltol = 1.0e-5, abstol = 1.0e-8)

function oderkf(F, x0, tspan, a, b4, b5; reltol = 1.0e-5, abstol = 1.0e-8)
# see p.91 in the Ascher & Petzold reference for more infomation.
pow = 1/5

Expand All @@ -196,25 +195,26 @@ function oderkf{T}(F::Function, tspan::AbstractVector, x0::AbstractVector{T}, a,
h = (tfinal - t)/100 # initial guess at a step size
x = x0
tout = t # first output time
xout = x.' # first output solution
xout = Array(typeof(x0), 1)
xout[1] = x # first output solution

k = zeros(eltype(x), length(c), length(x))
k[1,:] = F(t,x) # first stage
k = Array(typeof(x0), length(c))
k[1] = F(t,x) # first stage

while t < tfinal && h >= hmin
if t + h > tfinal
h = tfinal - t
end

for j = 2:length(c)
k[j,:] = F(t + h.*c[j], x + h.*(a[j,1:j-1]*k[1:j-1,:]).')
k[j] = F(t + h.*c[j], x + h.*(a[j,1:j-1]*k[1:j-1])[1])
end

# compute the 4th order estimate
x4 = x + h.*(b4*k).'
x4 = x + h.*(b4*k)[1]

# compute the 5th order estimate
x5 = x + h.*(b5*k).'
x5 = x + h.*(b5*k)[1]

# estimate the local truncation error
gamma1 = x5 - x4
Expand All @@ -228,7 +228,7 @@ function oderkf{T}(F::Function, tspan::AbstractVector, x0::AbstractVector{T}, a,
t = t + h
x = x5 # <-- using the higher order estimate is called 'local extrapolation'
tout = [tout; t]
xout = [xout; x.']
push!(xout, x)

# 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),
Expand All @@ -238,9 +238,9 @@ function oderkf{T}(F::Function, tspan::AbstractVector, x0::AbstractVector{T}, a,
# 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,:]
k[1] = k[end]
else
k[1,:] = F(t,x) # first stage
k[1] = F(t,x) # first stage
end
end

Expand All @@ -252,7 +252,7 @@ function oderkf{T}(F::Function, tspan::AbstractVector, x0::AbstractVector{T}, a,
println("Step size grew too small. t=", t, ", h=", h, ", x=", x)
end

return (tout, xout)
return tout, xout
end

# Both the Dormand-Prince and Fehlberg 4(5) coefficients are from a tableau in
Expand All @@ -273,7 +273,7 @@ const dp_coefficients = ([ 0 0 0 0 0
# 5th order b-coefficients
[35/384 0 500/1113 125/192 -2187/6784 11/84 0],
)
ode45_dp(F, tspan, x0; kwargs...) = oderkf(F, tspan, x0, dp_coefficients...; kwargs...)
ode45_dp(F, x0, tspan; kwargs...) = oderkf(F, x0, tspan, dp_coefficients...; kwargs...)

# Fehlberg coefficients
const fb_coefficients = ([ 0 0 0 0 0
Expand All @@ -287,7 +287,7 @@ const fb_coefficients = ([ 0 0 0 0 0
# 5th order b-coefficients
[16/135 0 6656/12825 28561/56430 -9/50 2/55],
)
ode45_fb(F, tspan, x0; kwargs...) = oderkf(F, tspan, x0, fb_coefficients...; kwargs...)
ode45_fb(F, x0, tspan; kwargs...) = oderkf(F, x0, tspan, fb_coefficients...; kwargs...)

# Cash-Karp coefficients
# Numerical Recipes in Fortran 77
Expand All @@ -302,7 +302,7 @@ const ck_coefficients = ([ 0 0 0 0 0
# 5th order b-coefficients
[2825/27648 0 18575/48384 13525/55296 277/14336 1/4],
)
ode45_ck(F, tspan, x0; kwargs...) = oderkf(F, tspan, x0, ck_coefficients...; kwargs...)
ode45_ck(F, x0, tspan; kwargs...) = oderkf(F, x0, tspan, ck_coefficients...; kwargs...)

# Use Dormand Prince version of ode45 by default
const ode45 = ode45_dp
Expand All @@ -316,67 +316,87 @@ const ode45 = ode45_dp
# 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})
function ode4(F, x0, tspan)
h = diff(tspan)
x = Array(T, length(tspan), length(x0))
x[1,:] = x0
x = Array(typeof(x0), length(tspan))
x[1] = x0

midxdot = Array(T, 4, length(x0))
midxdot = Array(typeof(x0), 4)
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])
midxdot[1] = F(tspan[i], x[i])
midxdot[2] = 2*F(tspan[i]+h[i]./2, x[i] + midxdot[1].*h[i]./2)
midxdot[3] = 2*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
x[i+1] = x[i] + 1/6 .*h[i].*sum(midxdot)
end
return (tspan, x)
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)
function oderosenbrock(F, x0, tspan, gamma, a, b, c; jacobian=nothing)
# Crude forward finite differences estimator as fallback
# FIXME: This doesn't really work if x is anything but a Vector or a scalar
function fdjacobian(F, x::Number, t)
ftx = F(t, x)

# The 100 below is heuristic
dx = (x .+ (x==0))./100
dFdx = (F(t,x+dx)-ftx)./dx

return dFdx
end

function fdjacobian(F, x, t)
ftx = F(t, x)
lx = max(length(x),1)
dFdx = zeros(eltype(x), lx, lx)
for j = 1:lx
# The 100 below is heuristic
dx = zeros(eltype(x), lx)
dx[j] = (x[j] .+ (x[j]==0))./100
dFdx[:,j] = (F(t,x+dx)-ftx)./dx[j]
end
return dFdx
end

if typeof(jacobian) == Function
G = jacobian
else
G = (t, x)->fdjacobian(F, x, t)
end

h = diff(tspan)
x = Array(T, length(tspan), length(x0))
x[1,:] = x0
x = Array(typeof(x0), length(tspan))
x[1] = x0

solstep = 1
while tspan[solstep] < maximum(tspan)
ts = tspan[solstep]
hs = h[solstep]
xs = reshape(x[solstep,:], size(x0))
xs = x[solstep]
dFdx = G(ts, xs)
jac = eye(size(dFdx,1))./gamma./hs-dFdx
# FIXME
if size(dFdx,1) == 1
jac = 1/gamma/hs - dFdx[1]
else
jac = eye(dFdx)./gamma./hs - dFdx
end

g = zeros(size(a,1), length(x0))
g[1,:] = jac \ F(ts + b[1].*hs, xs)
g = Array(typeof(x0), size(a,1))
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)
g[i] = (jac \ (F(ts + b[i].*hs, xs + (a[i,1:i-1]*g[1:i-1])[1]) + (c[i,1:i-1]*g[1:i-1])[1]./hs))
end

x[solstep+1,:] = x[solstep,:] + b*g
x[solstep+1] = x[solstep] + (b*g)[1]
solstep += 1
end
return (tspan, x)
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
const kr4_coefficients = (0.231,
Expand All @@ -389,8 +409,9 @@ const kr4_coefficients = (0.231,
-5.07167533877 0 0 0
6.02015272865 0.1597500684673 0 0
-1.856343618677 -8.50538085819 -2.08407513602 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...)

ode4s_kr(F, x0, tspan; jacobian=nothing) = oderosenbrock(F, x0, tspan, kr4_coefficients...; jacobian=jacobian)

# Shampine coefficients
const s4_coefficients = (0.5,
[ 0 0 0 0
Expand All @@ -403,19 +424,16 @@ const s4_coefficients = (0.5,
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...)
ode4s_s(F, x0, tspan; jacobian=nothing) = oderosenbrock(F, x0, tspan, s4_coefficients...; jacobian=jacobian)

# 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)
function ode_ms(F, x0, tspan, order::Integer)
h = diff(tspan)
x = zeros(T, length(tspan), length(x0))
x[1,:] = x0
x = Array(typeof(x0), length(tspan))
x[1] = x0

if 1 <= order <= 4
b = [ 1 0 0 0
Expand All @@ -438,13 +456,13 @@ function ode_ms{T}(F::Function, tspan::AbstractVector, x0::AbstractVector{T}, or
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]
xdot[i] = F(tspan[i], x[i])
x[i+1] = x[i] + (b[steporder,1:steporder]*xdot[i-(steporder-1):i])[1].*h[i]
end
return (tspan, x)
return [tspan], x
end

# Use order 4 by default
ode4ms(F, tspan, x0) = ode_ms(F, tspan, x0, 4)
ode4ms(F, x0, tspan) = ode_ms(F, x0, tspan, 4)

end # module ODE
11 changes: 6 additions & 5 deletions test/tests.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,28 +19,29 @@ for solver in solvers
# dy
# -- = 6 ==> y = 6t
# dt
t,y=solver((t,y)->6, [0:.1:1], [0.])
t,y=solver((t,y)->6, 0., [0:.1:1])
@test maximum(abs(y-6t)) < tol

# dy
# -- = 2t ==> y = t.^2
# dt
t,y=solver((t,y)->2t, [0:.001:1], [0.])
t,y=solver((t,y)->2t, 0., [0:.001:1])
@test maximum(abs(y-t.^2)) < tol

# dy
# -- = y ==> y = y0*e.^t
# dt
t,y=solver((t,y)->y, [0:.001:1], [1.])
t,y=solver((t,y)->y, 1., [0:.001:1])
@test maximum(abs(y-e.^t)) < tol

# dv dw
# -- = -w, -- = v ==> v = v0*cos(t) - w0*sin(t), w = w0*cos(t) + v0*sin(t)
# dt dt
#
# y = [v, w]
t,y=solver((t,y)->[-y[2], y[1]], [0:.001:2*pi], [1., 2.])
@test maximum(abs(y-[cos(t)-2*sin(t) 2*cos(t)+sin(t)])) < tol
t,y=solver((t,y)->[-y[2], y[1]], [1., 2.], [0:.001:2*pi])
ys = hcat(y...).' # convert Vector{Vector{Float}} to Matrix{Float}
@test maximum(abs(ys-[cos(t)-2*sin(t) 2*cos(t)+sin(t)])) < tol
end

println("All looks OK")