From 38db7cb81389af0b44a396f427ade18c7f9a4f0f Mon Sep 17 00:00:00 2001
From: Anna Jaruga <ajaruga@caltech.edu>
Date: Fri, 21 Jun 2024 18:17:58 -0700
Subject: [PATCH] Add sgs cloud diagnostics for diagnostic edmf

---
 config/default_configs/default_config.yml     |  2 +-
 .../diagnostic_edmfx_bomex_box.yml            |  1 +
 src/cache/cloud_fraction.jl                   | 53 ++++++++++++++++++-
 src/solver/model_getters.jl                   |  2 +
 src/solver/types.jl                           |  1 +
 5 files changed, 57 insertions(+), 2 deletions(-)

diff --git a/config/default_configs/default_config.yml b/config/default_configs/default_config.yml
index 5bf0989c10e..3f2b5d365c2 100644
--- a/config/default_configs/default_config.yml
+++ b/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"
diff --git a/config/model_configs/diagnostic_edmfx_bomex_box.yml b/config/model_configs/diagnostic_edmfx_bomex_box.yml
index 5ab746b408a..29b5dcc00c5 100644
--- a/config/model_configs/diagnostic_edmfx_bomex_box.yml
+++ b/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
diff --git a/src/cache/cloud_fraction.jl b/src/cache/cloud_fraction.jl
index 6eefe0c476e..ed357f251d7 100644
--- a/src/cache/cloud_fraction.jl
+++ b/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,57 @@ 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, ᶜ∇θ,
diff --git a/src/solver/model_getters.jl b/src/solver/model_getters.jl
index 8a604bf44ae..307a3178a57 100644
--- a/src/solver/model_getters.jl
+++ b/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
diff --git a/src/solver/types.jl b/src/solver/types.jl
index afbdd6f44b8..c0db60d3d02 100644
--- a/src/solver/types.jl
+++ b/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