remove ᶜspecific_tke

Author: costachris

src/cache/diagnostic_edmf_precomputed_quantities.jl

@@ -359,7 +359,8 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!(
             specific(Y.c.ρe_tot, Y.c.ρ),
         ),
     )
-    ᶜtke⁰ = ᶜspecific_tke(Y, p)
+    ᶜρa⁰ = @. lazy(ρa⁰(Y.c.ρ, ᶜρaʲs, turbconv_model))
+    ᶜtke⁰ = @. lazy(specific_tke(Y.c.ρ, Y.c.sgs⁰.ρatke, ᶜρa⁰, turbconv_model))
 
     for i in 2:Spaces.nlevels(axes(Y.c))
         ρ_level = Fields.field_values(Fields.level(Y.c.ρ, i))
@@ -992,9 +993,11 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_env_closures!
 
     if p.atmos.turbconv_model isa DiagnosticEDMFX
         (; ᶜρaʲs, ᶠu³ʲs, ᶜdetrʲs, ᶠu³⁰, ᶜu⁰) = p.precomputed
+        ᶜρa⁰ = @. lazy(ρa⁰(Y.c.ρ, ᶜρaʲs, turbconv_model))
     elseif p.atmos.turbconv_model isa EDOnlyEDMFX
         ᶠu³⁰ = p.precomputed.ᶠu³
         ᶜu⁰ = ᶜu
+        ᶜρa⁰ = Y.c.ρ
     end
     @. ᶜu⁰ = C123(Y.c.uₕ) + ᶜinterp(C123(ᶠu³⁰))  # Set here, but used in a different function
 
@@ -1021,7 +1024,7 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_env_closures!
     @. ᶜprandtl_nvec =
         turbulent_prandtl_number(params, ᶜlinear_buoygrad, ᶜstrain_rate_norm)
 
-    ᶜtke⁰ = ᶜspecific_tke(Y, p)
+    ᶜtke⁰ = @. lazy(specific_tke(Y.c.ρ, Y.c.sgs⁰.ρatke, ᶜρa⁰, turbconv_model))
 
     ᶜtke_exch = p.scratch.ᶜtemp_scalar_2
     @. ᶜtke_exch = 0

src/cache/prognostic_edmf_precomputed_quantities.jl

@@ -24,7 +24,9 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_environment!(
     (; ᶜp, ᶜK) = p.precomputed
     (; ᶠu₃⁰, ᶜu⁰, ᶠu³⁰, ᶜK⁰, ᶜts⁰) = p.precomputed
 
-    ᶜtke⁰ = ᶜspecific_tke(Y, p)
+    ᶜρa⁰ = @. lazy(ρa⁰(Y.c.ρ, Y.c.sgsʲs, turbconv_model))
+    ᶜtke⁰ = @. lazy(specific_tke(Y.c.ρ, Y.c.sgs⁰.ρatke, ᶜρa⁰, turbconv_model))
+
     set_sgs_ᶠu₃!(u₃⁰, ᶠu₃⁰, Y, turbconv_model)
     set_velocity_quantities!(ᶜu⁰, ᶠu³⁰, ᶜK⁰, ᶠu₃⁰, Y.c.uₕ, ᶠuₕ³)
     # @. ᶜK⁰ += ᶜtke⁰
@@ -361,8 +363,8 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_explicit_clos
     ᶜdz = Fields.Δz_field(axes(Y.c))
     ᶜlg = Fields.local_geometry_field(Y.c)
     ᶠlg = Fields.local_geometry_field(Y.f)
-    ᶜtke⁰ = ᶜspecific_tke(Y, p)
     ᶜρa⁰ = @. lazy(ρa⁰(Y.c.ρ, Y.c.sgsʲs, turbconv_model))
+    ᶜtke⁰ = @. lazy(specific_tke(Y.c.ρ, Y.c.sgs⁰.ρatke, ᶜρa⁰, turbconv_model))
 
     ᶜvert_div = p.scratch.ᶜtemp_scalar
     ᶜmassflux_vert_div = p.scratch.ᶜtemp_scalar_2

src/diagnostics/Diagnostics.jl

@@ -59,7 +59,7 @@ import ..draft_area
 import ..compute_gm_mixing_length!
 import ..horizontal_integral_at_boundary
 import ..ρa⁰
-import ..ᶜspecific_tke
+import ..specific_tke
 import ..ᶜspecific_env_value
 
 # We need the abbreviations for symbols like curl, grad, and so on

src/diagnostics/edmfx_diagnostics.jl

@@ -1150,7 +1150,17 @@ function compute_tke!(
     time,
     turbconv_model::Union{EDOnlyEDMFX, PrognosticEDMFX, DiagnosticEDMFX},
 )
-    ᶜtke = ᶜspecific_tke(state, cache)
+    if turbconv_model isa PrognosticEDMFX
+        sgsʲs = state.c.sgsʲs
+    elseif turbconv_model isa DiagnosticEDMFX
+        (; ᶜρaʲs) = cache.precomputed
+        sgsʲs = ᶜρaʲs
+    else
+        sgsʲs = nothing
+    end
+
+    ᶜρa⁰ = @. lazy(ρa⁰(state.c.ρ, sgsʲs, turbconv_model))
+    ᶜtke = @. lazy(specific_tke(state.c.ρ, state.c.sgs⁰.ρatke, ᶜρa⁰, turbconv_model))
     if isnothing(out)
         return Base.materialize(ᶜtke)
     else

src/prognostic_equations/advection.jl

@@ -194,18 +194,18 @@ NVTX.@annotate function explicit_vertical_advection_tendency!(Yₜ, Y, p, t)
             turbconv_model isa EDOnlyEDMFX ? p.precomputed.ᶠu³ :
             p.precomputed.ᶠu³⁰
         ) : nothing
