Add implict solver for diagnostic EDMF

.buildkite/pipeline.yml

@@ -590,6 +590,14 @@ steps:
         agents:
           slurm_mem: 20GB
 
+      - label: ":genie: Diagnostic EDMFX DYCOMS_RF01 in a box (implicit)"
+        command: >
+          julia --color=yes --project=examples examples/hybrid/driver.jl
+          --config_file $CONFIG_PATH/diagnostic_edmfx_dycoms_rf01_box_implicit.yml
+        artifact_paths: "diagnostic_edmfx_dycoms_rf01_box_implicit/*"
+        agents:
+          slurm_mem: 20GB
+
       - label: ":genie: Diagnostic EDMFX Rico in a box"
         command: >
           julia --color=yes --project=examples examples/hybrid/driver.jl

config/default_configs/default_config.yml

@@ -246,6 +246,9 @@ prognostic_surface:
 implicit_diffusion:
   help: "Whether to treat the vertical diffusion tendency implicitly [`false` (default), `true`]"
   value: false
+approximate_linear_solve_iters:
+  help: "Number of iterations for the approximate linear solve (used when `implicit_diffusion` is true)"
+  value: 1
 override_τ_precip:
   help: "If true, sets τ_precip to dt. Otherwise, τ_precip is set to the value in the toml dictionary"
   value: true

config/model_configs/diagnostic_edmfx_dycoms_rf01_box.yml

@@ -25,5 +25,5 @@ z_max: 1500
 z_stretch: false
 dt: 40secs 
 t_end: 4hours 
-dt_save_to_disk: 2mins
+dt_save_to_disk: 10mins
 toml: [toml/diagnostic_edmfx_box.toml]

config/model_configs/diagnostic_edmfx_dycoms_rf01_box_implicit.yml

@@ -0,0 +1,31 @@
+job_id: diagnostic_edmfx_dycoms_rf01_box_implicit
+initial_condition: DYCOMS_RF01 
+subsidence: DYCOMS 
+edmf_coriolis: DYCOMS_RF01 
+rad: DYCOMS_RF01 
+surface_setup: DYCOMS_RF01
+turbconv: diagnostic_edmfx 
+implicit_diffusion: true
+approximate_linear_solve_iters: 2
+prognostic_tke: true
+edmfx_upwinding: first_order 
+edmfx_entr_model: "Generalized" 
+edmfx_detr_model: "Generalized" 
+edmfx_nh_pressure: true 
+edmfx_sgs_mass_flux: true
+edmfx_sgs_diffusive_flux: true
+moist: equil
+config: box 
+hyperdiff: "true" 
+kappa_4: 1e21
+x_max: 1e8 
+y_max: 1e8 
+x_elem: 2 
+y_elem: 2 
+z_elem: 30 
+z_max: 1500 
+z_stretch: false
+dt: 200secs
+t_end: 4hours 
+dt_save_to_disk: 10mins
+toml: [toml/diagnostic_edmfx_box.toml]

src/cache/precomputed_quantities.jl

@@ -40,6 +40,7 @@ function precomputed_quantities(Y, atmos)
             !(atmos.turbconv_model isa PrognosticEDMFX)
     @assert !(atmos.edmfx_detr_model isa ConstantAreaDetrainment) ||
             !(atmos.turbconv_model isa DiagnosticEDMFX)
+    @assert isnothing(atmos.turbconv_model) || isnothing(atmos.vert_diff)
     TST = thermo_state_type(atmos.moisture_model, FT)
     SCT = SurfaceConditions.surface_conditions_type(atmos, FT)
     n = n_mass_flux_subdomains(atmos.turbconv_model)
@@ -109,13 +110,20 @@ function precomputed_quantities(Y, atmos)
             ᶜK_h = similar(Y.c, FT),
             ρatke_flux = similar(Fields.level(Y.f, half), C3{FT}),
         ) : (;)
+    vert_diff_quantities = if atmos.vert_diff isa VerticalDiffusion
+        ᶜK_h = similar(Y.c, FT)
+        (; ᶜK_u = ᶜK_h, ᶜK_h) # ᶜK_u aliases ᶜK_h because they are always equal.
+    else
+        (;)
+    end
     precipitation_quantities =
         atmos.precip_model isa Microphysics1Moment ?
         (; ᶜwᵣ = similar(Y.c, FT), ᶜwₛ = similar(Y.c, FT)) : (;)
     return (;
         gs_quantities...,
         advective_sgs_quantities...,
         diagnostic_sgs_quantities...,
+        vert_diff_quantities...,
         precipitation_quantities...,
     )
 end
@@ -272,6 +280,13 @@ ts_sgs(thermo_params, moisture_model, specific, K, Φ, p) = thermo_state(
     p,
 )
 
