ensure rho*a_env = rho for diagnostic edmf

Author: costachris

src/cache/diagnostic_edmf_precomputed_quantities.jl

@@ -359,8 +359,8 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!(
             specific(Y.c.ρe_tot, Y.c.ρ),
         ),
     )
-    ᶜρa⁰ = @. lazy(ρa⁰(Y.c.ρ, ᶜρaʲs, turbconv_model))
-    ᶜtke⁰ = @. lazy(specific_tke(Y.c.ρ, Y.c.sgs⁰.ρatke, ᶜρa⁰, turbconv_model))
+
+    ᶜtke⁰ = @. lazy(specific_tke(Y.c.ρ, Y.c.sgs⁰.ρatke, Y.c.ρ, turbconv_model))
 
     for i in 2:Spaces.nlevels(axes(Y.c))
         ρ_level = Fields.field_values(Fields.level(Y.c.ρ, i))
@@ -987,12 +987,10 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_env_closures!
     ᶜlg = Fields.local_geometry_field(Y.c)
 
     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))
+        (; ᶠu³⁰, ᶜu⁰) = p.precomputed
     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
 

src/cache/prognostic_edmf_precomputed_quantities.jl

@@ -8,7 +8,7 @@ import ClimaCore: Spaces, Fields
 """
     set_prognostic_edmf_precomputed_quantities!(Y, p, ᶠuₕ³, t)
 
-Updates the edmf environment precomputed quantities stored in `p` for edmfx.
+Updates the edmf environment precomputed quantities stored in `p` for prognostic edmfx.
 """
 NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_environment!(
     Y,
@@ -336,7 +336,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_explicit_clos
     FT = eltype(params)
     n = n_mass_flux_subdomains(turbconv_model)
 
-    (; ᶜtke⁰, ᶜu, ᶜp, ᶠu³⁰, ᶜts⁰) = p.precomputed
+    (; ᶜu, ᶜp, ᶠu³⁰, ᶜts⁰) = p.precomputed
     (; ᶜlinear_buoygrad, ᶜstrain_rate_norm, ρatke_flux) = p.precomputed
     (;
         ᶜuʲs,

src/diagnostics/edmfx_diagnostics.jl

@@ -1152,14 +1152,11 @@ function compute_tke!(
 )
     if turbconv_model isa PrognosticEDMFX
         sgsʲs = state.c.sgsʲs
-    elseif turbconv_model isa DiagnosticEDMFX
-        (; ᶜρaʲs) = cache.precomputed
-        sgsʲs = ᶜρaʲs
+        ᶜρa⁰ = @. lazy(ρa⁰(state.c.ρ, sgsʲs, turbconv_model))
     else
-        sgsʲs = nothing
+        ᶜρa⁰ = state.c.ρ
     end
 
-    ᶜρa⁰ = @. lazy(ρa⁰(state.c.ρ, sgsʲs, turbconv_model))
     ᶜtke = @. lazy(
         specific_tke(state.c.ρ, state.c.sgs⁰.ρatke, ᶜρa⁰, turbconv_model),
     )
@@ -1330,9 +1327,15 @@ function compute_edt!(
     turbconv_model::Union{PrognosticEDMFX, DiagnosticEDMFX},
 )
     turbconv_params = CAP.turbconv_params(cache.params)
-    (; ᶜlinear_buoygrad, ᶜstrain_rate_norm, ᶜtke⁰) = cache.precomputed
+    (; ᶜlinear_buoygrad, ᶜstrain_rate_norm) = cache.precomputed
     (; params) = cache
 
+    ᶜρa⁰ =
+        turbconv_model isa PrognosticEDMFX ?
+        (@. lazy(ρa⁰(state.c.ρ, state.c.sgsʲs, turbconv_model))) : state.c.ρ
+    ᶜtke⁰ = @. lazy(
+        specific_tke(state.c.ρ, state.c.sgs⁰.ρatke, ᶜρa⁰, turbconv_model),
+    )
     ᶜmixing_length_field = ᶜmixing_length(state, cache)
     ᶜK_u = @. lazy(eddy_viscosity(turbconv_params, ᶜtke⁰, ᶜmixing_length_field))
     ᶜprandtl_nvec = @. lazy(
@@ -1415,7 +1418,13 @@ function compute_evu!(
     turbconv_model::Union{PrognosticEDMFX, DiagnosticEDMFX},
 )
     turbconv_params = CAP.turbconv_params(cache.params)
-    (; ᶜtke⁰) = cache.precomputed
+
+    ᶜρa⁰ =
+        turbconv_model isa PrognosticEDMFX ?
+        (@. lazy(ρa⁰(state.c.ρ, state.c.sgsʲs, turbconv_model))) : state.c.ρ
+    ᶜtke⁰ = @. lazy(
+        specific_tke(state.c.ρ, state.c.sgs⁰.ρatke, ᶜρa⁰, turbconv_model),
+    )
     ᶜmixing_length_field = ᶜmixing_length(state, cache)
     ᶜK_u = @. lazy(eddy_viscosity(turbconv_params, ᶜtke⁰, ᶜmixing_length_field))
 

src/prognostic_equations/advection.jl

@@ -196,8 +196,10 @@ NVTX.@annotate function explicit_vertical_advection_tendency!(Yₜ, Y, p, t)
         ) : nothing
     ᶜρa⁰ =
         advect_tke ?
-        (n > 0 ? (@. lazy(ρa⁰(Y.c.ρ, Y.c.sgsʲs, turbconv_model))) : Y.c.ρ) :
-        nothing
+        (
+            turbconv_model isa PrognosticEDMFX ?
+            (@. lazy(ρa⁰(Y.c.ρ, Y.c.sgsʲs, turbconv_model))) : Y.c.ρ
+        ) : nothing
     ᶜρ⁰ = if advect_tke
         if n > 0
             (; ᶜts⁰) = p.precomputed

src/prognostic_equations/eddy_diffusion_closures.jl

@@ -541,11 +541,16 @@ end
 function ᶜmixing_length(Y, p, property::Val{P} = Val{:master}()) where {P}
     (; params) = p
     (; ustar, obukhov_length) = p.precomputed.sfc_conditions
-    (; ᶜtke⁰) = p.precomputed
     (; ᶜlinear_buoygrad, ᶜstrain_rate_norm) = p.precomputed
     ᶜz = Fields.coordinate_field(Y.c).z
     z_sfc = Fields.level(Fields.coordinate_field(Y.f).z, Fields.half)
     ᶜdz = Fields.Δz_field(axes(Y.c))
+
+    turbconv_model = p.atmos.turbconv_model
+    ᶜρa⁰ =
+        turbconv_model isa PrognosticEDMFX ?
+        (@. lazy(ρa⁰(Y.c.ρ, Y.c.sgsʲs, turbconv_model))) : Y.c.ρ
+    ᶜtke⁰ = @. lazy(specific_tke(Y.c.ρ, Y.c.sgs⁰.ρatke, ᶜρa⁰, turbconv_model))
     sfc_tke = Fields.level(ᶜtke⁰, 1)
 
     ᶜprandtl_nvec = p.scratch.ᶜtemp_scalar_5
@@ -677,11 +682,12 @@ function ᶜtke_exchange(Y, p)
     (; turbconv_model) = p.atmos
     n = n_mass_flux_subdomains(turbconv_model)
     ᶜρa⁰ =
-        p.atmos.turbconv_model isa PrognosticEDMFX ? p.precomputed.ᶜρa⁰ : Y.c.ρ
-
+        p.atmos.turbconv_model isa PrognosticEDMFX ?
+        (@. lazy(ρa⁰(Y.c.ρ, Y.c.sgsʲs, turbconv_model))) : Y.c.ρ
+    ᶜtke⁰ = @. lazy(specific_tke(Y.c.ρ, Y.c.sgs⁰.ρatke, ᶜρa⁰, turbconv_model))
 
     if p.atmos.turbconv_model isa PrognosticEDMFX
-        (; ᶜdetrʲs, ᶜtke⁰, ᶠu³⁰, ᶠu³ʲs) = p.precomputed
+        (; ᶜdetrʲs, ᶠu³⁰, ᶠu³ʲs) = p.precomputed
         ᶜtke_exch = p.scratch.ᶜtemp_scalar_2
         @. ᶜtke_exch = 0
         for j in 1:n
@@ -694,7 +700,7 @@ function ᶜtke_exchange(Y, p)
 
         return ᶜtke_exch
     elseif p.atmos.turbconv_model isa DiagnosticEDMFX
-        (; ᶜdetrʲs, ᶜtke⁰, ᶠu³⁰, ᶠu³ʲs, ᶜρaʲs) = p.precomputed
+        (; ᶜdetrʲs, ᶠu³⁰, ᶠu³ʲs, ᶜρaʲs) = p.precomputed
         ᶜtke_exch = p.scratch.ᶜtemp_scalar_2
         @. ᶜtke_exch = 0
         for j in 1:n

src/prognostic_equations/edmfx_sgs_flux.jl

@@ -318,7 +318,7 @@ function edmfx_sgs_mass_flux_tendency!(
         # TODO: the following adds the environment flux to the tendency
         # Make active and test later
         # @. ᶠu³_diff = p.precomputed.ᶠu³⁰ - ᶠu³
-        # ᶜρa⁰ = @.lazy(ρa⁰(Y.c))
+        # ρa⁰(Y.c.ρ, Y.c.sgsʲs, turbconv_model)
         # ᶜρ⁰ = p.scratch.ᶜtemp_scalar_2
         # @. ᶜρ⁰ = TD.air_density(
         #     CAP.thermodynamics_params(p.params),
@@ -393,7 +393,7 @@ function edmfx_sgs_diffusive_flux_tendency!(
     turbconv_params = CAP.turbconv_params(params)
     c_d = CAP.tke_diss_coeff(turbconv_params)
     (; ᶜu⁰, ᶜK⁰, ᶜlinear_buoygrad, ᶜstrain_rate_norm) = p.precomputed
-    (; ᶜmixing_length, ᶜK_u, ᶜK_h, ρatke_flux) = p.precomputed
+    (; ρatke_flux) = p.precomputed
     ᶠgradᵥ = Operators.GradientC2F()
     ᶜρa⁰ = @. lazy(ρa⁰(Y.c.ρ, Y.c.sgsʲs, turbconv_model))
     ᶜtke⁰ = @. lazy(specific_tke(Y.c.ρ, Y.c.sgs⁰.ρatke, ᶜρa⁰, turbconv_model))
@@ -520,13 +520,7 @@ function edmfx_sgs_diffusive_flux_tendency!(
     (; ᶜu, ᶜts) = p.precomputed
     (; ρatke_flux) = p.precomputed
     ᶠgradᵥ = Operators.GradientC2F()
-    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))
+    ᶜtke⁰ = @. lazy(specific_tke(Y.c.ρ, Y.c.sgs⁰.ρatke, Y.c.ρ, turbconv_model))
 
     if p.atmos.edmfx_model.sgs_diffusive_flux isa Val{true}
 

src/prognostic_equations/edmfx_tke.jl

@@ -27,9 +27,10 @@ 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, ᶜtke⁰) = p.precomputed
+    (; ᶜstrain_rate_norm, ᶜlinear_buoygrad) = p.precomputed
     turbconv_params = CAP.turbconv_params(p.params)
-
+    ᶜρa⁰ = Y.c.ρ # EDOnly
+    ᶜtke⁰ = @. lazy(specific_tke(Y.c.ρ, Y.c.sgs⁰.ρatke, ᶜρa⁰, turbconv_model))
     ᶜ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))
