Improve buoyancy_gradients design

CliMA · Jan 16, 2024 · e9ea81a · e9ea81a
1 parent 97bf370
commit e9ea81a
Show file tree

Hide file tree

Showing 6 changed files with 94 additions and 73 deletions.
diff --git a/src/ClimaAtmos.jl b/src/ClimaAtmos.jl
@@ -1,6 +1,7 @@
 module ClimaAtmos
 
 using NVTX, Colors
+import Thermodynamics as TD
 
 include(joinpath("parameters", "Parameters.jl"))
 import .Parameters as CAP

diff --git a/src/cache/diagnostic_edmf_precomputed_quantities.jl b/src/cache/diagnostic_edmf_precomputed_quantities.jl
@@ -759,17 +759,12 @@ 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.liquid_specific_humidity(thermo_params, ᶜts),
-            TD.ice_specific_humidity(thermo_params, ᶜts),
-            TD.dry_pottemp(thermo_params, ᶜts),                               # θ_sat
-            TD.liquid_ice_pottemp(thermo_params, ᶜts),                        # θ_liq_ice_sat
+        EnvBuoyGradVars(
+            thermo_params,
+            ᶜts,
             projected_vector_data(
                 C3,
                 ᶜgradᵥ(ᶠinterp(TD.virtual_pottemp(thermo_params, ᶜts))),
@@ -781,9 +776,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
@@ -239,17 +239,12 @@ 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.liquid_specific_humidity(thermo_params, ᶜts⁰),
-            TD.ice_specific_humidity(thermo_params, ᶜts⁰),
-            TD.dry_pottemp(thermo_params, ᶜts⁰),                               # θ_sat
-            TD.liquid_ice_pottemp(thermo_params, ᶜts⁰),                        # θ_liq_ice_sat
+        EnvBuoyGradVars(
+            thermo_params,
+            ᶜts⁰,
             projected_vector_data(
                 C3,
                 ᶜgradᵥ(ᶠinterp(TD.virtual_pottemp(thermo_params, ᶜts⁰))),
@@ -261,9 +256,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,26 +1,29 @@
+import Thermodynamics.Parameters as TDP
 """
     buoyancy_gradients(
-        params,
+        ::AbstractEnvBuoyGradClosure,
+        thermo_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(
-    params,
+    ebgc::AbstractEnvBuoyGradClosure,
+    thermo_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)
-    molmass_ratio = CAP.molmass_ratio(params)
-    R_d = CAP.R_d(params)
-    R_v = CAP.R_v(params)
+    g = TDP.grav(thermo_params)
+    molmass_ratio = TDP.molmass_ratio(thermo_params)
+    R_d = TDP.R_d(thermo_params)
+    R_v = TDP.R_v(thermo_params)
 
     phase_part = TD.PhasePartition(FT(0), FT(0), FT(0)) # assuming R_d = R_m
     Π = TD.exner_given_pressure(thermo_params, bg_model.p, phase_part)
@@ -85,27 +88,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/prognostic_equations/gm_sgs_closures.jl b/src/prognostic_equations/gm_sgs_closures.jl
@@ -33,35 +33,27 @@ function compute_gm_mixing_length!(ᶜmixing_length, Y, p)
 
     ᶜlinear_buoygrad = p.scratch.ᶜtemp_scalar
     @. ᶜlinear_buoygrad = buoyancy_gradients(
-        params,
+        BuoyGradMean(),
+        thermo_params,
         p.atmos.moisture_model,
-        EnvBuoyGrad(
-            BuoyGradMean(),
-            TD.air_temperature(thermo_params, ᶜts),           # t_sat
-            TD.vapor_specific_humidity(thermo_params, ᶜts),   # qv_sat
-            TD.total_specific_humidity(thermo_params, ᶜts),   # qt_sat
-            TD.liquid_specific_humidity(thermo_params, ᶜts),  # q_liq
-            TD.ice_specific_humidity(thermo_params, ᶜts),     # q_ice
-            TD.dry_pottemp(thermo_params, ᶜts),               # θ_sat
-            TD.liquid_ice_pottemp(thermo_params, ᶜts),        # θ_liq_ice_sat
+        EnvBuoyGradVars(
+            thermo_params,
+            ᶜts,
             projected_vector_data(
                 C3,
                 ᶜgradᵥ(ᶠinterp(TD.virtual_pottemp(thermo_params, ᶜts))),
                 ᶜlg,
-            ),                                               # ∂θv∂z_unsat
+            ),                                                                 # ∂θv∂z_unsat
             projected_vector_data(
                 C3,
                 ᶜgradᵥ(ᶠinterp(TD.total_specific_humidity(thermo_params, ᶜts))),
                 ᶜlg,
-            ),                                               # ∂qt∂z_sat
+            ),            # ∂qt∂z_sat
             projected_vector_data(
                 C3,
                 ᶜ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.ρ,                                           # ρ
+            ),                                                                 # ∂θl∂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
@@ -122,11 +123,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"
@@ -141,27 +142,59 @@ 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(
+    thermo_params,
+    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)
+    ql_sat = TD.liquid_specific_humidity(thermo_params, ts)
+    qi_sat = TD.ice_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,
+        ql_sat,
+        qi_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