Merge pull request #2404 from CliMA/zs/refactor_vertdiff

move diffusivity to precomputed quantities

config/gpu_configs/gpu_diagnostic_edmfx_aquaplanet.yml

@@ -1,5 +1,4 @@
 job_id: gpu_diagnostic_edmfx_aquaplanet
-vert_diff: "true" 
 surface_setup: DefaultExchangeCoefficients 
 rad: gray
 turbconv: diagnostic_edmfx 

config/perf_configs/flame_perf_target_diagnostic_edmfx.yml

@@ -1,9 +1,10 @@
 dt_save_to_disk: "600secs"
+vert_diff: "false"
+turbconv: "diagnostic_edmfx"
 edmfx_entr_model: "Generalized"
+edmfx_detr_model: "Generalized"
 edmfx_nh_pressure: true
 edmfx_upwinding: "first_order"
 prognostic_tke: true
 edmfx_sgs_flux: true
-turbconv: "diagnostic_edmfx"
 job_id: "flame_perf_target_diagnostic_edmfx"
-edmfx_detr_model: "Generalized"

config/perf_configs/flame_perf_target_prognostic_edmfx_aquaplanet.yml

@@ -1,7 +1,7 @@
 job_id: flame_perf_target_prognostic_edmfx_aquaplanet
-vert_diff: "true" 
 surface_setup: DefaultExchangeCoefficients 
 rad: gray
+vert_diff: "false"
 turbconv: prognostic_edmfx 
 prognostic_tke: true
 edmfx_upwinding: first_order 

regression_tests/ref_counter.jl

@@ -1 +1 @@
-139
+140

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/implicit/implicit_solver.jl

@@ -161,9 +161,10 @@ function ImplicitEquationJacobian(
             (@name(c.ρe_tot), available_tracer_names...),
         )...,
         (@name(c.uₕ), @name(c.uₕ)) =>
