Fix hard-coded ARS343, simplify code and add comments

Author: dennisYatunin

src/solvers/hard_coded_ars343.jl

@@ -10,11 +10,12 @@ function step_u!(integrator, cache::IMEXARKCache, ::ARS343)
     (; U, T_lim, T_exp, T_imp, temp, γ, newtons_method_cache) = cache
     T_lim! = !isnothing(f.T_lim!) ? f.T_lim! : (args...) -> nothing
     T_exp! = !isnothing(f.T_exp!) ? f.T_exp! : (args...) -> nothing
+    dtγ = dt * γ
 
     if !isnothing(newtons_method_cache)
         jacobian = newtons_method_cache.j
         if (!isnothing(jacobian)) && needs_update!(newtons_method.update_j, NewTimeStep(t))
-            T_imp!.Wfact(jacobian, u, p, dt * γ, t)
+            T_imp!.Wfact(jacobian, u, p, dtγ, t)
         end
     end
 
@@ -36,29 +37,24 @@ function step_u!(integrator, cache::IMEXARKCache, ::ARS343)
     let i = i
         t_imp = t + dt * c_imp[i]
         cache_imp!(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
+        implicit_equation_residual! = (residual, U′) -> begin
+            T_imp!(residual, U′, p, t_imp)
+            @. residual = temp + dtγ * residual - U′
         end
-        implicit_equation_jacobian! = (jacobian, Ui) -> begin
-            T_imp!.Wfact(jacobian, Ui, p, dt * a_imp[i, i], t_imp)
-        end
-        call_cache_imp! = Ui -> cache_imp!(Ui, p, t_imp)
         solve_newton!(
             newtons_method,
             newtons_method_cache,
             U,
             implicit_equation_residual!,
-            implicit_equation_jacobian!,
-            call_cache_imp!,
-            nothing,
+            (jacobian, U′) -> T_imp!.Wfact(jacobian, U′, p, dtγ, t_imp),
+            U′ -> cache_imp!(U′, p, t_imp),
         )
-        @. T_imp[i] = (U - temp) / (dt * a_imp[i, i])
         dss!(U, p, t_imp)
         cache!(U, p, t_imp)
     end
     T_lim!(T_lim[i], U, p, t_exp)
     T_exp!(T_exp[i], U, p, t_exp)
+    @. T_imp[i] = (U - temp) / dtγ
 
     i = 3
     t_exp = t + dt * c_exp[i]
@@ -70,29 +66,24 @@ function step_u!(integrator, cache::IMEXARKCache, ::ARS343)
     let i = i
         t_imp = t + dt * c_imp[i]
         cache_imp!(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)
+        implicit_equation_residual! = (residual, U′) -> begin
+            T_imp!(residual, U′, p, t_imp)
+            @. residual = temp + dtγ * residual - U′
         end
-        call_cache_imp! = Ui -> cache_imp!(Ui, p, t_imp)
         solve_newton!(
             newtons_method,
             newtons_method_cache,
             U,
             implicit_equation_residual!,
-            implicit_equation_jacobian!,
-            call_cache_imp!,
-            nothing,
+            (jacobian, U′) -> T_imp!.Wfact(jacobian, U′, p, dtγ, t_imp),
+            U′ -> cache_imp!(U′, p, t_imp),
         )
-        @. T_imp[i] = (U - temp) / (dt * a_imp[i, i])
         dss!(U, p, t_imp)
         cache!(U, p, t_imp)
     end
     T_lim!(T_lim[i], U, p, t_exp)
     T_exp!(T_exp[i], U, p, t_exp)
+    @. T_imp[i] = (U - temp) / dtγ
 
     i = 4
     t_exp = t + dt
@@ -109,29 +100,24 @@ function step_u!(integrator, cache::IMEXARKCache, ::ARS343)
     let i = i
         t_imp = t + dt * c_imp[i]
         cache_imp!(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)
+        implicit_equation_residual! = (residual, U′) -> begin
+            T_imp!(residual, U′, p, t_imp)
+            @. residual = temp + dtγ * residual - U′
         end
-        call_cache_imp! = Ui -> cache_imp!(Ui, p, t_imp)
         solve_newton!(
             newtons_method,
             newtons_method_cache,
             U,
             implicit_equation_residual!,
-            implicit_equation_jacobian!,
-            call_cache_imp!,
-            nothing,
+            (jacobian, U′) -> T_imp!.Wfact(jacobian, U′, p, dtγ, t_imp),
+            U′ -> cache_imp!(U′, p, t_imp),
         )
-        @. T_imp[i] = (U - temp) / (dt * a_imp[i, i])
         dss!(U, p, t_imp)
         cache!(U, p, t_imp)
     end
     T_lim!(T_lim[i], U, p, t_exp)
     T_exp!(T_exp[i], U, p, t_exp)
+    @. T_imp[i] = (U - temp) / dtγ
 
     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]

