rft: clean up diffusivity code

Author: haakon-e

src/cache/eddy_diffusivity_coefficient.jl

@@ -1,27 +1,70 @@
-function ᶜcompute_eddy_diffusivity_coefficient(
-    ᶜρ,
-    vert_diff::DecayWithHeightDiffusion,
-)
+
+"""
+    ᶜcompute_eddy_diffusivity_coefficient(ᶜρ, (; D₀, H)::DecayWithHeightDiffusion)
+
+Return lazy representation of the vertical profile of eddy diffusivity 
+for the `DecayWithHeightDiffusion` model.
+
+The profile is given by:
+```
+K(z) = D₀ ⋅ exp(-(z - z_sfc) / H)
+```
+
+# Arguments
+- `ᶜρ`: Cell-centered field whose axes provide vertical coordinates.
+- Instance of `DecayWithHeightDiffusion` model, with fields:
+    - `D₀`: Surface eddy diffusivity magnitude.
+    - `H`: E-folding height for the exponential decay.
+
+# See also
+- [`DecayWithHeightDiffusion`] for the model specification
+- [`vertical_diffusion_boundary_layer_tendency!`] where this coefficient is applied
+"""
+function ᶜcompute_eddy_diffusivity_coefficient(ᶜρ, (; D₀, H)::DecayWithHeightDiffusion)
     (; ᶜz, ᶠz) = z_coordinate_fields(axes(ᶜρ))
     ᶠz_sfc = Fields.level(ᶠz, Fields.half)
-    return @. lazy(
-        eddy_diffusivity_coefficient_H(vert_diff.D₀, vert_diff.H, ᶠz_sfc, ᶜz),
-    )
+    return @. lazy(eddy_diffusivity_coefficient_H(D₀, H, ᶠz_sfc, ᶜz))
 end
+eddy_diffusivity_coefficient_H(D₀, H, z_sfc, z) = D₀ * exp(-(z - z_sfc) / H)
 
-function ᶜcompute_eddy_diffusivity_coefficient(
-    ᶜuₕ,
-    ᶜp,
-    vert_diff::VerticalDiffusion,
-)
+"""
+    ᶜcompute_eddy_diffusivity_coefficient(ᶜuₕ, ᶜp, (; C_E)::VerticalDiffusion)
+
+Return lazy representation of the vertical profile of eddy diffusivity 
+for the `VerticalDiffusion` model.
+
+The profile is given by:
+```
+K(z) = K_E, if p > p_pbl
+     = K_E * exp(-((p_pbl - p) / p_strato)^2), otherwise
+```
+where `K_E` is given by:
+```
+K_E = C_E ⋅ norm(ᶜuₕ(z_bot)) ⋅ Δz_bot / 2
+```
+where `z_bot` is the first interior center level of the model,
+and `Δz_bot` is the thickness of the surface layer.
+
+# Arguments
+- `ᶜuₕ`: Cell-centered horizontal velocity field; its first interior level is used.
+- `ᶜp`: Cell-centered thermodynamic pressure field (or proxy) used by the closure.
+- Instance of `VerticalDiffusion` model, with field `C_E`:
+    - `C_E`: Dimensionless eddy-coefficient factor.
+
+# See also
+- [`VerticalDiffusion`] for the model specification
+"""
+function ᶜcompute_eddy_diffusivity_coefficient(ᶜuₕ, ᶜp, (; C_E)::VerticalDiffusion)
     interior_uₕ = Fields.level(ᶜuₕ, 1)
     ᶜΔz_surface = Fields.Δz_field(interior_uₕ)
     return @. lazy(
-        eddy_diffusivity_coefficient(
-            vert_diff.C_E,
-            norm(interior_uₕ),
-            ᶜΔz_surface / 2,
-            ᶜp,
-        ),
+        eddy_diffusivity_coefficient(C_E, norm(interior_uₕ), ᶜΔz_surface / 2, ᶜp),
     )
 end
+
+function eddy_diffusivity_coefficient(C_E, norm_uₕ_bottom, Δz_bottom, p)
+    p_pbl = 85000
+    p_strato = 10000
+    K_E = C_E * norm_uₕ_bottom * Δz_bottom
+    return p > p_pbl ? K_E : K_E * exp(-((p_pbl - p) / p_strato)^2)
+end

src/cache/precomputed_quantities.jl

@@ -431,16 +431,6 @@ ts_gs(thermo_params, moisture_model, microphysics_model, ᶜY, K, Φ, ρ) =
         ρ,
     )
 
-function eddy_diffusivity_coefficient_H(D₀, H, z_sfc, z)
-    return D₀ * exp(-(z - z_sfc) / H)
-end
-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_implicit_precomputed_quantities!(Y, p, t)
     set_implicit_precomputed_quantities_part1!(Y, p, t)

src/prognostic_equations/vertical_diffusion_boundary_layer.jl

