Merge pull request #3708 from CliMA/zs/implicit_entrainment

make entrainment and detrainment implicit

Author: szy21

config/default_configs/default_config.yml

@@ -318,8 +318,11 @@ zero_tendency:
 implicit_sgs_advection:
   help: "Whether to treat the subgrid-scale vertical advection tendency implicitly [`false` (default), `true`]"
   value: false
+implicit_sgs_entr_detr:
+  help: "Whether to treat the subgrid-scale entrainment and detrainment tendency implicitly [`false` (default), `true`]. Setting it to true only works if implicit_sgs_advection is set to true."
+  value: false
 implicit_sgs_mass_flux:
-  help: "Whether to treat the subgrid-scale mass flux tendency implicitly or explicitly in grid-mean equations. Currently updraft only with Jacobian terms 0. [`false` (default), `true`]"
+  help: "Whether to treat the subgrid-scale mass flux tendency implicitly or explicitly in grid-mean equations. Currently updraft only with Jacobian terms 0. [`false` (default), `true`]. Setting it to true only works if both implicit_sgs_advection and implicit_diffusion are set to true."
   value: false
 edmf_coriolis:
   help: "EDMF coriolis [`nothing` (default), `Bomex`,`LifeCycleTan2018`,`Rico`,`ARM_SGP`,`DYCOMS_RF01`,`DYCOMS_RF02`,`GABLS`]"

config/model_configs/prognostic_edmfx_bomex_implicit_column.yml

@@ -6,6 +6,7 @@ surface_setup: "Bomex"
 turbconv: "prognostic_edmfx"
 implicit_diffusion: true
 implicit_sgs_advection: true
+implicit_sgs_entr_detr: true
 implicit_sgs_mass_flux: true
 approximate_linear_solve_iters: 2
 max_newton_iters_ode: 3

config/perf_configs/bm_aquaplanet_progedmf.yml

@@ -9,6 +9,7 @@ rayleigh_sponge: true
 viscous_sponge: true
 implicit_diffusion: true
 implicit_sgs_advection: true
+implicit_sgs_entr_detr: true
 implicit_sgs_mass_flux: true
 approximate_linear_solve_iters: 2
 moist: equil

src/cache/temporary_quantities.jl

@@ -66,7 +66,6 @@ function temporary_quantities(Y, atmos)
         ),
         ᶜdiffusion_h_matrix = similar(Y.c, TridiagonalMatrixRow{FT}),
         ᶜdiffusion_h_matrix_scaled = similar(Y.c, TridiagonalMatrixRow{FT}),
-        ᶜtridiagonal_matrix_scalar = similar(Y.c, TridiagonalMatrixRow{FT}),
         ᶜdiffusion_u_matrix = similar(Y.c, TridiagonalMatrixRow{FT}),
         ᶠtridiagonal_matrix_c3 = similar(Y.f, TridiagonalMatrixRow{C3{FT}}),
     )

src/prognostic_equations/implicit/implicit_solver.jl

@@ -99,6 +99,7 @@ struct ImplicitEquationJacobian{
     F2 <: DerivativeFlag,
     F3 <: DerivativeFlag,
     F4 <: DerivativeFlag,
+    F5 <: DerivativeFlag,
     T <: Fields.FieldVector,
     R <: Base.RefValue,
 }
