Add implict solver for diagnostic EDMF

.buildkite/pipeline.yml

@@ -582,6 +582,14 @@ steps:
         agents:
           slurm_mem: 20GB
 
+      - label: ":genie: Diagnostic EDMFX DYCOMS_RF01 in a box (explicit)"
+        command: >
+          julia --color=yes --project=examples examples/hybrid/driver.jl
+          --config_file $CONFIG_PATH/diagnostic_edmfx_dycoms_rf01_explicit_box.yml
+        artifact_paths: "diagnostic_edmfx_dycoms_rf01_explicitbox/*"
+        agents:
+          slurm_mem: 20GB
+
       - label: ":genie: Diagnostic EDMFX DYCOMS_RF01 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_bomex_box.yml

@@ -1,11 +1,12 @@
-
 job_id: "diagnostic_edmfx_bomex_box"
 initial_condition: "Bomex" 
 subsidence: "Bomex" 
 edmf_coriolis: "Bomex" 
 ls_adv: "Bomex" 
 surface_setup: "Bomex"
 turbconv: "diagnostic_edmfx" 
+implicit_diffusion: true
+approximate_linear_solve_iters: 2
 prognostic_tke: true
 edmfx_upwinding: "first_order" 
 edmfx_entr_model: "Generalized" 
@@ -24,7 +25,7 @@ y_elem: 2
 z_elem: 60 
 z_max: 3e3 
 z_stretch: false
-dt: "50secs" 
+dt: "200secs" 
 t_end: "6hours" 
 dt_save_to_disk: "10mins"
 toml: [toml/diagnostic_edmfx_box.toml]

config/model_configs/diagnostic_edmfx_bomex_stretched_box.yml

@@ -1,11 +1,12 @@
-
 job_id: "diagnostic_edmfx_bomex_stretched_box"
 initial_condition: "Bomex" 
 subsidence: "Bomex" 
 edmf_coriolis: "Bomex" 
 ls_adv: "Bomex" 
 surface_setup: "Bomex"
 turbconv: "diagnostic_edmfx" 
+implicit_diffusion: true
+approximate_linear_solve_iters: 2
 prognostic_tke: true
 edmfx_upwinding: "first_order" 
 edmfx_entr_model: "Generalized" 

config/model_configs/diagnostic_edmfx_dycoms_rf01_box.yml

@@ -5,6 +5,8 @@ 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" 
@@ -23,7 +25,7 @@ y_elem: 2
 z_elem: 30 
 z_max: 1500 
 z_stretch: false
-dt: 40secs 
+dt: 200secs
 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_explicit_box.yml

@@ -0,0 +1,29 @@
+job_id: diagnostic_edmfx_dycoms_rf01_explicit_box
+initial_condition: DYCOMS_RF01 
+subsidence: DYCOMS 
+edmf_coriolis: DYCOMS_RF01 
+rad: DYCOMS_RF01 
+surface_setup: DYCOMS_RF01
+turbconv: diagnostic_edmfx 
+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: 40secs 
+t_end: 4hours 
+dt_save_to_disk: 10mins
+toml: [toml/diagnostic_edmfx_box.toml]

config/model_configs/diagnostic_edmfx_gabls_box.yml

@@ -3,6 +3,8 @@ initial_condition: GABLS
 edmf_coriolis: GABLS
 surface_setup: GABLS
 turbconv: diagnostic_edmfx
+implicit_diffusion: true
+approximate_linear_solve_iters: 2
 prognostic_tke: true
 edmfx_upwinding: first_order
 edmfx_entr_model: "Generalized" 

config/model_configs/diagnostic_edmfx_rico_box.yml

@@ -1,4 +1,3 @@
-
 job_id: diagnostic_edmfx_rico_box
 initial_condition: Rico
 subsidence: Rico

config/model_configs/diagnostic_edmfx_trmm_box.yml

@@ -1,4 +1,3 @@
-
 job_id: diagnostic_edmfx_trmm_box
 initial_condition: TRMM_LBA
 rad: TRMM_LBA

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]

src/prognostic_equations/implicit/implicit_solver.jl

