Merge pull request #356 from CliMA/dy/move_T_imp

Swap DSS and T_imp approximation
CliMA · Jan 27, 2025 · 73d15e1 · 73d15e1 · dennisYatunin · Jan 27, 2025
2 parents 7b367ab + 6e84816
commit 73d15e1
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 32 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.8.0"
+version = "0.8.1"
 
 [deps]
 ClimaComms = "3a4d1b5c-c61d-41fd-a00a-5873ba7a1b0d"

diff --git a/src/solvers/imex_ark.jl b/src/solvers/imex_ark.jl
@@ -102,11 +102,12 @@ end
     (; T_exp_T_lim!, 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
-    (; U, T_lim, T_exp, T_imp, temp, γ, newtons_method_cache) = cache
+    (; U, T_lim, T_exp, T_imp, temp, newtons_method_cache) = cache
     s = length(b_exp)
 
     t_exp = t + dt * c_exp[i]
     t_imp = t + dt * c_imp[i]
+    dtγ = dt * a_imp[i, i]
 
     if has_T_lim(f) # Update based on limited tendencies from previous stages
         assign_fused_increment!(U, u, dt, a_exp, T_lim, Val(i))
@@ -123,20 +124,16 @@ end
 
     if isnothing(T_imp!) || iszero(a_imp[i, i])
         i ≠ 1 && cache!(U, p, t_exp)
-        if !all(iszero, a_imp[:, i]) || !iszero(b_imp[i])
-            # If its coefficient is 0, T_imp[i] is being treated explicitly.
-            isnothing(T_imp!) || T_imp!(T_imp[i], U, p, t_imp)
-        end
     else # Implicit solve
         @assert !isnothing(newtons_method)
         i ≠ 1 && cache_imp!(U, p, t_imp)
         @. temp = U
         implicit_equation_residual! = (residual, Ui) -> begin
             T_imp!(residual, Ui, p, t_imp)
-            @. residual = temp + dt * a_imp[i, i] * residual - Ui
+            @. residual = temp + dtγ * residual - Ui
         end
         implicit_equation_jacobian! = (jacobian, Ui) -> begin
-            T_imp!.Wfact(jacobian, Ui, p, dt * a_imp[i, i], t_imp)
+            T_imp!.Wfact(jacobian, Ui, p, dtγ, t_imp)
         end
         implicit_equation_cache! = Ui -> cache_imp!(Ui, p, t_imp)
         solve_newton!(
@@ -147,21 +144,23 @@ end
             implicit_equation_jacobian!,
             implicit_equation_cache!,
         )
-        if !all(iszero, a_imp[:, i]) || !iszero(b_imp[i])
-            # If T_imp[i] is being treated implicitly, ensure that it
-            # exactly satisfies the implicit equation before applying DSS.
-            @. T_imp[i] = (U - temp) / (dt * a_imp[i, i])
-        end
         dss!(U, p, t_imp)
         cache!(U, p, t_imp)
     end
 
     if !all(iszero, a_exp[:, i]) || !iszero(b_exp[i])
-        if !isnothing(T_exp_T_lim!)
-            T_exp_T_lim!(T_exp[i], T_lim[i], U, p, t_exp)
+        isnothing(T_exp_T_lim!) || T_exp_T_lim!(T_exp[i], T_lim[i], U, p, t_exp)
+        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
+    if !all(iszero, a_imp[:, i]) || !iszero(b_imp[i])
+        if iszero(a_imp[i, i])
+            # If its coefficient is 0, T_imp[i] is being treated explicitly.
+            isnothing(T_imp!) || T_imp!(T_imp[i], U, p, t_imp)
         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)
+            # If T_imp[i] is being treated implicitly, ensure that it
+            # exactly satisfies the implicit equation after applying DSS.
+            isnothing(T_imp!) || @. T_imp[i] = (U - temp) / dtγ
         end
     end
 

diff --git a/src/solvers/imex_ssprk.jl b/src/solvers/imex_ssprk.jl
@@ -79,6 +79,7 @@ function step_u!(integrator, cache::IMEXSSPRKCache)
     for i in 1:s
         t_exp = t + dt * c_exp[i]
         t_imp = t + dt * c_imp[i]
+        dtγ = dt * a_imp[i, i]
 
         if i == 1
             @. U_exp = u
@@ -106,20 +107,16 @@ function step_u!(integrator, cache::IMEXSSPRKCache)
 
         if isnothing(T_imp!) || iszero(a_imp[i, i])
             i ≠ 1 && cache!(U, p, t_exp)
-            if !all(iszero, a_imp[:, i]) || !iszero(b_imp[i])
-                # If its coefficient is 0, T_imp[i] is being treated explicitly.
-                isnothing(T_imp!) || T_imp!(T_imp[i], U, p, t_imp)
-            end
         else # Implicit solve
             @assert !isnothing(newtons_method)
             i ≠ 1 && cache_imp!(U, p, t_imp)
             @. temp = U
             implicit_equation_residual! = (residual, Ui) -> begin
                 T_imp!(residual, Ui, p, t_imp)
-                @. residual = temp + dt * a_imp[i, i] * residual - Ui
+                @. residual = temp + dtγ * residual - Ui
             end
             implicit_equation_jacobian! = (jacobian, Ui) -> begin
-                T_imp!.Wfact(jacobian, Ui, p, dt * a_imp[i, i], t_imp)
+                T_imp!.Wfact(jacobian, Ui, p, dtγ, t_imp)
             end
             implicit_equation_cache! = Ui -> cache_imp!(Ui, p, t_imp)
             solve_newton!(
@@ -130,21 +127,23 @@ function step_u!(integrator, cache::IMEXSSPRKCache)
                 implicit_equation_jacobian!,
                 implicit_equation_cache!,
             )
-            if !all(iszero, a_imp[:, i]) || !iszero(b_imp[i])
-                # If T_imp[i] is being treated implicitly, ensure that it
-                # exactly satisfies the implicit equation before applying DSS.
-                @. T_imp[i] = (U - temp) / (dt * a_imp[i, i])
-            end
             dss!(U, p, t_imp)
             cache!(U, p, t_imp)
         end
 
         if !iszero(β[i])
-            if !isnothing(T_exp_T_lim!)
-                T_exp_T_lim!(T_lim, T_exp, U, p, t_exp)
+            isnothing(T_exp_T_lim!) || T_exp_T_lim!(T_lim, T_exp, U, p, t_exp)
+            isnothing(T_lim!) || T_lim!(T_lim, U, p, t_exp)
+            isnothing(T_exp!) || T_exp!(T_exp, U, p, t_exp)
+        end
+        if !all(iszero, a_imp[:, i]) || !iszero(b_imp[i])
+            if iszero(a_imp[i, i])
+                # If its coefficient is 0, T_imp[i] is being treated explicitly.
+                isnothing(T_imp!) || T_imp!(T_imp[i], U, p, t_imp)
             else
-                isnothing(T_lim!) || T_lim!(T_lim, U, p, t_exp)
-                isnothing(T_exp!) || T_exp!(T_exp, U, p, t_exp)
+                # If T_imp[i] is being treated implicitly, ensure that it
+                # exactly satisfies the implicit equation after applying DSS.
+                isnothing(T_imp!) || @. T_imp[i] = (U - temp) / dtγ
             end
         end
     end