cleanup in cloud fraction

config/model_configs/sphere_held_suarez_rhoe_equilmoist_topography_dcmip.yml

@@ -7,7 +7,6 @@ vert_diff: "true"
 use_reference_state: false
 forcing: "held_suarez"
 precip_model: "0M"
-cloud_diagnostics: "diag_covariance"
 topography: "DCMIP200"
 job_id: "sphere_held_suarez_rhoe_equilmoist_topography_dcmip"
 moist: "equil"

examples/hybrid/driver.jl

@@ -174,5 +174,10 @@ if config.parsed_args["check_precipitation"]
         @test all(
             ClimaCore.isapprox(Yₜ_ρ[colidx], Yₜ_ρqₜ[colidx], rtol = eps(FT)),
         )
+
+        # cloud fraction diagnostics
+        @assert !any(isnan, sol.prob.p.precomputed.cloud_fraction[colidx])
+        @test minimum(sol.prob.p.precomputed.cloud_fraction[colidx]) >= FT(0)
+        @test maximum(sol.prob.p.precomputed.cloud_fraction[colidx]) <= FT(1)
     end
 end

src/ClimaAtmos.jl

@@ -13,7 +13,6 @@ include(joinpath("utils", "utilities.jl"))
 include(joinpath("utils", "debug_utils.jl"))
 include(joinpath("topography", "topography.jl"))
 include(joinpath("utils", "variable_manipulations.jl"))
