try adding better cloud diagnostics for diagnostic edmf

config/default_configs/default_config.yml

@@ -191,7 +191,7 @@ precip_model:
   help: "Precipitation model [`nothing` (default), `0M`]"
   value: ~
 cloud_model:
-  help: "Cloud model [`grid_scale`, `quadrature` (default)]"
+  help: "Cloud model [`grid_scale`, `quadrature` (default), `diagnostic_edmf`]"
   value: "quadrature"
 perf_summary:
   help: "Flag for collecting performance summary information"

config/model_configs/diagnostic_edmfx_bomex_box.yml

@@ -16,6 +16,7 @@ edmfx_sgs_diffusive_flux: true
 moist: "equil"
 call_cloud_diagnostics_per_stage: true
 config: "box"
+cloud_model: "diagnostic_edmf"
 x_max: 1e8
 y_max: 1e8
 x_elem: 2

src/cache/cloud_fraction.jl

@@ -66,7 +66,7 @@ NVTX.@annotate function set_cloud_fraction!(
 
     FT = eltype(params)
     thermo_params = CAP.thermodynamics_params(params)
-    (; ᶜts, ᶜp, ᶜmixing_length, cloud_diagnostics_tuple) = p.precomputed
+    (; ᶜts, ᶜmixing_length, cloud_diagnostics_tuple) = p.precomputed
     (; turbconv_model) = p.atmos
     if isnothing(turbconv_model)
         if p.atmos.call_cloud_diagnostics_per_stage isa
@@ -94,6 +94,55 @@ NVTX.@annotate function set_cloud_fraction!(
         thermo_params,
     )
 end
+NVTX.@annotate function set_cloud_fraction!(
+    Y,
+    p,
+    ::Union{EquilMoistModel, NonEquilMoistModel},
+    ::DiagnosticEDMFCloud,
+)
+    (; SG_quad, params) = p
+
+    FT = eltype(params)
+    thermo_params = CAP.thermodynamics_params(params)
+    (; ᶜts, ᶜmixing_length, cloud_diagnostics_tuple) = p.precomputed
+    (; ᶜρaʲs, ᶜρʲs, ᶜtsʲs) = p.precomputed
+    (; turbconv_model) = p.atmos
+
+    if p.atmos.call_cloud_diagnostics_per_stage isa CallCloudDiagnosticsPerStage
+        (; ᶜgradᵥ_θ_virt, ᶜgradᵥ_q_tot, ᶜgradᵥ_θ_liq_ice) = p.precomputed
+        thermo_params = CAP.thermodynamics_params(p.params)
+        @. ᶜgradᵥ_θ_virt =
+            ᶜgradᵥ(ᶠinterp(TD.virtual_pottemp(thermo_params, ᶜts)))
+        @. ᶜgradᵥ_q_tot =
+            ᶜgradᵥ(ᶠinterp(TD.total_specific_humidity(thermo_params, ᶜts)))
+        @. ᶜgradᵥ_θ_liq_ice =
+            ᶜgradᵥ(ᶠinterp(TD.liquid_ice_pottemp(thermo_params, ᶜts)))
+    end
+
+    # environment
+    coeff = FT(2.1) # TODO - move to parameters
+    @. cloud_diagnostics_tuple = quad_loop(
+        SG_quad,
+        ᶜts,
+        Geometry.WVector(p.precomputed.ᶜgradᵥ_q_tot),
+        Geometry.WVector(p.precomputed.ᶜgradᵥ_θ_liq_ice),
+        coeff,
+        ᶜmixing_length,
+        thermo_params,
+    )
+    # updrafts
+    n = n_mass_flux_subdomains(p.atmos.turbconv_model)
+
+    for j in 1:n
+        @. cloud_diagnostics_tuple += NamedTuple{(:cf, :q_liq, :q_ice)}(
+            tuple(
+                ifelse(TD.has_condensate(thermo_params, ᶜtsʲs.:($$j)), 1, 0),
+                draft_area(ᶜρaʲs.:($$j), ᶜρʲs.:($$j)) * TD.PhasePartition(thermo_params, ᶜtsʲs.:($$j)).liq,
+                draft_area(ᶜρaʲs.:($$j), ᶜρʲs.:($$j)) * TD.PhasePartition(thermo_params, ᶜtsʲs.:($$j)).ice,
+            ),
+        )
+    end
+end
 
 """
     function quad_loop(SG_quad, ts, ᶜ∇q, ᶜ∇θ,

src/solver/model_getters.jl

@@ -272,6 +272,8 @@ function get_cloud_model(parsed_args)
         GridScaleCloud()
     elseif cloud_model == "quadrature"
         QuadratureCloud()
+    elseif cloud_model == "diagnostic_edmf"
+        DiagnosticEDMFCloud()
     else
         error("Invalid cloud_model $(cloud_model)")
     end

src/solver/types.jl

@@ -15,6 +15,7 @@ struct Microphysics1Moment <: AbstractPrecipitationModel end
 abstract type AbstractCloudModel end
 struct GridScaleCloud <: AbstractCloudModel end
 struct QuadratureCloud <: AbstractCloudModel end
+struct DiagnosticEDMFCloud <: AbstractCloudModel end
 
 abstract type AbstractModelConfig end
 struct SingleColumnModel <: AbstractModelConfig end