Skip to content

Commit

Permalink
Add cloud fraction diagnostics
Browse files Browse the repository at this point in the history
  • Loading branch information
trontrytel committed Dec 2, 2023
1 parent 4d36738 commit 84a1573
Show file tree
Hide file tree
Showing 9 changed files with 197 additions and 161 deletions.
5 changes: 5 additions & 0 deletions examples/hybrid/driver.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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
2 changes: 1 addition & 1 deletion src/ClimaAtmos.jl
Original file line number Diff line number Diff line change
Expand Up @@ -25,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(
Expand Down Expand Up @@ -80,7 +81,6 @@ include(joinpath("prognostic_equations", "edmfx_entr_detr.jl"))
include(joinpath("prognostic_equations", "edmfx_tke.jl"))
include(joinpath("prognostic_equations", "edmfx_sgs_flux.jl"))
include(joinpath("prognostic_equations", "edmfx_boundary_condition.jl"))
include(joinpath("prognostic_equations", "cloud_fraction.jl"))
include(
joinpath("parameterized_tendencies", "microphysics", "precipitation.jl"),
)
Expand Down
21 changes: 15 additions & 6 deletions src/cache/cache.jl
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@ struct AtmosCache{
COR,
SFC,
GHOST,
ENV,
SGQ,
PREC,
SCRA,
HYPE,
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -146,11 +147,19 @@ 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)
precomputing_arguments = (;
atmos,
core,
params,
sfc_setup,
precomputed,
scratch,
simulation,
SG_quad,
)

# Coupler compatibility
isnothing(precomputing_arguments.sfc_setup) &&
Expand Down Expand Up @@ -183,7 +192,7 @@ function build_cache(Y, atmos, params, surface_setup, simulation)
core,
sfc_setup,
ghost_buffer,
env_thermo_quad,
SG_quad,
precomputed,
scratch,
hyperdiff,
Expand Down
152 changes: 152 additions & 0 deletions src/cache/cloud_fraction.jl
Original file line number Diff line number Diff line change
@@ -0,0 +1,152 @@
import StaticArrays as SA
import ClimaCore.RecursiveApply: rzero, ,

# TODO: write a test with scalars that are linear with z
"""
Diagnose horizntal covariances based on vertical gradients
(i.e. taking turbulence production as the only term)
"""
function covariance_from_grad(coeff, ᶜmixing_length, ᶜ∇Φ, ᶜ∇Ψ)
return 2 * coeff * ᶜmixing_length^2 * dot(ᶜ∇Φ, ᶜ∇Ψ)
end

"""
Compute the grid scale cloud fraction based on sub-grid scale properties
"""
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, 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)

coeff = FT(2.1) # TODO - move to parameters

@. cloud_fraction = quad_loop(
SG_quad,
ᶜp,
ᶜqₜ,
ᶜθ,
covariance_from_grad(
coeff,
ᶜdz, # TODO - replace dz with mixinglength when using EDMF SGS
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,
)
end

"""
function quad_loop(SG_quad, p_c, qt_mean, θl_mean, qt′qt′, θl′θl′, θl′qt′, thermo_params)
where:
- SG_quad is a struct containing information about quadrature type and order
- p_c is the atmospheic pressure
- qt_mean, θl_mean is the grid mean q_tot and liquid ice potential temperature
- qt′qt′, θl′θl′, θl′qt′ are the (co)variances of q_tot and liquid ice potential temperature
- thermo params are the thermodynamics parameters
The function trnasforms and imposes additional limits on the quadrature points.
It returns cloud fraction computed as a sum over quadrature points.
"""
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))

# Epsilon defined per typical variable fluctuation
eps_q = FT(eps(FT)) * max(FT(eps(FT)), qt_mean)
eps_θ = FT(eps(FT))

# limit σ_q to prevent negative q_tot_hat
σ_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′)

# Enforce Cauchy-Schwarz inequality, numerically stable compute
_corr::FT = (θl′qt′ / max(σ_q, eps_q))
corr::FT = max(min(_corr / max(σ_θ, eps_θ), 1), -1)

# Conditionals
σ_c = sqrt(max(1 - corr * corr, 0)) * σ_θ

function get_x_hat::Tuple{<:Real, <:Real})
μ_c = θl_mean + sqrt2 * corr * σ_θ * χ[1]
θ_hat = μ_c + sqrt2 * σ_c * χ[2]
q_tot_hat = qt_mean + sqrt2 * σ_q * χ[1]
return (θ_hat, q_tot_hat)
end

# cloudy/dry categories for buoyancy in TKE
f_q_tot_sat(x_hat::Tuple{<:Real, <:Real}, hc) =
hc ? x_hat[2] : eltype(x_hat)(0)

get_ts(x_hat::Tuple{<:Real, <:Real}) =
thermo_state(thermo_params; p = p_c, θ = x_hat[1], q_tot = x_hat[2])
f_cf(x_hat::Tuple{<:Real, <:Real}, hc) =
hc ? eltype(x_hat)(1) : eltype(x_hat)(0)
function f(x_hat::Tuple{<:Real, <:Real})
ts = get_ts(x_hat)
hc = TD.has_condensate(thermo_params, ts)
return (; cf = f_cf(x_hat, hc), q_tot_sat = f_q_tot_sat(x_hat, hc))
end

return quad(f, get_x_hat, SG_quad).cf
end

"""
Compute f(A, B) as a sum over inner and outer quadrature points
that approximate the sub-grid scale variability of A and B
"""
function quad(f::F, get_x_hat::F1, quad) where {F <: Function, F1 <: Function}
χ = quad.a
weights = quad.w
quad_order = quadrature_order(quad)
FT = eltype(χ)
# zero outer quadrature points
T = typeof(f(get_x_hat((χ[1], χ[1]))))
outer_env = rzero(T)
@inbounds for m_q in 1:quad_order
# zero inner quadrature points
inner_env = rzero(T)
for m_h in 1:quad_order
x_hat = get_x_hat((χ[m_q], χ[m_h]))
inner_env = inner_env f(x_hat) weights[m_h] FT(1 / sqrt(π))
end
outer_env = outer_env inner_env weights[m_q] FT(1 / sqrt(π))
end
return outer_env
end
4 changes: 4 additions & 0 deletions src/cache/precomputed_quantities.jl
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,7 @@ function precomputed_quantities(Y, atmos)
ᶜh_tot = similar(Y.c, FT),
sfc_conditions = Fields.Field(SCT, Spaces.level(axes(Y.f), half)),
)
cloud_diagnostics = (; cloud_fraction = similar(Y.c, FT),)
advective_sgs_quantities =
atmos.turbconv_model isa PrognosticEDMFX ?
(;
Expand Down Expand Up @@ -125,6 +126,7 @@ function precomputed_quantities(Y, atmos)
diagnostic_sgs_quantities...,
vert_diff_quantities...,
precipitation_quantities...,
cloud_diagnostics...,
)
end

Expand Down Expand Up @@ -375,6 +377,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

Expand Down
49 changes: 14 additions & 35 deletions src/callbacks/callbacks.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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 = (;
Expand All @@ -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 = (;
Expand All @@ -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
Expand All @@ -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)

Expand Down Expand Up @@ -270,58 +267,40 @@ 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⁺),
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()
Expand Down
Loading

0 comments on commit 84a1573

Please sign in to comment.