WIP: Implement the ROCK2 method

src/OrdinaryDiffEq.jl

@@ -66,6 +66,7 @@ module OrdinaryDiffEq
   include("caches/adams_bashforth_moulton_caches.jl")
   include("caches/nordsieck_caches.jl")
   include("caches/bdf_caches.jl")
+  include("caches/rkc_caches.jl")
   include("caches/euler_imex_caches.jl")
 
   include("alg_utils.jl")
@@ -78,6 +79,7 @@ module OrdinaryDiffEq
   include("tableaus/rosenbrock_tableaus.jl")
   include("tableaus/sdirk_tableaus.jl")
   include("tableaus/rkn_tableaus.jl")
+  include("tableaus/rkc_tableaus.jl")
 
   include("integrators/type.jl")
   include("integrators/integrator_utils.jl")
@@ -105,6 +107,7 @@ module OrdinaryDiffEq
   include("perform_step/adams_bashforth_moulton_perform_step.jl")
   include("perform_step/nordsieck_perform_step.jl")
   include("perform_step/bdf_perform_step.jl")
+  include("perform_step/rkc_perform_step.jl")
   include("perform_step/euler_imex_perform_step.jl")
 
   include("dense/generic_dense.jl")
@@ -121,6 +124,7 @@ module OrdinaryDiffEq
   include("adams_utils.jl")
   include("bdf_utils.jl")
   include("exponential_utils.jl")
+  include("rkc_utils.jl")
   include("derivative_wrappers.jl")
   include("iterator_interface.jl")
   include("constants.jl")
@@ -176,6 +180,8 @@ module OrdinaryDiffEq
   export Nystrom4, Nystrom4VelocityIndependent, Nystrom5VelocityIndependent,
          IRKN3, IRKN4, DPRKN6, DPRKN8, DPRKN12, ERKN4, ERKN5
 
+  export ROCK2
+
   export AB3, AB4, AB5, ABM32, ABM43, ABM54
 
   export VCAB3, VCAB4, VCAB5, VCABM3, VCABM4, VCABM5

src/alg_utils.jl

@@ -246,6 +246,8 @@ alg_order(alg::JVODE) = 1  #dummy value
 alg_order(alg::ABDF2) = 2
 alg_order(alg::QNDF1) = 1
 
+alg_order(alg::ROCK2) = 2
+
 alg_maximum_order(alg) = alg_order(alg)
 alg_maximum_order(alg::CompositeAlgorithm) = maximum(alg_order(x) for x in alg.algs)
 

src/algorithms.jl

