Skip to content

Commit

Permalink
Merge pull request #612 from JuliaDiffEq/myb/radau
Browse files Browse the repository at this point in the history
WIP: implement fifth order Radau IIA
  • Loading branch information
ChrisRackauckas committed Jan 21, 2019
2 parents aa0d5e0 + 8947511 commit 09f48d9
Show file tree
Hide file tree
Showing 9 changed files with 334 additions and 16 deletions.
5 changes: 5 additions & 0 deletions src/OrdinaryDiffEq.jl
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,7 @@ module OrdinaryDiffEq
include("caches/feagin_caches.jl")
include("caches/verner_caches.jl")
include("caches/sdirk_caches.jl")
include("caches/firk_caches.jl")
include("caches/kencarp_kvaerno_caches.jl")
include("caches/generic_implicit_caches.jl")
include("caches/linear_caches.jl")
Expand All @@ -85,6 +86,7 @@ module OrdinaryDiffEq
include("tableaus/feagin_tableaus.jl")
include("tableaus/rosenbrock_tableaus.jl")
include("tableaus/sdirk_tableaus.jl")
include("tableaus/firk_tableaus.jl")
include("tableaus/rkn_tableaus.jl")
include("tableaus/rkc_tableaus.jl")

Expand All @@ -107,6 +109,7 @@ module OrdinaryDiffEq
include("perform_step/ssprk_perform_step.jl")
include("perform_step/sdirk_perform_step.jl")
include("perform_step/kencarp_kvaerno_perform_step.jl")
include("perform_step/firk_perform_step.jl")
include("perform_step/generic_implicit_perform_step.jl")
include("perform_step/rosenbrock_perform_step.jl")
include("perform_step/threaded_rk_perform_step.jl")
Expand Down Expand Up @@ -158,6 +161,8 @@ module OrdinaryDiffEq
ORK256, RK46NL, DP5, DP5Threaded, Tsit5, DP8, Vern6, Vern7, Vern8, TanYam7, TsitPap8,
Vern9,Feagin10, Feagin12, Feagin14, CompositeAlgorithm, Anas5

export RadauIIA5

export ImplicitEuler, ImplicitMidpoint, Trapezoid, TRBDF2, SDIRK2, Kvaerno3,
KenCarp3, Cash4, Hairer4, Hairer42, SSPSDIRK2, Kvaerno4, Kvaerno5,
KenCarp4, KenCarp5
Expand Down
6 changes: 5 additions & 1 deletion src/alg_utils.jl
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,7 @@ qmin_default(alg::DP8) = 1//3
qmax_default(alg::OrdinaryDiffEqAlgorithm) = 10
qmax_default(alg::CompositeAlgorithm) = minimum(qmax_default.(alg.algs))
qmax_default(alg::DP8) = 6
qmax_default(alg::RadauIIA5) = 8

get_chunksize(alg::OrdinaryDiffEqAlgorithm) = error("This algorithm does not have a chunk size defined.")
get_chunksize(alg::OrdinaryDiffEqAdaptiveImplicitAlgorithm{CS,AD}) where {CS,AD} = CS
Expand Down Expand Up @@ -101,7 +102,7 @@ get_current_alg_order(alg::JVODE,cache) = get_current_adaptive_order(alg,cache)
get_current_alg_order(alg::QNDF,cache) = cache.order
get_current_adaptive_order(alg::QNDF,cache) = cache.order

alg_adaptive_order(alg::OrdinaryDiffEqAdaptiveAlgorithm) = error("Algorithm is adaptive with no order")
#alg_adaptive_order(alg::OrdinaryDiffEqAdaptiveAlgorithm) = error("Algorithm is adaptive with no order")
get_current_adaptive_order(alg::OrdinaryDiffEqAlgorithm,cache) = alg_adaptive_order(alg)
get_current_adaptive_order(alg::CompositeAlgorithm,cache) = alg_adaptive_order(alg.algs[cache.current])