@@ -112,7 +113,8 @@ struct ImplicitEquationJacobian{
     diffusion_flag::F1
     topography_flag::F2
     sgs_advection_flag::F3
-    sgs_mass_flux_flag::F4
+    sgs_entr_detr_flag::F4
+    sgs_mass_flux_flag::F5
 
     # required by Krylov.jl to evaluate ldiv! with AbstractVector inputs
     temp_b::T
@@ -130,6 +132,7 @@ function Base.zero(jac::ImplicitEquationJacobian)
         jac.diffusion_flag,
         jac.topography_flag,
         jac.sgs_advection_flag,
+        jac.sgs_entr_detr_flag,
         jac.sgs_mass_flux_flag,
         jac.temp_b,
         jac.temp_x,
@@ -148,6 +151,8 @@ function ImplicitEquationJacobian(
                       IgnoreDerivative(),
     sgs_advection_flag = atmos.sgs_adv_mode == Implicit() ? UseDerivative() :
                          IgnoreDerivative(),
+    sgs_entr_detr_flag = atmos.sgs_entr_detr_mode == Implicit() ?
+                         UseDerivative() : IgnoreDerivative(),
     sgs_mass_flux_flag = atmos.sgs_mf_mode == Implicit() ? UseDerivative() :
                          IgnoreDerivative(),
     transform_flag = false,
@@ -403,6 +408,7 @@ function ImplicitEquationJacobian(
         diffusion_flag,
         topography_flag,
         sgs_advection_flag,
+        sgs_entr_detr_flag,
         sgs_mass_flux_flag,
         similar(Y),
         similar(Y),
@@ -491,6 +497,9 @@ NVTX.@annotate function Wfact!(A, Y, p, dtγ, t)
                 p.precomputed.bdmr_l,
                 p.precomputed.bdmr_r,
                 p.precomputed.bdmr,
+                p.precomputed.ᶜentrʲs,
+                p.precomputed.ᶜdetrʲs,
+                p.precomputed.ᶜturb_entrʲs,
             ) : (;)
         )...,
         p.core.ᶜΦ,
@@ -504,7 +513,6 @@ NVTX.@annotate function Wfact!(A, Y, p, dtγ, t)
         p.scratch.ᶜadvection_matrix,
         p.scratch.ᶜdiffusion_h_matrix,
         p.scratch.ᶜdiffusion_h_matrix_scaled,
-        p.scratch.ᶜtridiagonal_matrix_scalar,
         p.scratch.ᶜdiffusion_u_matrix,
         p.scratch.ᶠbidiagonal_matrix_ct3,
         p.scratch.ᶠbidiagonal_matrix_ct3_2,
@@ -528,6 +536,7 @@ function update_implicit_equation_jacobian!(A, Y, p, dtγ)
         matrix,
         diffusion_flag,
         sgs_advection_flag,
+        sgs_entr_detr_flag,
         topography_flag,
         sgs_mass_flux_flag,
     ) = A
@@ -545,7 +554,6 @@ function update_implicit_equation_jacobian!(A, Y, p, dtγ)
     (;
         ᶜdiffusion_h_matrix,
         ᶜdiffusion_h_matrix_scaled,
-        ᶜtridiagonal_matrix_scalar,
         ᶜdiffusion_u_matrix,
         params,
     ) = p
@@ -851,6 +859,7 @@ function update_implicit_equation_jacobian!(A, Y, p, dtγ)
             @. ᶜkappa_mʲ =
                 TD.gas_constant_air(thermo_params, ᶜtsʲs.:(1)) /
                 TD.cv_m(thermo_params, ᶜtsʲs.:(1))
+
             ∂ᶜq_totʲ_err_∂ᶜq_totʲ =
                 matrix[@name(c.sgsʲs.:(1).q_tot), @name(c.sgsʲs.:(1).q_tot)]
             @. ∂ᶜq_totʲ_err_∂ᶜq_totʲ =
@@ -995,6 +1004,22 @@ function update_implicit_equation_jacobian!(A, Y, p, dtγ)
                     dtγ * ᶠtridiagonal_matrix_c3 ⋅
                     DiagonalMatrixRow(adjoint(CT3(Y.f.sgsʲs.:(1).u₃))) - (I_u₃,)
             end
+
+            # entrainment and detrainment
+            (; ᶜentrʲs, ᶜdetrʲs, ᶜturb_entrʲs) = p
+            # This assumes entrainment and detrainment rates are constant in the Jacobian
+            @. ∂ᶜq_totʲ_err_∂ᶜq_totʲ -=
+                dtγ * DiagonalMatrixRow(ᶜentrʲs.:(1) + ᶜturb_entrʲs.:(1))
+            @. ∂ᶜmseʲ_err_∂ᶜmseʲ -=
+                dtγ * DiagonalMatrixRow(ᶜentrʲs.:(1) + ᶜturb_entrʲs.:(1))
+            @. ∂ᶜρaʲ_err_∂ᶜρaʲ +=
+                dtγ * DiagonalMatrixRow(ᶜentrʲs.:(1) - ᶜdetrʲs.:(1))
+            @. ∂ᶠu₃ʲ_err_∂ᶠu₃ʲ -=
+                dtγ * (DiagonalMatrixRow(
+                    (ᶠinterp(ᶜentrʲs.:(1) + ᶜturb_entrʲs.:(1))) *
+                    (one_C3xACT3,),
+                ))
+
             # add updraft mass flux contributions to grid-mean
             if use_derivative(sgs_mass_flux_flag)
 

