Merge #208

208: Define AbstractClimaODEFunction r=charleskawczynski a=charleskawczynski This PR defines `AbstractClimaODEFunction`, so that we can, from ClimaAtmos, define our own `AtmosODEFunction` with custom hooks (post stage callback etc.) and rearrange them in a custom `step_u!`. Co-authored-by: Charles Kawczynski <[email protected]>
CliMA · Sep 27, 2023 · 585daa1 · 585daa1 · charleskawczynski · Sep 27, 2023
2 parents 270d06a + fb34440
commit 585daa1
Show file tree

Hide file tree

Showing 5 changed files with 15 additions and 14 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "ClimaTimeSteppers"
 uuid = "595c0a79-7f3d-439a-bc5a-b232dc3bde79"
 authors = ["Climate Modeling Alliance"]
-version = "0.7.8"
+version = "0.7.9"
 
 [deps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"

diff --git a/docs/src/plotting_utils.jl b/docs/src/plotting_utils.jl
@@ -11,7 +11,7 @@ using PrettyTables: pretty_table, ft_printf
 Return the predicted convergence order of the algorithm for the given ODE
 function (assuming that the algorithm converges).
 """
-function predicted_convergence_order(algorithm_name::AbstractAlgorithmName, ode_function::ClimaODEFunction)
+function predicted_convergence_order(algorithm_name::AbstractAlgorithmName, ode_function::AbstractClimaODEFunction)
     (imp_order, exp_order, combined_order) = imex_convergence_orders(algorithm_name)
     has_imp = !isnothing(ode_function.T_imp!)
     has_exp = !isnothing(ode_function.T_exp!) || !isnothing(ode_function.T_lim!)

diff --git a/src/functions.jl b/src/functions.jl
@@ -1,7 +1,10 @@
 import DiffEqBase
+export AbstractClimaODEFunction
 export ClimaODEFunction, ForwardEulerODEFunction
 
-Base.@kwdef struct ClimaODEFunction{TL, TE, TI, L, D, PE, PI} <: DiffEqBase.AbstractODEFunction{true}
+abstract type AbstractClimaODEFunction <: DiffEqBase.AbstractODEFunction{true} end
+
+Base.@kwdef struct ClimaODEFunction{TL, TE, TI, L, D, PE, PI} <: AbstractClimaODEFunction
     T_lim!::TL = nothing # nothing or (uₜ, u, p, t) -> ...
     T_exp!::TE = nothing # nothing or (uₜ, u, p, t) -> ...
     T_imp!::TI = nothing # nothing or (uₜ, u, p, t) -> ...
@@ -11,9 +14,9 @@ Base.@kwdef struct ClimaODEFunction{TL, TE, TI, L, D, PE, PI} <: DiffEqBase.Abst
     post_implicit!::PI = (u, p, t) -> nothing
 end
 
-# Don't wrap a ClimaODEFunction in an ODEFunction (makes ODEProblem work).
-DiffEqBase.ODEFunction{iip}(f::ClimaODEFunction) where {iip} = f
-DiffEqBase.ODEFunction(f::ClimaODEFunction) = f
+# Don't wrap a AbstractClimaODEFunction in an ODEFunction (makes ODEProblem work).
+DiffEqBase.ODEFunction{iip}(f::AbstractClimaODEFunction) where {iip} = f
+DiffEqBase.ODEFunction(f::AbstractClimaODEFunction) = f
 
 """
     ForwardEulerODEFunction(f; jac_prototype, Wfact, tgrad)

diff --git a/src/solvers/imex_ark.jl b/src/solvers/imex_ark.jl
@@ -45,13 +45,12 @@ function init_cache(prob::DiffEqBase.AbstractODEProblem, alg::IMEXAlgorithm{Unco
     return IMEXARKCache(U, T_lim, T_exp, T_imp, temp, γ, newtons_method_cache)
 end
 
-step_u!(integrator, cache::IMEXARKCache) = step_u!(integrator, cache, integrator.alg.name)
+step_u!(integrator, cache::IMEXARKCache) = step_u!(integrator, cache, integrator.sol.prob.f, integrator.alg.name)
 
 include("hard_coded_ars343.jl")
 # generic fallback
-function step_u!(integrator, cache::IMEXARKCache, name)
-    (; u, p, t, dt, sol, alg) = integrator
-    (; f) = sol.prob
+function step_u!(integrator, cache::IMEXARKCache, f, name)
+    (; u, p, t, dt, alg) = integrator
     (; T_lim!, T_exp!, T_imp!, lim!, dss!) = f
     (; tableau, newtons_method) = alg
     (; a_exp, b_exp, a_imp, b_imp, c_exp, c_imp) = tableau

diff --git a/src/solvers/imex_ssprk.jl b/src/solvers/imex_ssprk.jl
@@ -52,11 +52,10 @@ function init_cache(prob::DiffEqBase.AbstractODEProblem, alg::IMEXAlgorithm{SSP}
     return IMEXSSPRKCache(U, U_exp, U_lim, T_lim, T_exp, T_imp, temp, β, γ, newtons_method_cache)
 end
 
-step_u!(integrator, cache::IMEXSSPRKCache) = step_u!(integrator, cache, integrator.alg.name)
+step_u!(integrator, cache::IMEXSSPRKCache) = step_u!(integrator, cache, integrator.sol.prob.f, integrator.alg.name)
 
-function step_u!(integrator, cache::IMEXSSPRKCache, name)
-    (; u, p, t, dt, sol, alg) = integrator
-    (; f) = sol.prob
+function step_u!(integrator, cache::IMEXSSPRKCache, f, name)
+    (; u, p, t, dt, alg) = integrator
     (; T_lim!, T_exp!, T_imp!, lim!, dss!) = f
     (; tableau, newtons_method) = alg
     (; a_imp, b_imp, c_exp, c_imp) = tableau