@@ -53,8 +54,7 @@ function edmfx_tke_tendency!(
 )
     n = n_mass_flux_subdomains(turbconv_model)
     (; ᶜturb_entrʲs, ᶜentrʲs, ᶜdetrʲs, ᶠu³ʲs) = p.precomputed
-    (; ᶠu³⁰, ᶠu³, ᶜstrain_rate_norm, ᶜlinear_buoygrad, ᶜK_u, ᶜK_h) =
-        p.precomputed
+    (; ᶠu³⁰, ᶠu³, ᶜstrain_rate_norm, ᶜlinear_buoygrad) = p.precomputed
     turbconv_params = CAP.turbconv_params(p.params)
     FT = eltype(p.params)
 
@@ -86,6 +86,11 @@ function edmfx_tke_tendency!(
 
         ᶜmixing_length_field = p.scratch.ᶜtemp_scalar_2
         ᶜmixing_length_field .= ᶜmixing_length(Y, p)
+        ᶜρa⁰ =
+            turbconv_model isa PrognosticEDMFX ?
+            (@. lazy(ρa⁰(Y.c.ρ, Y.c.sgsʲs, turbconv_model))) : Y.c.ρ
+        ᶜtke⁰ =
+            @. lazy(specific_tke(Y.c.ρ, Y.c.sgs⁰.ρatke, ᶜρa⁰, turbconv_model))
         ᶜK_u = @. lazy(
             eddy_viscosity(turbconv_params, ᶜtke⁰, ᶜmixing_length_field),
         )

src/prognostic_equations/hyperdiffusion.jl

@@ -117,15 +117,9 @@ 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
-        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))
+        ᶜρa⁰ =
+            turbconv_model isa PrognosticEDMFX ?
+            (@. lazy(ρa⁰(Y.c.ρ, Y.c.sgsʲs, turbconv_model))) : Y.c.ρ
         ᶜtke⁰ =
             @. lazy(specific_tke(Y.c.ρ, Y.c.sgs⁰.ρatke, ᶜρa⁰, turbconv_model))
         (; ᶜ∇²tke⁰) = p.hyperdiff

