Merge pull request #3713 from CliMA/zs/implicit_nh_pressure

make nonhydrostatic pressure drag implicit

Author: szy21

.buildkite/pipeline.yml

@@ -134,6 +134,8 @@ steps:
           --job_id single_column_hydrostatic_balance_ft64
           --out_dir single_column_hydrostatic_balance_ft64/output_active
         artifact_paths: "single_column_hydrostatic_balance_ft64/output_active/*"
+        agents:
+          slurm_constraint: icelake|cascadelake|skylake|epyc
 
   - group: "Box Examples"
     steps:
@@ -363,6 +365,7 @@ steps:
         artifact_paths: "aquaplanet_nonequil_allsky_gw_res/output_active/*"
         agents:
           slurm_mem: 20GB
+          slurm_constraint: icelake|cascadelake|skylake|epyc
 
       - label: ":computer: aquaplanet equil allsky monin_obukhov varying insol gravity wave (raw_topo) high top zonally asymmetric"
         command: >

config/default_configs/default_config.yml

@@ -324,6 +324,9 @@ implicit_sgs_advection:
 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_nh_pressure:
+  help: "Whether to treat the subgrid-scale nonhydrostatic pressure drag 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`]. Setting it to true only works if both implicit_sgs_advection and implicit_diffusion are set to true."
   value: false

config/model_configs/prognostic_edmfx_bomex_implicit_column.yml

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

config/perf_configs/bm_aquaplanet_progedmf.yml

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

src/prognostic_equations/implicit/implicit_solver.jl

@@ -100,6 +100,7 @@ struct ImplicitEquationJacobian{
     F3 <: DerivativeFlag,
     F4 <: DerivativeFlag,
     F5 <: DerivativeFlag,
+    F6 <: DerivativeFlag,
     T <: Fields.FieldVector,
     R <: Base.RefValue,
 }
@@ -114,7 +115,8 @@ struct ImplicitEquationJacobian{
     topography_flag::F2
     sgs_advection_flag::F3
     sgs_entr_detr_flag::F4
-    sgs_mass_flux_flag::F5
+    sgs_nh_pressure_flag::F5
+    sgs_mass_flux_flag::F6
 
     # required by Krylov.jl to evaluate ldiv! with AbstractVector inputs
     temp_b::T
@@ -133,6 +135,7 @@ function Base.zero(jac::ImplicitEquationJacobian)
         jac.topography_flag,
         jac.sgs_advection_flag,
         jac.sgs_entr_detr_flag,
+        jac.sgs_nh_pressure_flag,
         jac.sgs_mass_flux_flag,
         jac.temp_b,
         jac.temp_x,
@@ -153,6 +156,8 @@ function ImplicitEquationJacobian(
                          IgnoreDerivative(),
     sgs_entr_detr_flag = atmos.sgs_entr_detr_mode == Implicit() ?
                          UseDerivative() : IgnoreDerivative(),
+    sgs_nh_pressure_flag = atmos.sgs_nh_pressure_mode == Implicit() ?
+                           UseDerivative() : IgnoreDerivative(),
     sgs_mass_flux_flag = atmos.sgs_mf_mode == Implicit() ? UseDerivative() :
                          IgnoreDerivative(),
     transform_flag = false,
@@ -409,6 +414,7 @@ function ImplicitEquationJacobian(
         topography_flag,
         sgs_advection_flag,
         sgs_entr_detr_flag,
+        sgs_nh_pressure_flag,
         sgs_mass_flux_flag,
         similar(Y),
         similar(Y),
@@ -497,11 +503,22 @@ NVTX.@annotate function Wfact!(A, Y, p, dtγ, t)
                 p.precomputed.bdmr_l,
                 p.precomputed.bdmr_r,
                 p.precomputed.bdmr,
+            ) : (;)
+        )...,
+        (
+            use_derivative(A.sgs_entr_detr_flag) &&
+            p.atmos.turbconv_model isa PrognosticEDMFX ?
+            (;
                 p.precomputed.ᶜentrʲs,
                 p.precomputed.ᶜdetrʲs,
                 p.precomputed.ᶜturb_entrʲs,
             ) : (;)
         )...,
+        (
+            use_derivative(A.sgs_nh_pressure_flag) &&
+            p.atmos.turbconv_model isa PrognosticEDMFX ?
+            (; p.precomputed.ᶠu₃⁰,) : (;)
+        )...,
         p.core.ᶜΦ,
         p.core.ᶠgradᵥ_ᶜΦ,
         p.scratch.ᶜtemp_scalar,
@@ -537,6 +554,7 @@ function update_implicit_equation_jacobian!(A, Y, p, dtγ)
         diffusion_flag,
         sgs_advection_flag,
         sgs_entr_detr_flag,
+        sgs_nh_pressure_flag,
         topography_flag,
         sgs_mass_flux_flag,
     ) = A
