Merge pull request #4124 from CliMA/zs/buoyancy_gradient2

use cloud fraction in buoyancy gradient

Author: szy21

docs/src/restarts.md

@@ -24,6 +24,11 @@ particularly useful for
     if that is not the case. When the warning is produced, it is your responsability
     to ensure that what you are doing makes sense.
 
+!!! note
+
+    The simulation can only be restarted in a reproducible way when using `GridScaleCloud`
+    for `cloud_model`.
+
 ### How Restarts Work
 
 `ClimaAtmos` periodically saves the simulation state to a file called a "restart

reproducibility_tests/ref_counter.jl

@@ -1,4 +1,4 @@
-278
+279
 
 # **README**
 #
@@ -20,6 +20,9 @@
 
 
 #=
+279
+- Use partial cloud fraction in buoyancy gradient calculation
+
 278
 - Add ∂/∂q elements to Jacobian
 

src/cache/cloud_fraction.jl

@@ -22,7 +22,6 @@ end
    Compute the grid scale cloud fraction based on sub-grid scale properties
 """
 NVTX.@annotate function set_cloud_fraction!(Y, p, ::DryModel, _)
-    (; turbconv_model) = p.atmos
     FT = eltype(p.params)
 
     p.precomputed.cloud_diagnostics_tuple .=
@@ -35,23 +34,9 @@ NVTX.@annotate function set_cloud_fraction!(
     ::GridScaleCloud,
 )
     (; params) = p
-    (; turbconv_model) = p.atmos
     (; ᶜts, cloud_diagnostics_tuple) = p.precomputed
     thermo_params = CAP.thermodynamics_params(params)
 
-    if isnothing(turbconv_model)
-        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
-    end
     if moist_model isa EquilMoistModel
         @. cloud_diagnostics_tuple = make_named_tuple(
             ifelse(TD.has_condensate(thermo_params, ᶜts), 1, 0),
@@ -83,10 +68,8 @@ NVTX.@annotate function set_cloud_fraction!(
     if isnothing(turbconv_model)
         if p.atmos.call_cloud_diagnostics_per_stage isa
            CallCloudDiagnosticsPerStage
-            (; ᶜgradᵥ_θ_virt, ᶜgradᵥ_q_tot, ᶜgradᵥ_θ_liq_ice) = p.precomputed
+            (; ᶜ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 =

src/cache/diagnostic_edmf_precomputed_quantities.jl

@@ -1339,14 +1339,9 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_env_closures!
     p,
     t,
 )
-    (; moisture_model, turbconv_model, microphysics_model) = p.atmos
-    n = n_mass_flux_subdomains(turbconv_model)
-    ᶜz = Fields.coordinate_field(Y.c).z
-    ᶜdz = Fields.Δz_field(axes(Y.c))
     (; params) = p
-    (; dt) = p
-    (; ᶜp, ᶜu, ᶜts, ᶠu³) = p.precomputed
-    (; ustar, obukhov_length) = p.precomputed.sfc_conditions
+    (; ᶜu, ᶜts, ᶠu³) = p.precomputed
+    (; ustar) = p.precomputed.sfc_conditions
     (; ᶜlinear_buoygrad, ᶜstrain_rate_norm) = p.precomputed
     (; ρatke_flux) = p.precomputed
     turbconv_params = CAP.turbconv_params(params)
@@ -1364,12 +1359,11 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_env_closures!
     @. ᶜlinear_buoygrad = buoyancy_gradients(
         BuoyGradMean(),
         thermo_params,
-        moisture_model,
         ᶜts,
+        p.precomputed.cloud_diagnostics_tuple.cf,
         C3,
-        p.precomputed.ᶜgradᵥ_θ_virt,    # ∂θv∂z_unsat
-        p.precomputed.ᶜgradᵥ_q_tot,     # ∂qt∂z_sat
-        p.precomputed.ᶜgradᵥ_θ_liq_ice, # ∂θl∂z_sat
+        p.precomputed.ᶜgradᵥ_q_tot,
+        p.precomputed.ᶜgradᵥ_θ_liq_ice,
         ᶜlg,
     )
 
@@ -1381,10 +1375,7 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_env_closures!
 
     ρatke_flux_values = Fields.field_values(ρatke_flux)
     ρ_int_values = Fields.field_values(Fields.level(Y.c.ρ, 1))
-    u_int_values = Fields.field_values(Fields.level(ᶜu, 1))
     ustar_values = Fields.field_values(ustar)
-    int_local_geometry_values =
-        Fields.field_values(Fields.level(Fields.local_geometry_field(Y.c), 1))
     sfc_local_geometry_values = Fields.field_values(
         Fields.level(Fields.local_geometry_field(Y.f), half),
     )

src/cache/precomputed_quantities.jl

@@ -85,7 +85,6 @@ function precomputed_quantities(Y, atmos)
     TST = thermo_state_type(atmos.moisture_model, FT)
     SCT = SurfaceConditions.surface_conditions_type(atmos, FT)
     cspace = axes(Y.c)
-    fspace = axes(Y.f)
     n = n_mass_flux_subdomains(atmos.turbconv_model)
     n_prog = n_prognostic_mass_flux_subdomains(atmos.turbconv_model)
     @assert !(atmos.turbconv_model isa PrognosticEDMFX) || n_prog == 1
@@ -98,6 +97,8 @@ function precomputed_quantities(Y, atmos)
     )
     cloud_diagnostics_tuple =
         similar(Y.c, @NamedTuple{cf::FT, q_liq::FT, q_ice::FT})
+    # Cloud fraction is used to calculate buoyancy gradient, so we initialize it to 0 here.
+    @. cloud_diagnostics_tuple.cf = FT(0)
     surface_precip_fluxes = (;
         surface_rain_flux = zeros(axes(Fields.level(Y.f, half))),
         surface_snow_flux = zeros(axes(Fields.level(Y.f, half))),
@@ -200,7 +201,6 @@ function precomputed_quantities(Y, atmos)
             ᶜentrʲs = similar(Y.c, NTuple{n, FT}),
             ᶜdetrʲs = similar(Y.c, NTuple{n, FT}),
             ᶜturb_entrʲs = similar(Y.c, NTuple{n, FT}),
-            ᶜgradᵥ_θ_virt = Fields.Field(C3{FT}, cspace),
             ᶜgradᵥ_q_tot = Fields.Field(C3{FT}, cspace),
             ᶜgradᵥ_θ_liq_ice = Fields.Field(C3{FT}, cspace),
             ᶠnh_pressure₃_buoyʲs = similar(Y.f, NTuple{n, C3{FT}}),
@@ -212,7 +212,6 @@ function precomputed_quantities(Y, atmos)
         (; ρatke_flux = similar(Fields.level(Y.f, half), C3{FT}),) : (;)
 
     sgs_quantities = (;
-        ᶜgradᵥ_θ_virt = Fields.Field(C3{FT}, cspace),
         ᶜgradᵥ_q_tot = Fields.Field(C3{FT}, cspace),
         ᶜgradᵥ_θ_liq_ice = Fields.Field(C3{FT}, cspace),
     )
@@ -548,8 +547,6 @@ NVTX.@annotate function set_explicit_precomputed_quantities_part1!(Y, p, t)
     end
 
     if turbconv_model isa AbstractEDMF
-        @. p.precomputed.ᶜgradᵥ_θ_virt =
-            ᶜgradᵥ(ᶠinterp(TD.virtual_pottemp(thermo_params, ᶜts)))
         @. p.precomputed.ᶜgradᵥ_q_tot =
             ᶜgradᵥ(ᶠinterp(TD.total_specific_humidity(thermo_params, ᶜts)))
         @. p.precomputed.ᶜgradᵥ_θ_liq_ice =
@@ -568,6 +565,15 @@ NVTX.@annotate function set_explicit_precomputed_quantities_part2!(Y, p, t)
     (; turbconv_model, moisture_model, cloud_model) = p.atmos
     (; call_cloud_diagnostics_per_stage) = p.atmos
 
+    # The buoyancy gradient depends on the cloud fraction, and the cloud fraction 
+    # depends on the mixing length, which depends on the buoyancy gradient. 
+    # We break this circular dependency by using cloud fraction from the previous time step in the 
+    # buoyancy gradient calculation. This breaks reproducible restart in general, 
+    # but we support reproducible restart with grid-scale cloud by recalculating the cloud fraction here.
+    if p.atmos.cloud_model isa GridScaleCloud
+        set_cloud_fraction!(Y, p, p.atmos.moisture_model, p.atmos.cloud_model)
+    end
+
     if turbconv_model isa PrognosticEDMFX
         set_prognostic_edmf_precomputed_quantities_explicit_closures!(Y, p, t)
         set_prognostic_edmf_precomputed_quantities_precipitation!(

src/cache/prognostic_edmf_precomputed_quantities.jl

@@ -164,7 +164,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_explicit_clos
     t,
 )
 
-    (; moisture_model, turbconv_model) = p.atmos
+    (; turbconv_model) = p.atmos
 
     (; params) = p
     (; dt) = p
@@ -299,14 +299,14 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_explicit_clos
         )
     end
 
-    (; ᶜgradᵥ_θ_virt, ᶜgradᵥ_q_tot, ᶜgradᵥ_θ_liq_ice) = p.precomputed
+    (; ᶜgradᵥ_q_tot, ᶜgradᵥ_θ_liq_ice, cloud_diagnostics_tuple) =
+        p.precomputed
     @. ᶜlinear_buoygrad = buoyancy_gradients( # TODO - do we need to modify buoyancy gradients based on NonEq + 1M tracers?
         BuoyGradMean(),
         thermo_params,
-        moisture_model,
         ᶜts,
+        cloud_diagnostics_tuple.cf,
         C3,
-        ᶜgradᵥ_θ_virt,
         ᶜgradᵥ_q_tot,
         ᶜgradᵥ_θ_liq_ice,
         ᶜlg,

src/callbacks/callbacks.jl

@@ -106,16 +106,12 @@ function external_driven_single_column!(integrator)
     evaluate!(ᶜls_subsidence, wa, t)
 end
 
-
-
 NVTX.@annotate function cloud_fraction_model_callback!(integrator)
     Y = integrator.u
     p = integrator.p
-    (; ᶜts, ᶜgradᵥ_θ_virt, ᶜgradᵥ_q_tot, ᶜgradᵥ_θ_liq_ice) = p.precomputed
+    (; ᶜts, ᶜgradᵥ_q_tot, ᶜgradᵥ_θ_liq_ice) = p.precomputed
     thermo_params = CAP.thermodynamics_params(p.params)
     if isnothing(p.atmos.turbconv_model)
-        @. ᶜ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 =

src/prognostic_equations/eddy_diffusion_closures.jl

@@ -11,13 +11,12 @@ import ClimaCore.Fields as Fields
     buoyancy_gradients(
         closure::AbstractEnvBuoyGradClosure,
         thermo_params,
-        moisture_model,
+
         # Arguments for the first method (most commonly called):
         ts::TD.ThermodynamicState,
         ::Type{C3}, # Covariant3 vector type, for projecting gradients
-        ∂θv∂z_unsat::AbstractField, # Vertical gradient of virtual potential temperature in unsaturated part
-        ∂qt∂z_sat::AbstractField,   # Vertical gradient of total specific humidity in saturated part
-        ∂θli∂z_sat::AbstractField,   # Vertical gradient of liquid-ice potential temperature in saturated part
+        ∂qt∂z::AbstractField,   # Vertical gradient of total specific humidity
+        ∂θli∂z::AbstractField,   # Vertical gradient of liquid-ice potential temperature
         local_geometry::Fields.LocalGeometry,
         # Argument for the second method (internal use with precomputed EnvBuoyGradVars):
         # bg_model::EnvBuoyGradVars
@@ -47,12 +46,10 @@ variables.
 Arguments:
 - `closure`: The environmental buoyancy gradient closure type (e.g., `BuoyGradMean`).
 - `thermo_params`: Thermodynamic parameters from `CLIMAParameters`.
-- `moisture_model`: Moisture model (e.g., `EquilMoistModel`, `NonEquilMoistModel`).
 - `ts`: Center-level thermodynamic state of the environment.
 - `C3`: The `ClimaCore.Geometry.Covariant3Vector` type, used for projecting input vertical gradients.
-- `∂θv∂z_unsat`: Field of vertical gradients of virtual potential temperature in the unsaturated part.
-- `∂qt∂z_sat`: Field of vertical gradients of total specific humidity in the saturated part.
-- `∂θli∂z_sat`: Field of vertical gradients of liquid-ice potential temperature in the saturated part.
+- `∂qt∂z`: Field of vertical gradients of total specific humidity.
+- `∂θli∂z`: Field of vertical gradients of liquid-ice potential temperature.
 - `local_geometry`: Field of local geometry at cell centers, used for gradient projection.
 The second method takes a precomputed `EnvBuoyGradVars` object instead of `ts` and gradient fields.
 
@@ -64,25 +61,23 @@ function buoyancy_gradients end
 function buoyancy_gradients(
     ebgc::AbstractEnvBuoyGradClosure,
     thermo_params,
-    moisture_model,
     ts,
+    cf,
     ::Type{C3},
-    ∂θv∂z_unsat,
-    ∂qt∂z_sat,
-    ∂θli∂z_sat,
+    ∂qt∂z,
+    ∂θli∂z,
     ᶜlg,
 ) where {C3}
     return buoyancy_gradients(
         ebgc,
         thermo_params,
-        moisture_model,
         EnvBuoyGradVars(
             ts,
+            cf,
             projected_vector_buoy_grad_vars(
                 C3,
-                ∂θv∂z_unsat,
-                ∂qt∂z_sat,
-                ∂θli∂z_sat,
+                ∂qt∂z,
+                ∂θli∂z,
                 ᶜlg,
             ),
         ),
@@ -92,43 +87,42 @@ end
 function buoyancy_gradients(
     ebgc::AbstractEnvBuoyGradClosure,
     thermo_params,
-    moisture_model,
     bg_model::EnvBuoyGradVars,
 )
     FT = eltype(bg_model)
 
     g = TDP.grav(thermo_params)
     Rv_over_Rd = TDP.Rv_over_Rd(thermo_params)
-    R_d = TDP.R_d(thermo_params)
     R_v = TDP.R_v(thermo_params)
 
     ts = bg_model.ts
-    p = TD.air_pressure(thermo_params, ts)
-    Π = TD.exner_given_pressure(thermo_params, p)
-    ∂b∂θv = g * (R_d * TD.air_density(thermo_params, ts) / p) * Π
+    ∂b∂θv = g / TD.virtual_pottemp(thermo_params, ts)
 
-    t_sat = TD.air_temperature(thermo_params, ts)
+    T = TD.air_temperature(thermo_params, ts)
     λ = TD.liquid_fraction(thermo_params, ts)
     lh =
-        λ * TD.latent_heat_vapor(thermo_params, t_sat) +
-        (1 - λ) * TD.latent_heat_sublim(thermo_params, t_sat)
+        λ * TD.latent_heat_vapor(thermo_params, T) +
+        (1 - λ) * TD.latent_heat_sublim(thermo_params, T)
     cp_m = TD.cp_m(thermo_params, ts)
-    qv_sat = TD.vapor_specific_humidity(thermo_params, ts)
-    qt_sat = TD.total_specific_humidity(thermo_params, ts)
+    q_sat = TD.q_vap_saturation(thermo_params, ts)
+    q_tot = TD.total_specific_humidity(thermo_params, ts)
+    θ = TD.dry_pottemp(thermo_params, ts)
+    ∂b∂θli_unsat = ∂b∂θv * (1 + (Rv_over_Rd - 1) * q_tot)
+    ∂b∂qt_unsat = ∂b∂θv * (Rv_over_Rd - 1) * θ
     ∂b∂θli_sat = (
         ∂b∂θv *
-        (1 + Rv_over_Rd * (1 + lh / R_v / t_sat) * qv_sat - qt_sat) /
-        (1 + lh * lh / cp_m / R_v / t_sat / t_sat * qv_sat)
+        (1 + Rv_over_Rd * (1 + lh / R_v / T) * q_sat - q_tot) /
+        (1 + lh^2 / cp_m / R_v / T^2 * q_sat)
     )
     ∂b∂qt_sat =
-        (lh / cp_m / t_sat * ∂b∂θli_sat - ∂b∂θv) *
-        TD.dry_pottemp(thermo_params, ts)
+        (lh / cp_m / T * ∂b∂θli_sat - ∂b∂θv) * θ
 
     ∂b∂z = buoyancy_gradient_chain_rule(
         ebgc,
         bg_model,
         thermo_params,
-        ∂b∂θv,
+        ∂b∂θli_unsat,
+        ∂b∂qt_unsat,
         ∂b∂θli_sat,
         ∂b∂qt_sat,
     )
@@ -140,7 +134,6 @@ end
         closure::AbstractEnvBuoyGradClosure,
         bg_model::EnvBuoyGradVars,
         thermo_params,
-        ∂b∂θv::FT,
         ∂b∂θli_sat::FT,
         ∂b∂qt_sat::FT,
     ) where {FT}
@@ -165,7 +158,6 @@ Arguments:
 - `closure`: The environmental buoyancy gradient closure type.
 - `bg_model`: Precomputed environmental buoyancy gradient variables (`EnvBuoyGradVars`).
 - `thermo_params`: Thermodynamic parameters from `CLIMAParameters`.
-- `∂b∂θv`: Partial derivative of buoyancy w.r.t. virtual potential temperature (unsaturated part).
 - `∂b∂θli_sat`: Partial derivative of buoyancy w.r.t. liquid-ice potential temperature (saturated part).
 - `∂b∂qt_sat`: Partial derivative of buoyancy w.r.t. total specific humidity (saturated part).
 
@@ -176,18 +168,19 @@ function buoyancy_gradient_chain_rule(
     ::AbstractEnvBuoyGradClosure,
     bg_model::EnvBuoyGradVars,
     thermo_params,
-    ∂b∂θv,
+    ∂b∂θli_unsat,
+    ∂b∂qt_unsat,
     ∂b∂θli_sat,
     ∂b∂qt_sat,
 )
-    en_cld_frac = ifelse(TD.has_condensate(thermo_params, bg_model.ts), 1, 0)
-
-    ∂b∂z_θl_sat = ∂b∂θli_sat * bg_model.∂θli∂z_sat
-    ∂b∂z_qt_sat = ∂b∂qt_sat * bg_model.∂qt∂z_sat
+    ∂b∂z_θli_unsat = ∂b∂θli_unsat * bg_model.∂θli∂z
+    ∂b∂z_qt_unsat = ∂b∂qt_unsat * bg_model.∂qt∂z
+    ∂b∂z_unsat = ∂b∂z_θli_unsat + ∂b∂z_qt_unsat
+    ∂b∂z_θl_sat = ∂b∂θli_sat * bg_model.∂θli∂z
+    ∂b∂z_qt_sat = ∂b∂qt_sat * bg_model.∂qt∂z
     ∂b∂z_sat = ∂b∂z_θl_sat + ∂b∂z_qt_sat
-    ∂b∂z_unsat = ∂b∂θv * bg_model.∂θv∂z_unsat
 
-    ∂b∂z = (1 - en_cld_frac) * ∂b∂z_unsat + en_cld_frac * ∂b∂z_sat
+    ∂b∂z = (1 - bg_model.cf) * ∂b∂z_unsat + bg_model.cf * ∂b∂z_sat
 
     return ∂b∂z
 end