@@ -67,26 +67,14 @@ Arguments:
 Modifies components of tendency vector `Yₜ.c` (e.g., `Yₜ.c.uₕ`, `Yₜ.c.ρe_tot`, `Yₜ.c.ρ`, and
 various tracer fields such as `Yₜ.c.ρq_tot`).
 """
-
 vertical_diffusion_boundary_layer_tendency!(Yₜ, Y, p, t) =
-    vertical_diffusion_boundary_layer_tendency!(
-        Yₜ,
-        Y,
-        p,
-        t,
-        p.atmos.vertical_diffusion,
-    )
-
-vertical_diffusion_boundary_layer_tendency!(Yₜ, Y, p, t, ::Nothing) = nothing
-
-function vertical_diffusion_boundary_layer_tendency!(
-    Yₜ,
-    Y,
-    p,
-    t,
+    vertical_diffusion_boundary_layer_tendency!(Yₜ, Y, p, t, p.atmos.vertical_diffusion)
+
+vertical_diffusion_boundary_layer_tendency!(_, _, _, _, ::Nothing) = nothing
+
+function vertical_diffusion_boundary_layer_tendency!(Yₜ, Y, p, _,
     ::Union{VerticalDiffusion, DecayWithHeightDiffusion},
 )
-    FT = eltype(Y)
     (; vertical_diffusion) = p.atmos
     α_vert_diff_tracer = CAP.α_vert_diff_tracer(p.params)
     thermo_params = CAP.thermodynamics_params(p.params)
@@ -96,33 +84,20 @@ function vertical_diffusion_boundary_layer_tendency!(
     if vertical_diffusion isa DecayWithHeightDiffusion
         ᶜK_h .= ᶜcompute_eddy_diffusivity_coefficient(Y.c.ρ, vertical_diffusion)
     elseif vertical_diffusion isa VerticalDiffusion
-        ᶜK_h .= ᶜcompute_eddy_diffusivity_coefficient(
-            Y.c.uₕ,
-            ᶜp,
-            vertical_diffusion,
-        )
+        ᶜK_h .= ᶜcompute_eddy_diffusivity_coefficient(Y.c.uₕ, ᶜp, vertical_diffusion)
     end
 
     if !disable_momentum_vertical_diffusion(p.atmos.vertical_diffusion)
         ᶠstrain_rate = compute_strain_rate_face_vertical(ᶜu)
-        @. Yₜ.c.uₕ -= C12(
-            ᶜdivᵥ(-2 * ᶠinterp(Y.c.ρ) * ᶠinterp(ᶜK_h) * ᶠstrain_rate) / Y.c.ρ,
-        ) # assumes ᶜK_u = ᶜK_h
+        @. Yₜ.c.uₕ -= C12(ᶜdivᵥ(-2 * ᶠinterp(Y.c.ρ) * ᶠinterp(ᶜK_h) * ᶠstrain_rate) / Y.c.ρ)  # assumes ᶜK_u = ᶜK_h
     end
 
-    ᶜdivᵥ_ρe_tot = Operators.DivergenceF2C(
-        top = Operators.SetValue(C3(0)),
-        bottom = Operators.SetValue(C3(0)),
-    )
-    ᶜh_tot = @. lazy(
-        TD.total_specific_enthalpy(
-            thermo_params,
-            ᶜts,
-            specific(Y.c.ρe_tot, Y.c.ρ),
-        ),
-    )
-    @. Yₜ.c.ρe_tot -=
-        ᶜdivᵥ_ρe_tot(-(ᶠinterp(Y.c.ρ) * ᶠinterp(ᶜK_h) * ᶠgradᵥ(ᶜh_tot)))
+    top = bottom = Operators.SetValue(C3(0))
+    ᶜdivᵥ_ρχ = Operators.DivergenceF2C(; top, bottom)
+
+    e_tot = @. lazy(specific(Y.c.ρe_tot, Y.c.ρ))
+    ᶜh_tot = @. lazy(TD.total_specific_enthalpy(thermo_params, ᶜts, e_tot))
+    @. Yₜ.c.ρe_tot -= ᶜdivᵥ_ρχ(-(ᶠinterp(Y.c.ρ) * ᶠinterp(ᶜK_h) * ᶠgradᵥ(ᶜh_tot)))
 
     ᶜρχₜ_diffusion = p.scratch.ᶜtemp_scalar_2
     ᶜK_h_scaled = p.scratch.ᶜtemp_scalar_3
@@ -133,17 +108,8 @@ function vertical_diffusion_boundary_layer_tendency!(
         else
             @. ᶜK_h_scaled = ᶜK_h
         end
-        ᶜdivᵥ_ρχ = Operators.DivergenceF2C(
-            top = Operators.SetValue(C3(0)),
-            bottom = Operators.SetValue(C3(0)),
-        )
-        @. ᶜρχₜ_diffusion = ᶜdivᵥ_ρχ(
-            -(
-                ᶠinterp(Y.c.ρ) *
-                ᶠinterp(ᶜK_h_scaled) *
-                ᶠgradᵥ(specific(ᶜρχ, Y.c.ρ))
-            ),
-        )
+        ᶜχ = @. lazy(specific(ᶜρχ, Y.c.ρ))
+        @. ᶜρχₜ_diffusion = ᶜdivᵥ_ρχ(-(ᶠinterp(Y.c.ρ) * ᶠinterp(ᶜK_h_scaled) * ᶠgradᵥ(ᶜχ)))
         @. ᶜρχₜ -= ᶜρχₜ_diffusion
         # Only add contribution from total water diffusion to mass tendency
         # (exclude contributions from diffusion of condensate, precipitation)