src/prognostic_equations/implicit/implicit_tendency.jl

@@ -30,6 +30,11 @@ NVTX.@annotate function implicit_tendency!(Yₜ, Y, p, t)
         edmfx_sgs_diffusive_flux_tendency!(Yₜ, Y, p, t, p.atmos.turbconv_model)
     end
 
+
+    if p.atmos.sgs_entr_detr_mode == Implicit()
+        edmfx_entr_detr_tendency!(Yₜ, Y, p, t, p.atmos.turbconv_model)
+    end
+
     if p.atmos.sgs_mf_mode == Implicit()
         edmfx_sgs_mass_flux_tendency!(Yₜ, Y, p, t, p.atmos.turbconv_model)
     end

src/prognostic_equations/remaining_tendency.jl

@@ -112,7 +112,9 @@ NVTX.@annotate function additional_tendency!(Yₜ, Y, p, t)
     surface_flux_tendency!(Yₜ, Y, p, t)
 
     radiation_tendency!(Yₜ, Y, p, t, p.atmos.radiation_mode)
-    edmfx_entr_detr_tendency!(Yₜ, Y, p, t, p.atmos.turbconv_model)
+    if p.atmos.sgs_entr_detr_mode == Explicit()
+        edmfx_entr_detr_tendency!(Yₜ, Y, p, t, p.atmos.turbconv_model)
+    end
     if p.atmos.sgs_mf_mode == Explicit()
         edmfx_sgs_mass_flux_tendency!(Yₜ, Y, p, t, p.atmos.turbconv_model)
     end

src/solver/type_getters.jl

@@ -55,6 +55,9 @@ function get_atmos(config::AtmosConfig, params)
     implicit_sgs_advection = parsed_args["implicit_sgs_advection"]
     @assert implicit_sgs_advection in (true, false)
 
+    implicit_sgs_entr_detr = parsed_args["implicit_sgs_entr_detr"]
+    @assert implicit_sgs_entr_detr in (true, false)
+
     implicit_sgs_mass_flux = parsed_args["implicit_sgs_mass_flux"]
     @assert implicit_sgs_mass_flux in (true, false)
 
@@ -103,6 +106,7 @@ function get_atmos(config::AtmosConfig, params)
         vert_diff,
         diff_mode = implicit_diffusion ? Implicit() : Explicit(),
         sgs_adv_mode = implicit_sgs_advection ? Implicit() : Explicit(),
+        sgs_entr_detr_mode = implicit_sgs_entr_detr ? Implicit() : Explicit(),
         sgs_mf_mode = implicit_sgs_mass_flux ? Implicit() : Explicit(),
         viscous_sponge = get_viscous_sponge_model(parsed_args, params, FT),
         smagorinsky_lilly = get_smagorinsky_lilly_model(parsed_args),

src/solver/types.jl

@@ -519,6 +519,7 @@ Base.@kwdef struct AtmosModel{
     VD,
     DM,
     SAM,
+    SEDM,
     SMM,
     VS,
     SL,
@@ -557,6 +558,9 @@ Base.@kwdef struct AtmosModel{
     vert_diff::VD = nothing
     diff_mode::DM = nothing
     sgs_adv_mode::SAM = nothing
+    """sgs_entr_detr_mode == Implicit() only works if sgs_adv_mode == Implicit()"""
+    sgs_entr_detr_mode::SEDM = nothing
+    """sgs_mf_mode == Implicit() only works if sgs_adv_mode == Implicit() and diff_mode == Implicit()"""
     sgs_mf_mode::SMM = nothing
     viscous_sponge::VS = nothing
     smagorinsky_lilly::SL = nothing