-    ᶜρa⁰ = advect_tke ? (n > 0 ? (@. lazy(ρa⁰(Y.c.ρ, Y.c.sgsʲs, turbconv_model))) : @. lazy(Y.c.ρ)) : nothing
+    ᶜρa⁰ = advect_tke ? (n > 0 ? (@. lazy(ρa⁰(Y.c.ρ, Y.c.sgsʲs, turbconv_model))) : Y.c.ρ) : nothing
     ᶜρ⁰ = if advect_tke
         if n > 0
             (; ᶜts⁰) = p.precomputed
             @. lazy(TD.air_density(thermo_params, ᶜts⁰))
         else
-            @. lazy(Y.c.ρ)
+            Y.c.ρ
         end
     else
         nothing
     end
-    ᶜtke⁰ = advect_tke ? (ᶜspecific_tke(Y, p)) : nothing
+    ᶜtke⁰ = advect_tke ? (@. lazy(specific_tke(Y.c.ρ, Y.c.sgs⁰.ρatke, ᶜρa⁰, turbconv_model))) : nothing
     ᶜa_scalar = p.scratch.ᶜtemp_scalar
     ᶜω³ = p.scratch.ᶜtemp_CT3
     ᶠω¹² = p.scratch.ᶠtemp_CT12

src/prognostic_equations/edmfx_sgs_flux.jl

