refactor(diffusion): consolidate the TKE eddy-diffusivity pipeline

Extract the TKE -> mixing length -> Prandtl -> eddy viscosity/diffusivity
chain, previously duplicated at seven sites, into `eddy_diffusivities!`.
Callers supply the mixing-length output field and optional `K_u`/`K_h`
output fields, preserving each site's scratch slots and lazy or materialized
policy; the closure inputs (`grid_scale`, `buoyancy_gradient`) pass through
unchanged. Migrated sites: the vertical and horizontal EDMFX SGS diffusive
fluxes, the vertical and horizontal updraft diffusion, the TKE tendency for
both EDMF models, and the diffusion Jacobian coefficients.

Author: haakon-e

src/prognostic_equations/eddy_diffusion_closures.jl

@@ -468,6 +468,53 @@ function ᶜmixing_length(
     return @. lazy(get_mixing_length_field(ᶜmixing_length_tuple, property))
 end
 
+"""
+    eddy_diffusivities!(ᶜmixing_length_field, Y, p; grid_scale, buoyancy_gradient, ᶜK_u_field, ᶜK_h_field)
+
+Compute the eddy viscosity `K_u` and eddy diffusivity `K_h` of the TKE-based
+closure, materializing the mixing length into `ᶜmixing_length_field`.
+
+# Keyword Arguments
+
+  - `grid_scale`: grid-scale limit on the mixing length. By default, the
+    vertical cell thickness.
+  - `buoyancy_gradient`: stability input of the mixing length and the turbulent
+    Prandtl number. By default, the centered `ᶜlinear_buoygrad`.
+  - `ᶜK_u_field`, `ᶜK_h_field`: output fields for `K_u` and `K_h`. By default
+    `nothing`, in which case the corresponding value is lazy.
+
+# Returns
+
+  - `(; ᶜK_u, ᶜK_h)`.
+"""
+function eddy_diffusivities!(
+    ᶜmixing_length_field, Y, p;
+    grid_scale = Fields.Δz_field(axes(Y.c)),
+    buoyancy_gradient = p.precomputed.ᶜlinear_buoygrad,
+    ᶜK_u_field = nothing,
+    ᶜK_h_field = nothing,
+)
+    (; params) = p
+    (; ᶜstrain_rate_norm) = p.precomputed
+    turbconv_params = CAP.turbconv_params(params)
+    ᶜtke = @. lazy(specific(Y.c.ρtke, Y.c.ρ))
+    ᶜmixing_length_field .= ᶜmixing_length(Y, p; grid_scale, buoyancy_gradient)
+    ᶜK_u = if isnothing(ᶜK_u_field)
+        @. lazy(eddy_viscosity(turbconv_params, ᶜtke, ᶜmixing_length_field))
+    else
+        @. ᶜK_u_field =
+            eddy_viscosity(turbconv_params, ᶜtke, ᶜmixing_length_field)
+    end
+    ᶜprandtl_nvec =
+        @. lazy(turbulent_prandtl_number(params, buoyancy_gradient, ᶜstrain_rate_norm))
+    ᶜK_h = if isnothing(ᶜK_h_field)
+        @. lazy(eddy_diffusivity(ᶜK_u, ᶜprandtl_nvec))
+    else
+        @. ᶜK_h_field = eddy_diffusivity(ᶜK_u, ᶜprandtl_nvec)
+    end
+    return (; ᶜK_u, ᶜK_h)
+end
+
 """
     gradient_richardson_number(params, ᶜN²_eff, ᶜstrain_rate_norm)
 

src/prognostic_equations/edmfx_sgs_flux.jl

@@ -231,21 +231,8 @@ function edmfx_sgs_diffusive_flux_tendency!(
 
     if p.atmos.edmfx_model.sgs_diffusive_flux isa Val{true}
 
-        (; ᶜlinear_buoygrad, ᶜstrain_rate_norm) = p.precomputed
-        # scratch to prevent GPU Kernel parameter memory error
         ᶜmixing_length_field = p.scratch.ᶜtemp_scalar_2
-        ᶜmixing_length_field .= ᶜmixing_length(Y, p)
-        ᶜK_u = @. lazy(
-            eddy_viscosity(turbconv_params, ᶜtke, ᶜmixing_length_field),
-        )
-        ᶜprandtl_nvec = @. lazy(
-            turbulent_prandtl_number(
-                params,
-                ᶜlinear_buoygrad,
-                ᶜstrain_rate_norm,
-            ),
-        )
-        ᶜK_h = @. lazy(eddy_diffusivity(ᶜK_u, ᶜprandtl_nvec))
+        (; ᶜK_u, ᶜK_h) = eddy_diffusivities!(ᶜmixing_length_field, Y, p)
 
         ᶠρaK_h = p.scratch.ᶠtemp_scalar
         @. ᶠρaK_h = ᶠinterp(Y.c.ρ) * ᶠinterp(ᶜK_h)
@@ -382,25 +369,20 @@ function edmfx_sgs_horizontal_diffusive_flux_tendency!(
         return nothing
     iscolumn(axes(Y.c)) && return nothing
     (; params) = p
-    turbconv_params = CAP.turbconv_params(params)
-    (; ᶜlinear_buoygrad, ᶜstrain_rate_norm) = p.precomputed
+    (; ᶜlinear_buoygrad) = p.precomputed
     ᶜρ = Y.c.ρ
     ᶜtke = @. lazy(specific(Y.c.ρtke, ᶜρ))
 
     # Mixing length limited by the horizontal node spacing, with the centered
     # buoyancy gradient as the stability input
     Δx = Spaces.node_horizontal_length_scale(Spaces.horizontal_space(axes(Y.c)))
-    ᶜmixing_length_h = p.scratch.ᶜtemp_scalar_2
-    ᶜmixing_length_h .= ᶜmixing_length(
-        Y, p; grid_scale = Δx, buoyancy_gradient = ᶜlinear_buoygrad,
+    ᶜK = eddy_diffusivities!(
+        p.scratch.ᶜtemp_scalar_2, Y, p;
+        grid_scale = Δx, buoyancy_gradient = ᶜlinear_buoygrad,
+        ᶜK_u_field = p.scratch.ᶜtemp_scalar_4, ᶜK_h_field = p.scratch.ᶜtemp_scalar,
     )
-
-    ᶜK_u_h = p.scratch.ᶜtemp_scalar_4
-    @. ᶜK_u_h = eddy_viscosity(turbconv_params, ᶜtke, ᶜmixing_length_h)
-    ᶜprandtl_nvec =
-        @. lazy(turbulent_prandtl_number(params, ᶜlinear_buoygrad, ᶜstrain_rate_norm))
-    ᶜK_h_h = p.scratch.ᶜtemp_scalar
-    @. ᶜK_h_h = eddy_diffusivity(ᶜK_u_h, ᶜprandtl_nvec)
+    ᶜK_u_h = ᶜK.ᶜK_u
+    ᶜK_h_h = ᶜK.ᶜK_h
 
     # Total enthalpy, using the dry-static-energy + water-enthalpy
     # decomposition; see the matching vertical term in

src/prognostic_equations/edmfx_tke.jl

@@ -26,17 +26,8 @@
 edmfx_tke_tendency!(Yₜ, Y, p, t, turbconv_model) = nothing
 
 function edmfx_tke_tendency!(Yₜ, Y, p, t, turbconv_model::EDOnlyEDMFX)
-    (; params) = p
     (; ᶜstrain_rate_norm, ᶜlinear_buoygrad) = p.precomputed
-    turbconv_params = CAP.turbconv_params(p.params)
-    ᶜtke = @. lazy(specific(Y.c.ρtke, Y.c.ρ))
-    ᶜmixing_length_field = p.scratch.ᶜtemp_scalar
-    ᶜmixing_length_field .= ᶜmixing_length(Y, p)
-    ᶜK_u = @. lazy(eddy_viscosity(turbconv_params, ᶜtke, ᶜmixing_length_field))
-    ᶜprandtl_nvec = @. lazy(
-        turbulent_prandtl_number(params, ᶜlinear_buoygrad, ᶜstrain_rate_norm),
-    )
-    ᶜK_h = @. lazy(eddy_diffusivity(ᶜK_u, ᶜprandtl_nvec))
+    (; ᶜK_u, ᶜK_h) = eddy_diffusivities!(p.scratch.ᶜtemp_scalar, Y, p)
 
     # shear production
     @. Yₜ.c.ρtke += 2 * Y.c.ρ * ᶜK_u * ᶜstrain_rate_norm
@@ -52,23 +43,9 @@ function edmfx_tke_tendency!(
     turbconv_model::PrognosticEDMFX,
 )
     (; ᶜstrain_rate_norm, ᶜlinear_buoygrad) = p.precomputed
-    turbconv_params = CAP.turbconv_params(p.params)
 
     if use_prognostic_tke(turbconv_model)
-        ᶜmixing_length_field = p.scratch.ᶜtemp_scalar_2
-        ᶜmixing_length_field .= ᶜmixing_length(Y, p)
-        ᶜtke = @. lazy(specific(Y.c.ρtke, Y.c.ρ))
-        ᶜK_u = @. lazy(
-            eddy_viscosity(turbconv_params, ᶜtke, ᶜmixing_length_field),
-        )
-        ᶜprandtl_nvec = @. lazy(
-            turbulent_prandtl_number(
-                p.params,
-                ᶜlinear_buoygrad,
-                ᶜstrain_rate_norm,
-            ),
-        )
-        ᶜK_h = @. lazy(eddy_diffusivity(ᶜK_u, ᶜprandtl_nvec))
+        (; ᶜK_u, ᶜK_h) = eddy_diffusivities!(p.scratch.ᶜtemp_scalar_2, Y, p)
 
         # shear production
         @. Yₜ.c.ρtke += 2 * Y.c.ρ * ᶜK_u * ᶜstrain_rate_norm

src/prognostic_equations/implicit/manual_sparse_jacobian.jl

@@ -796,9 +796,7 @@ end
 # Eddy diffusivity and viscosity used by both grid-scale and SGS implicit
 # diffusion. May write to ᶜtemp_scalar_3, ᶜtemp_scalar_4, and ᶜtemp_scalar_6.
 function eddy_diffusivity_coefficients!(Y, p)
-    (; params) = p
     (; turbconv_model, vertical_diffusion, smagorinsky_lilly) = p.atmos
-    turbconv_params = CAP.turbconv_params(params)
     (; ᶜp) = p.precomputed
     ᶜK_u = p.scratch.ᶜtemp_scalar_4
     ᶜK_h = p.scratch.ᶜtemp_scalar_6
@@ -813,17 +811,12 @@ function eddy_diffusivity_coefficients!(Y, p)
         ᶜK_u = p.precomputed.ᶜνₜ_v
         ᶜK_h = p.precomputed.ᶜD_v
     elseif turbconv_model isa AbstractEDMF
-        (; ᶜlinear_buoygrad, ᶜstrain_rate_norm) = p.precomputed
-        ᶜtke = @. lazy(specific(Y.c.ρtke, Y.c.ρ))
-        ᶜmixing_length_field = p.scratch.ᶜtemp_scalar_3
-        ᶜmixing_length_field .= ᶜmixing_length(Y, p)
-        ᶜK_u = p.scratch.ᶜtemp_scalar_4
-        @. ᶜK_u = eddy_viscosity(turbconv_params, ᶜtke, ᶜmixing_length_field)
-        ᶜprandtl_nvec = @. lazy(
-            turbulent_prandtl_number(params, ᶜlinear_buoygrad, ᶜstrain_rate_norm),
+        ᶜK = eddy_diffusivities!(
+            p.scratch.ᶜtemp_scalar_3, Y, p;
+            ᶜK_u_field = p.scratch.ᶜtemp_scalar_4, ᶜK_h_field = p.scratch.ᶜtemp_scalar_6,
         )
-        ᶜK_h = p.scratch.ᶜtemp_scalar_6
-        @. ᶜK_h = eddy_diffusivity(ᶜK_u, ᶜprandtl_nvec)
+        ᶜK_u = ᶜK.ᶜK_u
+        ᶜK_h = ᶜK.ᶜK_h
     end
     return (; ᶜK_u, ᶜK_h)
 end

src/prognostic_equations/mass_flux_closures.jl

@@ -175,7 +175,6 @@ function edmfx_vertical_diffusion_tendency!(
         (; params) = p
         (; ᶜρʲs) = p.precomputed
         FT = eltype(p.params)
-        turbconv_params = CAP.turbconv_params(params)
         n = n_mass_flux_subdomains(turbconv_model)
         ᶜdivᵥ_mse = Operators.DivergenceF2C(
             top = Operators.SetValue(C3(0)),
@@ -186,16 +185,7 @@ function edmfx_vertical_diffusion_tendency!(
             bottom = Operators.SetValue(C3(0)),
         )
 
-        (; ᶜlinear_buoygrad, ᶜstrain_rate_norm) = p.precomputed
-        ᶜtke = @. lazy(specific(Y.c.ρtke, Y.c.ρ))
-        # scratch to prevent GPU Kernel parameter memory error
-        ᶜmixing_length_field = p.scratch.ᶜtemp_scalar
-        ᶜmixing_length_field .= ᶜmixing_length(Y, p)
-        ᶜK_u = @. lazy(eddy_viscosity(turbconv_params, ᶜtke, ᶜmixing_length_field))
-        ᶜprandtl_nvec = @. lazy(
-            turbulent_prandtl_number(params, ᶜlinear_buoygrad, ᶜstrain_rate_norm),
-        )
-        ᶜK_h = @. lazy(eddy_diffusivity(ᶜK_u, ᶜprandtl_nvec))
+        (; ᶜK_h) = eddy_diffusivities!(p.scratch.ᶜtemp_scalar, Y, p)
 
         for j in 1:n
             ᶜρʲ = ᶜρʲs.:($j)
@@ -244,21 +234,16 @@ function edmfx_horizontal_diffusion_tendency!(
     p.atmos.edmfx_model.horizontal_diffusion isa Val{true} || return nothing
     iscolumn(axes(Y.c)) && return nothing
     (; params) = p
-    (; ᶜρʲs, ᶜlinear_buoygrad, ᶜstrain_rate_norm) = p.precomputed
-    turbconv_params = CAP.turbconv_params(params)
+    (; ᶜρʲs, ᶜlinear_buoygrad) = p.precomputed
     n = n_mass_flux_subdomains(turbconv_model)
 
-    ᶜtke = @. lazy(specific(Y.c.ρtke, Y.c.ρ))
     Δx = Spaces.node_horizontal_length_scale(Spaces.horizontal_space(axes(Y.c)))
-    ᶜmixing_length_h = p.scratch.ᶜtemp_scalar
-    ᶜmixing_length_h .= ᶜmixing_length(
-        Y, p; grid_scale = Δx, buoyancy_gradient = ᶜlinear_buoygrad,
+    ᶜK = eddy_diffusivities!(
+        p.scratch.ᶜtemp_scalar, Y, p;
+        grid_scale = Δx, buoyancy_gradient = ᶜlinear_buoygrad,
+        ᶜK_h_field = p.scratch.ᶜtemp_scalar_2,
     )
-    ᶜK_u_h = @. lazy(eddy_viscosity(turbconv_params, ᶜtke, ᶜmixing_length_h))
-    ᶜprandtl_nvec =
-        @. lazy(turbulent_prandtl_number(params, ᶜlinear_buoygrad, ᶜstrain_rate_norm))
-    ᶜK_h_h = p.scratch.ᶜtemp_scalar_2
-    @. ᶜK_h_h = eddy_diffusivity(ᶜK_u_h, ᶜprandtl_nvec)
+    ᶜK_h_h = ᶜK.ᶜK_h
 
     ᶜq_totʲₜ_diffusion = p.scratch.ᶜtemp_scalar_3
     for j in 1:n