-            use_derivative(diffusion_flag) &&
-            diffuse_momentum(atmos.vert_diff) ?
-            similar(Y.c, TridiagonalRow) : FT(-1) * I,
+            use_derivative(diffusion_flag) && (
+                !isnothing(atmos.turbconv_model) ||
+                diffuse_momentum(atmos.vert_diff)
+            ) ? similar(Y.c, TridiagonalRow) : FT(-1) * I,
         MatrixFields.unrolled_map(
             name -> (name, name) => FT(-1) * I,
             other_available_names,
@@ -247,6 +248,11 @@ function Wfact!(A, Y, p, dtγ, t)
             p.precomputed.ᶠu³,
             p.precomputed.ᶜK,
             p.precomputed.ᶜp,
+            p.precomputed.ᶜh_tot,
+            (
+                use_derivative(A.diffusion_flag) ?
+                (; p.precomputed.ᶜK_u, p.precomputed.ᶜK_h) : (;)
+            )...,
             p.core.∂ᶜK_∂ᶠu₃,
             p.core.ᶜΦ,
             p.core.ᶠgradᵥ_ᶜΦ,
@@ -256,7 +262,6 @@ function Wfact!(A, Y, p, dtγ, t)
             p.scratch.ᶠtemp_scalar,
             p.params,
             p.atmos,
-            p.precomputed.ᶜh_tot,
             (
                 rayleigh_sponge isa RayleighSponge ?
                 (; p.rayleigh_sponge.ᶠβ_rayleigh_w) : (;)
@@ -511,41 +516,43 @@ function update_implicit_equation_jacobian!(A, Y, p, dtγ, colidx)
             ∂ᶜK_∂ᶠu₃[colidx] - (I_u₃,)
     end
 
-    # We can express the implicit diffusion tendency for scalars and horizontal
-    # velocity as
-    # ᶜρχₜ = ᶜadvdivᵥ(ᶠρK_E * ᶠgradᵥ(ᶜχ)) and
-    # ᶜuₕₜ = ᶜadvdivᵥ(ᶠρK_E * ᶠgradᵥ(ᶜuₕ)) / ᶜρ.
+    # We can express the implicit diffusion tendency for horizontal velocity and
+    # scalars as
+    # ᶜuₕₜ = ᶜadvdivᵥ(ᶠinterp(ᶜρ) * ᶠinterp(ᶜK_u) * ᶠgradᵥ(ᶜuₕ)) / ᶜρ and
+    # ᶜρχₜ = ᶜadvdivᵥ(ᶠinterp(ᶜρ) * ᶠinterp(ᶜK_h) * ᶠgradᵥ(ᶜχ)).
+    # The implicit diffusion tendency for horizontal velocity actually uses
+    # 2 * ᶠstrain_rate instead of ᶠgradᵥ(ᶜuₕ), but these are roughly equivalent.
     # The implicit diffusion tendency for density is the sum of the ᶜρχₜ's, but
     # we approximate the derivative of this sum with respect to ᶜρ as 0.
     # This means that
-    # ∂(ᶜρχₜ)/∂(ᶜρχ) =
-    #     ᶜadvdivᵥ_matrix() ⋅ DiagonalMatrixRow(ᶠρK_E) ⋅ ᶠgradᵥ_matrix() ⋅
-    #     ∂(ᶜχ)/∂(ᶜρχ) and
     # ∂(ᶜuₕₜ)/∂(ᶜuₕ) =
     #     DiagonalMatrixRow(1 / ᶜρ) ⋅ ᶜadvdivᵥ_matrix() ⋅
-    #     DiagonalMatrixRow(ᶠρK_E) ⋅ ᶠgradᵥ_matrix().
+    #     DiagonalMatrixRow(ᶠinterp(ᶜρ) * ᶠinterp(ᶜK_u)) ⋅ ᶠgradᵥ_matrix() and
+    # ∂(ᶜρχₜ)/∂(ᶜρχ) =
+    #     ᶜadvdivᵥ_matrix() ⋅ DiagonalMatrixRow(ᶠinterp(ᶜρ) * ᶠinterp(ᶜK_h)) ⋅
+    #     ᶠgradᵥ_matrix() ⋅ ∂(ᶜχ)/∂(ᶜρχ).
     # In general, ∂(ᶜχ)/∂(ᶜρχ) = DiagonalMatrixRow(1 / ᶜρ), except for the case
     # ∂(ᶜh_tot)/∂(ᶜρe_tot) =
     #     ∂((ᶜρe_tot + ᶜp) / ᶜρ)/∂(ᶜρe_tot) =
     #     (I + ∂(ᶜp)/∂(ᶜρe_tot)) ⋅ DiagonalMatrixRow(1 / ᶜρ) ≈
     #     DiagonalMatrixRow((1 + R_d / cv_d) / ᶜρ).
     if use_derivative(diffusion_flag)
-        (; C_E) = p.atmos.vert_diff
-        ᶠp = ᶠρK_E = p.ᶠtemp_scalar
-        interior_uₕ = Fields.level(Y.c.uₕ, 1)
-        ᶜΔz_surface = Fields.Δz_field(interior_uₕ)
-        @. ᶠp[colidx] = ᶠinterp(ᶜp[colidx])
-        @. ᶠρK_E[colidx] =
-            ᶠinterp(Y.c.ρ[colidx]) * eddy_diffusivity_coefficient(
-                C_E,
-                norm(interior_uₕ[colidx]),
-                ᶜΔz_surface[colidx] / 2,
-                ᶠp[colidx],
-            )
+        if (
+            !isnothing(p.atmos.turbconv_model) ||
+            diffuse_momentum(p.atmos.vert_diff)
+        )
+            ∂ᶜuₕ_err_∂ᶜuₕ = matrix[@name(c.uₕ), @name(c.uₕ)]
+            @. ∂ᶜuₕ_err_∂ᶜuₕ[colidx] =
+                dtγ * DiagonalMatrixRow(1 / ᶜρ[colidx]) ⋅ ᶜadvdivᵥ_matrix() ⋅
+                DiagonalMatrixRow(
+                    ᶠinterp(ᶜρ[colidx]) * ᶠinterp(p.ᶜK_u[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(ᶠρK_E[colidx]) ⋅
+            dtγ * ᶜadvdivᵥ_matrix() ⋅
+            DiagonalMatrixRow(ᶠinterp(ᶜρ[colidx]) * ᶠinterp(p.ᶜK_h[colidx])) ⋅
             ᶠgradᵥ_matrix() ⋅ DiagonalMatrixRow((1 + R_d / cv_d) / ᶜρ[colidx]) -
             (I,)
 
@@ -560,15 +567,9 @@ function update_implicit_equation_jacobian!(A, Y, p, dtγ, colidx)
             MatrixFields.has_field(Y, ρχ_name) || return
             ∂ᶜρχ_err_∂ᶜρχ = matrix[ρχ_name, ρχ_name]
             @. ∂ᶜρχ_err_∂ᶜρχ[colidx] =
-                dtγ * ᶜadvdivᵥ_matrix() ⋅ DiagonalMatrixRow(ᶠρK_E[colidx]) ⋅
-                ᶠgradᵥ_matrix() ⋅ DiagonalMatrixRow(1 / ᶜρ[colidx]) - (I,)
-        end
-
-        if diffuse_momentum(p.atmos.vert_diff)
-            ∂ᶜuₕ_err_∂ᶜuₕ = matrix[@name(c.uₕ), @name(c.uₕ)]
-            @. ∂ᶜuₕ_err_∂ᶜuₕ[colidx] =
-                dtγ * DiagonalMatrixRow(1 / ᶜρ[colidx]) ⋅ ᶜadvdivᵥ_matrix() ⋅
-                DiagonalMatrixRow(ᶠρK_E[colidx]) ⋅ ᶠgradᵥ_matrix() - (I,)
+                dtγ * ᶜadvdivᵥ_matrix() ⋅ DiagonalMatrixRow(
+                    ᶠinterp(ᶜρ[colidx]) * ᶠinterp(p.ᶜK_h[colidx]),
+                ) ⋅ ᶠgradᵥ_matrix() ⋅ DiagonalMatrixRow(1 / ᶜρ[colidx]) - (I,)
         end
     end
 end

src/prognostic_equations/vertical_diffusion_boundary_layer.jl

@@ -13,43 +13,6 @@ import ClimaCore.Geometry as Geometry
 import ClimaCore.Fields as Fields
 import ClimaCore.Operators as Operators
 
-# Apply on potential temperature and moisture
-# 1) turn the liquid_theta into theta version
-# 2) have a total energy version (primary goal)
-
-function eddy_diffusivity_coefficient(C_E::FT, norm_v_a, z_a, p) where {FT}
-    p_pbl = FT(85000)
-    p_strato = FT(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
-
-function surface_thermo_state(
-    ::GCMSurfaceThermoState,
-    thermo_params,
-    T_sfc,
-    ts_int,
-    t,
-)
-    ρ_sfc =
-        TD.air_density(thermo_params, ts_int) *
-        (
-            T_sfc / TD.air_temperature(thermo_params, ts_int)
-        )^(
-            TD.cv_m(thermo_params, ts_int) /
-            TD.gas_constant_air(thermo_params, ts_int)
-        )
-    q_sfc =
-        TD.q_vap_saturation_generic(thermo_params, T_sfc, ρ_sfc, TD.Liquid())
-    if ts_int isa TD.PhaseDry
-        return TD.PhaseDry_ρT(thermo_params, ρ_sfc, T_sfc)
-    elseif ts_int isa TD.PhaseEquil
-        return TD.PhaseEquil_ρTq(thermo_params, ρ_sfc, T_sfc, q_sfc)
-    else
-        error("Unsupported thermo option")
-    end
-end
-
 function vertical_diffusion_boundary_layer_tendency!(Yₜ, Y, p, t)
     Fields.bycolumn(axes(Y.c.uₕ)) do colidx
         (; vert_diff) = p.atmos
@@ -75,45 +38,43 @@ function vertical_diffusion_boundary_layer_tendency!(
     colidx,
     ::VerticalDiffusion,
 )
-    ᶜρ = Y.c.ρ
-    FT = Spaces.undertype(axes(ᶜρ))
-    (; ᶜp, ᶜspecific, sfc_conditions) = p.precomputed # assume ᶜts and ᶜp have been updated
-    (; C_E) = p.atmos.vert_diff
-
+    FT = eltype(Y)
+    (; ᶜu, ᶜh_tot, ᶜspecific, ᶜK_u, ᶜK_h, sfc_conditions) = p.precomputed
     ᶠgradᵥ = Operators.GradientC2F() # apply BCs to ᶜdivᵥ, which wraps ᶠgradᵥ
 
-    FT = eltype(Y)
-    interior_uₕ = Fields.level(Y.c.uₕ, 1)
-    ᶠp = ᶠρK_E = p.scratch.ᶠtemp_scalar
-    @. ᶠp[colidx] = ᶠinterp(ᶜp[colidx])
-    ᶜΔz_surface = Fields.Δz_field(interior_uₕ)
-    @. ᶠρK_E[colidx] =
-        ᶠinterp(Y.c.ρ[colidx]) * eddy_diffusivity_coefficient(
-            C_E,
-            norm(interior_uₕ[colidx]),
-            ᶜΔz_surface[colidx] / 2,
-            ᶠp[colidx],
+    if diffuse_momentum(p.atmos.vert_diff)
+        ᶠstrain_rate = p.scratch.ᶠtemp_UVWxUVW
+        compute_strain_rate_face!(ᶠstrain_rate[colidx], ᶜu[colidx])
+        @. Yₜ.c.uₕ[colidx] -= C12(
+            ᶜdivᵥ(
+                -2 *
+                ᶠinterp(Y.c.ρ[colidx]) *
+                ᶠinterp(ᶜK_u[colidx]) *
+                ᶠstrain_rate[colidx],
+            ) / Y.c.ρ[colidx],
         )
 
-    if diffuse_momentum(p.atmos.vert_diff)
+        # apply boundary condition for momentum flux
         ᶜdivᵥ_uₕ = Operators.DivergenceF2C(
             top = Operators.SetValue(C3(FT(0)) ⊗ C12(FT(0), FT(0))),
             bottom = Operators.SetValue(sfc_conditions.ρ_flux_uₕ[colidx]),
         )
         @. Yₜ.c.uₕ[colidx] -=
-            ᶜdivᵥ_uₕ(-(ᶠρK_E[colidx] * ᶠgradᵥ(Y.c.uₕ[colidx]))) / Y.c.ρ[colidx]
+            ᶜdivᵥ_uₕ(-(FT(0) * ᶠgradᵥ(Y.c.uₕ[colidx]))) / Y.c.ρ[colidx]
     end
 
-    if :ρe_tot in propertynames(Y.c)
-        (; ᶜh_tot) = p.precomputed
+    ᶜ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(
+        -(
+            ᶠinterp(Y.c.ρ[colidx]) *
+            ᶠinterp(ᶜK_h[colidx]) *
+            ᶠgradᵥ(ᶜh_tot[colidx])
+        ),
+    )
 
-        ᶜ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(-(ᶠρK_E[colidx] * ᶠgradᵥ(ᶜh_tot[colidx])))
-    end
     ᶜρχₜ_diffusion = p.scratch.ᶜtemp_scalar
     ρ_flux_χ = p.scratch.sfc_temp_C3
     for (ᶜρχₜ, ᶜχ, χ_name) in matching_subfields(Yₜ.c, ᶜspecific)
@@ -127,8 +88,13 @@ function vertical_diffusion_boundary_layer_tendency!(
             top = Operators.SetValue(C3(FT(0))),
             bottom = Operators.SetValue(ρ_flux_χ[colidx]),
         )
-        @. ᶜρχₜ_diffusion[colidx] =
-            ᶜdivᵥ_ρχ(-(ᶠρK_E[colidx] * ᶠgradᵥ(ᶜχ[colidx])))
+        @. ᶜρχₜ_diffusion[colidx] = ᶜdivᵥ_ρχ(
+            -(
+                ᶠinterp(Y.c.ρ[colidx]) *
+                ᶠinterp(ᶜK_h[colidx]) *
+                ᶠgradᵥ(ᶜχ[colidx])
+            ),
+        )
         @. ᶜρχₜ[colidx] -= ᶜρχₜ_diffusion[colidx]
         @. Yₜ.c.ρ[colidx] -= ᶜρχₜ_diffusion[colidx]
     end