Merge pull request #2548 from Shreyas-Ekanathan/master

Small Radau Optimizations

lib/OrdinaryDiffEqFIRK/src/controllers.jl

@@ -1,7 +1,7 @@
 function step_accept_controller!(integrator, controller::PredictiveController, alg::AdaptiveRadau, q)
     @unpack qmin, qmax, gamma, qsteady_min, qsteady_max = integrator.opts
     @unpack cache = integrator
-    @unpack num_stages, step, iter, hist_iter = cache
+    @unpack num_stages, step, iter, hist_iter, index = cache
 
     EEst = DiffEqBase.value(integrator.EEst)
 
@@ -25,12 +25,14 @@ function step_accept_controller!(integrator, controller::PredictiveController, a
     max_stages = (alg.max_order - 1) ÷ 4 * 2 + 1
     min_stages = (alg.min_order - 1) ÷ 4 * 2 + 1
     if (step > 10)
-        if (hist_iter < 2.6 && num_stages <= max_stages)
+        if (hist_iter < 2.6 && num_stages < max_stages)
             cache.num_stages += 2
+            cache.index += 1
             cache.step = 1
             cache.hist_iter = iter
-        elseif ((hist_iter > 8 || cache.status == VerySlowConvergence || cache.status == Divergence) && num_stages >= min_stages)
+        elseif ((hist_iter > 8 || cache.status == VerySlowConvergence || cache.status == Divergence) && num_stages > min_stages)
             cache.num_stages -= 2 
+            cache.index -= 1
             cache.step = 1
             cache.hist_iter = iter
         end
@@ -48,8 +50,9 @@ function step_reject_controller!(integrator, controller::PredictiveController, a
     cache.hist_iter = hist_iter
     min_stages = (alg.min_order - 1) ÷ 4 * 2 + 1
     if (step > 10)
-        if ((hist_iter > 8 || cache.status == VerySlowConvergence || cache.status == Divergence) && num_stages >= min_stages)
+        if ((hist_iter > 8 || cache.status == VerySlowConvergence || cache.status == Divergence) && num_stages > min_stages)
             cache.num_stages -= 2 
+            cache.index -= 1
             cache.step = 1
             cache.hist_iter = iter
         end

lib/OrdinaryDiffEqFIRK/src/firk_caches.jl

@@ -497,6 +497,7 @@ mutable struct AdaptiveRadauConstantCache{F, Tab, Tol, Dt, U, JType} <:
     num_stages::Int
     step::Int
     hist_iter::Float64
+    index::Int
 end
 
 function alg_cache(alg::AdaptiveRadau, u, rate_prototype, ::Type{uEltypeNoUnits},
@@ -508,30 +509,32 @@ function alg_cache(alg::AdaptiveRadau, u, rate_prototype, ::Type{uEltypeNoUnits}
 
     max_order = alg.max_order
     min_order = alg.min_order
-    max = (max_order - 1) ÷ 4 * 2 + 1
-    min = (min_order - 1) ÷ 4 * 2 + 1
+    max_stages = (max_order - 1) ÷ 4 * 2 + 1
+    min_stages = (min_order - 1) ÷ 4 * 2 + 1
     if (alg.min_order < 5)
         error("min_order choice $min_order below 5 is not compatible with the algorithm")
-    elseif (max < min)
+    elseif (max_stages < min_stages)
         error("max_order $max_order is below min_order $min_order")
     end
-    num_stages = min
+    num_stages = min_stages
 
     tabs = [RadauIIATableau5(uToltype, constvalue(tTypeNoUnits)), RadauIIATableau9(uToltype, constvalue(tTypeNoUnits)), RadauIIATableau13(uToltype, constvalue(tTypeNoUnits))]
-    i = 9
-    while i <= max
+    i = max(min_stages, 9)
+    while i <= max_stages
         push!(tabs, RadauIIATableau(uToltype, constvalue(tTypeNoUnits), i))
         i += 2
     end
-    cont = Vector{typeof(u)}(undef, max)
-    for i in 1:max
+    cont = Vector{typeof(u)}(undef, max_stages)
+    for i in 1:max_stages
         cont[i] = zero(u)
     end
 
+    index = min((min_stages - 1) ÷ 2, 4)
+
     κ = alg.κ !== nothing ? convert(uToltype, alg.κ) : convert(uToltype, 1 // 100)
     J = false .* _vec(rate_prototype) .* _vec(rate_prototype)'
     AdaptiveRadauConstantCache(uf, tabs, κ, one(uToltype), 10000, cont, dt, dt,
-        Convergence, J, num_stages, 1, 0.0)
+        Convergence, J, num_stages, 1, 0.0, index)
 end
 
 mutable struct AdaptiveRadauCache{uType, cuType, tType, uNoUnitsType, rateType, JType, W1Type, W2Type,
@@ -578,6 +581,7 @@ mutable struct AdaptiveRadauCache{uType, cuType, tType, uNoUnitsType, rateType,
     num_stages::Int
     step::Int
     hist_iter::Float64
+    index::Int
 end
 
 function alg_cache(alg::AdaptiveRadau, u, rate_prototype, ::Type{uEltypeNoUnits},
@@ -589,56 +593,58 @@ function alg_cache(alg::AdaptiveRadau, u, rate_prototype, ::Type{uEltypeNoUnits}
 
     max_order = alg.max_order
     min_order = alg.min_order
-    max = (max_order - 1) ÷ 4 * 2 + 1
-    min = (min_order - 1) ÷ 4 * 2 + 1
+    max_stages = (max_order - 1) ÷ 4 * 2 + 1
+    min_stages = (min_order - 1) ÷ 4 * 2 + 1
     if (alg.min_order < 5)
         error("min_order choice $min_order below 5 is not compatible with the algorithm")
-    elseif (max < min)
+    elseif (max_stages < min_stages)
         error("max_order $max_order is below min_order $min_order")
     end
-    num_stages = min
+    num_stages = min_stages
 
     tabs = [RadauIIATableau5(uToltype, constvalue(tTypeNoUnits)), RadauIIATableau9(uToltype, constvalue(tTypeNoUnits)), RadauIIATableau13(uToltype, constvalue(tTypeNoUnits))]
-    i = 9
-    while i <= max
+    i = max(min_stages, 9)
+    while i <= max_stages
         push!(tabs, RadauIIATableau(uToltype, constvalue(tTypeNoUnits), i))
         i += 2
     end
 
+    index = min((min_stages - 1) ÷ 2, 4)
+
     κ = alg.κ !== nothing ? convert(uToltype, alg.κ) : convert(uToltype, 1 // 100)
 
-    z = Vector{typeof(u)}(undef, max)
-    w = Vector{typeof(u)}(undef, max)
-    for i in 1 : max
+    z = Vector{typeof(u)}(undef, max_stages)
+    w = Vector{typeof(u)}(undef, max_stages)
+    for i in 1 : max_stages
         z[i] = zero(u)
         w[i] = zero(u)
     end
 
-    αdt = [zero(t) for i in 1:max]
-    βdt = [zero(t) for i in 1:max]
-    c_prime = Vector{typeof(t)}(undef, max) #time stepping
-    for i in 1 : max
+    αdt = [zero(t) for i in 1:max_stages]
+    βdt = [zero(t) for i in 1:max_stages]
+    c_prime = Vector{typeof(t)}(undef, max_stages) #time stepping
+    for i in 1 : max_stages
         c_prime[i] = zero(t)
     end
 
     dw1 = zero(u)
     ubuff = zero(u)
-    dw2 = [similar(u, Complex{eltype(u)}) for _ in 1 : (max - 1) ÷ 2]
+    dw2 = [similar(u, Complex{eltype(u)}) for _ in 1 : (max_stages - 1) ÷ 2]
     recursivefill!.(dw2, false)
-    cubuff = [similar(u, Complex{eltype(u)}) for _ in 1 : (max - 1) ÷ 2]
+    cubuff = [similar(u, Complex{eltype(u)}) for _ in 1 : (max_stages - 1) ÷ 2]
     recursivefill!.(cubuff, false)
-    dw = [zero(u) for i in 1 : max]
+    dw = [zero(u) for i in 1:max_stages]
 
-    cont = [zero(u) for i in 1:max]
+    cont = [zero(u) for i in 1:max_stages]
 
-    derivatives = Matrix{typeof(u)}(undef, max, max)
-    for i in 1 : max, j in 1 : max
+    derivatives = Matrix{typeof(u)}(undef, max_stages, max_stages)
+    for i in 1 : max_stages, j in 1 : max_stages
         derivatives[i, j] = zero(u)
     end
 
     fsalfirst = zero(rate_prototype)
-    fw = [zero(rate_prototype) for i in 1 : max]
-    ks = [zero(rate_prototype) for i in 1 : max]
+    fw = [zero(rate_prototype) for i in 1 : max_stages]
+    ks = [zero(rate_prototype) for i in 1 : max_stages]
 
     k = ks[1]
 
@@ -647,7 +653,7 @@ function alg_cache(alg::AdaptiveRadau, u, rate_prototype, ::Type{uEltypeNoUnits}
         error("Non-concrete Jacobian not yet supported by AdaptiveRadau.")
     end
 
-    W2 = [similar(J, Complex{eltype(W1)}) for _ in 1 : (max - 1) ÷ 2]
+    W2 = [similar(J, Complex{eltype(W1)}) for _ in 1 : (max_stages - 1) ÷ 2]
     recursivefill!.(W2, false)
 
     du1 = zero(rate_prototype)
@@ -665,7 +671,7 @@ function alg_cache(alg::AdaptiveRadau, u, rate_prototype, ::Type{uEltypeNoUnits}
 
     linsolve2 = [
         init(LinearProblem(W2[i], _vec(cubuff[i]); u0 = _vec(dw2[i])), alg.linsolve, alias_A = true, alias_b = true,
-            assumptions = LinearSolve.OperatorAssumptions(true)) for i in 1 : (max - 1) ÷ 2]
+            assumptions = LinearSolve.OperatorAssumptions(true)) for i in 1 : (max_stages - 1) ÷ 2]
 
     rtol = reltol isa Number ? reltol : zero(reltol)
     atol = reltol isa Number ? reltol : zero(reltol)
@@ -677,6 +683,6 @@ function alg_cache(alg::AdaptiveRadau, u, rate_prototype, ::Type{uEltypeNoUnits}
         uf, tabs, κ, one(uToltype), 10000, tmp,
         atmp, jac_config,
         linsolve1, linsolve2, rtol, atol, dt, dt,
-        Convergence, alg.step_limiter!, num_stages, 1, 0.0)
+        Convergence, alg.step_limiter!, num_stages, 1, 0.0, index)
 end
 

lib/OrdinaryDiffEqFIRK/src/firk_perform_step.jl

@@ -1354,8 +1354,8 @@ end
 @muladd function perform_step!(integrator, cache::AdaptiveRadauConstantCache,
     repeat_step = false)
     @unpack t, dt, uprev, u, f, p = integrator
-    @unpack tabs, num_stages = cache
-    tab = tabs[(num_stages - 1) ÷ 2]
+    @unpack tabs, num_stages, index = cache
+    tab = tabs[index]
     @unpack T, TI, γ, α, β, c, e = tab
     @unpack κ, cont = cache
     @unpack internalnorm, abstol, reltol, adaptive = integrator.opts
@@ -1595,8 +1595,8 @@ end
 
 @muladd function perform_step!(integrator, cache::AdaptiveRadauCache, repeat_step = false)
     @unpack t, dt, uprev, u, f, p, fsallast, fsalfirst = integrator
-    @unpack num_stages, tabs = cache
-    tab = tabs[(num_stages - 1) ÷ 2]
+    @unpack num_stages, tabs, index = cache
+    tab = tabs[index]
     @unpack T, TI, γ, α, β, c, e = tab
     @unpack κ, cont, derivatives, z, w, c_prime, αdt, βdt= cache
     @unpack dw1, ubuff, dw2, cubuff, dw = cache

lib/OrdinaryDiffEqFIRK/test/ode_high_order_firk_tests.jl

@@ -7,7 +7,7 @@ testTol = 0.5
 prob_ode_linear_big = remake(prob_ode_linear, u0 = big.(prob_ode_linear.u0), tspan = big.(prob_ode_linear.tspan))
 prob_ode_2Dlinear_big = remake(prob_ode_2Dlinear, u0 = big.(prob_ode_2Dlinear.u0), tspan = big.(prob_ode_2Dlinear.tspan))
 
-for i in [17, 21], prob in [prob_ode_linear_big, prob_ode_2Dlinear_big]
+for i in [17, 21, 25], prob in [prob_ode_linear_big, prob_ode_2Dlinear_big]
     dts = 1 ./ 2 .^ (4.25:-1:0.25)
     sim21 = test_convergence(dts, prob, AdaptiveRadau(min_order = i, max_order = i))
     @test sim21.𝒪est[:final]≈ i atol=testTol