Merge pull request #348 from CliMA/dy/move_dss

dennisYatunin · web-flow · commit 965ee2a8ee16 · 2025-01-16T16:41:44.000-08:00
Move DSS to the end of each implicit stage
diff --git a/.buildkite/Manifest.toml b/.buildkite/Manifest.toml
@@ -283,7 +283,7 @@ version = "0.2.11"
 deps = ["ClimaComms", "Colors", "DataStructures", "DiffEqBase", "KernelAbstractions", "Krylov", "LinearAlgebra", "LinearOperators", "NVTX", "SciMLBase", "StaticArrays"]
 path = ".."
 uuid = "595c0a79-7f3d-439a-bc5a-b232dc3bde79"
-version = "0.7.38"
+version = "0.7.39"
 weakdeps = ["BenchmarkTools", "CUDA", "OrderedCollections", "PrettyTables", "StatsBase"]
 
     [deps.ClimaTimeSteppers.extensions]
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.38"
+version = "0.7.39"
 
 [deps]
 ClimaComms = "3a4d1b5c-c61d-41fd-a00a-5873ba7a1b0d"
diff --git a/docs/Manifest.toml b/docs/Manifest.toml
@@ -241,7 +241,7 @@ version = "0.2.11"
 deps = ["ClimaComms", "Colors", "DataStructures", "DiffEqBase", "KernelAbstractions", "Krylov", "LinearAlgebra", "LinearOperators", "NVTX", "SciMLBase", "StaticArrays"]
 path = ".."
 uuid = "595c0a79-7f3d-439a-bc5a-b232dc3bde79"
-version = "0.7.38"
+version = "0.7.39"
 
     [deps.ClimaTimeSteppers.extensions]
     ClimaTimeSteppersBenchmarkToolsExt = ["CUDA", "BenchmarkTools", "OrderedCollections", "StatsBase", "PrettyTables"]
diff --git a/src/solvers/hard_coded_ars343.jl b/src/solvers/hard_coded_ars343.jl
@@ -33,7 +33,6 @@ function step_u!(integrator, cache::IMEXARKCache, ::ARS343)
     @. U = u + dt * a_exp[i, 1] * T_lim[1]
     lim!(U, p, t_exp, u)
     @. U += dt * a_exp[i, 1] * T_exp[1]
-    dss!(U, p, t_exp)
     post_explicit!(U, p, t_exp)
 
     @. temp = U # used in closures
@@ -49,13 +48,18 @@ function step_u!(integrator, cache::IMEXARKCache, ::ARS343)
         call_post_implicit! = Ui -> begin
             post_implicit!(Ui, p, t_imp)
         end
+        call_post_implicit_last! = Ui -> begin
+            dss!(Ui, p, t_imp)
+            post_implicit!(Ui, p, t_imp)
+        end
         solve_newton!(
             newtons_method,
             newtons_method_cache,
             U,
             implicit_equation_residual!,
             implicit_equation_jacobian!,
             call_post_implicit!,
+            call_post_implicit_last!,
         )
     end
 
@@ -69,7 +73,6 @@ function step_u!(integrator, cache::IMEXARKCache, ::ARS343)
     @. U = u + dt * a_exp[i, 1] * T_lim[1] + dt * a_exp[i, 2] * T_lim[2]
     lim!(U, p, t_exp, u)
     @. U += dt * a_exp[i, 1] * T_exp[1] + dt * a_exp[i, 2] * T_exp[2] + dt * a_imp[i, 2] * T_imp[2]
-    dss!(U, p, t_exp)
     post_explicit!(U, p, t_exp)
 
     @. temp = U # used in closures
@@ -85,13 +88,18 @@ function step_u!(integrator, cache::IMEXARKCache, ::ARS343)
         call_post_implicit! = Ui -> begin
             post_implicit!(Ui, p, t_imp)
         end
+        call_post_implicit_last! = Ui -> begin
+            dss!(Ui, p, t_imp)
+            post_implicit!(Ui, p, t_imp)
+        end
         solve_newton!(
             newtons_method,
             newtons_method_cache,
             U,
             implicit_equation_residual!,
             implicit_equation_jacobian!,
             call_post_implicit!,
+            call_post_implicit_last!,
         )
     end
 
