Merge #85

85: Revamp imex ark r=charleskawczynski a=charleskawczynski

This PR revamps the imex ark time stepper, as outlined in the google doc. It's still a WIP.

Related issue: #83.

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

Author: 3 people

Project.toml

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

docs/src/algorithms.md

@@ -15,8 +15,8 @@ ForwardEulerODEFunction
 ## IMEX ARK methods
 
 ```@docs
-IMEXARKAlgorithm
-make_IMEXARKTableau
+OldIMEXARKAlgorithm
+Oldmake_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.
@@ -90,7 +90,7 @@ ARK548L2SA2KennedyCarpenter
 
 ### ARS
 
-TODO: `ARS111`, `ARS121`, and `ARS343` should probably be types, and not constants, so that we can properly document them.
+TODO: add `ARS111`, `ARS121`, and `ARS343` docs.
 
 ## Multirate
 

perf/flame.jl

@@ -30,7 +30,7 @@ elseif problem_str=="fe"
 else
     error("Bad option")
 end
-algorithm = CTS.IMEXARKAlgorithm(ARS343(), NewtonsMethod(; max_iters = 2))
+algorithm = CTS.OldIMEXARKAlgorithm(OldARS343(), 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 = CTS.IMEXARKAlgorithm(ARS343(), NewtonsMethod(; linsolve = linsolve_direct, max_iters = 2))
+    algorithm = CTS.OldIMEXARKAlgorithm(OldARS343(), 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

@@ -67,7 +67,9 @@ include("algorithms.jl")
 abstract type DistributedODEAlgorithm <: DiffEqBase.AbstractODEAlgorithm end
 
 abstract type AbstractIMEXARKAlgorithm <: DistributedODEAlgorithm end
-abstract type AbstractIMEXARKTableau end
+
+abstract type AbstractTableau end
+abstract type AbstractIMEXARKTableau <: AbstractTableau end
 
 """
     tableau(::DistributedODEAlgorithm)
@@ -93,6 +95,7 @@ include("solvers/convergence_condition.jl")
 include("solvers/convergence_checker.jl")
 include("solvers/newtons_method.jl")
 include("solvers/imex_ark.jl")
+include("solvers/imex_ark2.jl")
 
 # Include concrete implementations
 include("solvers/imex_ark_tableaus.jl")

src/convergence_orders.jl

@@ -36,6 +36,7 @@ const second_order_tableau = [
 #####
 const third_order_tableau = [
     ARS233,
+    OldARS343,
     ARS343,
     ARS443,
     IMKG342a,

src/solvers/imex_ark.jl

@@ -112,44 +112,44 @@ Is this too messy to do in the general case?
 
 Don't forget about the possible memory optimizations!
 =#
-export IMEXARKAlgorithm, make_IMEXARKTableau
+export OldIMEXARKAlgorithm, Oldmake_IMEXARKTableau
 
 using Base: broadcasted, materialize!
 using StaticArrays: SMatrix, SVector
 
 """
-    IMEXARKAlgorithm <: DistributedODEAlgorithm
+    OldIMEXARKAlgorithm <: DistributedODEAlgorithm
 
 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} <: AbstractIMEXARKAlgorithm
+struct OldIMEXARKAlgorithm{as, cs, N} <: AbstractIMEXARKAlgorithm
     newtons_method::N
 end
 
-IMEXARKAlgorithm{as, cs}(newtons_method::N) where {as, cs, N} =
-    IMEXARKAlgorithm{as, cs, N}(newtons_method)
+OldIMEXARKAlgorithm{as, cs}(newtons_method::N) where {as, cs, N} =
+    OldIMEXARKAlgorithm{as, cs, N}(newtons_method)
 
 """
-    IMEXARKAlgorithm(::AbstractIMEXARKAlgorithm, newtons_method)
+    OldIMEXARKAlgorithm(::AbstractIMEXARKAlgorithm, newtons_method)
 
 Returns the imex ARK algorithm for a particular algorithm.
 """
-function IMEXARKAlgorithm(tab::AbstractIMEXARKTableau, newtons_method)
+function OldIMEXARKAlgorithm(tab::AbstractIMEXARKTableau, newtons_method)
     tableau(tab)(newtons_method)
 end
 
 """
-    make_IMEXARKTableau(; a_exp, b_exp, c_exp, a_imp, b_imp, c_imp)
+    Oldmake_IMEXARKTableau(; a_exp, b_exp, c_exp, a_imp, b_imp, c_imp)
 
-Generates an `IMEXARKAlgorithm` type from an IMEX ARK Butcher tableau. Only
+Generates an `OldIMEXARKAlgorithm` 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_IMEXARKTableau(;
+function Oldmake_IMEXARKTableau(;
     a_exp::SMatrix{s, s},
     b_exp::SVector{s} = vec(a_exp[end, :]),
     c_exp::SVector{s} = vec(sum(a_exp; dims = 2)),
@@ -165,7 +165,7 @@ function make_IMEXARKTableau(;
         as = (vcat(a_exp, b_exp'), vcat(a_imp, b_imp'))
     end
     cs = (c_exp, c_imp)
-    return IMEXARKAlgorithm{as, cs}
+    return OldIMEXARKAlgorithm{as, cs}
 end
 
 # General helper functions
@@ -208,7 +208,7 @@ end
 #       interval graph for all required cached values.
 function cache(
     prob::DiffEqBase.AbstractODEProblem,
-    alg::IMEXARKAlgorithm{as, cs};
+    alg::OldIMEXARKAlgorithm{as, cs};
     kwargs...
 ) where {as, cs}
     f_cache(χ, a, f_type) = is_increment(f_type) ?
@@ -440,7 +440,7 @@ step_u!(integrator, cache::IMEXARKCache) =
 
 function not_generated_cache(
     prob::DiffEqBase.AbstractODEProblem,
-    alg::IMEXARKAlgorithm{as, cs};
+    alg::OldIMEXARKAlgorithm{as, cs};
     kwargs...
 ) where {as, cs}
     f_cache(χ, a, f_type) = is_increment(f_type) ?