src/prognostic_equations/implicit/manual_sparse_jacobian.jl

@@ -559,6 +559,7 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t)
     end
 
     if use_derivative(diffusion_flag)
+        (; turbconv_model) = p.atmos
         turbconv_params = CAP.turbconv_params(params)
         FT = eltype(params)
         (; vertical_diffusion) = p.atmos
@@ -575,7 +576,13 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t)
             )
             ᶜK_u = ᶜK_h
         else
-            (; ᶜtke⁰, ᶜlinear_buoygrad, ᶜstrain_rate_norm) = p.precomputed
+            (; ᶜlinear_buoygrad, ᶜstrain_rate_norm) = p.precomputed
+            ᶜρa⁰ =
+                p.atmos.turbconv_model isa PrognosticEDMFX ?
+                (@. lazy(ρa⁰(Y.c.ρ, Y.c.sgsʲs, turbconv_model))) : Y.c.ρ
+            ᶜtke⁰ = @. lazy(
+                specific_tke(Y.c.ρ, Y.c.sgs⁰.ρatke, ᶜρa⁰, turbconv_model),
+            )
             ᶜmixing_length_field = p.scratch.ᶜtemp_scalar_3
             ᶜmixing_length_field .= ᶜmixing_length(Y, p)
             ᶜK_u = @. lazy(
@@ -658,11 +665,15 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t)
             turbconv_model = p.atmos.turbconv_model
             ᶜρa⁰ =
                 p.atmos.turbconv_model isa PrognosticEDMFX ?