Expand Down Expand Up @@ -200,6 +201,7 @@ alg_order(alg::TanYam7) = 7
alg_order(alg::TsitPap8) = 8
alg_order(alg::GenericImplicitEuler) = 1
alg_order(alg::GenericTrapezoid) = 2
alg_order(alg::RadauIIA5) = 5
alg_order(alg::ImplicitEuler) = 1
alg_order(alg::MidpointSplitting) = 2
alg_order(alg::LinearExponential) = 1
Expand Down Expand Up @@ -281,6 +283,8 @@ alg_adaptive_order(alg::Feagin14) = 12
alg_adaptive_order(alg::Rosenbrock23) = 3
alg_adaptive_order(alg::Rosenbrock32) = 2

alg_adaptive_order(alg::RadauIIA5) = 3

alg_adaptive_order(alg::GenericImplicitEuler) = 0
alg_adaptive_order(alg::GenericTrapezoid) = 1
alg_adaptive_order(alg::ImplicitEuler) = 0
Expand Down
28 changes: 28 additions & 0 deletions src/algorithms.jl
Original file line number Diff line number Diff line change
Expand Up @@ -363,6 +363,34 @@ LinearExponential(;krylov=:off, m=10, iop=0) = LinearExponential(krylov, m, iop)

################################################################################

# FIRK Methods

struct RadauIIA5{CS,AD,F,FDT,T2,Tol,Controller} <: OrdinaryDiffEqNewtonAdaptiveAlgorithm{CS,AD,Controller}
linsolve::F
diff_type::FDT
smooth_est::Bool
extrapolant::Symbol
new_jac_conv_bound::T2
κ::Tol
tol::Tol
max_iter::Int
min_iter::Int
end
RadauIIA5(;chunk_size=0,autodiff=true,diff_type=Val{:central},
linsolve=DEFAULT_LINSOLVE,
extrapolant=:constant,new_jac_conv_bound=1e-3,
controller=:Predictive=nothing,
tol=nothing,max_iter=10,min_iter=1,smooth_est=true) =
RadauIIA5{chunk_size,autodiff,typeof(linsolve),
typeof(diff_type),
typeof(new_jac_conv_bound),typeof(tol),
controller}(linsolve,
diff_type,smooth_est,extrapolant,
new_jac_conv_bound,κ,tol,
max_iter,min_iter)

################################################################################

# SDIRK Methods

struct ImplicitEuler{CS,AD,F,F2,FDT,T2,Controller} <: OrdinaryDiffEqNewtonAdaptiveAlgorithm{CS,AD,Controller}
Expand Down
25 changes: 25 additions & 0 deletions src/caches/firk_caches.jl
Original file line number Diff line number Diff line change
@@ -0,0 +1,25 @@
mutable struct RadauIIA5ConstantCache{F,Tab,Tol,Dt,U} <: OrdinaryDiffEqConstantCache
uf::F
tab::Tab
κ::Tol
tol::Tol
ηold::Tol
nl_iters::Int
dtprev::Dt
cont1::U
cont2::U
cont3::U
end

function alg_cache(alg::RadauIIA5,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,
tTypeNoUnits,uprev,uprev2,f,t,dt,reltol,p,calck,::Type{Val{false}})

uf = DiffEqDiffTools.UDerivativeWrapper(f, t, p)
uToltype = real(uBottomEltypeNoUnits)
tab = RadauIIA5Tableau(uToltype, real(tTypeNoUnits))

