Factor out implicit and wfact pcq calls

CliMA · Sep 30, 2023 · 68889f5 · 68889f5
1 parent 8519aa6
commit 68889f5
Show file tree

Hide file tree

Showing 5 changed files with 15 additions and 4 deletions.
diff --git a/src/prognostic_equations/implicit/implicit_tendency.jl b/src/prognostic_equations/implicit/implicit_tendency.jl
@@ -7,7 +7,6 @@ import ClimaCore: Fields, Geometry
 NVTX.@annotate function implicit_tendency!(Yₜ, Y, p, t)
     fill_with_nans!(p)
     Yₜ .= zero(eltype(Yₜ))
-    set_precomputed_quantities!(Y, p, t)
     Fields.bycolumn(axes(Y.c)) do colidx
         implicit_vertical_advection_tendency!(Yₜ, Y, p, t, colidx)
         if p.turbconv_model isa TurbulenceConvection.EDMFModel

diff --git a/src/prognostic_equations/implicit/wfact.jl b/src/prognostic_equations/implicit/wfact.jl
@@ -60,7 +60,6 @@ end
 
 NVTX.@annotate function Wfact!(W, Y, p, dtγ, t)
     fill_with_nans!(p)
-    set_precomputed_quantities!(Y, p, t)
     Fields.bycolumn(axes(Y.c)) do colidx
         Wfact!(W, Y, p, dtγ, t, colidx)
     end

diff --git a/src/time_stepper/hc_ars343.jl b/src/time_stepper/hc_ars343.jl
@@ -16,6 +16,7 @@ function CTS.step_u!(
         jacobian = newtons_method_cache.j
         if (!isnothing(jacobian)) &&
            CTS.needs_update!(newtons_method.update_j, CTS.NewTimeStep(t))
+            post_implicit!(u, p, t)
             T_imp!.Wfact(jacobian, u, p, dt * γ, t)
         end
     end
@@ -45,11 +46,13 @@ function CTS.step_u!(
         t_imp = t + dt * c_imp[i]
         implicit_equation_residual! =
             (residual, Ui) -> begin
+                post_implicit!(Ui, p, t_imp)
                 T_imp!(residual, Ui, p, t_imp)
                 @. residual = temp + dt * a_imp[i, i] * residual - Ui
             end
         implicit_equation_jacobian! =
             (jacobian, Ui) -> begin
+                post_implicit!(Ui, p, t_imp)
                 T_imp!.Wfact(jacobian, Ui, p, dt * a_imp[i, i], t_imp)
             end
         call_post_implicit! = Ui -> begin

diff --git a/src/time_stepper/imex_ark.jl b/src/time_stepper/imex_ark.jl
@@ -20,6 +20,7 @@ function CTS.step_u!(
             if γ isa Nothing
                 sdirk_error(name)
             else
+                post_implicit!(u, p, t)
                 T_imp!.Wfact(jacobian, u, p, dt * γ, t)
             end
         end
@@ -56,13 +57,16 @@ function CTS.step_u!(
                 # TODO: can/should we remove these closures?
                 implicit_equation_residual! =
                     (residual, Ui) -> begin
+                        post_implicit!(Ui, p, t_imp)
                         T_imp!(residual, Ui, p, t_imp)
                         @. residual =
                             temp + dt * a_imp[i, i] * residual - Ui
                     end
                 implicit_equation_jacobian! =
-                    (jacobian, Ui) ->
+                    (jacobian, Ui) -> begin
+                        post_implicit!(Ui, p, t_imp)
                         T_imp!.Wfact(jacobian, Ui, p, dt * a_imp[i, i], t_imp)
+                    end
 
                 CTS.solve_newton!(
                     newtons_method,
@@ -81,6 +85,7 @@ function CTS.step_u!(
                 if iszero(a_imp[i, i])
                     # If its coefficient is 0, T_imp[i] is effectively being
                     # treated explicitly.
+                    post_implicit!(Ui, p, t_imp)
                     T_imp!(T_imp[i], U, p, t_imp)
                 else
                     # If T_imp[i] is being treated implicitly, ensure that it

diff --git a/src/time_stepper/imex_ssprk.jl b/src/time_stepper/imex_ssprk.jl
@@ -21,6 +21,7 @@ function CTS.step_u!(
             if γ isa Nothing
                 sdirk_error(name)
             else
+                post_implicit!(u, p, t)
                 T_imp!.Wfact(jacobian, u, p, dt * γ, t)
             end
         end
@@ -56,12 +57,15 @@ function CTS.step_u!(
             # TODO: can/should we remove these closures?
             implicit_equation_residual! =
                 (residual, Ui) -> begin
+                    post_implicit!(Ui, p, t_imp)
                     T_imp!(residual, Ui, p, t_imp)
                     @. residual = temp + dt * a_imp[i, i] * residual - Ui
                 end
             implicit_equation_jacobian! =
-                (jacobian, Ui) ->
+                (jacobian, Ui) -> begin
+                    post_implicit!(Ui, p, t_imp)
                     T_imp!.Wfact(jacobian, Ui, p, dt * a_imp[i, i], t_imp)
+                end
             CTS.solve_newton!(
                 newtons_method,
                 newtons_method_cache,
@@ -79,6 +83,7 @@ function CTS.step_u!(
             if iszero(a_imp[i, i])
                 # If its coefficient is 0, T_imp[i] is effectively being
                 # treated explicitly.
+                post_implicit!(U, p, t_imp)
                 T_imp!(T_imp[i], U, p, t_imp)
             else
                 # If T_imp[i] is being treated implicitly, ensure that it