WIP: implement fifth order Radau IIA

src/OrdinaryDiffEq.jl

@@ -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")
@@ -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")
 
@@ -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")
@@ -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

src/alg_utils.jl

@@ -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
@@ -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])
 
@@ -199,6 +200,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
@@ -280,6 +282,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

src/algorithms.jl

@@ -362,6 +362,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}

src/caches/firk_caches.jl

@@ -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

src/integrators/integrator_utils.jl

@@ -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)
@@ -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))
+    integrator.qold = q
     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 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)
@@ -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

src/perform_step/firk_perform_step.jl

@@ -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 nw, η
+  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
+
+    # transform `u'` to `z`
+    z1 = @. TI11 * ff1 + TI12 * ff2 + TI13 * ff3
+    z2 = @. TI21 * ff1 + TI22 * ff2 + TI23 * ff3
+    z3 = @. 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
+
+    z1 = LU1 \ (@. z1 - γdt*Mw1)
+    z2 = @. z2 - αdt*Mw2 + βdt*Mw3
+    z3 = @. z3 - βdt*Mw2 - αdt*Mw3
+    z23 = LU2 \ (@. z2 + z3*im)
+    z2 = real(z23)
+    z3 = imag(z23)
+
+    iter != 1 && (nwprev = nw)
+    nw = internalnorm(z1) + internalnorm(z2) + internalnorm(z3)
+
+    # check early stopping criterion
+    if iter != 1
+      θ = nw/nwprev
+      θ ≥ 1 || nw * θ^(max_iter - iter) > κtol * (1 - θ) && break
+    end
+
+    w1 = @. w1 + z1
+    w2 = @. w2 + z2
+    w3 = @. w3 + z3
+
+    # 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 || η * nw > κtol
+    else
+      η = θ / (1 - θ) # calculated for possible early stopping
+      do_newton = iter < min_iter || η * nw > κ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) || integrator.success_iter > 0
+    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