κ = alg.κ !== nothing ? uToltype(alg.κ) : uToltype(1//100)
tol = alg.tol !== nothing ? uToltype(alg.tol) : uToltype(min(0.03,first(reltol)^(0.5)))

RadauIIA5ConstantCache(uf, tab, κ, tol, zero(tol), 10000, dt, u, u, u)
end
36 changes: 21 additions & 15 deletions src/integrators/integrator_utils.jl
Original file line number Diff line number Diff line change
Expand Up @@ -163,14 +163,14 @@ function stepsize_controller!(integrator,alg)
@fastmath q = q11/(qold^beta2)
integrator.q11 = q11
@fastmath q = max(inv(integrator.opts.qmax),min(inv(integrator.opts.qmin),q/integrator.opts.gamma))
if q <= integrator.opts.qsteady_max && q >= integrator.opts.qsteady_min
q = one(q)
end
end
q
end

function step_accept_controller!(integrator,alg,q)
if q <= integrator.opts.qsteady_max && q >= integrator.opts.qsteady_min
q = one(q)
end
integrator.qold = max(integrator.EEst,integrator.opts.qoldinit)
integrator.dt/q #dtnew
end
Expand Down Expand Up @@ -225,7 +225,7 @@ end


function stepsize_controller!(integrator,alg::Union{StandardControllerAlgs,
OrdinaryDiffEqNewtonAdaptiveAlgorithm{:Standard}})
OrdinaryDiffEqNewtonAdaptiveAlgorithm{CS,AD,:Standard}}) where {CS, AD}
# Standard stepsize controller
if iszero(integrator.EEst)
q = inv(integrator.opts.qmax)
Expand All @@ -237,37 +237,43 @@ function stepsize_controller!(integrator,alg::Union{StandardControllerAlgs,
q
end
function step_accept_controller!(integrator,alg::Union{StandardControllerAlgs,
OrdinaryDiffEqNewtonAdaptiveAlgorithm{:Standard}},q)
OrdinaryDiffEqNewtonAdaptiveAlgorithm{CS,AD,:Standard}},q) where {CS, AD}
integrator.dt/q # dtnew
end
function step_reject_controller!(integrator,alg::Union{StandardControllerAlgs,
OrdinaryDiffEqNewtonAdaptiveAlgorithm{:Standard}})
OrdinaryDiffEqNewtonAdaptiveAlgorithm{CS,AD,:Standard}}) where {CS, AD}
integrator.dt = integrator.qold
end

function stepsize_controller!(integrator,
alg::OrdinaryDiffEqNewtonAdaptiveAlgorithm{:Predictive})
alg::OrdinaryDiffEqNewtonAdaptiveAlgorithm{CS,AD,:Predictive}) where {CS, AD}

# Gustafsson predictive stepsize controller

if iszero(integrator.EEst)
q = inv(integrator.opts.qmax)
else
gamma = integrator.opts.gamma
niters = integrator.cache.newton_iters
fac = min(gamma,(1+2*integrator.alg.max_newton_iter)*gamma/(niters+2*integrator.alg.max_newton_iter))
expo = 1/(get_current_alg_order(integrator.alg,integrator.cache)+1)
if alg isa RadauIIA5
@unpack nl_iters = integrator.cache
@unpack max_iter = alg
else
@unpack nl_iters, max_iter = integrator.cache.nlsolve.cache
end
niters = nl_iters
fac = min(gamma,(1+2*max_iter)*gamma/(niters+2*max_iter))
expo = 1/(get_current_adaptive_order(integrator.alg,integrator.cache)+1)
qtmp = (integrator.EEst^expo)/fac
@fastmath q = max(inv(integrator.opts.qmax),min(inv(integrator.opts.qmin),qtmp))
if q <= integrator.opts.qsteady_max && q >= integrator.opts.qsteady_min
q = one(q)
end
integrator.qold = q
end
q
end
function step_accept_controller!(integrator,
alg::OrdinaryDiffEqNewtonAdaptiveAlgorithm{:Predictive},q)
alg::OrdinaryDiffEqNewtonAdaptiveAlgorithm{CS,AD,:Predictive},q) where {CS, AD}
if q <= integrator.opts.qsteady_max && q >= integrator.opts.qsteady_min
q = one(q)
end
if integrator.success_iter > 0
expo = 1/(get_current_adaptive_order(integrator.alg,integrator.cache)+1)
qgus=(integrator.dtacc/integrator.dt)*(((integrator.EEst^2)/integrator.erracc)^expo)
Expand All @@ -281,7 +287,7 @@ function step_accept_controller!(integrator,
integrator.dt/qacc
end
function step_reject_controller!(integrator,
alg::OrdinaryDiffEqNewtonAdaptiveAlgorithm{:Predictive})
alg::OrdinaryDiffEqNewtonAdaptiveAlgorithm{CS,AD,:Predictive}) where {CS, AD}
if integrator.success_iter == 0
integrator.dt *= 0.1
else
Expand Down
158 changes: 158 additions & 0 deletions src/perform_step/firk_perform_step.jl
Original file line number Diff line number Diff line change
@@ -0,0 +1,158 @@
function initialize!(integrator, cache::RadauIIA5ConstantCache)
integrator.kshortsize = 2
integrator.k = typeof(integrator.k)(undef, integrator.kshortsize)
integrator.fsalfirst = integrator.f(integrator.uprev, integrator.p, integrator.t) # Pre-start fsal