+function eddy_diffusivity_coefficient(C_E, norm_v_a, z_a, p)
+    p_pbl = 85000
+    p_strato = 10000
+    K_E = C_E * norm_v_a * z_a
+    return p > p_pbl ? K_E : K_E * exp(-((p_pbl - p) / p_strato)^2)
+end
+
 """
     set_precomputed_quantities!(Y, p, t)
 
@@ -289,7 +304,7 @@ function instead of recomputing the value yourself. Otherwise, it will be
 difficult to ensure that the duplicated computations are consistent.
 """
 NVTX.@annotate function set_precomputed_quantities!(Y, p, t)
-    (; moisture_model, turbconv_model, precip_model) = p.atmos
+    (; moisture_model, turbconv_model, vert_diff, precip_model) = p.atmos
     thermo_params = CAP.thermodynamics_params(p.params)
     n = n_mass_flux_subdomains(turbconv_model)
     thermo_args = (thermo_params, moisture_model)
@@ -344,6 +359,18 @@ NVTX.@annotate function set_precomputed_quantities!(Y, p, t)
         set_diagnostic_edmf_precomputed_quantities_env_closures!(Y, p, t)
     end
 
+    if vert_diff isa VerticalDiffusion
+        (; ᶜK_h) = p.precomputed
+        interior_uₕ = Fields.level(Y.c.uₕ, 1)
+        ᶜΔz_surface = Fields.Δz_field(interior_uₕ)
+        @. ᶜK_h = eddy_diffusivity_coefficient(
+            p.atmos.vert_diff.C_E,
+            norm(interior_uₕ),
+            ᶜΔz_surface / 2,
+            ᶜp,
+        )
+    end
+
     if precip_model isa Microphysics1Moment
         set_precipitation_precomputed_quantities!(Y, p, t)
     end

src/prognostic_equations/edmfx_sgs_flux.jl

