Merge #100

100: Make user-facing types r=charleskawczynski a=charleskawczynski

This PR:
 - Adds a new abstract type `AbstractImexARKAlgorithm <: DistributedODEAlgorithm`
 - Adds user-facing types for all of the imex ark algorithms, all of which subtype `AbstractImexARKAlgorithm`
 - Defines `theoretical_convergence_order` for these new types
 - Adjusts the tests accordingly

This should all help with two things
 - Initializing algorithms. One issue we've had in ClimaAtmos is that initializing an algorithm requires first inspecting if an algorithm is a function, in which case assumptions are made about how that function is called. This is really clunky, and usage requires understanding ClimaTimeStepper internals.
 - Documentation. Several of these existing "algorithms", are just constants pointing to a generic tableau specific to that algorithm. Documenting these algorithms is part of our requirements, so this must be redesigned.

This should help us more easily peel off parts from the ClimaAtmos PR, and the revamp PR.

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

Author: bors[bot]

Project.toml

@@ -1,7 +1,7 @@
 name = "ClimaTimeSteppers"
 uuid = "595c0a79-7f3d-439a-bc5a-b232dc3bde79"
 authors = ["Climate Modeling Alliance"]
-version = "0.2.6"
+version = "0.3.0"
 
 [deps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"

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)),