src/solvers/imex_ark.jl

@@ -120,29 +120,28 @@ 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))
 
+    # Run the implicit solver, apply DSS, and update the cache. When γ == 0,
+    # the implicit solver does not need to be run. On stage i == 1, we do not
+    # need to apply DSS and update the cache because we did that at the end of
+    # the previous timestep.
     i ≠ 1 && dss!(U, p, t_exp)
-
     if isnothing(T_imp!) || iszero(a_imp[i, i])
         i ≠ 1 && cache!(U, p, t_exp)
-    else # Implicit solve
+    else
         @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γ * residual - Ui
-        end
-        implicit_equation_jacobian! = (jacobian, Ui) -> begin
-            T_imp!.Wfact(jacobian, Ui, p, dtγ, t_imp)
+        implicit_equation_residual! = (residual, U′) -> begin
+            T_imp!(residual, U′, p, t_imp)
+            @. residual = temp + dtγ * residual - U′
         end
-        implicit_equation_cache! = Ui -> cache_imp!(Ui, p, t_imp)
         solve_newton!(
             newtons_method,
             newtons_method_cache,
             U,
             implicit_equation_residual!,
-            implicit_equation_jacobian!,
-            implicit_equation_cache!,
+            (jacobian, U′) -> T_imp!.Wfact(jacobian, U′, p, dtγ, t_imp),
+            U′ -> cache_imp!(U′, p, t_imp),
         )
         dss!(U, p, t_imp)
         cache!(U, p, t_imp)
@@ -155,11 +154,12 @@ end
     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.
+            # When γ == 0, T_imp[i] is treated explicitly.
             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 after applying DSS.
+            # When γ != 0, T_imp[i] is treated implicitly, so it must satisfy
+            # the implicit equation. To ensure that T_imp[i] only includes the
+            # effect of applying DSS to T_imp!, U and temp must be DSSed.
             isnothing(T_imp!) || @. T_imp[i] = (U - temp) / dtγ
         end
     end

src/solvers/imex_ssprk.jl

@@ -103,29 +103,28 @@ function step_u!(integrator, cache::IMEXSSPRKCache)
             end
         end
 
+        # Run the implicit solver, apply DSS, and update the cache. When γ == 0,
+        # the implicit solver does not need to be run. On stage i == 1, we do
+        # not need to apply DSS and update the cache because we did that at the
+        # end of the previous timestep.
         i ≠ 1 && dss!(U, p, t_exp)
-
         if isnothing(T_imp!) || iszero(a_imp[i, i])
             i ≠ 1 && cache!(U, p, t_exp)
-        else # Implicit solve
+        else
             @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γ * residual - Ui
-            end
-            implicit_equation_jacobian! = (jacobian, Ui) -> begin
-                T_imp!.Wfact(jacobian, Ui, p, dtγ, t_imp)
+            implicit_equation_residual! = (residual, U′) -> begin
+                T_imp!(residual, U′, p, t_imp)
+                @. residual = temp + dtγ * residual - U′
             end
-            implicit_equation_cache! = Ui -> cache_imp!(Ui, p, t_imp)
             solve_newton!(
                 newtons_method,
                 newtons_method_cache,
                 U,
                 implicit_equation_residual!,
-                implicit_equation_jacobian!,
-                implicit_equation_cache!,
+                (jacobian, U′) -> T_imp!.Wfact(jacobian, U′, p, dtγ, t_imp),
+                U′ -> cache_imp!(U′, p, t_imp),
             )
             dss!(U, p, t_imp)
             cache!(U, p, t_imp)
@@ -138,11 +137,12 @@ function step_u!(integrator, cache::IMEXSSPRKCache)
         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.
+                # When γ == 0, T_imp[i] is treated explicitly.
                 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 after applying DSS.
+                # When γ != 0, T_imp[i] is treated implicitly, so it must satisfy
+                # the implicit equation. To ensure that T_imp[i] only includes the
+                # effect of applying DSS to T_imp!, U and temp must be DSSed.
                 isnothing(T_imp!) || @. T_imp[i] = (U - temp) / dtγ
             end
         end