Make user-facing types

Author: charleskawczynski

docs/src/algorithms.md

@@ -16,7 +16,7 @@ ForwardEulerODEFunction
 
 ```@docs
 IMEXARKAlgorithm
-make_IMEXARKAlgorithm
+make_IMEXARKTableau
 ```
 
 The convergence orders of the provided methods are verified using test cases from [ARKode](http://runge.math.smu.edu/ARKode_example.pdf). Plots of the solutions to these test cases, the errors of these solutions, and the convergence orders with respect to `dt` are shown below.

perf/flame.jl

@@ -30,7 +30,7 @@ elseif problem_str=="fe"
 else
     error("Bad option")
 end
-algorithm = ARS343(NewtonsMethod(; max_iters = 2))
+algorithm = CTS.IMEXARKAlgorithm(ARS343(), NewtonsMethod(; max_iters = 2))
 dt = 0.01
 integrator = DiffEqBase.init(prob, algorithm; dt)
 not_generated_cache = CTS.not_generated_cache(prob, algorithm)

perf/jet.jl

@@ -15,7 +15,7 @@ end
 cts = joinpath(dirname(@__DIR__));
 include(joinpath(cts, "test", "problems.jl"))
 function config_integrators(problem)
-    algorithm = ARS343(NewtonsMethod(; linsolve = linsolve_direct, max_iters = 2))
+    algorithm = CTS.IMEXARKAlgorithm(ARS343(), NewtonsMethod(; linsolve = linsolve_direct, max_iters = 2))
     dt = 0.01
     integrator = DiffEqBase.init(problem, algorithm; dt)
     not_generated_integrator = DiffEqBase.init(problem, algorithm; dt)

src/ClimaTimeSteppers.jl

@@ -66,14 +66,24 @@ include("algorithms.jl")
 
 abstract type DistributedODEAlgorithm <: DiffEqBase.AbstractODEAlgorithm end
 
+abstract type AbstractIMEXARKAlgorithm <: DistributedODEAlgorithm end
+abstract type AbstractIMEXARKTableau end
+
+"""
+    tableau(::DistributedODEAlgorithm)
+
+Returns the tableau for a particular algorithm.
+"""
+function tableau end
+
 """
     theoretical_convergence_order
 
 Returns the theoretical convergence order of an ODE algorithm
 """
 function theoretical_convergence_order end
-theoretical_convergence_order(alg::DistributedODEAlgorithm) =
-    error("No convergence order found for algo $alg, please open an issue or PR.")
+theoretical_convergence_order(tab) =
+    error("No convergence order found for tableau $tab, please open an issue or PR.")
 
 SciMLBase.allowscomplex(alg::DistributedODEAlgorithm) = true
 include("integrators.jl")

src/convergence_orders.jl

@@ -2,51 +2,53 @@
 ##### 1st order
 #####
 
-const first_order_methods = [
-    # ARS111,
-    # ARS121,
+# TODO: is it better to use `first_order_tableau = Union{ARS111,ARS121}`? to
+#       reduce the number of methods?
+const first_order_tableau = [
+    ARS111,
+    ARS121,
 ]
 
 #####
 ##### 2nd order
 #####
 
-const second_order_methods = [
-    # ARS122,
-    # ARS232,
-    # ARS222,
-    # IMKG232a,
-    # IMKG232b,
-    # IMKG242a,
-    # IMKG242b,
-    # IMKG252a,
-    # IMKG252b,
-    # IMKG253a,
-    # IMKG253b,
-    # IMKG254a,
-    # IMKG254b,
-    # IMKG254c,
-    # HOMMEM1,
+const second_order_tableau = [
+    ARS122,
+    ARS232,
+    ARS222,
+    IMKG232a,
+    IMKG232b,
+    IMKG242a,
+    IMKG242b,
+    IMKG252a,
+    IMKG252b,
+    IMKG253a,
+    IMKG253b,
+    IMKG254a,
+    IMKG254b,
+    IMKG254c,
+    HOMMEM1,
 ]
 
 #####
 ##### 3rd order
 #####
-const third_order_methods = [
-    # ARS233,
-    # ARS343,
-    # ARS443,
-    # IMKG342a,
-    # IMKG343a,
-    # DBM453,
+const third_order_tableau = [
+    ARS233,
+    ARS343,
+    ARS443,
+    IMKG342a,
+    IMKG343a,
+    DBM453,
 ]
 
-for m in first_order_methods
+for m in first_order_tableau
     @eval theoretical_convergence_order(::$m) = 1
 end
-for m in second_order_methods
+for m in second_order_tableau
     @eval theoretical_convergence_order(::$m) = 2
 end
-for m in third_order_methods
+for m in third_order_tableau
     @eval theoretical_convergence_order(::$m) = 3
 end

src/integrators.jl

@@ -50,7 +50,7 @@ function tstops_and_saveat_heaps(t0, tf, tstops, saveat)
     tstops = DataStructures.BinaryHeap{FT, ordering}(tstops)
 
     if isnothing(saveat)
-        saveat = [t0, tf]        
+        saveat = [t0, tf]
     elseif saveat isa Number
         saveat > zero(saveat) || error("saveat value must be positive")
         saveat = tf > t0 ? saveat : -saveat
@@ -101,7 +101,7 @@ function DiffEqBase.__init(
     callback = DiffEqBase.CallbackSet(callback, saving_callback)
     isempty(callback.continuous_callbacks) ||
         error("Continuous callbacks are not supported")
-    
+
     integrator = DistributedODEIntegrator(
         alg,
         u0,

src/solvers/imex_ark.jl

@@ -112,7 +112,7 @@ Is this too messy to do in the general case?
 
 Don't forget about the possible memory optimizations!
 =#
-export IMEXARKAlgorithm, make_IMEXARKAlgorithm
+export IMEXARKAlgorithm, make_IMEXARKTableau
 
 using Base: broadcasted, materialize!
 using StaticArrays: SMatrix, SVector
@@ -124,23 +124,32 @@ A generic implementation of an IMEX ARK algorithm that can handle arbitrary
 Butcher tableaus and problems specified using either `ForwardEulerODEFunction`s
 or regular `ODEFunction`s.
 """
-struct IMEXARKAlgorithm{as, cs, N} <: DistributedODEAlgorithm
+struct IMEXARKAlgorithm{as, cs, N} <: AbstractIMEXARKAlgorithm
     newtons_method::N
 end
 
 IMEXARKAlgorithm{as, cs}(newtons_method::N) where {as, cs, N} =
     IMEXARKAlgorithm{as, cs, N}(newtons_method)
 
 """
-    make_IMEXARKAlgorithm(; a_exp, b_exp, c_exp, a_imp, b_imp, c_imp)
+    IMEXARKAlgorithm(::AbstractIMEXARKAlgorithm, newtons_method)
+
+Returns the imex ARK algorithm for a particular algorithm.
+"""
+function IMEXARKAlgorithm(tab::AbstractIMEXARKTableau, newtons_method)
+    tableau(tab)(newtons_method)
+end
+
+"""
+    make_IMEXARKTableau(; a_exp, b_exp, c_exp, a_imp, b_imp, c_imp)
 
 Generates an `IMEXARKAlgorithm` type from an IMEX ARK Butcher tableau. Only
 `a_exp` and `a_imp` are required arguments; the default values for `b_exp` and
 `b_imp` assume that the algorithm is FSAL (first same as last), and the default
 values for `c_exp` and `c_imp` assume that the algorithm is internally
 consistent.
 """
-function make_IMEXARKAlgorithm(;
+function make_IMEXARKTableau(;
     a_exp::SMatrix{s, s},
     b_exp::SVector{s} = vec(a_exp[end, :]),
     c_exp::SVector{s} = vec(sum(a_exp; dims = 2)),