@@ -13,7 +13,7 @@ struct IgnoreEnthalpyDerivative <: EnthalpyDerivativeFlag end
 use_derivative(flag::EnthalpyDerivativeFlag) = flag == UseEnthalpyDerivative()
 
 """
-    ImplicitEquationJacobian(Y, atmos, diffusion_flag, enthalpy_flag, transform_flag)
+    ImplicitEquationJacobian(Y, atmos, diffusion_flag, enthalpy_flag, transform_flag, approximate_solve_iters)
 
 A wrapper for the matrix ``∂E/∂Y``, where ``E(Y)`` is the "error" of the
 implicit step with the state ``Y``.
@@ -86,6 +86,8 @@ is reached.
 - `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
+- `approximate_solve_iters::Int`: number of iterations to take for the
+  approximate linear solve required by `UseDiffusionDerivative()`
 """
 struct ImplicitEquationJacobian{
     M <: MatrixFields.FieldMatrix,
@@ -120,6 +122,7 @@ function ImplicitEquationJacobian(
     diffusion_flag,
     enthalpy_flag,
     transform_flag,
+    approximate_solve_iters,
 )
     FT = Spaces.undertype(axes(Y.c))
     one_C3xACT3 = C3(FT(1)) * CT3(FT(1))'
@@ -137,10 +140,10 @@ function ImplicitEquationJacobian(
         @name(c.ρq_ice),
         @name(c.ρq_rai),
         @name(c.ρq_sno),
+        @name(c.sgs⁰.ρatke),
     )
     available_tracer_names = MatrixFields.unrolled_filter(is_in_Y, tracer_names)
     other_names = (
-        @name(c.sgs⁰.ρatke),
         @name(c.sgsʲs),
         @name(f.sgsʲs),
         @name(c.turbconv),
@@ -185,8 +188,10 @@ function ImplicitEquationJacobian(
 
     alg =
         use_derivative(diffusion_flag) ?
-        MatrixFields.ApproximateBlockArrowheadIterativeSolve(@name(c)) :
-        MatrixFields.BlockArrowheadSolve(@name(c))
+        MatrixFields.ApproximateBlockArrowheadIterativeSolve(
+            @name(c);
+            n_iters = approximate_solve_iters,
+        ) : MatrixFields.BlockArrowheadSolve(@name(c))
 
     # By default, the ApproximateBlockArrowheadIterativeSolve takes 1 iteration
     # and approximates the A[c, c] block with a MainDiagonalPreconditioner.
@@ -242,7 +247,6 @@ _linsolve!(x, A, b, update_matrix = false; kwargs...) = ldiv!(x, A, b)
 function Wfact!(A, Y, p, dtγ, t)
     NVTX.@range "Wfact!" color = colorant"green" begin
         # Remove unnecessary values from p to avoid allocations in bycolumn.
-        (; rayleigh_sponge) = p.atmos
         p′ = (;
             p.precomputed.ᶜspecific,
             p.precomputed.ᶠu³,
@@ -253,6 +257,11 @@ function Wfact!(A, Y, p, dtγ, t)
                 use_derivative(A.diffusion_flag) ?
                 (; p.precomputed.ᶜK_u, p.precomputed.ᶜK_h) : (;)
             )...,
+            (
+                use_derivative(A.diffusion_flag) &&
+                p.atmos.turbconv_model isa PrognosticEDMFX ?
+                (; p.precomputed.ᶜρa⁰) : (;)
+            )...,
             p.core.∂ᶜK_∂ᶠu₃,
             p.core.ᶜΦ,
             p.core.ᶠgradᵥ_ᶜΦ,
@@ -263,7 +272,7 @@ function Wfact!(A, Y, p, dtγ, t)
             p.params,
             p.atmos,
             (
-                rayleigh_sponge isa RayleighSponge ?
+                p.atmos.rayleigh_sponge isa RayleighSponge ?
                 (; p.rayleigh_sponge.ᶠβ_rayleigh_w) : (;)
             )...,
         )
@@ -444,6 +453,12 @@ function update_implicit_equation_jacobian!(A, Y, p, dtγ, colidx)
         end
     end
 
+    # TODO: Move the vertical advection of ρatke into the implicit tendency.
+    if MatrixFields.has_field(Y, @name(c.sgs⁰.ρatke))
+        ∂ᶜρatke_err_∂ᶠu₃ = matrix[@name(c.sgs⁰.ρatke), @name(f.u₃)]
+        ∂ᶜρatke_err_∂ᶠu₃[colidx] .= (zero(eltype(∂ᶜρatke_err_∂ᶠu₃)),)
+    end
+
     # We can express the implicit tendency for vertical velocity as
     # ᶠu₃ₜ =
     #     -(ᶠgradᵥ(ᶜp - ᶜp_ref) + ᶠinterp(ᶜρ - ᶜρ_ref) * ᶠgradᵥ_ᶜΦ) /
@@ -549,27 +564,24 @@ function update_implicit_equation_jacobian!(A, Y, p, dtγ, colidx)
                 ) ⋅ ᶠgradᵥ_matrix() - (I,)
         end
 
-        ∂ᶜρe_tot_err_∂ᶜρe_tot = matrix[@name(c.ρe_tot), @name(c.ρe_tot)]
-        @. ∂ᶜρe_tot_err_∂ᶜρe_tot[colidx] =
-            dtγ * ᶜadvdivᵥ_matrix() ⋅
-            DiagonalMatrixRow(ᶠinterp(ᶜρ[colidx]) * ᶠinterp(p.ᶜK_h[colidx])) ⋅
-            ᶠgradᵥ_matrix() ⋅ DiagonalMatrixRow((1 + R_d / cv_d) / ᶜρ[colidx]) -
-            (I,)
-
-        tracer_names = (
-            @name(c.ρq_tot),
-            @name(c.ρq_liq),
-            @name(c.ρq_ice),
-            @name(c.ρq_rai),
-            @name(c.ρq_sno),
+        ᶜρ⁰ = p.atmos.turbconv_model isa PrognosticEDMFX ? p.ᶜρa⁰ : ᶜρ
+        scalar_info = (
+            (@name(c.ρe_tot), ᶜρ, p.ᶜK_h, 1 + R_d / cv_d),
+            (@name(c.ρq_tot), ᶜρ, p.ᶜK_h, 1),
+            (@name(c.ρq_liq), ᶜρ, p.ᶜK_h, 1),
+            (@name(c.ρq_ice), ᶜρ, p.ᶜK_h, 1),
+            (@name(c.ρq_rai), ᶜρ, p.ᶜK_h, 1),
+            (@name(c.ρq_sno), ᶜρ, p.ᶜK_h, 1),
+            (@name(c.sgs⁰.ρatke), ᶜρ⁰, p.ᶜK_u, 1),
         )
-        MatrixFields.unrolled_foreach(tracer_names) do ρχ_name
+        MatrixFields.unrolled_foreach(scalar_info) do (ρχ_name, ᶜρ_χ, ᶜK_χ, s_χ)
             MatrixFields.has_field(Y, ρχ_name) || return
             ∂ᶜρχ_err_∂ᶜρχ = matrix[ρχ_name, ρχ_name]
             @. ∂ᶜρχ_err_∂ᶜρχ[colidx] =
                 dtγ * ᶜadvdivᵥ_matrix() ⋅ DiagonalMatrixRow(
-                    ᶠinterp(ᶜρ[colidx]) * ᶠinterp(p.ᶜK_h[colidx]),
-                ) ⋅ ᶠgradᵥ_matrix() ⋅ DiagonalMatrixRow(1 / ᶜρ[colidx]) - (I,)
+                    ᶠinterp(ᶜρ_χ[colidx]) * ᶠinterp(ᶜK_χ[colidx]),
+                ) ⋅ ᶠgradᵥ_matrix() ⋅ DiagonalMatrixRow(s_χ / ᶜρ_χ[colidx]) -
+                (I,)
         end
     end
 end

src/prognostic_equations/implicit/implicit_tendency.jl

@@ -18,6 +18,14 @@ NVTX.@annotate function implicit_tendency!(Yₜ, Y, p, t)
                 colidx,
                 p.atmos.vert_diff,
             )
+            edmfx_sgs_diffusive_flux_tendency!(
+                Yₜ,
+                Y,
+                p,
+                t,
+                colidx,
+                p.atmos.turbconv_model,
+            )
         end
         # NOTE: All ρa tendencies should be applied before calling this function
         pressure_work_tendency!(Yₜ, Y, p, t, colidx, p.atmos.turbconv_model)