Merge pull request #226 from CliMA/ck/dispatch

Re-simplify dispatch, remove hard coded ARS

Project.toml

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

perf/flame.jl

@@ -35,7 +35,7 @@ end
 p = @allocated do_work!(integrator, cache)
 using Test
 @testset "Allocations" begin
-    @test_broken p == 0
+    @test p == 0
 end
 
 import ProfileCanvas

src/solvers/imex_ark.jl

@@ -45,12 +45,10 @@ 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.sol.prob.f, integrator.alg.name)
-
-# include("hard_coded_ars343.jl")
 # generic fallback
-function step_u!(integrator, cache::IMEXARKCache, f, name)
+function step_u!(integrator, cache::IMEXARKCache)
     (; u, p, t, dt, alg) = integrator
+    (; f) = integrator.sol.prob
     (; post_explicit!, post_implicit!) = f
     (; T_lim!, T_exp!, T_imp!, lim!, dss!) = f
     (; tableau, newtons_method) = alg
@@ -107,11 +105,10 @@ function step_u!(integrator, cache::IMEXARKCache, f, name)
             @. temp = U
             post_explicit!(U, p, t_imp)
             # TODO: can/should we remove these closures?
-            implicit_equation_residual! =
-                (residual, Ui) -> begin
-                    T_imp!(residual, Ui, p, t_imp)
-                    @. residual = temp + dt * a_imp[i, i] * residual - Ui
-                end
+            implicit_equation_residual! = (residual, Ui) -> begin
+                T_imp!(residual, Ui, p, t_imp)
+                @. residual = temp + dt * a_imp[i, i] * residual - Ui
+            end
             implicit_equation_jacobian! = (jacobian, Ui) -> T_imp!.Wfact(jacobian, Ui, p, dt * a_imp[i, i], t_imp)
             call_post_implicit! = Ui -> begin
                 post_implicit!(Ui, p, t_imp)

src/solvers/imex_ssprk.jl

@@ -52,10 +52,9 @@ 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.sol.prob.f, integrator.alg.name)
-
-function step_u!(integrator, cache::IMEXSSPRKCache, f, name)
+function step_u!(integrator, cache::IMEXSSPRKCache)
     (; u, p, t, dt, alg) = integrator
+    (; f) = integrator.sol.prob
     (; post_explicit!, post_implicit!) = f
     (; T_lim!, T_exp!, T_imp!, lim!, dss!) = f
     (; tableau, newtons_method) = alg