separate advection and diffusion jacobian blocks

config/default_configs/default_edmf_config.yml

@@ -8,6 +8,9 @@ advection_test:
 gs_tendency:
   help: "Turns on all grid-scale tendencies [`true` (default), `false`]"
   value: true
+implicit_sgs_advection:
+  help: "Whether to treat the subgrid-scale vertical advection tendency implicitly [`false` (default), `true`]"
+  value: false
 edmf_coriolis:
   help: "EDMF coriolis [`nothing` (default), `Bomex`,`LifeCycleTan2018`,`Rico`,`ARM_SGP`,`DYCOMS_RF01`,`DYCOMS_RF02`,`GABLS`]"
   value: ~

config/model_configs/prognostic_edmfx_bomex_box.yml

@@ -6,6 +6,7 @@ ls_adv: "Bomex"
 surface_setup: "Bomex"
 turbconv: "prognostic_edmfx"
 implicit_diffusion: true
+implicit_sgs_advection: true
 approximate_linear_solve_iters: 2
 max_newton_iters_ode: 3
 edmfx_upwinding: first_order

config/model_configs/prognostic_edmfx_bomex_stretched_box.yml

@@ -6,6 +6,7 @@ ls_adv: "Bomex"
 surface_setup: "Bomex"
 turbconv: "prognostic_edmfx"
 implicit_diffusion: true
+implicit_sgs_advection: true
 approximate_linear_solve_iters: 2
 max_newton_iters_ode: 3
 edmfx_upwinding: first_order

config/model_configs/prognostic_edmfx_dycoms_rf01_box.yml

@@ -6,6 +6,7 @@ rad: DYCOMS_RF01
 surface_setup: DYCOMS_RF01
 turbconv: prognostic_edmfx
 implicit_diffusion: true
+implicit_sgs_advection: true
 approximate_linear_solve_iters: 2
 max_newton_iters_ode: 3
 edmfx_upwinding: first_order

config/model_configs/prognostic_edmfx_gabls_box.yml

@@ -4,6 +4,7 @@ edmf_coriolis: GABLS
 surface_setup: GABLS
 turbconv: "prognostic_edmfx"
 implicit_diffusion: true
+implicit_sgs_advection: true
 approximate_linear_solve_iters: 2
 max_newton_iters_ode: 3
 edmfx_upwinding: "first_order"

config/model_configs/prognostic_edmfx_simpleplume_column.yml

@@ -3,6 +3,7 @@ initial_condition: "SimplePlume"
 surface_setup: "SimplePlume"
 turbconv: "prognostic_edmfx"
 implicit_diffusion: true
+implicit_sgs_advection: true
 approximate_linear_solve_iters: 2
 max_newton_iters_ode: 3
 gs_tendency: false

config/perf_configs/flame_perf_target_prognostic_edmfx_aquaplanet.yml

@@ -3,6 +3,8 @@ surface_setup: DefaultExchangeCoefficients
 rad: gray
 vert_diff: "false"
 turbconv: prognostic_edmfx 
+implicit_diffusion: true
+implicit_sgs_advection: true
 prognostic_tke: true
 edmfx_upwinding: first_order 
 edmfx_entr_model: "Generalized" 

src/prognostic_equations/implicit/implicit_solver.jl