@@ -287,6 +287,9 @@ Base.@pure JVODE(algorithm=:Adams;bias1=6, bias2=6,bias3=10,
                  addon=1//10^6) = JVODE(algorithm,bias1,bias2,bias3,addon)
 Base.@pure JVODE_Adams(;kwargs...) = JVODE(:Adams;kwargs...)
 
+# ROCK methods
+struct ROCK2 <: OrdinaryDiffEqAdaptiveAlgorithm end
+
 ################################################################################
 
 # Generic implicit methods

src/caches/rkc_caches.jl

@@ -0,0 +1,40 @@
+mutable struct ROCK2ConstantCache{T,T2,zType} <: OrdinaryDiffEqConstantCache
+  ms::SVector{46, Int}
+  fp1::SVector{46, T}
+  fp2::SVector{46, T}
+  recf::Vector{T2}
+  zprev::zType
+  mdegprev::Int
+  mdeg::Int
+  recind::Int
+end
+struct ROCK2Cache{uType,rateType,uEltypeNoUnits} <: OrdinaryDiffEqMutableCache # WIP
+  u::uType
+  uprev::uType
+  gprev::uType
+  gprev2::uType
+  tmp::uType
+  atmp::uEltypeNoUnits
+  fsalfirst::rateType
+  k::rateType
+  k2::rateType
+  constantcache::ROCK2ConstantCache
+end
+u_cache(c::ROCK2Cache)    = (c.atmp,)
+du_cache(c::ROCK2Cache)   = (c.fsalfirst,c.k,c.k2)
+
+function alg_cache(alg::ROCK2,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,uprev2,f,t,dt,reltol,p,calck,::Type{Val{true}})
+  constantcache = ROCK2ConstantCache(uEltypeNoUnits, uEltypeNoUnits, u) # WIP: not sure about what type to use in here
+  gprev = similar(u)
+  gprev2 = similar(u)
+  tmp = similar(u)
+  atmp = similar(u,uEltypeNoUnits,axes(u))
+  fsalfirst = zero(rate_prototype)
+  k = zero(rate_prototype)
+  k2 = zero(rate_prototype)
+  ROCK2Cache(u, uprev, gprev, gprev2, tmp, atmp, fsalfirst, k, k2, constantcache)
+end
+
+function alg_cache(alg::ROCK2,u,rate_prototype,uEltypeNoUnits,uBottomEltypeNoUnits,tTypeNoUnits,uprev,uprev2,f,t,dt,reltol,p,calck,::Type{Val{false}})
+  ROCK2ConstantCache(uEltypeNoUnits, uEltypeNoUnits, u) # WIP: not sure about what type to use in here
+end

src/perform_step/rkc_perform_step.jl

@@ -0,0 +1,121 @@
+function initialize!(integrator, cache::ROCK2ConstantCache)
+  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::ROCK2ConstantCache, repeat_step=false)
+  @unpack t, dt, uprev, u, f, p, fsalfirst = integrator
+  @unpack ms, fp1, fp2, recf = cache
+  # The number of stage.
+  mdeg = Int(floor(sqrt((1.5 + dt * integrator.eigen_est)/0.811) + 1))
+  if mdeg >= 200
+    mdeg = 200
+  end
+  cache.mdeg = max(mdeg, 3) - 2
+  cache.mdeg != cache.mdegprev && choosedeg!(cache)
+  # recurrence
+  # for the first stage
+  temp1 = dt * recf[cache.recind]
+  ci1 = t + temp1
+  ci2 = t + temp1
+  ci3 = t
+  gprev2 = copy(uprev)
+  gprev = uprev + temp1 * fsalfirst
+  ms[cache.mdeg] < 2 && ( u = gprev )
+  # for the second to the ms[cache.mdeg] th stages
+  for i in 2:ms[cache.mdeg]
+    temp1 = dt * recf[cache.recind + 2 * (i - 2) + 1]
+    temp3 = -recf[cache.recind + 2 * (i - 2) + 2]
+    temp2 = 1 - temp3
+    ci1 = temp1 + temp2 * ci2 + temp3 * ci3
+    u = temp1 * u + temp2 * gprev + temp3 * gprev2
+    i < ms[cache.mdeg] && (gprev2 = gprev; gprev = u)
+    ci3 = ci2
+    ci2 = ci1
+  end # end if
+  # two-stage finishing procedure.
+  temp1 = dt * fp1[cache.mdeg]
+  temp2 = dt * fp2[cache.mdeg]
+  gprev2 = f(u, p, ci1)
+  gprev = u + temp1 * gprev2
+  ci1 += temp1
+  u = f(gprev, p, ci1)
+  temp3 = temp2 * (u - gprev2)
+  u = gprev + temp1 * u + temp3
+  # error estimate
+  if integrator.opts.adaptive
+    atmp = calculate_residuals(temp3, uprev, u, integrator.opts.abstol, integrator.opts.reltol,integrator.opts.internalnorm)
+    integrator.EEst = integrator.opts.internalnorm(atmp)
+  end
+  integrator.k[1] = integrator.fsalfirst
+  integrator.k[2] = integrator.fsallast = f(u, p, t+dt)
+  integrator.u = u
+end
+
+function initialize!(integrator, cache::ROCK2Cache)
+  integrator.kshortsize = 2
+  resize!(integrator.k, integrator.kshortsize)
+  integrator.fsalfirst = cache.fsalfirst  # done by pointers, no copying
+  integrator.fsallast = cache.k
+  integrator.k[1] = integrator.fsalfirst
+  integrator.k[2] = integrator.fsallast
+  integrator.f(integrator.fsalfirst, integrator.uprev, integrator.p, integrator.t) # Pre-start fsal
+end
+
+@muladd function perform_step!(integrator, cache::ROCK2Cache, repeat_step=false)
+  @unpack t, dt, uprev, u, f, p, fsalfirst = integrator
+  @unpack k, k2, tmp, gprev2, gprev, atmp = cache
+  @unpack ms, fp1, fp2, recf = cache.constantcache
+  ccache = cache.constantcache
+  # The number of stage.
+  mdeg = Int(floor(sqrt((1.5 + dt * integrator.eigen_est)/0.811) + 1))
+  if mdeg >= 200
+    mdeg = 200
+  end
+  ccache.mdeg = max(mdeg, 3) - 2
+  ccache.mdeg != ccache.mdegprev && choosedeg!(cache)
+  # recurrence
+  # for the first stage
+  temp1 = dt * recf[ccache.recind]
+  ci1 = t + temp1
+  ci2 = t + temp1
+  ci3 = t
+  @. gprev2 = uprev
+  @. gprev = uprev + temp1 * fsalfirst
+  ms[ccache.mdeg] < 2 && ( @. u = gprev )
+  # for the second to the ms[ccache.mdeg] th stages
+  for i in 2:ms[ccache.mdeg]
+    temp1 = dt * recf[ccache.recind + 2 * (i - 2) + 1]
+    temp3 = -recf[ccache.recind + 2 * (i - 2) + 2]
+    temp2 = 1 - temp3
+    ci1 = temp1 + temp2 * ci2 + temp3 * ci3
+    @. u = temp1 * u + temp2 * gprev + temp3 * gprev2
+    i < ms[ccache.mdeg] && (gprev2 .= gprev; gprev .= u)
+    ci3 = ci2
+    ci2 = ci1
+  end # end if
+  # two-stage finishing procedure.
+  temp1 = dt * fp1[ccache.mdeg]
+  temp2 = dt * fp2[ccache.mdeg]
+  f(k, u, p, ci1)
+  @. gprev = u + temp1 * k
+  ci1 += temp1
+  f(k2, gprev, p, ci1)
+  @. tmp = temp2 * (k2 - k)
+  @. u = gprev + temp1 * k2 + tmp
+  # error estimate
+  if integrator.opts.adaptive
+    calculate_residuals!(atmp, tmp, uprev, u, integrator.opts.abstol, integrator.opts.reltol,integrator.opts.internalnorm)
+    integrator.EEst = integrator.opts.internalnorm(atmp)
+  end
+  integrator.k[1] = integrator.fsalfirst
+  f(integrator.fsallast, u, p, t+dt)
+  integrator.k[2] = integrator.fsallast
+  integrator.u = u
+end