@@ -405,7 +405,7 @@ function edmfx_sgs_diffusive_flux_tendency!(
     (; ᶜmixing_length, ᶜK_u, ᶜK_h, ρatke_flux) = p.precomputed
     ᶠgradᵥ = Operators.GradientC2F()
     ᶜρa⁰ = @. lazy(ρa⁰(Y.c.ρ, Y.c.sgsʲs, turbconv_model))
-    ᶜtke⁰ = ᶜspecific_tke(Y, p)
+    ᶜtke⁰ = @. lazy(specific_tke(Y.c.ρ, Y.c.sgs⁰.ρatke, ᶜρa⁰, turbconv_model))
 
     if p.atmos.edmfx_model.sgs_diffusive_flux isa Val{true}
         ᶠρaK_h = p.scratch.ᶠtemp_scalar
@@ -513,7 +513,13 @@ function edmfx_sgs_diffusive_flux_tendency!(
     (; ᶜu, ᶜmixing_length, ᶜts) = p.precomputed
     (; ᶜK_u, ᶜK_h, ρatke_flux) = p.precomputed
     ᶠgradᵥ = Operators.GradientC2F()
-    ᶜtke⁰ = ᶜspecific_tke(Y, p)
+    if turbconv_model isa DiagnosticEDMFX
+        (; ᶜρaʲs) = p.precomputed
+        ᶜρa⁰ = @. lazy(ρa⁰(Y.c.ρ, ᶜρaʲs, turbconv_model))
+    elseif turbconv_model isa EDOnlyEDMFX
+        ᶜρa⁰ = Y.c.ρ
+    end
+    ᶜtke⁰ = @. lazy(specific_tke(Y.c.ρ, Y.c.sgs⁰.ρatke, ᶜρa⁰, turbconv_model))
 
     if p.atmos.edmfx_model.sgs_diffusive_flux isa Val{true}
         ᶠρaK_h = p.scratch.ᶠtemp_scalar

src/prognostic_equations/edmfx_tke.jl

@@ -50,7 +50,7 @@ function edmfx_tke_tendency!(
     FT = eltype(p.params)
 
 
-    ᶜρa⁰ = turbconv_model isa PrognosticEDMFX ? (@. lazy(ρa⁰(Y.c.ρ, Y.c.sgsʲs, turbconv_model))) : @. lazy(Y.c.ρ)
+    ᶜρa⁰ = turbconv_model isa PrognosticEDMFX ? (@. lazy(ρa⁰(Y.c.ρ, Y.c.sgsʲs, turbconv_model))) : Y.c.ρ
     nh_pressure3_buoyʲs =
         turbconv_model isa PrognosticEDMFX ?
         p.precomputed.ᶠnh_pressure₃_buoyʲs : p.precomputed.ᶠnh_pressure³_buoyʲs
@@ -77,7 +77,7 @@ function edmfx_tke_tendency!(
         # buoyancy production
         @. Yₜ.c.sgs⁰.ρatke -= ᶜρa⁰ * ᶜK_h * ᶜlinear_buoygrad
 
-        ᶜtke⁰ = ᶜspecific_tke(Y, p)
+        ᶜtke⁰ = @. lazy(specific_tke(Y.c.ρ, Y.c.sgs⁰.ρatke, ᶜρa⁰, turbconv_model))
 
         # entrainment and detraiment
         # using ᶜu⁰ and local geometry results in allocation

src/prognostic_equations/hyperdiffusion.jl

@@ -117,7 +117,16 @@ NVTX.@annotate function prep_hyperdiffusion_tendency!(Yₜ, Y, p, t)
         wdivₕ(gradₕ(specific(Y.c.ρe_tot, Y.c.ρ) + ᶜp / Y.c.ρ - ᶜh_ref))
 
     if diffuse_tke
-        ᶜtke⁰ = ᶜspecific_tke(Y, p)
+        if turbconv_model isa PrognosticEDMFX
+            sgsʲs = Y.c.sgsʲs
+        elseif turbconv_model isa DiagnosticEDMFX
+            (; ᶜρaʲs) = p.precomputed
+            sgsʲs = ᶜρaʲs
+        else
+            sgsʲs = nothing
+        end
+        ᶜρa⁰ = @. lazy(ρa⁰(Y.c.ρ, sgsʲs, turbconv_model))
+        ᶜtke⁰ = @. lazy(specific_tke(Y.c.ρ, Y.c.sgs⁰.ρatke, ᶜρa⁰, turbconv_model))
         (; ᶜ∇²tke⁰) = p.hyperdiff
         @. ᶜ∇²tke⁰ = wdivₕ(gradₕ(ᶜtke⁰))
     end
@@ -154,7 +163,6 @@ NVTX.@annotate function apply_hyperdiffusion_tendency!(Yₜ, Y, p, t)
         (; ᶜ∇²uₕʲs, ᶜ∇²uᵥʲs, ᶜ∇²uʲs, ᶜ∇²mseʲs) = p.hyperdiff
     end
     if use_prognostic_tke(turbconv_model)
-        ᶜtke⁰ = ᶜspecific_tke(Y, p)
         (; ᶜ∇²tke⁰) = p.hyperdiff
     end
 

src/prognostic_equations/implicit/manual_sparse_jacobian.jl

@@ -625,11 +625,12 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t)
             (; dt) = p
             turbconv_model = p.atmos.turbconv_model
             ᶜmixing_length = p.precomputed.ᶜmixing_length
-            ᶜtke⁰ = ᶜspecific_tke(Y, p)
             ᶜρa⁰ =
-                p.atmos.turbconv_model isa PrognosticEDMFX ? (@. lazy(ρa⁰(Y.c.ρ, Y.c.sgsʲs, turbconv_model))) : @. lazy(Y.c.ρ)
+                p.atmos.turbconv_model isa PrognosticEDMFX ? (@. lazy(ρa⁰(Y.c.ρ, Y.c.sgsʲs, turbconv_model))) : Y.c.ρ
             ᶜρatke⁰ = Y.c.sgs⁰.ρatke
 
+            ᶜtke⁰ = @. lazy(specific_tke(Y.c.ρ, ᶜρatke⁰, ᶜρa⁰, turbconv_model))
+
             @inline tke_dissipation_rate_tendency(tke⁰, mixing_length) =
                 tke⁰ >= 0 ? c_d * sqrt(tke⁰) / mixing_length : 1 / float(dt)
             @inline ∂tke_dissipation_rate_tendency_∂tke⁰(tke⁰, mixing_length) =

src/utils/variable_manipulations.jl

@@ -375,23 +375,12 @@ function ρa⁰(ρ, sgsʲs, turbconv_model)
         return env_value(ρ, sgsʲ -> sgsʲ.ρa, sgsʲs)
 
     elseif turbconv_model isa DiagnosticEDMFX
-        (; ᶜρaʲs) = p.precomputed
         return env_value(ρ, ᶜρaʲ -> ᶜρaʲ, sgsʲs)
     else
-        return Y.c.ρ
+        return ρ
     end
 end
 
-"""
-Arguments:
-- `ρ`: The model state, which contains the grid-mean density `Y.c.ρ` and
-        the draft subdomain states `Y.c.sgsʲs` (for PrognosticEDMFX).
-- `sgsʲs`: Y.c.sgsʲs for prognostic of The cache, containing precomputed quantities and turbconv_model.
-"""
-function ρa⁰(ρ, sgsʲs)
-    # ρ - Σ ρaʲ
-    return env_value(ρ, sgsʲ -> sgsʲ.ρa, sgsʲs)
-end
 
 """
     ᶜspecific_tke(Y, p)
@@ -404,47 +393,28 @@ fallback value (`ρχ_fallback`) in the limit of small environmental area
 fraction.
 
 Arguments:
-- `Y`: The state, containing the grid-mean density `ρ` and the environment SGS state `Y.c.sgs⁰`.
-- `p`: The cache, containing precomputed quantities and turbconv_model.
+- `ρ`: The grid-mean density `ρ` and the environment SGS state `Y.c.sgs⁰`.
+- `ρatke`: The cache, containing precomputed quantities and turbconv_model.
+- `ρa⁰`: The environment area-weighted density.
 
 Returns:
 - The specific TKE of the environment (`tke⁰`).
 """
-function ᶜspecific_tke(Y, p)
-    turbconv_model = p.atmos.turbconv_model
-    ᶜρa⁰ = @. lazy(ρa⁰(Y.c.ρ, Y.c.sgsʲs, turbconv_model))
-
-    sgs⁰ = Y.c.sgs⁰
+function specific_tke(ρ, ρatke, ρa⁰, turbconv_model)
 
-    # no sgs weighting function needed for EDOnlyEDMFX
-    if turbconv_model isa EDOnlyEDMFX
-        return @. lazy(specific(sgs⁰.ρatke, ᶜρa⁰))
-    else
-        return @. lazy(
-            specific(
-                sgs⁰.ρatke,     # ρaχ for environment TKE
-                ᶜρa⁰,        # ρa for environment, now computed internally
-                0,              # Fallback ρχ is zero for TKE
-                Y.c.ρ,           # Fallback ρ
-                turbconv_model,
-            ),
-        )
-    end
-end
-
-function specific_tke(ᶜρ, ᶜρatke, ᶜρa⁰, turbconv_model)
-    return specific(
-            ᶜρatke,    # ρaχ for environment TKE
-            ᶜρa⁰, # ρa for environment, now computed internally
+    if turbconv_model isa PrognosticEDMFX || turbconv_model isa DiagnosticEDMFX
+        return specific(
+            ρatke,    # ρaχ for environment TKE
+            ρa⁰, # ρa for environment, now computed internally
             0,         # Fallback ρχ is zero for TKE
-            ᶜρ,        # Fallback ρ
+            ρ,        # Fallback ρ
             turbconv_model,
         )
+    else
+        return specific(ρatke, ρa⁰)
+    end
 end
 
-function specific_tke(ᶜρatke, ᶜρa⁰, ::EDOnlyEDMFX)
-    return specific(ᶜρatke, ᶜρa⁰)
-end
 
 """
     ᶜspecific_env_mse(Y, p)