@@ -169,23 +169,34 @@ function edmfx_sgs_diffusive_flux_tendency!(
 
     FT = Spaces.undertype(axes(Y.c))
     (; sfc_conditions) = p.precomputed
-    (; ᶜρa⁰, ᶜu⁰, ᶜK⁰, ᶜmse⁰, ᶜq_tot⁰) = p.precomputed
-    (; ᶜK_u, ᶜK_h) = p.precomputed
+    (; ᶜρa⁰, ᶜu⁰, ᶜK⁰, ᶜmse⁰, ᶜq_tot⁰, ᶜtke⁰) = p.precomputed
+    (; ᶜK_u, ᶜK_h, ρatke_flux) = p.precomputed
     ᶠgradᵥ = Operators.GradientC2F()
 
     if p.atmos.edmfx_sgs_diffusive_flux
-        # energy
         ᶠρaK_h = p.scratch.ᶠtemp_scalar
         @. ᶠρaK_h[colidx] = ᶠinterp(ᶜρa⁰[colidx]) * ᶠinterp(ᶜK_h[colidx])
+        ᶠρaK_u = p.scratch.ᶠtemp_scalar
+        @. ᶠρaK_u[colidx] = ᶠinterp(ᶜρa⁰[colidx]) * ᶠinterp(ᶜK_u[colidx])
 
+        # energy
         ᶜdivᵥ_ρe_tot = Operators.DivergenceF2C(
             top = Operators.SetValue(C3(FT(0))),
             bottom = Operators.SetValue(sfc_conditions.ρ_flux_h_tot[colidx]),
         )
         @. Yₜ.c.ρe_tot[colidx] -= ᶜdivᵥ_ρe_tot(
             -(ᶠρaK_h[colidx] * ᶠgradᵥ(ᶜmse⁰[colidx] + ᶜK⁰[colidx])),
         )
-
+        if use_prognostic_tke(turbconv_model)
+            # turbulent transport (diffusive flux)
+            # boundary condition for the diffusive flux
+            ᶜdivᵥ_ρatke = Operators.DivergenceF2C(
+                top = Operators.SetValue(C3(FT(0))),
+                bottom = Operators.SetValue(ρatke_flux[colidx]),
+            )
+            @. Yₜ.c.sgs⁰.ρatke[colidx] -=
+                ᶜdivᵥ_ρatke(-(ᶠρaK_u[colidx] * ᶠgradᵥ(ᶜtke⁰[colidx])))
+        end
         if !(p.atmos.moisture_model isa DryModel)
             # specific humidity
             ᶜρχₜ_diffusion = p.scratch.ᶜtemp_scalar
@@ -202,8 +213,6 @@ function edmfx_sgs_diffusive_flux_tendency!(
         end
 
         # momentum
-        ᶠρaK_u = p.scratch.ᶠtemp_scalar
-        @. ᶠρaK_u[colidx] = ᶠinterp(ᶜρa⁰[colidx]) * ᶠinterp(ᶜK_u[colidx])
         ᶠstrain_rate = p.scratch.ᶠtemp_UVWxUVW
         compute_strain_rate_face!(ᶠstrain_rate[colidx], ᶜu⁰[colidx])
         @. Yₜ.c.uₕ[colidx] -= C12(
@@ -234,22 +243,35 @@ function edmfx_sgs_diffusive_flux_tendency!(
 
     FT = Spaces.undertype(axes(Y.c))
     (; sfc_conditions) = p.precomputed
-    (; ᶜu, ᶜh_tot, ᶜspecific) = p.precomputed
-    (; ᶜK_u, ᶜK_h) = p.precomputed
+    (; ᶜu, ᶜh_tot, ᶜspecific, ᶜtke⁰) = p.precomputed
+    (; ᶜK_u, ᶜK_h, ρatke_flux) = p.precomputed
     ᶠgradᵥ = Operators.GradientC2F()
 
     if p.atmos.edmfx_sgs_diffusive_flux
-        # energy
         ᶠρaK_h = p.scratch.ᶠtemp_scalar
         @. ᶠρaK_h[colidx] = ᶠinterp(Y.c.ρ[colidx]) * ᶠinterp(ᶜK_h[colidx])
+        ᶠρaK_u = p.scratch.ᶠtemp_scalar
+        @. ᶠρaK_u[colidx] = ᶠinterp(Y.c.ρ[colidx]) * ᶠinterp(ᶜK_u[colidx])
 
+        # energy
         ᶜdivᵥ_ρe_tot = Operators.DivergenceF2C(
             top = Operators.SetValue(C3(FT(0))),
             bottom = Operators.SetValue(sfc_conditions.ρ_flux_h_tot[colidx]),
         )
         @. Yₜ.c.ρe_tot[colidx] -=
             ᶜdivᵥ_ρe_tot(-(ᶠρaK_h[colidx] * ᶠgradᵥ(ᶜh_tot[colidx])))
 
+        if use_prognostic_tke(turbconv_model)
+            # turbulent transport (diffusive flux)
+            # boundary condition for the diffusive flux
+            ᶜdivᵥ_ρatke = Operators.DivergenceF2C(
+                top = Operators.SetValue(C3(FT(0))),
+                bottom = Operators.SetValue(ρatke_flux[colidx]),
+            )
+            @. Yₜ.c.sgs⁰.ρatke[colidx] -=
+                ᶜdivᵥ_ρatke(-(ᶠρaK_u[colidx] * ᶠgradᵥ(ᶜtke⁰[colidx])))
+        end
+
         if !(p.atmos.moisture_model isa DryModel)
             # specific humidity
             ᶜρχₜ_diffusion = p.scratch.ᶜtemp_scalar
@@ -267,8 +289,6 @@ function edmfx_sgs_diffusive_flux_tendency!(
         end
 
         # momentum
-        ᶠρaK_u = p.scratch.ᶠtemp_scalar
-        @. ᶠρaK_u[colidx] = ᶠinterp(Y.c.ρ[colidx]) * ᶠinterp(ᶜK_u[colidx])
         ᶠstrain_rate = p.scratch.ᶠtemp_UVWxUVW
         compute_strain_rate_face!(ᶠstrain_rate[colidx], ᶜu[colidx])
         @. Yₜ.c.uₕ[colidx] -= C12(

src/prognostic_equations/edmfx_tke.jl

@@ -21,22 +21,11 @@ function edmfx_tke_tendency!(
     (; ᶜentrʲs, ᶜdetrʲs, ᶠu³ʲs) = p.precomputed
     (; ᶠu³⁰, ᶜstrain_rate_norm, ᶜlinear_buoygrad, ᶜtke⁰, ᶜmixing_length) =
         p.precomputed
-    (; ᶜK_u, ᶜK_h, ρatke_flux) = p.precomputed
+    (; ᶜK_u, ᶜK_h) = p.precomputed
     (; dt) = p.simulation
     ᶜρa⁰ = turbconv_model isa PrognosticEDMFX ? p.precomputed.ᶜρa⁰ : Y.c.ρ
-    ᶠgradᵥ = Operators.GradientC2F()
 
-    ᶠρaK_u = p.scratch.ᶠtemp_scalar
     if use_prognostic_tke(turbconv_model)
-        # turbulent transport (diffusive flux)
-        @. ᶠρaK_u[colidx] = ᶠinterp(ᶜρa⁰[colidx]) * ᶠinterp(ᶜK_u[colidx])
-        # boundary condition for the diffusive flux
-        ᶜdivᵥ_ρatke = Operators.DivergenceF2C(
-            top = Operators.SetValue(C3(FT(0))),
-            bottom = Operators.SetValue(ρatke_flux[colidx]),
-        )
-        @. Yₜ.c.sgs⁰.ρatke[colidx] -=
-            ᶜdivᵥ_ρatke(-(ᶠρaK_u[colidx] * ᶠgradᵥ(ᶜtke⁰[colidx])))
         # shear production
         @. Yₜ.c.sgs⁰.ρatke[colidx] +=
             2 * ᶜρa⁰[colidx] * ᶜK_u[colidx] * ᶜstrain_rate_norm[colidx]