Refactor tests, prep for tab-aware alg

Author: charleskawczynski

docs/src/algo_comparisons.jl

@@ -15,7 +15,8 @@ include(joinpath(cts_dir, "test", "problems.jl"))
     tab2 = (tab2..., IMKG254b, IMKG254c, HOMMEM1)
     tab3 = (ARS233, ARS343, ARS443, IMKG342a, IMKG343a, DBM453)
     tabs = [tab1..., tab2..., tab3...]
+    tabs = map(t -> t(), tabs)
     test_algs("IMEX ARK", tabs, ark_analytic_nonlin_test_cts(Float64), 400)
     test_algs("IMEX ARK", tabs, ark_analytic_sys_test_cts(Float64), 60)
-    test_algs("IMEX ARK", tabs, ark_analytic_test_cts(Float64), 16000; super_convergence = ARS121)
+    test_algs("IMEX ARK", tabs, ark_analytic_test_cts(Float64), 16000; super_convergence = ARS121())
 end

docs/src/plotting_utils.jl

@@ -86,14 +86,15 @@ function test_algs(
     analytic_end_sol = [analytic_sols[end]]
 
     for tab in tableaus
-        (prob, alg) = problem_algo(test_case, tab)
+        prob = problem(test_case, tab)
+        alg = algorithm(test_case, tab)
         predicted_order = if super_convergence == tab
-            CTS.theoretical_convergence_order(tab()) + 1
+            CTS.theoretical_convergence_order(tab) + 1
         else
-            CTS.theoretical_convergence_order(tab())
+            CTS.theoretical_convergence_order(tab)
         end
         linestyle = linestyles[(predicted_order - 1) % length(linestyles) + 1]
-        alg_name = string(nameof(tab))
+        alg_name = string(nameof(typeof(tab)))
 
         # Use tstops to fix saving issues due to machine precision (e.g. if the
         # integrator needs to save at t but it stops at t - eps(), it will skip

test/convergence.jl

@@ -72,6 +72,7 @@ function tabulate_convergence_orders()
         IMKG343a,
         DBM453,
     ]
+    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)

test/convergence_utils.jl

@@ -59,8 +59,9 @@ function default_expected_order(alg, tab)
 end
 
 function test_convergence_order!(test_case, tab, results = Dict(); refinement_range)
-    prob, alg = problem_algo(test_case, tab)
-    expected_order = default_expected_order(alg, tab())
+    prob = problem(test_case, tab)
+    alg = algorithm(test_case, tab)
+    expected_order = default_expected_order(alg, tab)
     cr = OCT.refinement_study(
         prob,
         alg;
@@ -96,8 +97,8 @@ function tabulate_convergence_orders(test_cases, tabs, results)
     columns = map(test_cases) do test_case
         map(tab -> results[tab, test_case.test_name], tabs)
     end
-    expected_order = map(tab -> default_expected_order(nothing, tab()), tabs)
-    tab_names = map(tab -> "$tab ($(default_expected_order(nothing, tab())))", tabs)
+    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
     data_summary = map(d -> summary(d), data)

test/problems.jl

@@ -274,11 +274,11 @@ function IntegratorTestCase(;
     jac_prototype = Matrix{FT}(undef, length(Y₀), length(Y₀))
     func_args = (; jac_prototype, Wfact = Wfact!, tgrad = tgrad!)
     tendency_func = ODEFunction(tendency!; func_args...)
-    if isnothing(implicit_tendency!) # assume that related args are also nothing
+    split_tendency_func = if isnothing(implicit_tendency!) # assume that related args are also nothing
         no_tendency!(Yₜ, Y, _, t) = Yₜ .= 0
-        split_tendency_func = SplitFunction(tendency_func, no_tendency!)
+        SplitFunction(tendency_func, no_tendency!)
     else
-        split_tendency_func = SplitFunction(ODEFunction(implicit_tendency!; func_args...), explicit_tendency!)
+        SplitFunction(ODEFunction(implicit_tendency!; func_args...), explicit_tendency!)
     end
     make_prob(func) = ODEProblem(func, Y₀, (FT(0), t_end), nothing)
     IntegratorTestCase(
@@ -307,12 +307,14 @@ function ClimaIntegratorTestCase(;
     jac_prototype = Matrix{FT}(undef, length(Y₀), length(Y₀))
     func_args = (; jac_prototype, Wfact = Wfact!, tgrad = tgrad!)
     tendency_func = ClimaODEFunction(; T_imp! = ODEFunction(tendency!; func_args...))
-    if isnothing(implicit_tendency!) # assume that related args are also nothing
-        split_tendency_func = ClimaODEFunction(; T_imp! = ODEFunction(tendency!; func_args...))
+
+    T_imp! = if isnothing(implicit_tendency!)
+        # assume that related args are also nothing
+        ODEFunction(tendency!; func_args...)
     else
-        split_tendency_func =
-            ClimaODEFunction(; T_exp! = explicit_tendency!, T_imp! = ODEFunction(implicit_tendency!; func_args...))
+        ODEFunction(implicit_tendency!; func_args...)
     end
+    split_tendency_func = ClimaODEFunction(; T_exp! = explicit_tendency!, T_imp!)
     make_prob(func) = ODEProblem(func, Y₀, (FT(0), t_end), nothing)
     IntegratorTestCase(
         test_name,

test/utils.jl

@@ -1,14 +1,14 @@
 import ClimaTimeSteppers as CTS
 using Test
 
-function problem_algo(test_case, tab)
-    if tab() isa CTS.AbstractIMEXARKTableau
+problem(test_case, tab::CTS.AbstractIMEXARKTableau) = test_case.split_prob
+problem(test_case, tab) = test_case.prob
+
+function algorithm(test_case, tab)
+    return if tab isa CTS.AbstractIMEXARKTableau
         max_iters = test_case.linear_implicit ? 1 : 2 # TODO: is 2 enough?
-        alg = CTS.IMEXARKAlgorithm(tab(), NewtonsMethod(; max_iters))
-        prob = test_case.split_prob
+        CTS.IMEXARKAlgorithm(tab, NewtonsMethod(; max_iters))
     else
-        alg = tab()
-        prob = test_case.prob
+        tab
     end
-    return (prob, alg)
 end