Add DSS before implicit solve, fix T_imp approximation and callbacks

Project.toml

@@ -1,7 +1,7 @@
 name = "ClimaTimeSteppers"
 uuid = "595c0a79-7f3d-439a-bc5a-b232dc3bde79"
 authors = ["Climate Modeling Alliance"]
-version = "0.7.39"
+version = "0.7.40"
 
 [deps]
 ClimaComms = "3a4d1b5c-c61d-41fd-a00a-5873ba7a1b0d"

src/solvers/hard_coded_ars343.jl

@@ -22,9 +22,7 @@ function step_u!(integrator, cache::IMEXARKCache, ::ARS343)
 
     i::Int = 1
     t_exp = t
-    @. U = u
-    lim!(U, p, t_exp, u)
-    dss!(U, p, t_exp)
+    @. U = u # TODO: This is unnecessary; we can just pass u to T_exp and T_lim
     T_lim!(T_lim[i], U, p, t_exp)
     T_exp!(T_exp[i], U, p, t_exp)
 
@@ -33,38 +31,32 @@ 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]
-    post_explicit!(U, p, t_exp)
-
+    dss!(U, p, t_exp)
     @. temp = U # used in closures
     let i = i
         t_imp = t + dt * c_imp[i]
+        post_implicit!(U, p, t_imp)
         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) -> begin
             T_imp!.Wfact(jacobian, Ui, p, dt * a_imp[i, i], t_imp)
         end
-        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
+        call_post_implicit! = Ui -> post_implicit!(Ui, p, t_imp)
         solve_newton!(
             newtons_method,
             newtons_method_cache,
             U,
             implicit_equation_residual!,
             implicit_equation_jacobian!,
             call_post_implicit!,
-            call_post_implicit_last!,
+            nothing,
         )
+        @. T_imp[i] = (U - temp) / (dt * a_imp[i, i])
+        dss!(U, p, t_imp)
+        post_explicit!(U, p, t_imp)
     end
-
-    @. T_imp[i] = (U - temp) / (dt * a_imp[i, i])
-
     T_lim!(T_lim[i], U, p, t_exp)
     T_exp!(T_exp[i], U, p, t_exp)
 
@@ -73,40 +65,35 @@ 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]
-    post_explicit!(U, p, t_exp)
-
+    dss!(U, p, t_exp)
     @. temp = U # used in closures
     let i = i
         t_imp = t + dt * c_imp[i]
+        post_implicit!(U, p, t_imp)
         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) -> begin
             T_imp!.Wfact(jacobian, Ui, p, dt * a_imp[i, i], t_imp)
         end
-        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
+        call_post_implicit! = Ui -> post_implicit!(Ui, p, t_imp)
         solve_newton!(
             newtons_method,
             newtons_method_cache,
             U,
             implicit_equation_residual!,
             implicit_equation_jacobian!,
             call_post_implicit!,
-            call_post_implicit_last!,
+            nothing,
         )
+        @. T_imp[i] = (U - temp) / (dt * a_imp[i, i])
+        dss!(U, p, t_imp)
+        post_explicit!(U, p, t_imp)
     end
-
-    @. T_imp[i] = (U - temp) / (dt * a_imp[i, i])
-
     T_lim!(T_lim[i], U, p, t_exp)
     T_exp!(T_exp[i], U, p, t_exp)
+
     i = 4
     t_exp = t + dt
     @. U = u + dt * a_exp[i, 1] * T_lim[1] + dt * a_exp[i, 2] * T_lim[2] + dt * a_exp[i, 3] * T_lim[3]
@@ -117,44 +104,35 @@ 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]
-    post_explicit!(U, p, t_exp)
-
+    dss!(U, p, t_exp)
     @. temp = U # used in closures
     let i = i
         t_imp = t + dt * c_imp[i]
+        post_implicit!(U, p, t_imp)
         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) -> begin
             T_imp!.Wfact(jacobian, Ui, p, dt * a_imp[i, i], t_imp)
         end
-        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
+        call_post_implicit! = Ui -> post_implicit!(Ui, p, t_imp)
         solve_newton!(
             newtons_method,
             newtons_method_cache,
             U,
             implicit_equation_residual!,
             implicit_equation_jacobian!,
             call_post_implicit!,
-            call_post_implicit_last!,
+            nothing,
         )
+        @. T_imp[i] = (U - temp) / (dt * a_imp[i, i])
+        dss!(U, p, t_imp)
+        post_explicit!(U, p, t_imp)
     end
-
-    @. T_imp[i] = (U - temp) / (dt * a_imp[i, i])
-
     T_lim!(T_lim[i], U, p, t_exp)
     T_exp!(T_exp[i], U, p, t_exp)
 
-    # final
-    i = -1
-
     t_final = t + dt
     @. temp = u + dt * b_exp[2] * T_lim[2] + dt * b_exp[3] * T_lim[3] + dt * b_exp[4] * T_lim[4]
     lim!(temp, p, t_final, u)