# Avoid undefined entries if k is an array of arrays
integrator.fsallast = zero(integrator.fsalfirst)
integrator.k[1] = integrator.fsalfirst
integrator.k[2] = integrator.fsallast
end

@muladd function perform_step!(integrator, cache::RadauIIA5ConstantCache, repeat_step=false)
@unpack t,dt,uprev,u,f,p = integrator
@unpack T11, T12, T13, T21, T22, T23, T31, TI11, TI12, TI13, TI21, TI22, TI23, TI31, TI32, TI33 = cache.tab
@unpack c1, c2, γ, α, β, e1, e2, e3 = cache.tab
@unpack tol, κ, cont1, cont2, cont3 = cache
@unpack internalnorm, abstol, reltol, adaptive = integrator.opts
alg = unwrap_alg(integrator, true)
@unpack min_iter, max_iter = alg
mass_matrix = integrator.f.mass_matrix
is_compos = integrator.alg isa CompositeAlgorithm

# precalculations
c1m1 = c1-1
c2m1 = c2-1
c1mc2= c1-c2
κtol = κ*tol # used in Newton iteration
γdt, αdt, βdt = γ/dt, α/dt, β/dt
J = calc_J(integrator, cache, is_compos)
if u isa Number
LU1 = γdt*mass_matrix - J
LU2 = (αdt + βdt*im)*mass_matrix - J
else
LU1 = lu(γdt*mass_matrix - J)
LU2 = lu((αdt + βdt*im)*mass_matrix - J)
end

# TODO better initial guess
if integrator.iter == 1 || integrator.u_modified || alg.extrapolant == :constant
cache.dtprev = one(cache.dtprev)
z1 = w1 = map(zero, u)
z2 = w2 = map(zero, u)
z3 = w3 = map(zero, u)
cache.cont1 = map(zero, u)
cache.cont2 = map(zero, u)
cache.cont3 = map(zero, u)
else
c3′ = dt/cache.dtprev
c1′ = c1*c3′
c2′ = c2*c3′
z1 = @. c1′ * (cont1 + (c1′-c2m1) * (cont2 + (c1′-c1m1) * cont3))
z2 = @. c2′ * (cont1 + (c2′-c2m1) * (cont2 + (c2′-c1m1) * cont3))
z3 = @. c3′ * (cont1 + (c3′-c2m1) * (cont2 + (c3′-c1m1) * cont3))
w1 = @. TI11 * z1 + TI12 * z2 + TI13 * z3
w2 = @. TI21 * z1 + TI22 * z2 + TI23 * z3
w3 = @. TI31 * z1 + TI32 * z2 + TI33 * z3
end

# Newton iteration
local ndw, η
do_newton = true
iter = 0
while do_newton && iter < max_iter
iter += 1
ff1 = f(uprev+z1, p, t+c1*dt)
ff2 = f(uprev+z2, p, t+c2*dt)
ff3 = f(uprev+z3, p, t+ dt) # c3 = 1