-                (@.lazy(ρa⁰(Y.c))) : Y.c.ρ
-            ᶜtke⁰ = @. lazy(specific_tke(Y.c.sgs⁰, Y.c, turbconv_model))
+                (@. lazy(ρa⁰(Y.c.ρ, Y.c.sgsʲs, turbconv_model))) : Y.c.ρ
+            ᶜtke⁰ = @. lazy(
+                specific_tke(Y.c.ρ, Y.c.sgs⁰.ρatke, ᶜρa⁰, turbconv_model),
+            )
             ᶜρatke⁰ = Y.c.sgs⁰.ρatke
 
-            ᶜmixing_length_field = ᶜmixing_length(Y, p)
+            # scratch to prevent GPU Kernel parameter memory error
+            ᶜmixing_length_field = p.scratch.ᶜtemp_scalar_3
+            ᶜmixing_length_field .= ᶜmixing_length(Y, p)
 
             @inline tke_dissipation_rate_tendency(tke⁰, mixing_length) =
                 tke⁰ >= 0 ? c_d * sqrt(tke⁰) / mixing_length : 1 / float(dt)

src/utils/variable_manipulations.jl

@@ -328,8 +328,8 @@ function ᶜspecific_env_value(::Val{χ_name}, Y, p) where {χ_name}
         end
 
         ᶜρaχ⁰ = @. lazy(ᶜρχ - ᶜρaχʲs_sum)