@@ -109,7 +117,6 @@ function step_u!(integrator, cache::IMEXARKCache, ::ARS343)
         dt * a_exp[i, 3] * T_exp[3] +
         dt * a_imp[i, 2] * T_imp[2] +
         dt * a_imp[i, 3] * T_imp[3]
-    dss!(U, p, t_exp)
     post_explicit!(U, p, t_exp)
 
     @. temp = U # used in closures
@@ -125,13 +132,18 @@ function step_u!(integrator, cache::IMEXARKCache, ::ARS343)
         call_post_implicit! = Ui -> begin
             post_implicit!(Ui, p, t_imp)
         end
+        call_post_implicit_last! = Ui -> begin
+            dss!(Ui, p, t_imp)
+            post_implicit!(Ui, p, t_imp)
+        end
         solve_newton!(
             newtons_method,
             newtons_method_cache,
             U,
             implicit_equation_residual!,
             implicit_equation_jacobian!,
             call_post_implicit!,
+            call_post_implicit_last!,
         )
     end
 
diff --git a/src/solvers/imex_ark.jl b/src/solvers/imex_ark.jl
@@ -119,13 +119,14 @@ end
     has_T_exp(f) && fused_increment!(U, dt, a_exp, T_exp, Val(i))
     isnothing(T_imp!) || fused_increment!(U, dt, a_imp, T_imp, Val(i))
 
-    i ≠ 1 && dss!(U, p, t_exp)
-
-    if !(!isnothing(T_imp!) && !iszero(a_imp[i, i]))
+    if isnothing(T_imp!) || iszero(a_imp[i, i])
+        i ≠ 1 && dss!(U, p, t_imp)
         i ≠ 1 && post_explicit!(U, p, t_imp)
     else # Implicit solve
         @assert !isnothing(newtons_method)
         @. temp = U
+        # We do not need to apply DSS yet because the implicit solve does not
+        # involve any horizontal derivatives.
         i ≠ 1 && post_explicit!(U, p, t_imp)
         # TODO: can/should we remove these closures?
         implicit_equation_residual! = (residual, Ui) -> begin
@@ -137,11 +138,7 @@ end
             post_implicit!(Ui, p, t_imp)
         end
         call_post_implicit_last! = Ui -> begin
-            if (!all(iszero, a_imp[:, i]) || !iszero(b_imp[i])) && !iszero(a_imp[i, i])
-                # If T_imp[i] is being treated implicitly, ensure that it
-                # exactly satisfies the implicit equation.
-                @. T_imp[i] = (Ui - temp) / (dt * a_imp[i, i])
-            end
+            dss!(Ui, p, t_imp)
             post_implicit!(Ui, p, t_imp)
         end
 
@@ -156,15 +153,15 @@ end
         )
     end
 
-    # We do not need to DSS U again because the implicit solve should
-    # give the same results for redundant columns (as long as the implicit
-    # tendency only acts in the vertical direction).
-
     if !all(iszero, a_imp[:, i]) || !iszero(b_imp[i])
-        if iszero(a_imp[i, i]) && !isnothing(T_imp!)
+        if iszero(a_imp[i, i])
             # If its coefficient is 0, T_imp[i] is effectively being
             # treated explicitly.
-            T_imp!(T_imp[i], U, p, t_imp)
+            isnothing(T_imp!) || T_imp!(T_imp[i], U, p, t_imp)
+        else
+            # If T_imp[i] is being treated implicitly, ensure that it
+            # exactly satisfies the implicit equation.
+            isnothing(T_imp!) || @. T_imp[i] = (U - temp) / (dt * a_imp[i, i])
         end
     end
 
diff --git a/src/solvers/imex_ssprk.jl b/src/solvers/imex_ssprk.jl
@@ -94,8 +94,6 @@ function step_u!(integrator, cache::IMEXSSPRKCache)
             @. U_exp = (1 - β[i - 1]) * u + β[i - 1] * U_exp
         end
 
-        i ≠ 1 && dss!(U_exp, p, t_exp)
-
         @. U = U_exp
         if !isnothing(T_imp!) # Update based on implicit tendencies from previous stages
             for j in 1:(i - 1)