fw1 = @. TI11 * ff1 + TI12 * ff2 + TI13 * ff3
fw2 = @. TI21 * ff1 + TI22 * ff2 + TI23 * ff3
fw3 = @. TI31 * ff1 + TI32 * ff2 + TI33 * ff3

if mass_matrix isa UniformScaling # `UniformScaling` doesn't play nicely with broadcast
Mw1 = @. mass_matrix.λ * w1
Mw2 = @. mass_matrix.λ * w2
Mw3 = @. mass_matrix.λ * w3
else
Mw1 = mass_matrix*w1
Mw2 = mass_matrix*w2
Mw3 = mass_matrix*w3
end

rhs1 = @. fw1 - γdt*Mw1
rhs2 = @. fw2 - αdt*Mw2 + βdt*Mw3
rhs3 = @. fw3 - βdt*Mw2 - αdt*Mw3
dw1 = LU1 \ rhs1
dw23 = LU2 \ (@. rhs2 + rhs3*im)
dw2 = real(dw23)
dw3 = imag(dw23)

iter != 1 && (ndwprev = ndw)
ndw = internalnorm(dw1) + internalnorm(dw2) + internalnorm(dw3)

# check early stopping criterion
if iter != 1
θ = ndw/ndwprev
θ 1 || ndw * θ^(max_iter - iter) > κtol * (1 - θ) && break
end

w1 = @. w1 + dw1
w2 = @. w2 + dw2
w3 = @. w3 + dw3

# transform `w` to `z`
z1 = @. T11 * w1 + T12 * w2 + T13 * w3
z2 = @. T21 * w1 + T22 * w2 + T23 * w3
z3 = @. T31 * w1 + w2 # T32 = 1, T33 = 0

# check stopping criterion
if iter == 1
η = max(cache.ηold, eps(eltype(reltol)))^(0.8)
do_newton = iszero(integrator.success_iter) || iter < min_iter || η * ndw > κtol
else
η = θ / (1 - θ) # calculated for possible early stopping
do_newton = iter < min_iter || η * ndw > κtol
end
end
cache.ηold = η
integrator.force_stepfail = do_newton
cache.nl_iters = iter
do_newton && return

u = @. uprev + z3

if adaptive
e1dt, e2dt, e3dt = e1/dt, e2/dt, e3/dt
tmp = @. e1dt*z1 + e2dt*z2 + e3dt*z3
mass_matrix != I && (tmp = mass_matrix*tmp)
utilde = @. integrator.fsalfirst + tmp
alg.smooth_est && (utilde = LU1 \ utilde)
atmp = calculate_residuals(utilde, uprev, u, abstol, reltol, internalnorm)
integrator.EEst = internalnorm(atmp)

if integrator.iter == 1 || integrator.u_modified || integrator.EEst > oneunit(integrator.EEst)
f0 = f(uprev .+ utilde, p, t)
utilde = @. f0 + tmp
alg.smooth_est && (utilde = LU1 \ utilde)
atmp = calculate_residuals(utilde, uprev, u, abstol, reltol, internalnorm)
integrator.EEst = internalnorm(atmp)
end
end

if integrator.EEst <= oneunit(integrator.EEst)
cache.dtprev = dt
if alg.extrapolant != :constant
cache.cont1 = @. (z2 - z3)/c2m1
tmp = @. (z1 - z2)/c1mc2
cache.cont2 = @. (tmp - cache.cont1) / c1m1
cache.cont3 = @. cache.cont2 - (tmp - z1/c1)/c2
end
end

integrator.fsallast = f(u, p, t+dt)
integrator.k[1] = integrator.fsalfirst
integrator.k[2] = integrator.fsallast
integrator.u = u
return
end
Loading

0 comments on commit 09f48d9

Please sign in to comment.