@@ -580,6 +598,7 @@ function update_implicit_equation_jacobian!(A, Y, p, dtγ)
     ᶠJ = Fields.local_geometry_field(Y.f).J
     ᶜgⁱʲ = Fields.local_geometry_field(Y.c).gⁱʲ
     ᶠgⁱʲ = Fields.local_geometry_field(Y.f).gⁱʲ
+    ᶠlg = Fields.local_geometry_field(Y.f)
 
     ᶜkappa_m = p.ᶜtemp_scalar
     @. ᶜkappa_m =
@@ -1006,19 +1025,39 @@ function update_implicit_equation_jacobian!(A, Y, p, dtγ)
             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,),
-                ))
+            if use_derivative(sgs_entr_detr_flag)
+                (; ᶜ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,),
+                    ))
+            end
+
+            # non-hydrostatic pressure drag
+            # Only the quadratic drag term is considered in the Jacobian, the buoyancy term is ignored
+            if use_derivative(sgs_nh_pressure_flag)
+                (; ᶠu₃⁰) = p
+                turbconv_params = CAP.turbconv_params(params)
+                α_d = CAP.pressure_normalmode_drag_coeff(turbconv_params)
+                scale_height =
+                    CAP.R_d(params) * CAP.T_surf_ref(params) / CAP.grav(params)
+                H_up_min = CAP.min_updraft_top(turbconv_params)
+                @. ∂ᶠu₃ʲ_err_∂ᶠu₃ʲ -=
+                    dtγ * (DiagonalMatrixRow(
+                        2 *
+                        α_d *
+                        Geometry._norm((Y.f.sgsʲs.:(1).u₃ - ᶠu₃⁰), ᶠlg) /
+                        max(scale_height, H_up_min) * (one_C3xACT3,),
+                    ))
+            end
 
             # add updraft mass flux contributions to grid-mean
             if use_derivative(sgs_mass_flux_flag)

src/prognostic_equations/implicit/implicit_tendency.jl

@@ -39,6 +39,10 @@ NVTX.@annotate function implicit_tendency!(Yₜ, Y, p, t)
         edmfx_sgs_mass_flux_tendency!(Yₜ, Y, p, t, p.atmos.turbconv_model)
     end
 
+    if p.atmos.sgs_nh_pressure_mode == Implicit()
+        edmfx_nh_pressure_tendency!(Yₜ, Y, p, t, p.atmos.turbconv_model)
+    end
+
     # NOTE: All ρa tendencies should be applied before calling this function
     pressure_work_tendency!(Yₜ, Y, p, t, p.atmos.turbconv_model)
 

src/prognostic_equations/remaining_tendency.jl

@@ -118,7 +118,9 @@ NVTX.@annotate function additional_tendency!(Yₜ, Y, p, t)
     if p.atmos.sgs_mf_mode == Explicit()
         edmfx_sgs_mass_flux_tendency!(Yₜ, Y, p, t, p.atmos.turbconv_model)
     end
-    edmfx_nh_pressure_tendency!(Yₜ, Y, p, t, p.atmos.turbconv_model)
+    if p.atmos.sgs_nh_pressure_mode == Explicit()
+        edmfx_nh_pressure_tendency!(Yₜ, Y, p, t, p.atmos.turbconv_model)
+    end
     edmfx_filter_tendency!(Yₜ, Y, p, t, p.atmos.turbconv_model)
     edmfx_tke_tendency!(Yₜ, Y, p, t, p.atmos.turbconv_model)
     # Non-equilibrium cloud formation

src/solver/type_getters.jl

@@ -58,6 +58,9 @@ function get_atmos(config::AtmosConfig, params)
     implicit_sgs_entr_detr = parsed_args["implicit_sgs_entr_detr"]
     @assert implicit_sgs_entr_detr in (true, false)
 
+    implicit_sgs_nh_pressure = parsed_args["implicit_sgs_nh_pressure"]
+    @assert implicit_sgs_nh_pressure in (true, false)
+
     implicit_sgs_mass_flux = parsed_args["implicit_sgs_mass_flux"]
     @assert implicit_sgs_mass_flux in (true, false)
 
@@ -107,6 +110,8 @@ function get_atmos(config::AtmosConfig, params)
         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_nh_pressure_mode = implicit_sgs_nh_pressure ? 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

@@ -560,6 +560,8 @@ Base.@kwdef struct AtmosModel{
     sgs_adv_mode::SAM = nothing
     """sgs_entr_detr_mode == Implicit() only works if sgs_adv_mode == Implicit()"""
     sgs_entr_detr_mode::SEDM = nothing
+    """sgs_nh_pressure_mode == Implicit() only works if sgs_adv_mode == Implicit()"""
+    sgs_nh_pressure_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