-include(joinpath("utils", "cloud_fraction.jl"))
 
 include(
     joinpath("parameterized_tendencies", "radiation", "radiation_utilities.jl"),
@@ -26,6 +25,7 @@ include(joinpath("cache", "prognostic_edmf_precomputed_quantities.jl"))
 include(joinpath("cache", "diagnostic_edmf_precomputed_quantities.jl"))
 include(joinpath("cache", "precipitation_precomputed_quantities.jl"))
 include(joinpath("cache", "precomputed_quantities.jl"))
+include(joinpath("cache", "cloud_fraction.jl"))
 
 include(joinpath("initial_conditions", "InitialConditions.jl"))
 include(

src/cache/cache.jl

@@ -7,7 +7,7 @@ struct AtmosCache{
     COR,
     SFC,
     GHOST,
-    ENV,
+    SGQ,
     PREC,
     SCRA,
     HYPE,
@@ -49,7 +49,8 @@ struct AtmosCache{
     """Center and face ghost buffers used by DSS"""
     ghost_buffer::GHOST
 
-    env_thermo_quad::ENV
+    """Struct with sub-grid sampling quadrature"""
+    SG_quad::SGQ
 
     """Quantities that are updated with set_precomputed_quantities!"""
     precomputed::PREC
@@ -146,11 +147,11 @@ function build_cache(Y, atmos, params, surface_setup, simulation)
 
     sfc_setup = surface_setup(params)
     scratch = temporary_quantities(Y, atmos)
-    env_thermo_quad = SGSQuadrature(FT)
+    SG_quad = SGSQuadrature(FT)
 
     precomputed = precomputed_quantities(Y, atmos)
     precomputing_arguments =
-        (; atmos, core, params, sfc_setup, precomputed, scratch, simulation)
+        (; atmos, core, params, sfc_setup, precomputed, scratch, simulation, SG_quad)
 
     # Coupler compatibility
     isnothing(precomputing_arguments.sfc_setup) &&
@@ -183,7 +184,7 @@ function build_cache(Y, atmos, params, surface_setup, simulation)
         core,
         sfc_setup,
         ghost_buffer,
-        env_thermo_quad,
+        SG_quad,
         precomputed,
         scratch,
         hyperdiff,

src/utils/cloud_fraction.jl → src/cache/cloud_fraction.jl

@@ -2,54 +2,71 @@ import StaticArrays as SA
 import ClimaCore.RecursiveApply: rzero, ⊞, ⊠
 
 # TODO: write a test with scalars that are linear with z
-function get_covariance(coeff, ᶜmixing_length, ᶜscalar1_grad, ᶜscalar2_grad)
-    return 2 * coeff * ᶜmixing_length^2 * dot(ᶜscalar1_grad, ᶜscalar2_grad)
+function covariance_from_grad(coeff, ᶜmixing_length, ᶜ∇Φ, ᶜ∇Ψ)
+    return 2 * coeff * ᶜmixing_length^2 * dot(ᶜ∇Φ, ᶜ∇Ψ)
 end
 
-function get_cloud_fraction(Y, p, colidx, ::SGSQuadrature)
-    (; params) = p
+function set_cloud_fraction!(Y, p, ::DryModel)
+    FT = eltype(Y)
+    (; cloud_fraction) = p.precomputed
+    @. cloud_fraction = FT(0)
+end
+function set_cloud_fraction!(Y, p, ::Union{EquilMoistModel, NonEquilMoistModel})
+    (; SG_quad, params) = p
+
     FT = eltype(params)
     thermo_params = CAP.thermodynamics_params(params)
     ᶜdz = Fields.Δz_field(axes(Y.c))
     (; ᶜts, ᶜp) = p.precomputed
 
-    (; qt′qt′, θl′θl′, θl′qt′) = p.precomputed.cloud_diagnostics_quantities
-    (; cloud_fraction) = p.precomputed.cloud_diagnostics_quantities
-    (; ᶜqₜ_grad, ᶜθ_grad) = p.precomputed.cloud_diagnostics_quantities
+    (; ᶜqₜ′qₜ′, ᶜθ′θ′, ᶜθ′qₜ′) = p.precomputed
+    (; cloud_fraction) = p.precomputed
 
     ᶜqₜ = p.scratch.ᶜtemp_scalar
     ᶜθ = p.scratch.ᶜtemp_scalar_2
     @. ᶜqₜ = TD.total_specific_humidity(thermo_params, ᶜts)
     @. ᶜθ = TD.liquid_ice_pottemp(thermo_params, ᶜts)
 
-    @. ᶜqₜ_grad = Geometry.WVector(ᶜgradᵥ(ᶠinterp(ᶜqₜ)))
-    @. ᶜθ_grad = Geometry.WVector(ᶜgradᵥ(ᶠinterp(ᶜθ)))
-
     coeff = FT(2.1) # TODO - move to parameters
+
     # TODO - replace dz with mixinglength when using EDMF SGS
-    @. qt′qt′ = get_covariance(coeff, ᶜdz, ᶜqₜ_grad, ᶜqₜ_grad)
-    @. θl′θl′ = get_covariance(coeff, ᶜdz, ᶜθ_grad, ᶜθ_grad)
-    @. θl′qt′ = get_covariance(coeff, ᶜdz, ᶜθ_grad, ᶜqₜ_grad)
-
-    vars = (;
-        p_c = ᶜp,
-        qt_mean = q_tot,
-        θl_mean = θ_liq_ice,
-        qt′qt′,
-        θl′θl′,
-        θl′qt′,
+    @. cloud_fraction = quad_loop(
+        SG_quad,
+        ᶜp,
+        ᶜqₜ,
+        ᶜθ,
+        covariance_from_grad(
+            coeff,
+            ᶜdz,
+            Geometry.WVector(ᶜgradᵥ(ᶠinterp(ᶜqₜ))),
+            Geometry.WVector(ᶜgradᵥ(ᶠinterp(ᶜqₜ))),
+        ),
+        covariance_from_grad(
+            coeff,
+            ᶜdz,
+            Geometry.WVector(ᶜgradᵥ(ᶠinterp(ᶜθ))),
+            Geometry.WVector(ᶜgradᵥ(ᶠinterp(ᶜθ))),
+        ),
+        covariance_from_grad(
+            coeff,
+            ᶜdz,
+            Geometry.WVector(ᶜgradᵥ(ᶠinterp(ᶜθ))),
+            Geometry.WVector(ᶜgradᵥ(ᶠinterp(ᶜqₜ)))
+        ),
+        thermo_params
     )
-
-    @. cloud_fraction = quad_loop(::SGSQuadrature, vars, thermo_params)::FT
 end
 
-function quad_loop(env_thermo_quad::SGSQuadrature, vars, thermo_params)
-
-    # qt - total water specific humidity
-    # θl - liquid ice potential temperature
-    # _mean and ′ - subdomain mean and (co)variances
-    (; qt′qt′, qt_mean, θl′θl′, θl_mean, θl′qt′, p_c) = vars
-
+function quad_loop(
+    SG_quad::SGSQuadrature,
+    p_c,
+    qt_mean,
+    θl_mean,
+    qt′qt′,
+    θl′θl′,
+    θl′qt′,
+    thermo_params
+)
     FT = eltype(qt_mean)
 
     sqrt2 = FT(sqrt(2))
@@ -59,7 +76,7 @@ function quad_loop(env_thermo_quad::SGSQuadrature, vars, thermo_params)
     eps_θ = FT(eps(FT))
 
     # limit σ_q to prevent negative q_tot_hat
-    σ_q_lim = -qt_mean / (sqrt2 * env_thermo_quad.a[1]) # TODO: is this correct?
+    σ_q_lim = -qt_mean / (sqrt2 * SG_quad.a[1]) # TODO: is this correct?
     σ_q::FT = min(sqrt(qt′qt′), σ_q_lim)
     σ_θ::FT = sqrt(θl′θl′)
 
@@ -91,9 +108,10 @@ function quad_loop(env_thermo_quad::SGSQuadrature, vars, thermo_params)
         return (; cf = f_cf(x_hat, hc), q_tot_sat = f_q_tot_sat(x_hat, hc))
     end
 
-    return quad(f, get_x_hat, env_thermo_quad).cf
+    return quad(f, get_x_hat, SG_quad).cf
 end
 
+# TODO f::F1 <: Functon
 function quad(f, get_x_hat::F, quad_type) where {F <: Function}
     χ = quad_type.a
     weights = quad_type.w

src/cache/precomputed_quantities.jl

@@ -53,6 +53,12 @@ function precomputed_quantities(Y, atmos)
         ᶜh_tot = similar(Y.c, FT),
         sfc_conditions = Fields.Field(SCT, Spaces.level(axes(Y.f), half)),
     )
+    cloud_diagnostics = (;
+        ᶜqₜ′qₜ′ = similar(Y.c, FT),
+        ᶜθ′θ′ = similar(Y.c, FT),
+        ᶜθ′qₜ′ = similar(Y.c, FT),
+        cloud_fraction = similar(Y.c, FT),
+    )
     advective_sgs_quantities =
         atmos.turbconv_model isa PrognosticEDMFX ?
         (;
@@ -112,22 +118,12 @@ function precomputed_quantities(Y, atmos)
     precipitation_quantities =
         atmos.precip_model isa Microphysics1Moment ?
         (; ᶜwᵣ = similar(Y.c, FT), ᶜwₛ = similar(Y.c, FT)) : (;)
-    cloud_diagnostics_quantities =
-        atmos.cloud_diagnostics isa DiagnosticCovariance ?
-        (;
-            ᶜqₜ′qₜ′ = similar(Y.c, FT),
-            ᶜθ′θ′ = similar(Y.c, FT),
-            ᶜθ′qₜ′ = similar(Y.c, FT),
-            cloud_fraction = similar(Y.c, FT)
-            ᶜqₜ_grad = (Y.f, Geometry.WVector{FT}),
-            ᶜθ_grad = (Y.f, Geometry.WVector{FT}),
-        ) : (;)
     return (;
         gs_quantities...,
         advective_sgs_quantities...,
         diagnostic_sgs_quantities...,
         precipitation_quantities...,
-        cloud_diagnostics_quantities...,
+        cloud_diagnostics...,
     )
 end
 
@@ -359,6 +355,8 @@ NVTX.@annotate function set_precomputed_quantities!(Y, p, t)
         set_precipitation_precomputed_quantities!(Y, p, t)
     end
 
+    set_cloud_fraction!(Y, p, moisture_model)
+
     return nothing
 end
 

src/callbacks/callbacks.jl

@@ -174,7 +174,7 @@ NVTX.@annotate function rrtmgp_model_callback!(integrator)
 end
 
 function common_diagnostics(p, ᶜu, ᶜts)
-    (; env_thermo_quad, params) = p
+    (; params) = p
     thermo_params = CAP.thermodynamics_params(params)
     ᶜρ = TD.air_density.(thermo_params, ᶜts)
     diagnostics = (;
@@ -186,6 +186,7 @@ function common_diagnostics(p, ᶜu, ᶜts)
         specific_enthalpy = TD.specific_enthalpy.(thermo_params, ᶜts),
         buoyancy = CAP.grav(p.params) .* (p.core.ᶜρ_ref .- ᶜρ) ./ ᶜρ,
         density = TD.air_density.(thermo_params, ᶜts),
+        cloud_fraction_gm = p.precomputed.cloud_fraction
     )
     if !(p.atmos.moisture_model isa DryModel)
         diagnostics = (;
@@ -195,12 +196,6 @@ function common_diagnostics(p, ᶜu, ᶜts)
             q_ice = TD.ice_specific_humidity.(thermo_params, ᶜts),
             q_tot = TD.total_specific_humidity.(thermo_params, ᶜts),
             relative_humidity = TD.relative_humidity.(thermo_params, ᶜts),
-            cloud_fraction_gm = get_cloud_fraction.(
-                thermo_params,
-                env_thermo_quad,
-                p.precomputed.ᶜp,
-                ᶜts,
-            ),
         )
     end
     return diagnostics
@@ -211,13 +206,15 @@ end
 add_prefix(diagnostics::NamedTuple{names}, prefix) where {names} =
     NamedTuple{Symbol.(prefix, names)}(values(diagnostics))
 
-cloud_fraction(thermo_params, ts, area::FT) where {FT} =
+# TODO - temporary crutch untill we move this to the new diagnostics
+# and start using quadratures in EDMFX
+draft_cloud_fraction(thermo_params, ts, area::FT) where {FT} =
     TD.has_condensate(thermo_params, ts) && area > 1e-3 ? FT(1) : FT(0)
 
 NVTX.@annotate function compute_diagnostics(integrator)
     (; t, u, p) = integrator
     Y = u
-    (; params, env_thermo_quad) = p
+    (; params) = p
     FT = eltype(params)
     thermo_params = CAP.thermodynamics_params(params)
 
@@ -270,58 +267,41 @@ NVTX.@annotate function compute_diagnostics(integrator)
         env_diagnostics = (;
             common_diagnostics(p, ᶜu⁰, ᶜts⁰)...,
             area = ᶜa⁰,
-            cloud_fraction = get_cloud_fraction.(
-                thermo_params,
-                env_thermo_quad,
-                ᶜp,
-                ᶜts⁰,
-            ),
+            cloud_fraction = draft_cloud_fraction.(thermo_params, ᶜts⁰, ᶜa⁰),
             tke = ᶜtke⁰,
             mixing_length = ᶜmixing_length,
         )
         draft_diagnostics = (;
             common_diagnostics(p, ᶜu⁺, ᶜts⁺)...,
             area = ᶜa⁺,
-            cloud_fraction = cloud_fraction.(thermo_params, ᶜts⁺, ᶜa⁺),
+            cloud_fraction = draft_cloud_fraction.(thermo_params, ᶜts⁺, ᶜa⁺),
         )
         turbulence_convection_diagnostic = (;
             add_prefix(env_diagnostics, :env_)...,
             add_prefix(draft_diagnostics, :draft_)...,
-            cloud_fraction = ᶜa⁰ .*
-                             get_cloud_fraction.(
-                thermo_params,
-                env_thermo_quad,
-                ᶜp,
-                ᶜts⁰,
-            ) .+ ᶜa⁺ .* cloud_fraction.(thermo_params, ᶜts⁺, ᶜa⁺),
+            cloud_fraction =
+                ᶜa⁰ .* draft_cloud_fraction.(thermo_params, ᶜts⁰, ᶜa⁰) .+
+                ᶜa⁺ .* draft_cloud_fraction.(thermo_params, ᶜts⁺, ᶜa⁺),
         )
     elseif p.atmos.turbconv_model isa DiagnosticEDMFX
         (; ᶜtke⁰, ᶜmixing_length) = p.precomputed
         (; ᶜu⁺, ᶜts⁺, ᶜa⁺) = output_diagnostic_sgs_quantities(Y, p, t)
         env_diagnostics = (;
-            cloud_fraction = get_cloud_fraction.(
-                thermo_params,
-                env_thermo_quad,
-                ᶜp,
-                ᶜts,
-            ),
+            cloud_fraction = draft_cloud_fraction.(thermo_params, ᶜts, FT(1)),
             tke = ᶜtke⁰,
             mixing_length = ᶜmixing_length,
         )
         draft_diagnostics = (;
             common_diagnostics(p, ᶜu⁺, ᶜts⁺)...,
             area = ᶜa⁺,
-            cloud_fraction = cloud_fraction.(thermo_params, ᶜts⁺, ᶜa⁺),
+            cloud_fraction = draft_cloud_fraction.(thermo_params, ᶜts⁺, ᶜa⁺),
         )
         turbulence_convection_diagnostic = (;
             add_prefix(env_diagnostics, :env_)...,
             add_prefix(draft_diagnostics, :draft_)...,
-            cloud_fraction = get_cloud_fraction.(
-                thermo_params,
-                env_thermo_quad,
-                ᶜp,
-                ᶜts,
-            ) .+ ᶜa⁺ .* cloud_fraction.(thermo_params, ᶜts⁺, ᶜa⁺),
+            cloud_fraction =
+                draft_cloud_fraction.(thermo_params, ᶜts, FT(1)) .+
+                ᶜa⁺ .* draft_cloud_fraction.(thermo_params, ᶜts⁺, ᶜa⁺),
         )
     else
         turbulence_convection_diagnostic = NamedTuple()

src/solver/model_getters.jl

@@ -221,17 +221,6 @@ function get_precipitation_model(parsed_args)
     end
 end
 
-function get_cloud_diagnostics_type(parsed_args)
-    cloud_diag = parsed_args["cloud_diagnostics"]
-    return if cloud_diag == nothing
-        nothing
-    elseif cloud_diag = "diag_covariance"
-        DiagnosticCovariance()
-    else
-        error("Invalid cloud diagnostics type $(cloud_diagnostics)")
-    end
-end
-
 function get_forcing_type(parsed_args)
     forcing = parsed_args["forcing"]
     @assert forcing in (nothing, "held_suarez")