src/rkc_utils.jl

@@ -0,0 +1,80 @@
+# This function calculates the largest eigenvalue
+# (absolute value wise) by power iteration.
+
+function maxeig!(integrator, cache::OrdinaryDiffEqConstantCache)
+  isfirst = integrator.iter == 1 || integrator.u_modified
+  @unpack t, dt, uprev, u, f, p, fsalfirst = integrator
+  maxiter = 50
+  safe = 1.2
+  # Initial guess for eigenvector `z`
+  if isfirst
+    fz = fsalfirst
+    z = f(fz, p, t)
+  else
+    z = cache.zprev
+  end
+  # Perturbation
+  u_norm = integrator.opts.internalnorm(uprev)
+  z_norm = integrator.opts.internalnorm(z)
+  pert   = eps(u_norm)
+  sqrt_pert = sqrt(pert)
+  is_u_zero = u_norm == zero(u_norm)
+  is_z_zero = z_norm == zero(z_norm)
+  # Normalize `z` such that z-u lie in a circle
+  if ( !is_u_zero && !is_z_zero )
+    dz_u = u_norm * sqrt_pert
+    quot = dz_u/z_norm
+    z = uprev + quot*z
+  elseif !is_u_zero
+    dz_u = u_norm * sqrt_pert
+    z = uprev + uprev*dz_u
+  elseif !is_z_zero
+    dz_u = pert
+    quot = dz_u/z_norm
+    z *= quot
+  else
+    dz_u = pert
+    z = dz_u
+  end # endif
+  # Start power iteration
+  integrator.eigen_est = 0
+  for iter in 1:maxiter
+    fz = f(z, p, t)
+    tmp = fz - fsalfirst
+    Δ  = integrator.opts.internalnorm(tmp)
+    eig_prev = integrator.eigen_est
+    integrator.eigen_est = Δ/dz_u * safe
+    # Convergence
+    if iter >= 2 && abs(eig_prev - integrator.eigen_est) < integrator.eigen_est*0.05
+      # Store the eigenvector
+      cache.zprev = z
+      return true
+    end
+    # Next `z`
+    if Δ != zero(Δ)
+      quot = dz_u/Δ
+      z = uprev + quot*tmp
+    else
+      # An arbitrary change on `z`
+      nind = length(uprev)
+      ind = 1 + iter % nind
+      z[ind] = uprev[ind] - (z[ind] - uprev[ind])
+    end
+  end
+  return false
+end
+
+function choosedeg!(cache::T) where T
+  isconst = T <: OrdinaryDiffEqConstantCache
+  isconst || ( cache = cache.constantcache )
+  @unpack ms, fp1, fp2, recf, zprev = cache
+  recind = 0
+  @inbounds for i in 1:46
+    recind += ms[i]*2
+    if ms[i] > cache.mdeg
+      cache.mdeg = i
+      cache.recind = recind
+      return nothing
+    end # end if
+  end # end for
+end