-        # Denominator: ρa⁰ = ρ - Σ ρaʲ
-        ᶜρa⁰ = @. lazy(ρa⁰(Y.c.ρ, p.precomputed.ᶜρaʲs, turbconv_model))
+        # Denominator: ρa⁰ = ρ - Σ ρaʲ, assume ᶜρa⁰ = ρ
+        ᶜρa⁰ = Y.c.ρ
     end
 
     return @. lazy(specific(
@@ -342,23 +342,21 @@ function ᶜspecific_env_value(::Val{χ_name}, Y, p) where {χ_name}
 end
 
 """
-    ρa⁰(Y, p)
+    ρa⁰(ρ, sgsʲs, turbconv_model)
 
 Computes the environment area-weighted density (`ρa⁰`).
 
-This function uses the `ᶜenv_value` helper, which applies the domain
-decomposition principle (`GridMean = Environment + Sum(Drafts)`) to calculate
-the environment area-weighted density by subtracting the sum of all draft
-subdomain area-weighted densities (`ρaʲ`) from the grid-mean density (`ρ`).
+This function calculates the environment area-weighted density by subtracting the sum of all draft subdomain area-weighted densities (`ρaʲ`) from the grid-mean density (`ρ`), following the domain decomposition principle (`GridMean = Environment + Sum(Drafts)`).
 
 Arguments:
-- `ρ`: Grid-mean density
+- `ρ`: Grid-mean density.
 - `sgsʲs`: Iterable of draft subdomain quantities.
-            PrognosticEDMFX: Y.c.sgsʲs
-            DiagnosticEDMFX: p.precomputed.ᶜρaʲs
+    - For `PrognosticEDMFX`: typically `Y.c.sgsʲs`
+    - For `DiagnosticEDMFX`: typically `p.precomputed.ᶜρaʲs`
+- `turbconv_model`: The turbulence convection model (e.g., `PrognosticEDMFX`, `DiagnosticEDMFX`, or others).
 
 Returns:
-- The area-weighted density (`ρa⁰`).
+- The area-weighted density of the environment (`ρa⁰`).
 """
 
 function ρa⁰(ρ, sgsʲs, turbconv_model)
@@ -375,19 +373,20 @@ end
 
 
 """
-    ᶜspecific_tke(Y, p)
+    specific_tke(ρ, ρatke, ρa⁰, turbconv_model)
 
-Computes the specific turbulent kinetic energy (`tke`) in the environment (`tke⁰`).
+Computes the specific turbulent kinetic energy (TKE) in the environment.
 
-This is a specialized helper that encapsulates the call to the regularized
-`specific` function for the TKE variable. It provides `0` as the grid-scale
-fallback value (`ρχ_fallback`) in the limit of small environmental area
-fraction.
+This function returns the specific TKE of the environment by regularizing the
+area-weighted TKE density (`ρatke`) with the environment area-weighted density
+(`ρa⁰`). In the limit of vanishing environment area fraction, the fallback value
+for TKE is zero.
 
 Arguments:
-- `ρ`: 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.
+- `ρ`: Grid-mean density.
+- `ρatke`: Area-weighted TKE density in the environment.
+- `ρa⁰`: Area-weighted density of the environment.
+- `turbconv_model`: The turbulence convection model (e.g., `PrognosticEDMFX`, `DiagnosticEDMFX`, or others).
 
 Returns:
 - The specific TKE of the environment (`tke⁰`).
@@ -461,8 +460,8 @@ function ᶜspecific_env_mse(Y, p)
             @. ᶜρamseʲ_sum += ᶜρaʲ * ᶜmseʲ
         end
         ρa⁰mse⁰ = @. lazy(ᶜρmse - ᶜρamseʲ_sum)
-        # Denominator: ρa⁰ = ρ - Σ ρaʲ
-        ᶜρa⁰ = @. lazy(ρa⁰(Y.c.ρ, ᶜρaʲs, turbconv_model))
+        # Denominator: ρa⁰ = ρ - Σ ρaʲ, assume ᶜρa⁰ = ρ
+        ᶜρa⁰ = Y.c.ρ
     end
 
     return @. lazy(specific(ρa⁰mse⁰, ᶜρa⁰, ᶜρmse, Y.c.ρ, turbconv_model))