src/solvers/imex_ark.jl

@@ -119,50 +119,43 @@ 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])
-        i ≠ 1 && dss!(U, p, t_imp)
-        i ≠ 1 && post_explicit!(U, p, t_imp)
+        i ≠ 1 && post_explicit!(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 && post_implicit!(U, p, t_imp)
         @. 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
             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)
+        implicit_equation_jacobian! = (jacobian, Ui) -> begin
+            T_imp!.Wfact(jacobian, Ui, p, dt * a_imp[i, i], t_imp)
         end
-        call_post_implicit_last! = Ui -> begin
-            dss!(Ui, p, t_imp)
-            post_implicit!(Ui, p, t_imp)
-        end
-
+        call_post_implicit! = Ui -> post_implicit!(Ui, p, t_imp)
         solve_newton!(
             newtons_method,
             newtons_method_cache,
             U,
             implicit_equation_residual!,
             implicit_equation_jacobian!,
             call_post_implicit!,
-            call_post_implicit_last!,
+            nothing,
         )
-    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 effectively being
-            # treated explicitly.
-            isnothing(T_imp!) || T_imp!(T_imp[i], U, p, t_imp)
-        else
+        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.
-            isnothing(T_imp!) || @. T_imp[i] = (U - temp) / (dt * a_imp[i, i])
+            # exactly satisfies the implicit equation before applying DSS.
+            @. T_imp[i] = (U - temp) / (dt * a_imp[i, i])
         end
+        dss!(U, p, t_imp)
+        post_explicit!(U, p, t_imp)
     end
 
     if !all(iszero, a_exp[:, i]) || !iszero(b_exp[i])

src/solvers/imex_ssprk.jl

@@ -102,50 +102,43 @@ function step_u!(integrator, cache::IMEXSSPRKCache)
             end
         end
 
+        i ≠ 1 && dss!(U, p, t_exp)
+
         if isnothing(T_imp!) || iszero(a_imp[i, i])
-            i ≠ 1 && dss!(U, p, t_imp)
-            i ≠ 1 && post_explicit!(U, p, t_imp)
+            i ≠ 1 && post_explicit!(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 && post_implicit!(U, p, t_imp)
             @. 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
                 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)
-            end
-            call_post_implicit_last! = Ui -> begin
-                dss!(Ui, p, t_imp)
-                post_implicit!(Ui, p, t_imp)
+            implicit_equation_jacobian! = (jacobian, Ui) -> begin
+                T_imp!.Wfact(jacobian, Ui, p, dt * a_imp[i, i], t_imp)
             end
-
+            call_post_implicit! = Ui -> post_implicit!(Ui, p, t_imp)
             solve_newton!(
                 newtons_method,
                 newtons_method_cache,
                 U,
                 implicit_equation_residual!,
                 implicit_equation_jacobian!,
                 call_post_implicit!,
-                call_post_implicit_last!,
+                nothing,
             )
-        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 effectively being
-                # treated explicitly.
-                isnothing(T_imp!) || T_imp!(T_imp[i], U, p, t_imp)
-            else
+            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.
-                isnothing(T_imp!) || @. T_imp[i] = (U - temp) / (dt * a_imp[i, i])
+                # exactly satisfies the implicit equation before applying DSS.
+                @. T_imp[i] = (U - temp) / (dt * a_imp[i, i])
             end
+            dss!(U, p, t_imp)
+            post_explicit!(U, p, t_imp)
         end
 
         if !iszero(β[i])

src/solvers/rosenbrock.jl

@@ -123,7 +123,7 @@ function step_u!(int, cache::RosenbrockCache{Nstages}) where {Nstages}
     T_exp_lim! = int.sol.prob.f.T_exp_T_lim!
     tgrad! = isnothing(T_imp!) ? nothing : T_imp!.tgrad
 
-    (; post_explicit!, post_implicit!, dss!) = int.sol.prob.f
+    (; post_explicit!, dss!) = int.sol.prob.f
 
     # TODO: This is only valid when Γ[i, i] is constant, otherwise we have to
     # move this in the for loop
@@ -150,15 +150,15 @@ function step_u!(int, cache::RosenbrockCache{Nstages}) where {Nstages}
             U .+= A[i, j] .* k[j]
         end
 
-        # NOTE: post_implicit! is a misnomer
-        if !isnothing(post_implicit!)
+        # NOTE: post_explicit! is a misnomer; should be post_stage!
+        if !isnothing(post_explicit!)
             # 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)
+            (i != 1) && post_explicit!(U, p, t + αi * dt)
         end
 
         if !isnothing(T_imp!)
@@ -203,7 +203,7 @@ function step_u!(int, cache::RosenbrockCache{Nstages}) where {Nstages}
     end
 
     dss!(u, p, t + dt)
-    post_implicit!(u, p, t + dt)
+    post_explicit!(u, p, t + dt)
     return nothing
 end