Improve buoyancy_gradients design

CliMA · Nov 15, 2023 · 40c0e37 · 40c0e37
1 parent e293677
commit 40c0e37
Show file tree

Hide file tree

Showing 4 changed files with 71 additions and 46 deletions.
diff --git a/src/cache/diagnostic_edmf_precomputed_quantities.jl b/src/cache/diagnostic_edmf_precomputed_quantities.jl
@@ -682,15 +682,11 @@ function set_diagnostic_edmf_precomputed_quantities_env_closures!(Y, p, t)
     @. ᶜu⁰ = C123(Y.c.uₕ) + ᶜinterp(C123(ᶠu³⁰))
 
     @. ᶜlinear_buoygrad = buoyancy_gradients(
-        params,
+        BuoyGradMean(),
+        thermo_params,
         moisture_model,
-        EnvBuoyGrad(
-            BuoyGradMean(),
-            TD.air_temperature(thermo_params, ᶜts),                           # t_sat
-            TD.vapor_specific_humidity(thermo_params, ᶜts),                   # qv_sat
-            q_tot,                                                            # qt_sat
-            TD.dry_pottemp(thermo_params, ᶜts),                               # θ_sat
-            TD.liquid_ice_pottemp(thermo_params, ᶜts),                        # θ_liq_ice_sat
+        EnvBuoyGradVars(
+            ᶜts,
             projected_vector_data(
                 C3,
                 ᶜgradᵥ(ᶠinterp(TD.virtual_pottemp(thermo_params, ᶜts))),
@@ -702,9 +698,6 @@ function set_diagnostic_edmf_precomputed_quantities_env_closures!(Y, p, t)
                 ᶜgradᵥ(ᶠinterp(TD.liquid_ice_pottemp(thermo_params, ᶜts))),
                 ᶜlg,
             ),                                                                 # ∂θl∂z_sat
-            ᶜp,                                                                # p
-            ifelse(TD.has_condensate(thermo_params, ᶜts), 1, 0),               # en_cld_frac
-            Y.c.ρ,                                                             # ρ
         ),
     )
 

diff --git a/src/cache/prognostic_edmf_precomputed_quantities.jl b/src/cache/prognostic_edmf_precomputed_quantities.jl
@@ -233,15 +233,11 @@ function set_prognostic_edmf_precomputed_quantities_closures!(Y, p, t)
 
     # First order approximation: Use environmental mean fields.
     @. ᶜlinear_buoygrad = buoyancy_gradients(
-        params,
+        BuoyGradMean(),
+        thermo_params,
         moisture_model,
-        EnvBuoyGrad(
-            BuoyGradMean(),
-            TD.air_temperature(thermo_params, ᶜts⁰),                           # t_sat
-            TD.vapor_specific_humidity(thermo_params, ᶜts⁰),                   # qv_sat
-            ᶜq_tot⁰,                                                           # qt_sat
-            TD.dry_pottemp(thermo_params, ᶜts⁰),                               # θ_sat
-            TD.liquid_ice_pottemp(thermo_params, ᶜts⁰),                        # θ_liq_ice_sat
+        EnvBuoyGradVars(
+            ᶜts⁰,
             projected_vector_data(
                 C3,
                 ᶜgradᵥ(ᶠinterp(TD.virtual_pottemp(thermo_params, ᶜts⁰))),
@@ -253,9 +249,6 @@ function set_prognostic_edmf_precomputed_quantities_closures!(Y, p, t)
                 ᶜgradᵥ(ᶠinterp(TD.liquid_ice_pottemp(thermo_params, ᶜts⁰))),
                 ᶜlg,
             ),                                                                 # ∂θl∂z_sat
-            ᶜp,                                                                # p
-            ifelse(TD.has_condensate(thermo_params, ᶜts⁰), 1, 0),              # en_cld_frac
-            ᶜρ⁰,                                                               # ρ
         ),
     )
 

diff --git a/src/prognostic_equations/buoyancy_gradients.jl b/src/prognostic_equations/buoyancy_gradients.jl
@@ -1,20 +1,23 @@
 """
     buoyancy_gradients(
+        ::AbstractEnvBuoyGradClosure,
         params,
         moisture_model,
-        bg_model::EnvBuoyGrad{FT, EBG}
-    ) where {FT <: Real, EBG <: AbstractEnvBuoyGradClosure}
+        bg_model::EnvBuoyGradVars,
+    )
 
 Returns the vertical buoyancy gradients in the environment, as well as in its dry and cloudy volume fractions.
-The dispatch on EnvBuoyGrad type is performed at the EnvBuoyGrad construction time, and the analytical solutions
+The dispatch on EnvBuoyGradVars type is performed at the EnvBuoyGradVars construction time, and the analytical solutions
 used here are consistent for both mean fields and conditional fields obtained from assumed distributions
 over the conserved thermodynamic variables.
 """
 function buoyancy_gradients(
+    ebgc::AbstractEnvBuoyGradClosure,
     params,
     moisture_model,
-    bg_model::EnvBuoyGrad{FT, EBG},
-) where {FT <: Real, EBG <: AbstractEnvBuoyGradClosure}
+    bg_model::EnvBuoyGradVars,
+)
+    FT = eltype(bg_model)
 
     thermo_params = CAP.thermodynamics_params(params)
     g = CAP.grav(params)
@@ -63,27 +66,35 @@ function buoyancy_gradients(
         ∂b∂qt_sat = FT(0)
     end
 
-    ∂b∂z = buoyancy_gradient_chain_rule(bg_model, ∂b∂θv, ∂b∂θl_sat, ∂b∂qt_sat)
+    ∂b∂z = buoyancy_gradient_chain_rule(
+        ebgc,
+        bg_model,
+        ∂b∂θv,
+        ∂b∂θl_sat,
+        ∂b∂qt_sat,
+    )
     return ∂b∂z
 end
 
 """
     buoyancy_gradient_chain_rule(
-        bg_model::EnvBuoyGrad{FT, EBG},
+        ::AbstractEnvBuoyGradClosure,
+        bg_model::EnvBuoyGradVars{FT, EBG},
         ∂b∂θv::FT,
         ∂b∂θl_sat::FT,
         ∂b∂qt_sat::FT,
-    ) where {FT <: Real, EBG <: AbstractEnvBuoyGradClosure}
+    ) where {FT <: Real}
 
 Returns the vertical buoyancy gradients in the environment, as well as in its dry and cloudy volume fractions,
 from the partial derivatives with respect to thermodynamic variables in dry and cloudy volumes.
 """
 function buoyancy_gradient_chain_rule(
-    bg_model::EnvBuoyGrad{FT, EBG},
+    ::AbstractEnvBuoyGradClosure,
+    bg_model::EnvBuoyGradVars{FT},
     ∂b∂θv::FT,
     ∂b∂θl_sat::FT,
     ∂b∂qt_sat::FT,
-) where {FT <: Real, EBG <: AbstractEnvBuoyGradClosure}
+) where {FT <: Real}
     if bg_model.en_cld_frac > FT(0)
         ∂b∂z_θl_sat = ∂b∂θl_sat * bg_model.∂θl∂z_sat
         ∂b∂z_qt_sat = ∂b∂qt_sat * bg_model.∂qt∂z_sat

diff --git a/src/solver/types.jl b/src/solver/types.jl
@@ -1,5 +1,6 @@
 import FastGaussQuadrature
 import StaticArrays as SA
+import Thermodynamics as TD
 
 abstract type AbstractMoistureModel end
 struct DryModel <: AbstractMoistureModel end
@@ -110,11 +111,11 @@ struct BuoyGradMean <: AbstractEnvBuoyGradClosure end
 Base.broadcastable(x::BuoyGradMean) = tuple(x)
 
 """
-    EnvBuoyGrad
+    EnvBuoyGradVars
 
 Variables used in the environmental buoyancy gradient computation.
 """
-Base.@kwdef struct EnvBuoyGrad{FT, EBC <: AbstractEnvBuoyGradClosure}
+Base.@kwdef struct EnvBuoyGradVars{FT}
     "temperature in the saturated part"
     t_sat::FT
     "vapor specific humidity  in the saturated part"
@@ -125,27 +126,54 @@ Base.@kwdef struct EnvBuoyGrad{FT, EBC <: AbstractEnvBuoyGradClosure}
     θ_sat::FT
     "liquid ice potential temperature in the saturated part"
     θ_liq_ice_sat::FT
-    "virtual potential temperature gradient in the non saturated part"
-    ∂θv∂z_unsat::FT
-    "total specific humidity gradient in the saturated part"
-    ∂qt∂z_sat::FT
-    "liquid ice potential temperature gradient in the saturated part"
-    ∂θl∂z_sat::FT
     "reference pressure"
     p::FT
     "cloud fraction"
     en_cld_frac::FT
     "density"
     ρ::FT
+    "virtual potential temperature gradient in the non saturated part"
+    ∂θv∂z_unsat::FT
+    "total specific humidity gradient in the saturated part"
+    ∂qt∂z_sat::FT
+    "liquid ice potential temperature gradient in the saturated part"
+    ∂θl∂z_sat::FT
 end
-function EnvBuoyGrad(
-    ::EBG,
-    t_sat::FT,
-    args...,
-) where {FT <: Real, EBG <: AbstractEnvBuoyGradClosure}
-    return EnvBuoyGrad{FT, EBG}(t_sat, args...)
+
+
+function EnvBuoyGradVars(
+    ts::TD.ThermodynamicState,
+    ∂θv∂z_unsat::FT,
+    ∂qt∂z_sat::FT,
+    ∂θl∂z_sat::FT,
+) where {FT}
+    t_sat = TD.air_temperature(thermo_params, ᶜts)
+    qv_sat = TD.vapor_specific_humidity(thermo_params, ᶜts)
+    qt_sat = TD.total_specific_humidity(thermo_params, ᶜts)
+    θ_sat = TD.dry_pottemp(thermo_params, ᶜts)
+    θ_liq_ice_sat = TD.liquid_ice_pottemp(thermo_params, ᶜts)
+    p = TD.air_pressure(thermo_params, ᶜts)
+    en_cld_frac = ifelse(TD.has_condensate(thermo_params, ᶜts), 1, 0)
+    ρ = TD.air_density(thermo_params, ᶜts)
+    return EnvBuoyGradVars{FT}(
+        t_sat,
+        qv_sat,
+        qt_sat,
+        θ_sat,
+        θ_liq_ice_sat,
+        p,
+        en_cld_frac,
+        ρ,
+        ∂θv∂z_unsat,
+        ∂qt∂z_sat,
+        ∂θl∂z_sat,
+    )
 end
 
+Base.eltype(::EnvBuoyGradVars{FT}) where {FT} = FT
+Base.broadcastable(x::EnvBuoyGradVars) = tuple(x)
+
+
 abstract type AbstractEDMF end
 
 struct PrognosticEDMFX{N, TKE, FT} <: AbstractEDMF