Allow commscontext input, def compute_T_exp_T_lim

Author: charleskawczynski

src/functions.jl

@@ -4,14 +4,15 @@ export ClimaODEFunction, ForwardEulerODEFunction
 
 abstract type AbstractClimaODEFunction <: DiffEqBase.AbstractODEFunction{true} end
 
-Base.@kwdef struct ClimaODEFunction{TL, TE, TI, L, D, PE, PI} <: AbstractClimaODEFunction
+Base.@kwdef struct ClimaODEFunction{TL, TE, TI, L, D, PE, PI, CC} <: 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) -> ...
     lim!::L = (u, p, t, u_ref) -> nothing
     dss!::D = (u, p, t) -> nothing
     post_explicit!::PE = (u, p, t) -> nothing
     post_implicit!::PI = (u, p, t) -> nothing
+    comms_context::CC = nothing
 end
 
 # Don't wrap a AbstractClimaODEFunction in an ODEFunction (makes ODEProblem work).

src/solvers/imex_ark.jl

@@ -169,9 +169,28 @@ end
     end
 
     if !all(iszero, a_exp[:, i]) || !iszero(b_exp[i])
-        isnothing(T_lim!) || T_lim!(T_lim[i], U, p, t_exp)
-        isnothing(T_exp!) || T_exp!(T_exp[i], U, p, t_exp)
+        if !isnothing!(T_lim!) && !isnothing(T_exp!)
+            (; comms_context) = f
+            compute_T_exp_T_lim!(T_lim[i], T_exp[i], U, p, t, T_exp!, T_lim!, comms_context)
+        else
+            isnothing(T_lim!) || T_lim!(T_lim[i], U, p, t_exp)
+            isnothing(T_exp!) || T_exp!(T_exp[i], U, p, t_exp)
+        end
     end
 
     return nothing
 end
+
+@inline function compute_T_exp_T_lim!(
+    T_lim,
+    T_exp,
+    U,
+    p,
+    t,
+    T_exp!,
+    T_lim!,
+    ::Union{Nothing, ClimaComms.AbstractCommsContext},
+)
+    T_lim!(T_lim, U, p, t)
+    T_exp!(T_exp, U, p, t)
+end

src/solvers/imex_ssprk.jl

@@ -153,11 +153,12 @@ function step_u!(integrator, cache::IMEXSSPRKCache)
         end
 
         if !iszero(β[i])
-            if !isnothing(T_lim!)
-                T_lim!(T_lim, U, p, t_exp)
-            end
-            if !isnothing(T_exp!)
-                T_exp!(T_exp, U, p, t_exp)
+            if !isnothing!(T_lim!) && !isnothing(T_exp!)
+                (; comms_context) = f
+                compute_T_exp_T_lim!(T_lim[i], T_exp[i], U, p, t, T_exp!, T_lim!, comms_context)
+            else
+                isnothing(T_lim!) || T_lim!(T_lim[i], U, p, t_exp)
+                isnothing(T_exp!) || T_exp!(T_exp[i], U, p, t_exp)
             end
         end
     end