Merge #143

143: Increase test coverage r=charleskawczynski a=charleskawczynski

This PR:
 - Fixes some bugs, to get more code running
 - Un-comments some tests, to increase test coverage
 - Refactors, and fixes bugs in, the convergence tabulation

Co-authored-by: Charles Kawczynski <[email protected]>

Author: bors[bot]

src/ClimaTimeSteppers.jl

@@ -88,6 +88,7 @@ include("nl_solvers/newtons_method.jl")
 
 
 n_stages_ntuple(::Type{<:NTuple{Nstages}}) where {Nstages} = Nstages
+n_stages_ntuple(::Type{<:SVector{Nstages}}) where {Nstages} = Nstages
 
 # Include concrete implementations
 include("solvers/imex_ark_tableaus.jl")

src/solvers/ark.jl

@@ -117,7 +117,7 @@ function step_u!(int, cache::AdditiveRungeKuttaFullCache, f::DiffEqBase.SplitFun
 
     Nstages = n_stages(cache)
     (; C, Aimpl, Aexpl, B) = cache.tableau
-    (; U, R, L, W, linsolve) = cache
+    (; U, R, L, W, linsolve!) = cache
     (; u, p, t, dt) = int
 
     U = U
@@ -177,7 +177,7 @@ end
 function step_u!(int, cache::AdditiveRungeKuttaFullCache{Nstages}, f::DiffEqBase.ODEFunction) where {Nstages}
     (; C, Aimpl, Aexpl, B) = cache.tableau
     (; u, p, t, dt) = int
-    (; U, R, L, W, linsolve) = cache
+    (; U, R, L, W, linsolve!) = cache
     Uhat = R[end] # can be used as work array, as only used in last stage
 
     fL! = f.jvp # linear part
@@ -303,13 +303,15 @@ If the keyword `paperversion=true` is used, the coefficients from the paper are
 used. Otherwise it uses coefficients that make the scheme (much) more stable but less
 accurate
 """
-Base.@kwdef struct ARK2GiraldoKellyConstantinescu{L} <: AdditiveRungeKutta
+Base.@kwdef struct ARK2GiraldoKellyConstantinescu{L, PV} <: AdditiveRungeKutta
     linsolve::L
-    paperversion::Bool = false
 end
+ARK2GiraldoKellyConstantinescu(linsolve; paperversion::Bool = false) =
+    ARK2GiraldoKellyConstantinescu{typeof(linsolve), paperversion}(linsolve)
+paperversion(::ARK2GiraldoKellyConstantinescu{L, PV}) where {L, PV} = PV
 
 function tableau(ark::ARK2GiraldoKellyConstantinescu, RT)
-    a32 = RT(ark.paperversion ? (3 + 2 * sqrt(2)) / 6 : 1 // 2)
+    a32 = RT(paperversion(ark) ? (3 + 2 * sqrt(2)) / 6 : 1 // 2)
     RKA_explicit = @SArray [
         RT(0) RT(0) RT(0)
         RT(2 - sqrt(2)) RT(0) RT(0)

src/solvers/lsrk.jl

@@ -62,16 +62,15 @@ function init_inner(prob, outercache::LowStorageRungeKutta2NIncCache, dt)
 end
 function update_inner!(innerinteg, outercache::LowStorageRungeKutta2NIncCache, f_slow, u, p, t, dt, stage)
 
-    (; C, A, B) = cache.tableau
+    (; C, A, B) = outercache.tableau
     f_offset = innerinteg.sol.prob.f
-    tab = outercache.tableau
     N = n_stages(outercache)
 
     τ0 = t + C[stage] * dt
     τ1 = stage == N ? t + dt : t + C[stage + 1] * dt
     f_offset.α = τ0
     innerinteg.t = zero(τ0)
-    innerinteg.tstop = τ1 - τ0
+    DiffEqBase.add_tstop!(innerinteg, τ1 - τ0) # TODO: verify correctness
 
     #  du .= f(u, p, t + C[stage]*dt) .+ A[stage] .* du
     f_slow(f_offset.x, u, p, τ0, 1, A[stage])

src/solvers/mis.jl

@@ -102,7 +102,7 @@ function update_inner!(innerinteg, outercache::MultirateInfinitesimalStepCache,
         F,
         innerinteg.u,
         f_offset.x,
-        tab,
+        outercache.tableau, # TODO: verify correctness
         i,
         N,
         dt;
@@ -117,7 +117,7 @@ function update_inner!(innerinteg, outercache::MultirateInfinitesimalStepCache,
     f_offset.β = (c[i] - c̃[i]) / d[i]
 
     innerinteg.t = zero(t)
-    innerinteg.tstop = d[i] * dt
+    DiffEqBase.add_tstop!(innerinteg, d[i] * dt) # TODO: verify correctness
 end
 
 @kernel function mis_update!(u, ΔU, F, innerinteg_u, f_offset_x, tab, i, N, dt)

src/solvers/multirate.jl

@@ -61,11 +61,10 @@ function step_u!(int, cache::MultirateCache)
 
         # TODO: make this more generic
         # there are 2 strategies we can use here:
-        #  a. use same fast_dt for all slow stages, use `adjustfinal=true`
+        #  a. use same fast_dt for all slow stages
         #     - problems for ARK (e.g. requires expensive LU factorization)
         #  b. use different fast_dt, cache expensive ops
 
-        innerinteg.adjustfinal = true
         DiffEqBase.solve!(innerinteg)
         innerinteg.dt = fast_dt # reset
     end

src/solvers/wickerskamarock.jl

@@ -54,7 +54,8 @@ function update_inner!(innerinteg, outercache::WickerSkamarockRungeKuttaCache, f
     end
 
     innerinteg.t = t
-    innerinteg.tstop = i == N ? t + dt : t + c[i + 1] * dt
+    t_star = i == N ? t + dt : t + c[i + 1] * dt
+    DiffEqBase.add_tstop!(innerinteg, t_star) # TODO: verify correctness
 end
 
 

test/convergence.jl

@@ -45,9 +45,9 @@ end
 
 #=
 using Revise; include("test/convergence.jl")
-results = tabulate_convergence_orders();
+results = tabulate_convergence_orders_imex_ark();
 =#
-function tabulate_convergence_orders()
+function tabulate_convergence_orders_imex_ark()
     tabs = [
         ARS111,
         ARS121,
@@ -76,54 +76,78 @@ function tabulate_convergence_orders()
     tabs = map(t -> t(), tabs)
     test_cases = all_test_cases(Float64)
     results = convergence_order_results(tabs, test_cases)
-    tabulate_convergence_orders(test_cases, tabs, results)
+    algs = algorithm.(tabs)
+    prob_names = map(t -> t.test_name, test_cases)
+    expected_orders = ODE.alg_order.(tabs)
+    tabulate_convergence_orders(prob_names, algs, results, expected_orders; tabs)
     return results
 end
-tabulate_convergence_orders()
+tabulate_convergence_orders_imex_ark()
 
-#=
-if ArrayType == Array
-for (prob, sol) in [
-    imex_autonomous_prob => imex_autonomous_sol,
-    #imex_nonautonomous_prob => imex_nonautonomous_sol,
-]
-# IMEX
-    @test convergence_order(prob, sol, ARK1ForwardBackwardEuler(linsolve=DirectSolver), dts)       ≈ 1 atol=0.1
-    @test convergence_order(prob, sol, ARK2ImplicitExplicitMidpoint(linsolve=DirectSolver), dts)   ≈ 2 atol=0.05
-    @test convergence_order(prob, sol, ARK2GiraldoKellyConstantinescu(linsolve=DirectSolver), dts) ≈ 2 atol=0.05
-    @test convergence_order(prob, sol, ARK2GiraldoKellyConstantinescu(linsolve=DirectSolver; paperversion=true), dts) ≈ 2 atol=0.05
-    @test convergence_order(prob, sol, ARK437L2SA1KennedyCarpenter(linsolve=DirectSolver), dts)    ≈ 4 atol=0.05
-    @test convergence_order(prob, sol, ARK548L2SA2KennedyCarpenter(linsolve=DirectSolver), dts)    ≈ 5 atol=0.05
-end
+function tabulate_convergence_orders_ark()
+    # IMEX
+    co = Dict()
+    names_probs_sols = [
+        (:auto, imex_autonomous_prob(Array{Float64}), imex_autonomous_sol),
+        (:nonauto, imex_nonautonomous_prob(Array{Float64}), imex_nonautonomous_sol),
+    ]
+    algs_orders = [
+        (ARK1ForwardBackwardEuler(DirectSolver), 1),
+        (ARK2ImplicitExplicitMidpoint(DirectSolver), 2),
+        (ARK2GiraldoKellyConstantinescu(DirectSolver), 2),
+        (ARK2GiraldoKellyConstantinescu(DirectSolver; paperversion = true), 2),
+        (ARK437L2SA1KennedyCarpenter(DirectSolver), 4),
+        (ARK548L2SA2KennedyCarpenter(DirectSolver), 5),
+    ]
+    dts = 0.5 .^ (4:7)
+    for (name, prob, sol) in names_probs_sols
+        for (alg, ord) in algs_orders
+            co[name, typeof(alg)] = (ord, convergence_order(prob, sol, alg, dts))
+        end
+    end
+    prob_names = first.(names_probs_sols)
+    algs = first.(algs_orders)
+    expected_orders = last.(algs_orders)
+    tabulate_convergence_orders(prob_names, algs, co, expected_orders)
 end
-for (prob, sol) in [
-    imex_autonomous_prob => imex_autonomous_sol,
-    imex_nonautonomous_prob => imex_nonautonomous_sol,
-    # kpr_multirate_prob => kpr_sol,
-]
-    # Multirate
-    @test convergence_order(prob, sol, Multirate(LSRK54CarpenterKennedy(),LSRK54CarpenterKennedy()), dts;
-        fast_dt = 0.5^12, adjustfinal=true) ≈ 4 atol=0.05
-    # MIS
-    @test convergence_order(prob, sol, Multirate(LSRK54CarpenterKennedy(),MIS2()), dts;
-        fast_dt = 0.5^12, adjustfinal=true) ≈ 2 atol=0.05
-    @test convergence_order(prob, sol, Multirate(LSRK54CarpenterKennedy(),MIS3C()), dts;
-        fast_dt = 0.5^12, adjustfinal=true) ≈ 2 atol=0.05
-    @test convergence_order(prob, sol, Multirate(LSRK54CarpenterKennedy(),MIS4()), dts;
-        fast_dt = 0.5^12, adjustfinal=true) ≈ 3 atol=0.05
-    @test convergence_order(prob, sol, Multirate(LSRK54CarpenterKennedy(),MIS4a()), dts;
-        fast_dt = 0.5^12, adjustfinal=true) ≈ 3 atol=0.05
-    @test convergence_order(prob, sol, Multirate(LSRK54CarpenterKennedy(),TVDMISA()), dts;
-        fast_dt = 0.5^12, adjustfinal=true) ≈ 2 atol=0.05
-    @test convergence_order(prob, sol, Multirate(LSRK54CarpenterKennedy(),TVDMISB()), dts;
-        fast_dt = 0.5^12, adjustfinal=true) ≈ 2 atol=0.05
-
-    # Wicker Skamarock
-    @test convergence_order(prob, sol, Multirate(LSRK54CarpenterKennedy(),WSRK2()), dts;
-        fast_dt = 0.5^12, adjustfinal=true) ≈ 2 atol=0.05
-    @test convergence_order(prob, sol, Multirate(LSRK54CarpenterKennedy(),WSRK3()), dts;
-        fast_dt = 0.5^12, adjustfinal=true) ≈ 2 atol=0.05
 
+tabulate_convergence_orders_ark()
+
+function tabulate_convergence_orders_multirate()
+
+    co = Dict()
+    names_probs_sols = [
+        (:imex_auto, imex_autonomous_prob(Array{Float64}), imex_autonomous_sol),
+        (:imex_nonauto, imex_nonautonomous_prob(Array{Float64}), imex_nonautonomous_sol),
+        # (:kpr_multirate, kpr_multirate_prob(), kpr_sol),
+    ]
+    dts = 0.5 .^ (4:7)
+
+    algs_orders = [
+        # Multirate
+        (Multirate(LSRK54CarpenterKennedy(), LSRK54CarpenterKennedy()), 4),
+        # MIS
+        (Multirate(LSRK54CarpenterKennedy(), MIS2()), 2),
+        (Multirate(LSRK54CarpenterKennedy(), MIS3C()), 2),
+        (Multirate(LSRK54CarpenterKennedy(), MIS4()), 3),
+        (Multirate(LSRK54CarpenterKennedy(), MIS4a()), 3),
+        (Multirate(LSRK54CarpenterKennedy(), TVDMISA()), 2),
+        (Multirate(LSRK54CarpenterKennedy(), TVDMISB()), 2),
+        # Wicker Skamarock
+        (Multirate(LSRK54CarpenterKennedy(), WSRK2()), 2),
+        (Multirate(LSRK54CarpenterKennedy(), WSRK3()), 2),
+    ]
 
+    for (name, prob, sol) in names_probs_sols
+        for (alg, ord) in algs_orders
+            co[name, typeof(alg)] = (ord, convergence_order(prob, sol, alg, dts; fast_dt = 0.5^12))
+        end
     end
-=#
+
+    prob_names = first.(names_probs_sols)
+    algs = first.(algs_orders)
+    expected_orders = last.(algs_orders)
+    tabulate_convergence_orders(prob_names, algs, co, expected_orders)
+end
+
+tabulate_convergence_orders_multirate()

test/convergence_orders.jl

@@ -4,13 +4,13 @@
 
 # TODO: is it better to use `first_order_tableau = Union{ARS111,ARS121}`? to
 #       reduce the number of methods?
-const first_order_tableau = [ARS111, ARS121]
+first_order_tableau() = [ARS111, ARS121]
 
 #####
 ##### 2nd order
 #####
 
-const second_order_tableau = [
+second_order_tableau() = [
     ARS122,
     ARS232,
     ARS222,
@@ -31,18 +31,18 @@ const second_order_tableau = [
 #####
 ##### 3rd order
 #####
-const third_order_tableau = [ARS233, ARS343, ARS443, IMKG342a, IMKG343a, DBM453]
+third_order_tableau() = [ARS233, ARS343, ARS443, IMKG342a, IMKG343a, DBM453]
 
 import OrdinaryDiffEq as ODE
 import ClimaTimeSteppers as CTS
 ODE.alg_order(alg::CTS.IMEXARKAlgorithm) = ODE.alg_order(alg.tab)
 
-for m in first_order_tableau
+for m in first_order_tableau()
     @eval ODE.alg_order(::$m) = 1
 end
-for m in second_order_tableau
+for m in second_order_tableau()
     @eval ODE.alg_order(::$m) = 2
 end
-for m in third_order_tableau
+for m in third_order_tableau()
     @eval ODE.alg_order(::$m) = 3
 end

test/convergence_utils.jl

@@ -26,10 +26,11 @@ on the set of `dt` values in `dts`. Extra `kwargs` are passed to `solve`
 `solution` should be a function with a method `solution(u0, p, t)`.
 """
 function convergence_errors(prob, sol, method, dts; kwargs...)
+    hide_warning = (; kwargshandle = DiffEqBase.KeywordArgSilent)
     errs = map(dts) do dt
         # copy the problem so we don't mutate u0
         prob_copy = deepcopy(prob)
-        u = solve(prob_copy, method; dt = dt, saveat = (prob.tspan[2],), kwargs...)
+        u = solve(prob_copy, method; dt = dt, saveat = (prob.tspan[2],), kwargs..., hide_warning...)
         norm(u .- sol(prob.u0, prob.p, prob.tspan[end]))
     end
     return errs
@@ -66,13 +67,12 @@ function test_convergence_order!(test_case, tab, results = Dict(); refinement_ra
         refinement_range, # ::UnitRange, 2:4 is more fine than 1:3
     )
     computed_order = maximum(cr.computed_order)
-    results[tab, test_case.test_name] = (; expected_order, computed_order)
+    results[test_case.test_name, typeof(alg)] = (; expected_order, computed_order)
     return nothing
 end
 
-distance(theoretic, computed) = abs(computed - theoretic) / theoretic
-pass_conv(theoretic, computed) = distance(theoretic, computed) * 100 < 10
-fail_conv(theoretic, computed) = !pass_conv(computed, theoretic) && !super_conv(theoretic, computed)
+pass_conv(theoretic, computed) = abs(computed - theoretic) / theoretic * 100 < 10 && computed > 0
+fail_conv(theoretic, computed) = !pass_conv(theoretic, computed) && !super_conv(theoretic, computed)
 super_conv(theoretic, computed) = (computed - theoretic) / theoretic * 100 > 10
 
 #= Calls `test_convergence_order!` for each combination of test case
@@ -89,31 +89,32 @@ function convergence_order_results(tabs, test_cases)
     return results
 end
 
-function tabulate_convergence_orders(test_cases, tabs, results)
-    columns = map(test_cases) do test_case
-        map(tab -> results[tab, test_case.test_name], tabs)
+function tabulate_convergence_orders(prob_names, algs, results, expected_orders; tabs = nothing)
+    data = hcat(map(prob_names) do name
+        map(alg -> results[name, typeof(alg)], algs)
+    end...)
+    alg_names = if tabs ≠ nothing
+        @. string(nameof(typeof(tabs)))
+    else
+        @. string(typeof(algs))
     end
-    expected_order = map(tab -> default_expected_order(nothing, tab), tabs)
-    tab_names = map(tab -> "$tab ($(default_expected_order(nothing, tab)))", tabs)
-    data = hcat(columns...)
-    summary(result) = result.computed_order
+    summary(result) = last(result)
     data_summary = map(d -> summary(d), data)
 
-    table_data = hcat(tab_names, data_summary)
+    table_data = hcat(alg_names, data_summary)
     precentage_fail = sum(fail_conv.(getindex.(data, 1), getindex.(data, 2))) / length(data) * 100
     @info "Percentage of failed convergence order tests: $precentage_fail"
     fail_conv_hl = PrettyTables.Highlighter(
-        (data, i, j) -> j ≠ 1 && fail_conv(expected_order[i], data[i, j]),
+        (data, i, j) -> j ≠ 1 && fail_conv(expected_orders[i], data[i, j]),
         PrettyTables.crayon"red bold",
     )
     super_conv_hl = PrettyTables.Highlighter(
-        (data, i, j) -> j ≠ 1 && super_conv(expected_order[i], data[i, j]),
+        (data, i, j) -> j ≠ 1 && super_conv(expected_orders[i], data[i, j]),
         PrettyTables.crayon"yellow bold",
     )
     tab_column_hl = PrettyTables.Highlighter((data, i, j) -> j == 1, PrettyTables.crayon"green bold")
-    test_case_names = map(test_case -> test_case.test_name, test_cases)
 
-    header = (["Tableau (theoretic)", test_case_names...],
+    header = (["Tableau (theoretic)", prob_names...],
     # ["", ["" for tc in test_case_names]...],
     )
 

test/ode_tests_basic.jl

@@ -109,7 +109,7 @@ Qexact = exactsolution(finaltime, q0, t0)
                 for (n, dt) in enumerate(dts)
                     Q .= Qinit
                     prob = IncrementingODEProblem(rhs!, Q, (t0, finaltime))
-                    solve(prob, method; dt=dt, adjustfinal=true)
+                    solve(prob, method; dt=dt)
                     errors[n] = norm(Q - Qexact)
                 end
                 rates = log2.(errors[1:(end - 1)] ./ errors[2:end])
@@ -129,7 +129,7 @@ Qexact = exactsolution(finaltime, q0, t0)
                 # rhs_arg! =
                 #   split_explicit_implicit ? rhs_nonlinear! : rhs!
                 prob = SplitODEProblem(rhs_linear!, rhs_nonlinear!, Q, (t0, finaltime))
-                solve(prob, method(DirectSolver); dt = dt, adjustfinal = true)
+                solve(prob, method(DirectSolver); dt = dt)
                 errors[n] = norm(Q - Qexact)
                 @show (log2(dt), norm(Q - Qexact))
             end
@@ -160,7 +160,7 @@ Qexact = exactsolution(finaltime, q0, t0)
                 # rhs_arg! =
                 #   split_explicit_implicit ? rhs_nonlinear! : rhs!
                 prob = ODEProblem(ODEFunction(rhs!, jvp = rhs_linear!), Q, (t0, finaltime))
-                solve(prob, method(DirectSolver); dt = dt, adjustfinal = true)
+                solve(prob, method(DirectSolver); dt = dt)
                 errors[n] = norm(Q - Qexact)
                 @show (log2(dt), norm(Q - Qexact))
             end