@@ -104,11 +102,14 @@ function step_u!(integrator, cache::IMEXSSPRKCache)
             end
         end
 
-        if !(!isnothing(T_imp!) && !iszero(a_imp[i, i]))
+        if isnothing(T_imp!) || iszero(a_imp[i, i])
+            i ≠ 1 && dss!(U, p, t_imp)
             i ≠ 1 && post_explicit!(U, p, t_imp)
         else # Implicit solve
             @assert !isnothing(newtons_method)
             @. temp = U
+            # We do not need to apply DSS yet because the implicit solve does
+            # not involve any horizontal derivatives.
             post_explicit!(U, p, t_imp)
             # TODO: can/should we remove these closures?
             implicit_equation_residual! = (residual, Ui) -> begin
@@ -119,15 +120,10 @@ function step_u!(integrator, cache::IMEXSSPRKCache)
             call_post_implicit! = Ui -> begin
                 post_implicit!(Ui, p, t_imp)
             end
-            call_post_implicit_last! =
-                Ui -> begin
-                    if (!all(iszero, a_imp[:, i]) || !iszero(b_imp[i])) && !iszero(a_imp[i, i])
-                        # If T_imp[i] is being treated implicitly, ensure that it
-                        # exactly satisfies the implicit equation.
-                        @. T_imp[i] = (Ui - temp) / (dt * a_imp[i, i])
-                    end
-                    post_implicit!(Ui, p, t_imp)
-                end
+            call_post_implicit_last! = Ui -> begin
+                dss!(Ui, p, t_imp)
+                post_implicit!(Ui, p, t_imp)
+            end
 
             solve_newton!(
                 newtons_method,
@@ -140,15 +136,15 @@ function step_u!(integrator, cache::IMEXSSPRKCache)
             )
         end
 
-        # We do not need to DSS U again because the implicit solve should
-        # give the same results for redundant columns (as long as the implicit
-        # tendency only acts in the vertical direction).
-
         if !all(iszero, a_imp[:, i]) || !iszero(b_imp[i])
-            if iszero(a_imp[i, i]) && !isnothing(T_imp!)
+            if iszero(a_imp[i, i])
                 # If its coefficient is 0, T_imp[i] is effectively being
                 # treated explicitly.
-                T_imp!(T_imp[i], U, p, t_imp)
+                isnothing(T_imp!) || T_imp!(T_imp[i], U, p, t_imp)
+            else
+                # If T_imp[i] is being treated implicitly, ensure that it
+                # exactly satisfies the implicit equation.
+                isnothing(T_imp!) || @. T_imp[i] = (U - temp) / (dt * a_imp[i, i])
             end
         end
 
diff --git a/src/solvers/rosenbrock.jl b/src/solvers/rosenbrock.jl
@@ -152,10 +152,12 @@ function step_u!(int, cache::RosenbrockCache{Nstages}) where {Nstages}
 
         # NOTE: post_implicit! is a misnomer
         if !isnothing(post_implicit!)
-            # We update p on every stage but the first, and at the end of each
-            # timestep. Since the first stage is unchanged from the end of the
-            # previous timestep, this order of operations ensures that p is
-            # always consistent with the state, including between timesteps.
+            # We apply DSS and update p on every stage but the first, and at the
+            # end of each timestep. Since the first stage is unchanged from the
+            # end of the previous timestep, this order of operations ensures
+            # that the state is always continuous and that p is consistent with
+            # the state, including between timesteps.
+            (i != 1) && dss!(U, p, t + αi * dt)
             (i != 1) && post_implicit!(U, p, t + αi * dt)
         end
 
@@ -179,10 +181,6 @@ function step_u!(int, cache::RosenbrockCache{Nstages}) where {Nstages}
             fU .+= γi .* dt .* ∂Y∂t
         end
 
-        # We dss the tendency at every stage but the last. At the last stage, we
-        # dss the incremented state
-        (i != Nstages) && dss!(fU, p, t + αi * dt)
-
         for j in 1:(i - 1)
             fU .+= (C[i, j] / dt) .* k[j]
         end