@@ -9,7 +9,10 @@ use_derivative(::UseDerivative) = true
 use_derivative(::IgnoreDerivative) = false
 
 """
-    ImplicitEquationJacobian(Y, atmos; approximate_solve_iters, diffusion_flag, topography_flag, transform_flag)
+    ImplicitEquationJacobian(
+        Y, atmos; 
+        approximate_solve_iters, diffusion_flag, topography_flag, sgs_advection_flag, transform_flag
+    )
 
 A wrapper for the matrix ``∂E/∂Y``, where ``E(Y)`` is the "error" of the
 implicit step with the state ``Y``.
@@ -76,11 +79,15 @@ is reached.
 - `diffusion_flag::DerivativeFlag`: whether the derivative of the
   diffusion tendency with respect to the quantities being diffused should be
   computed or approximated as 0; must be either `UseDerivative()` or
-  `Ignoreerivative()` instead of a `Bool` to ensure type-stability
+  `IgnoreDerivative()` instead of a `Bool` to ensure type-stability
 - `topography_flag::DerivativeFlag`: whether the derivative of vertical
   contravariant velocity with respect to horizontal covariant velocity should be
   computed or approximated as 0; must be either `UseDerivative()` or
   `IgnoreDerivative()` instead of a `Bool` to ensure type-stability
+- `sgs_advection_flag::DerivativeFlag`: whether the derivative of the
+  subgrid-scale advection tendency with respect to the updraft quantities should be
+  computed or approximated as 0; must be either `UseDerivative()` or
+  `IgnoreDerivative()` instead of a `Bool` to ensure type-stability
 - `transform_flag::Bool`: whether the error should be transformed from ``E(Y)``
   to ``E(Y)/Δt``, which is required for non-Rosenbrock timestepping schemes from
   OrdinaryDiffEq.jl
@@ -90,6 +97,7 @@ struct ImplicitEquationJacobian{
     S <: MatrixFields.FieldMatrixSolver,
     F1 <: DerivativeFlag,
     F2 <: DerivativeFlag,
+    F3 <: DerivativeFlag,
     T <: Fields.FieldVector,
     R <: Base.RefValue,
 }
@@ -102,6 +110,7 @@ struct ImplicitEquationJacobian{
     # flags that determine how E'(Y) is approximated
     diffusion_flag::F1
     topography_flag::F2
+    sgs_advection_flag::F3
 
     # required by Krylov.jl to evaluate ldiv! with AbstractVector inputs
     temp_b::T
@@ -120,6 +129,8 @@ function ImplicitEquationJacobian(
                      IgnoreDerivative(),
     topography_flag = has_topography(axes(Y.c)) ? UseDerivative() :
                       IgnoreDerivative(),
+    sgs_advection_flag = atmos.sgs_adv_mode == Implicit() ? UseDerivative() :
+                         IgnoreDerivative(),
     transform_flag = false,
 )
     FT = Spaces.undertype(axes(Y.c))
@@ -214,14 +225,14 @@ function ImplicitEquationJacobian(
         )
     end
 
-    sgs_blocks = if atmos.turbconv_model isa PrognosticEDMFX
+    sgs_advection_blocks = if atmos.turbconv_model isa PrognosticEDMFX
         @assert n_prognostic_mass_flux_subdomains(atmos.turbconv_model) == 1
         sgs_scalar_names = (
             @name(c.sgsʲs.:(1).q_tot),
             @name(c.sgsʲs.:(1).mse),
             @name(c.sgsʲs.:(1).ρa)
         )
-        if use_derivative(diffusion_flag)
+        if use_derivative(sgs_advection_flag)
             (
                 MatrixFields.unrolled_map(
                     name -> (name, name) => similar(Y.c, TridiagonalRow),
@@ -261,7 +272,7 @@ function ImplicitEquationJacobian(
 
     matrix = MatrixFields.FieldMatrix(
         identity_blocks...,
-        sgs_blocks...,
+        sgs_advection_blocks...,
         advection_blocks...,
         diffusion_blocks...,
     )
@@ -287,49 +298,60 @@ function ImplicitEquationJacobian(
     )
 
     alg₂ = MatrixFields.BlockLowerTriangularSolve(@name(c.uₕ))
-    alg = if use_derivative(diffusion_flag)
-        alg₁_subalg₂ = if atmos.turbconv_model isa PrognosticEDMFX
-            (;
-                alg₂ = MatrixFields.BlockLowerTriangularSolve(
-                    @name(c.sgsʲs.:(1).q_tot);
-                    alg₂ = MatrixFields.BlockLowerTriangularSolve(
-                        @name(c.sgsʲs.:(1).mse);
-                        alg₂ = MatrixFields.BlockLowerTriangularSolve(
-                            @name(c.sgsʲs.:(1).ρa),
-                            @name(f.sgsʲs.:(1).u₃);
+    alg =
+        if use_derivative(diffusion_flag) || use_derivative(sgs_advection_flag)
+            alg₁_subalg₂ =
+                if atmos.turbconv_model isa PrognosticEDMFX &&
+                   use_derivative(sgs_advection_flag)
+                    diff_subalg =
+                        use_derivative(diffusion_flag) ?
+                        (;
                             alg₂ = MatrixFields.BlockLowerTriangularSolve(
                                 names₁_group₂...,
+                            )
+                        ) : (;)
+                    (;
+                        alg₂ = MatrixFields.BlockLowerTriangularSolve(
+                            @name(c.sgsʲs.:(1).q_tot);
+                            alg₂ = MatrixFields.BlockLowerTriangularSolve(
+                                @name(c.sgsʲs.:(1).mse);
+                                alg₂ = MatrixFields.BlockLowerTriangularSolve(
+                                    @name(c.sgsʲs.:(1).ρa),
+                                    @name(f.sgsʲs.:(1).u₃);
+                                    diff_subalg...,
+                                ),
                             ),
-                        ),
-                    ),
-                )
+                        )
+                    )
+                else
+                    is_in_Y(@name(c.ρq_tot)) ?
+                    (;
+                        alg₂ = MatrixFields.BlockLowerTriangularSolve(
+                            names₁_group₂...,
+                        )
+                    ) : (;)
+                end
+            alg₁ = MatrixFields.BlockLowerTriangularSolve(
+                names₁_group₁...;
+                alg₁_subalg₂...,
+            )
+            MatrixFields.ApproximateBlockArrowheadIterativeSolve(
+                names₁...;
+                alg₁,
+                alg₂,
+                P_alg₁ = MatrixFields.MainDiagonalPreconditioner(),
+                n_iters = approximate_solve_iters,
             )
         else
-            is_in_Y(@name(c.ρq_tot)) ?
-            (;
-                alg₂ = MatrixFields.BlockLowerTriangularSolve(names₁_group₂...)
-            ) : (;)
+            MatrixFields.BlockArrowheadSolve(names₁...; alg₂)
         end
-        alg₁ = MatrixFields.BlockLowerTriangularSolve(
-            names₁_group₁...;
-            alg₁_subalg₂...,
-        )
-        MatrixFields.ApproximateBlockArrowheadIterativeSolve(
-            names₁...;
-            alg₁,
-            alg₂,
-            P_alg₁ = MatrixFields.MainDiagonalPreconditioner(),
-            n_iters = approximate_solve_iters,
-        )
-    else
-        MatrixFields.BlockArrowheadSolve(names₁...; alg₂)
-    end
 
     return ImplicitEquationJacobian(
         matrix,
         MatrixFields.FieldMatrixSolver(alg, matrix, Y),
         diffusion_flag,
         topography_flag,
+        sgs_advection_flag,
         similar(Y),
         similar(Y),
         transform_flag,
@@ -397,9 +419,13 @@ NVTX.@annotate function Wfact!(A, Y, p, dtγ, t)
         (
             use_derivative(A.diffusion_flag) &&
             p.atmos.turbconv_model isa PrognosticEDMFX ?
+            (; p.precomputed.ᶜρa⁰) : (;)
+        )...,
+        (
+            use_derivative(A.sgs_advection_flag) &&
+            p.atmos.turbconv_model isa PrognosticEDMFX ?
             (;
                 p.core.ᶜgradᵥ_ᶠΦ,
-                p.precomputed.ᶜρa⁰,
                 p.precomputed.ᶜρʲs,
                 p.precomputed.ᶠu³ʲs,
                 p.precomputed.ᶜtsʲs,
@@ -441,7 +467,7 @@ NVTX.@annotate function Wfact!(A, Y, p, dtγ, t)
 end
 
 function update_implicit_equation_jacobian!(A, Y, p, dtγ, colidx)
-    (; matrix, diffusion_flag, topography_flag) = A
+    (; matrix, diffusion_flag, sgs_advection_flag, topography_flag) = A
     (; ᶜspecific, ᶜK, ᶜts, ᶜp, ᶜΦ, ᶠgradᵥ_ᶜΦ, ᶜρ_ref, ᶜp_ref) = p
     (;
         ᶜtemp_C3,
@@ -695,7 +721,7 @@ function update_implicit_equation_jacobian!(A, Y, p, dtγ, colidx)
     end
 
     if p.atmos.turbconv_model isa PrognosticEDMFX
-        if use_derivative(diffusion_flag)
+        if use_derivative(sgs_advection_flag)
             (; ᶜgradᵥ_ᶠΦ, ᶜρʲs, ᶠu³ʲs, ᶜtsʲs) = p
             is_third_order = edmfx_upwinding == Val(:third_order)
             ᶠupwind = is_third_order ? ᶠupwind3 : ᶠupwind1

src/prognostic_equations/implicit/implicit_tendency.jl

@@ -9,7 +9,7 @@ NVTX.@annotate function implicit_tendency!(Yₜ, Y, p, t)
     Yₜ .= zero(eltype(Yₜ))
     Fields.bycolumn(axes(Y.c)) do colidx
         implicit_vertical_advection_tendency!(Yₜ, Y, p, t, colidx)
-        if p.atmos.diff_mode == Implicit()
+        if p.atmos.sgs_adv_mode == Implicit()
             edmfx_sgs_vertical_advection_tendency!(
                 Yₜ,
                 Y,
@@ -18,6 +18,9 @@ NVTX.@annotate function implicit_tendency!(Yₜ, Y, p, t)
                 colidx,
                 p.atmos.turbconv_model,
             )
+        end
+
+        if p.atmos.diff_mode == Implicit()
             vertical_diffusion_boundary_layer_tendency!(
                 Yₜ,
                 Y,

src/prognostic_equations/remaining_tendency.jl

@@ -34,7 +34,7 @@ NVTX.@annotate function additional_tendency!(Yₜ, Y, p, t)
         edmf_coriolis_tendency!(Yₜ, Y, p, t, colidx, p.atmos.edmf_coriolis)
         large_scale_advection_tendency!(Yₜ, Y, p, t, colidx, p.atmos.ls_adv)
 
-        if p.atmos.diff_mode == Explicit()
+        if p.atmos.sgs_adv_mode == Explicit()
             edmfx_sgs_vertical_advection_tendency!(
                 Yₜ,
                 Y,
@@ -43,6 +43,9 @@ NVTX.@annotate function additional_tendency!(Yₜ, Y, p, t)
                 colidx,
                 p.atmos.turbconv_model,
             )
+        end
+
+        if p.atmos.diff_mode == Explicit()
             vertical_diffusion_boundary_layer_tendency!(
                 Yₜ,
                 Y,

src/solver/type_getters.jl

@@ -47,6 +47,9 @@ function get_atmos(config::AtmosConfig, params)
     implicit_diffusion = parsed_args["implicit_diffusion"]
     @assert implicit_diffusion in (true, false)
 
+    implicit_sgs_advection = parsed_args["implicit_sgs_advection"]
+    @assert implicit_sgs_advection in (true, false)
+
     model_config = get_model_config(parsed_args)
     vert_diff =
         get_vertical_diffusion_model(diffuse_momentum, parsed_args, params, FT)
@@ -82,6 +85,7 @@ function get_atmos(config::AtmosConfig, params)
         hyperdiff = get_hyperdiffusion_model(parsed_args, FT),
         vert_diff,
         diff_mode = implicit_diffusion ? Implicit() : Explicit(),
+        sgs_adv_mode = implicit_sgs_advection ? Implicit() : Explicit(),
         viscous_sponge = get_viscous_sponge_model(parsed_args, params, FT),
         rayleigh_sponge = get_rayleigh_sponge_model(parsed_args, params, FT),
         sfc_temperature = get_sfc_temperature_form(parsed_args),

src/solver/types.jl

@@ -340,6 +340,7 @@ Base.@kwdef struct AtmosModel{
     HD,
     VD,
     DM,
+    SAM,
     VS,
     RS,
     ST,
@@ -371,6 +372,7 @@ Base.@kwdef struct AtmosModel{
     hyperdiff::HD = nothing
     vert_diff::VD = nothing
     diff_mode::DM = nothing
+    sgs_adv_mode::SAM = nothing
     viscous_sponge::VS = nothing
     rayleigh_sponge::RS = nothing
     sfc_temperature::ST = nothing