From e70fbab035ea2a023a2a999deb384b45f91288a2 Mon Sep 17 00:00:00 2001 From: Haakon Ludvig Langeland Ervik <45243236+haakon-e@users.noreply.github.com> Date: Mon, 24 Nov 2025 16:15:31 -0800 Subject: [PATCH] clean up precip code --- .../precipitation_precomputed_quantities.jl | 861 +++++------------- .../microphysics/microphysics_wrappers.jl | 32 + src/utils/abbreviations.jl | 1 + src/utils/variable_manipulations.jl | 124 ++- 4 files changed, 386 insertions(+), 632 deletions(-) diff --git a/src/cache/precipitation_precomputed_quantities.jl b/src/cache/precipitation_precomputed_quantities.jl index 065637f1d8..d5f584561e 100644 --- a/src/cache/precipitation_precomputed_quantities.jl +++ b/src/cache/precipitation_precomputed_quantities.jl @@ -10,25 +10,37 @@ import CloudMicrophysics.P3Scheme as CMP3 import Thermodynamics as TD import ClimaCore.Operators as Operators import ClimaCore.Fields as Fields +import ClimaCore.MatrixFields: get_field const Iₗ = TD.internal_energy_liquid const Iᵢ = TD.internal_energy_ice """ - Kin(ᶜw_precip, ᶜu_air) + Kin(w_precip, u_air_uvw) - - ᶜw_precip - teminal velocity of cloud consensate or precipitation - - ᶜu_air - air velocity +Helper function to compute the kinetic energy of cloud condensate and precipitation. -Helper function to compute the kinetic energy of cloud condensate and -precipitation. +# Arguments +- `w_precip`: teminal velocity of cloud consensate or precipitation +- `u_air_uvw`: air velocity as a `UVW` field """ -function Kin(ᶜw_precip, ᶜu_air) - return @. lazy( - norm_sqr( - Geometry.UVWVector(0, 0, -(ᶜw_precip)) + Geometry.UVWVector(ᶜu_air), - ) / 2, - ) +Kin(w_precip, u_air_uvw) = norm_sqr(UVW(0, 0, -w_precip) + u_air_uvw) / 2 + +internal_energy_func(qᵪ_name) = + qᵪ_name == @name(q_liq) ? Iₗ : + qᵪ_name == @name(q_ice) ? Iᵢ : + qᵪ_name == @name(q_rai) ? Iₗ : + qᵪ_name == @name(q_sno) ? Iᵢ : error("Invalid tracer name $qᵪ_name") + +ρξw_div_ρξ_bounded(ρqw, ρq) = ifelse(ρq > zero(ρq), ρqw / ρq, zero(ρq)) + +function total_energy_func(qᵪ_name) + I = internal_energy_func(qᵪ_name) + total_energy(thp, ᶜts, ᶜΦ, ᶜwᵪ, ᶜu_uvw) = I(thp, ᶜts) + ᶜΦ + Kin(ᶜwᵪ, ᶜu_uvw) +end +function internal_plus_geopotential_energy_func(qᵪ_name) + I = internal_energy_func(qᵪ_name) + internal_plus_geopotential_energy(thp, ᶜts, ᶜΦ) = I(thp, ᶜts) + ᶜΦ end """ @@ -37,473 +49,203 @@ end Updates the precipitation terminal velocity, cloud sedimentation velocity, and their contribution to total water and energy advection. """ -function set_precipitation_velocities!(Y, p, _, _, _) +function set_precipitation_velocities!(_, p, _, _, _) (; ᶜwₜqₜ, ᶜwₕhₜ) = p.precomputed - @. ᶜwₜqₜ = Geometry.WVector(0) - @. ᶜwₕhₜ = Geometry.WVector(0) + @. ᶜwₜqₜ = WVec(0) + @. ᶜwₕhₜ = WVec(0) return nothing end -function set_precipitation_velocities!( - Y, - p, - moisture_model::NonEquilMoistModel, - microphysics_model::Microphysics1Moment, - _, -) - (; ᶜwₗ, ᶜwᵢ, ᶜwᵣ, ᶜwₛ, ᶜwₜqₜ, ᶜwₕhₜ, ᶜts, ᶜu) = p.precomputed +function set_precipitation_velocities!(Y, p, ::NonEquilMoistModel, ::Microphysics1Moment, _) + (; ᶜwₜqₜ, ᶜwₕhₜ, ᶜts, ᶜu) = p.precomputed (; ᶜΦ) = p.core cmc = CAP.microphysics_cloud_params(p.params) cmp = CAP.microphysics_1m_params(p.params) thp = CAP.thermodynamics_params(p.params) - - # compute the precipitation terminal velocity [m/s] - @. ᶜwᵣ = CM1.terminal_velocity( - cmp.pr, - cmp.tv.rain, - Y.c.ρ, - max(zero(Y.c.ρ), Y.c.ρq_rai / Y.c.ρ), - ) - @. ᶜwₛ = CM1.terminal_velocity( - cmp.ps, - cmp.tv.snow, - Y.c.ρ, - max(zero(Y.c.ρ), Y.c.ρq_sno / Y.c.ρ), - ) - # compute sedimentation velocity for cloud condensate [m/s] - @. ᶜwₗ = CMNe.terminal_velocity( - cmc.liquid, - cmc.Ch2022.rain, - Y.c.ρ, - max(zero(Y.c.ρ), Y.c.ρq_liq / Y.c.ρ), - ) - @. ᶜwᵢ = CMNe.terminal_velocity( - cmc.ice, - cmc.Ch2022.small_ice, - Y.c.ρ, - max(zero(Y.c.ρ), Y.c.ρq_ice / Y.c.ρ), - ) - - # compute their contributions to energy and total water advection - @. ᶜwₜqₜ = - Geometry.WVector( - ᶜwₗ * Y.c.ρq_liq + - ᶜwᵢ * Y.c.ρq_ice + - ᶜwᵣ * Y.c.ρq_rai + - ᶜwₛ * Y.c.ρq_sno, - ) / Y.c.ρ - @. ᶜwₕhₜ = - Geometry.WVector( - ᶜwₗ * Y.c.ρq_liq * (Iₗ(thp, ᶜts) + ᶜΦ + $(Kin(ᶜwₗ, ᶜu))) + - ᶜwᵢ * Y.c.ρq_ice * (Iᵢ(thp, ᶜts) + ᶜΦ + $(Kin(ᶜwᵢ, ᶜu))) + - ᶜwᵣ * Y.c.ρq_rai * (Iₗ(thp, ᶜts) + ᶜΦ + $(Kin(ᶜwᵣ, ᶜu))) + - ᶜwₛ * Y.c.ρq_sno * (Iᵢ(thp, ᶜts) + ᶜΦ + $(Kin(ᶜwₛ, ᶜu))), - ) / Y.c.ρ + ᶜu_uvw = @. lazy(UVW(ᶜu)) + + qᵪ_names = (@name(q_liq), @name(q_ice), @name(q_rai), @name(q_sno)) + @. ᶜwₜqₜ = WVec(0) + @. ᶜwₕhₜ = WVec(0) + for qᵪ_name in qᵪ_names + ᶜwᵪ = get_field(p.precomputed, get_ᶜwᵪ_name_from_qᵪ_name(qᵪ_name)) + ᶜqᵪ = ᶜspecific_gs_tracer(Y, qᵪ_name) + wᵪ = terminal_velocity_func_1M(cmc, cmp, qᵪ_name) + calc_hₜ = total_energy_func(qᵪ_name) + @. ᶜwᵪ = wᵪ(Y.c.ρ, clip(ᶜqᵪ)) + @. ᶜwₜqₜ += WVec(ᶜwᵪ * ᶜqᵪ) + @. ᶜwₕhₜ += WVec(ᶜwᵪ * ᶜqᵪ * calc_hₜ(thp, ᶜts, ᶜΦ, ᶜwᵪ, ᶜu_uvw)) + end return nothing end -function set_precipitation_velocities!( - Y, - p, - moisture_model::NonEquilMoistModel, - microphysics_model::Microphysics1Moment, - turbconv_model::PrognosticEDMFX, +function set_precipitation_velocities!(Y, p, + ::NonEquilMoistModel, ::Microphysics1Moment, turbconv_model::PrognosticEDMFX, ) - (; ᶜwₗ, ᶜwᵢ, ᶜwᵣ, ᶜwₛ, ᶜwₜqₜ, ᶜwₕhₜ) = p.precomputed (; ᶜΦ) = p.core - (; ᶜwₗʲs, ᶜwᵢʲs, ᶜwᵣʲs, ᶜwₛʲs) = p.precomputed - (; ᶜts⁰, ᶜtsʲs) = p.precomputed + (; ᶜwₜqₜ, ᶜwₕhₜ) = p.precomputed + (; ᶜts⁰, ᶜtsʲs, ᶜρʲs) = p.precomputed + ᶜsgsʲs = Y.c.sgsʲs + ᶜρ = Y.c.ρ cmc = CAP.microphysics_cloud_params(p.params) cmp = CAP.microphysics_1m_params(p.params) thp = CAP.thermodynamics_params(p.params) - FT = eltype(p.params) + # Environment quantities ᶜρ⁰ = @. lazy(TD.air_density(thp, ᶜts⁰)) - ᶜρa⁰ = @. lazy(ρa⁰(Y.c.ρ, Y.c.sgsʲs, turbconv_model)) - n = n_mass_flux_subdomains(turbconv_model) - - # scratch to compute env velocities - ᶜw⁰ = p.scratch.ᶜtemp_scalar - # scratch to add positive masses of subdomains - # TODO use Y.c.ρq instead of ᶜρχ for computing gs velocities by averaging velocities - # over subdomains once negative subdomain mass issues are resolved - ᶜρχ = p.scratch.ᶜtemp_scalar_2 - # scratch for adding energy fluxes over subdomains - ᶜρwₕhₜ = p.scratch.ᶜtemp_scalar_3 - - # Compute gs sedimentation velocity based on subdomain velocities (assuming gs flux - # equals sum of sgs fluxes) - - ᶜq_liq⁰ = ᶜspecific_env_value(@name(q_liq), Y, p) - ᶜq_ice⁰ = ᶜspecific_env_value(@name(q_ice), Y, p) - ᶜq_rai⁰ = ᶜspecific_env_value(@name(q_rai), Y, p) - ᶜq_sno⁰ = ᶜspecific_env_value(@name(q_sno), Y, p) - - # Cloud liquid - @. ᶜw⁰ = CMNe.terminal_velocity( - cmc.liquid, - cmc.Ch2022.rain, - ᶜρ⁰, - max(zero(Y.c.ρ), ᶜq_liq⁰), - ) - @. ᶜwₗ = ᶜρa⁰ * ᶜq_liq⁰ * ᶜw⁰ - @. ᶜρχ = max(zero(Y.c.ρ), ᶜρa⁰ * ᶜq_liq⁰) - for j in 1:n - @. ᶜwₗ += Y.c.sgsʲs.:($$j).ρa * Y.c.sgsʲs.:($$j).q_liq * ᶜwₗʲs.:($$j) - @. ᶜρχ += max(zero(Y.c.ρ), Y.c.sgsʲs.:($$j).ρa * Y.c.sgsʲs.:($$j).q_liq) - end - @. ᶜwₗ = ifelse(ᶜρχ > FT(0), ᶜwₗ / ᶜρχ, FT(0)) - - # contribution of env cloud liquid advection to htot advection - @. ᶜρwₕhₜ = ᶜρa⁰ * ᶜq_liq⁰ * ᶜw⁰ * (Iₗ(thp, ᶜts⁰) + ᶜΦ) - - # Cloud ice - @. ᶜw⁰ = CMNe.terminal_velocity( - cmc.ice, - cmc.Ch2022.small_ice, - ᶜρ⁰, - max(zero(Y.c.ρ), ᶜq_ice⁰), - ) - @. ᶜwᵢ = ᶜρa⁰ * ᶜq_ice⁰ * ᶜw⁰ - @. ᶜρχ = max(zero(Y.c.ρ), ᶜρa⁰ * ᶜq_ice⁰) - for j in 1:n - @. ᶜwᵢ += Y.c.sgsʲs.:($$j).ρa * Y.c.sgsʲs.:($$j).q_ice * ᶜwᵢʲs.:($$j) - @. ᶜρχ += max(zero(Y.c.ρ), Y.c.sgsʲs.:($$j).ρa * Y.c.sgsʲs.:($$j).q_ice) - end - @. ᶜwᵢ = ifelse(ᶜρχ > FT(0), ᶜwᵢ / ᶜρχ, FT(0)) - - # contribution of env cloud ice advection to htot advection - @. ᶜρwₕhₜ += ᶜρa⁰ * ᶜq_ice⁰ * ᶜw⁰ * (Iᵢ(thp, ᶜts⁰) + ᶜΦ) - - # Rain - @. ᶜw⁰ = CM1.terminal_velocity( - cmp.pr, - cmp.tv.rain, - ᶜρ⁰, - max(zero(Y.c.ρ), ᶜq_rai⁰), - ) - @. ᶜwᵣ = ᶜρa⁰ * ᶜq_rai⁰ * ᶜw⁰ - @. ᶜρχ = max(zero(Y.c.ρ), ᶜρa⁰ * ᶜq_rai⁰) - for j in 1:n - @. ᶜwᵣ += Y.c.sgsʲs.:($$j).ρa * Y.c.sgsʲs.:($$j).q_rai * ᶜwᵣʲs.:($$j) - @. ᶜρχ += max(zero(Y.c.ρ), Y.c.sgsʲs.:($$j).ρa * Y.c.sgsʲs.:($$j).q_rai) + ᶜρa⁰ = @. lazy(ρa⁰(ᶜρ, ᶜsgsʲs, turbconv_model)) + + @. ᶜwₕhₜ = WVec(0) + @. ᶜwₜqₜ = WVec(0) + + qᵪ_names = (@name(q_liq), @name(q_ice), @name(q_rai), @name(q_sno)) + for qᵪ_name in qᵪ_names + # Functions that depend on qᵪ_name + wᵪ_func = terminal_velocity_func_1M(cmc, cmp, qᵪ_name) # terminal velocity for tracer `q` + ρaqʲ(ᶜsgsʲ) = begin + # Calculate `ρa ⋅ qᵪ` for subdomain `j` + qᵪʲ = get_field(ᶜsgsʲ, qᵪ_name) + ᶜsgsʲ.ρa * qᵪʲ + end + ρaqwʲ(ᶜsgsʲ, ᶜρʲ) = begin + # Calculate `ρa ⋅ qᵪ ⋅ wᵪ` for subdomain `j` + qᵪʲ = get_field(ᶜsgsʲ, qᵪ_name) + wᵪʲ = wᵪ_func(ᶜρʲ, clip(qᵪʲ)) + ᶜsgsʲ.ρa * qᵪʲ * wᵪʲ + end + calc_hₜ = internal_plus_geopotential_energy_func(qᵪ_name) # TODO: do we need to add kinetic energy of subdomains? + # Get grid-scale quantities + ᶜqᵪ = ᶜspecific_gs_tracer(Y, qᵪ_name) + ᶜwᵪ = get_field(p.precomputed, get_ᶜwᵪ_name_from_qᵪ_name(qᵪ_name)) + # Get environment quantities + ᶜqᵪ⁰ = ᶜspecific_env_value(qᵪ_name, Y, p) + + # TODO: use Y.c.ρq instead of ᶜρqᵪ for computing gs velocities by averaging velocities + # over subdomains once negative subdomain mass issues are resolved + clip_ρaqʲ(ᶜsgsʲ) = clip(ρaqʲ(ᶜsgsʲ)) + ᶜρqᵪ = @. lazy(clip(ᶜρa⁰ * ᶜqᵪ⁰) + draft_sum(clip_ρaqʲ, ᶜsgsʲs)) + + ᶜwᵪ⁰ = @. lazy(wᵪ_func(ᶜρ⁰, clip(ᶜqᵪ⁰))) # environment terminal velocity + ᶜρaqwᵪ⁰ = @. lazy(ᶜρa⁰ * ᶜqᵪ⁰ * ᶜwᵪ⁰) # environment ρa ⋅ qᵪ ⋅ wᵪ + + # Compute gs sedimentation velocity based on subdomain velocities + # (assuming gs flux equals sum of sgs fluxes) + @. ᶜwᵪ = ρξw_div_ρξ_bounded(ᶜρaqwᵪ⁰ + draft_sum(ρaqwʲ, ᶜsgsʲs, ᶜρʲs), ᶜρqᵪ) + + # Calculate contribution of tracer `q` to total water and energy advection + ρaqwₕhₜʲ(Φ) = (ᶜsgsʲ, ᶜρʲ, ᶜtsʲ) -> ρaqwʲ(ᶜsgsʲ, ᶜρʲ) * calc_hₜ(thp, ᶜtsʲ, Φ) + ᶜρaqwᵪhₜ⁰ = @. lazy(ᶜρaqwᵪ⁰ * calc_hₜ(thp, ᶜts⁰, ᶜΦ)) + @. ᶜwₕhₜ += WVec(ᶜρaqwᵪhₜ⁰ + draft_sum(ρaqwₕhₜʲ(ᶜΦ), ᶜsgsʲs, ᶜρʲs, ᶜtsʲs)) / ᶜρ + @. ᶜwₜqₜ += WVec(ᶜwᵪ * ᶜqᵪ) end - @. ᶜwᵣ = ifelse(ᶜρχ > FT(0), ᶜwᵣ / ᶜρχ, FT(0)) - - # contribution of env rain advection to htot advection - @. ᶜρwₕhₜ += ᶜρa⁰ * ᶜq_rai⁰ * ᶜw⁰ * (Iₗ(thp, ᶜts⁰) + ᶜΦ) - - # Snow - @. ᶜw⁰ = CM1.terminal_velocity( - cmp.ps, - cmp.tv.snow, - ᶜρ⁰, - max(zero(Y.c.ρ), ᶜq_sno⁰), - ) - @. ᶜwₛ = ᶜρa⁰ * ᶜq_sno⁰ * ᶜw⁰ - @. ᶜρχ = max(zero(Y.c.ρ), ᶜρa⁰ * ᶜq_sno⁰) - for j in 1:n - @. ᶜwₛ += Y.c.sgsʲs.:($$j).ρa * Y.c.sgsʲs.:($$j).q_sno * ᶜwₛʲs.:($$j) - @. ᶜρχ += max(zero(Y.c.ρ), Y.c.sgsʲs.:($$j).ρa * Y.c.sgsʲs.:($$j).q_sno) - end - @. ᶜwₛ = ifelse(ᶜρχ > FT(0), ᶜwₛ / ᶜρχ, FT(0)) - - # contribution of env snow advection to htot advection - @. ᶜρwₕhₜ += ᶜρa⁰ * ᶜq_sno⁰ * ᶜw⁰ * (Iᵢ(thp, ᶜts⁰) + ᶜΦ) - - # Add contributions to energy and total water advection - # TODO do we need to add kinetic energy of subdomains? - for j in 1:n - @. ᶜρwₕhₜ += - Y.c.sgsʲs.:($$j).ρa * - ( - ᶜwₗʲs.:($$j) * Y.c.sgsʲs.:($$j).q_liq * (Iₗ(thp, ᶜtsʲs.:($$j)) + ᶜΦ) + - ᶜwᵢʲs.:($$j) * Y.c.sgsʲs.:($$j).q_ice * (Iᵢ(thp, ᶜtsʲs.:($$j)) + ᶜΦ) + - ᶜwᵣʲs.:($$j) * Y.c.sgsʲs.:($$j).q_rai * (Iₗ(thp, ᶜtsʲs.:($$j)) + ᶜΦ) + - ᶜwₛʲs.:($$j) * Y.c.sgsʲs.:($$j).q_sno * (Iᵢ(thp, ᶜtsʲs.:($$j)) + ᶜΦ) - ) - end - @. ᶜwₕhₜ = Geometry.WVector(ᶜρwₕhₜ) / Y.c.ρ - - @. ᶜwₜqₜ = - Geometry.WVector( - ᶜwₗ * Y.c.ρq_liq + - ᶜwᵢ * Y.c.ρq_ice + - ᶜwᵣ * Y.c.ρq_rai + - ᶜwₛ * Y.c.ρq_sno, - ) / Y.c.ρ - return nothing end -function set_precipitation_velocities!( - Y, - p, - moisture_model::NonEquilMoistModel, - microphysics_model::Microphysics2Moment, - _, -) - (; ᶜwₗ, ᶜwᵢ, ᶜwᵣ, ᶜwₛ, ᶜwₙₗ, ᶜwₙᵣ, ᶜwₜqₜ, ᶜwₕhₜ, ᶜts, ᶜu) = p.precomputed +function set_precipitation_velocities!(Y, p, ::NonEquilMoistModel, ::Microphysics2Moment, _) + (; ᶜwₜqₜ, ᶜwₕhₜ, ᶜts, ᶜu) = p.precomputed (; ᶜΦ) = p.core + ᶜu_uvw = @. lazy(UVW(ᶜu)) cmc = CAP.microphysics_cloud_params(p.params) cm1p = CAP.microphysics_1m_params(p.params) cm2p = CAP.microphysics_2m_params(p.params) thp = CAP.thermodynamics_params(p.params) - - # compute the precipitation terminal velocity [m/s] - # TODO sedimentation of snow is based on the 1M scheme - @. ᶜwₙᵣ = getindex( - CM2.rain_terminal_velocity( - cm2p.sb, - cm2p.rtv, - max(zero(Y.c.ρ), specific(Y.c.ρq_rai, Y.c.ρ)), - Y.c.ρ, - max(zero(Y.c.ρ), Y.c.ρn_rai), - ), - 1, - ) - @. ᶜwᵣ = getindex( - CM2.rain_terminal_velocity( - cm2p.sb, - cm2p.rtv, - max(zero(Y.c.ρ), specific(Y.c.ρq_rai, Y.c.ρ)), - Y.c.ρ, - max(zero(Y.c.ρ), Y.c.ρn_rai), - ), - 2, - ) - @. ᶜwₛ = CM1.terminal_velocity( - cm1p.ps, - cm1p.tv.snow, - Y.c.ρ, - max(zero(Y.c.ρ), specific(Y.c.ρq_sno, Y.c.ρ)), - ) - # compute sedimentation velocity for cloud condensate [m/s] - # TODO sedimentation of ice is based on the 1M scheme - @. ᶜwₙₗ = getindex( - CM2.cloud_terminal_velocity( - cm2p.sb.pdf_c, - cm2p.ctv, - max(zero(Y.c.ρ), specific(Y.c.ρq_liq, Y.c.ρ)), - Y.c.ρ, - max(zero(Y.c.ρ), Y.c.ρn_liq), - ), - 1, - ) - @. ᶜwₗ = getindex( - CM2.cloud_terminal_velocity( - cm2p.sb.pdf_c, - cm2p.ctv, - max(zero(Y.c.ρ), specific(Y.c.ρq_liq, Y.c.ρ)), - Y.c.ρ, - max(zero(Y.c.ρ), Y.c.ρn_liq), - ), - 2, - ) - @. ᶜwᵢ = CMNe.terminal_velocity( - cmc.ice, - cmc.Ch2022.small_ice, - Y.c.ρ, - max(zero(Y.c.ρ), specific(Y.c.ρq_ice, Y.c.ρ)), - ) - - # compute their contributions to energy and total water advection - @. ᶜwₜqₜ = - Geometry.WVector( - ᶜwₗ * Y.c.ρq_liq + - ᶜwᵢ * Y.c.ρq_ice + - ᶜwᵣ * Y.c.ρq_rai + - ᶜwₛ * Y.c.ρq_sno, - ) / Y.c.ρ - @. ᶜwₕhₜ = - Geometry.WVector( - ᶜwₗ * Y.c.ρq_liq * (Iₗ(thp, ᶜts) + ᶜΦ + $(Kin(ᶜwₗ, ᶜu))) + - ᶜwᵢ * Y.c.ρq_ice * (Iᵢ(thp, ᶜts) + ᶜΦ + $(Kin(ᶜwᵢ, ᶜu))) + - ᶜwᵣ * Y.c.ρq_rai * (Iₗ(thp, ᶜts) + ᶜΦ + $(Kin(ᶜwᵣ, ᶜu))) + - ᶜwₛ * Y.c.ρq_sno * (Iᵢ(thp, ᶜts) + ᶜΦ + $(Kin(ᶜwₛ, ᶜu))), - ) / Y.c.ρ + ᶜρ = Y.c.ρ + + qᵪ_names = (@name(q_liq), @name(q_ice), @name(q_rai), @name(q_sno)) + for qᵪ_name in qᵪ_names + # Get velocity fields and functions + ᶜwᵪ = get_field(p.precomputed, get_ᶜwᵪ_name_from_qᵪ_name(qᵪ_name)) + ᶜwₙᵪ = get_field(p.precomputed, get_ᶜwₙᵪ_name_from_qᵪ_name(qᵪ_name)) + wᵪ_func = terminal_velocity_mass_func_2M(cm2p, cmc, cm1p, qᵪ_name) + wₙᵪ_func = terminal_velocity_number_func_2M(cm2p, cmc, cm1p, qᵪ_name) + calc_hₜ = total_energy_func(qᵪ_name) + # Get tracer fields + ᶜqᵪ = ᶜspecific_gs_tracer(Y, qᵪ_name) + ρnᵪ_name = get_ρnᵪ_name_from_qᵪ_name(qᵪ_name) + ρnᵪ = get_field(Y.c, ρnᵪ_name) # this is `Y.c` if ρnᵪ is not tracked for qᵪ + # Compute velocity + @. ᶜwᵪ = wᵪ_func(ᶜρ, ᶜqᵪ, ρnᵪ) # `ρnᵪ` is not used if number is not tracked for qᵪ + if ρnᵪ_name != @name() # if number is tracked for qᵪ + @. ᶜwₙᵪ = wₙᵪ_func(ᶜρ, ᶜqᵪ, ρnᵪ) + end + ᶜwqᵪ = @. lazy(WVec(ᶜwᵪ * ᶜqᵪ)) + @. ᶜwₜqₜ += ᶜwqᵪ + @. ᶜwₕhₜ += ᶜwqᵪ * calc_hₜ(thp, ᶜts, ᶜΦ, ᶜwᵪ, ᶜu_uvw) + end return nothing end -function set_precipitation_velocities!( - Y, - p, - moisture_model::NonEquilMoistModel, - microphysics_model::Microphysics2Moment, - turbconv_model::PrognosticEDMFX, +function set_precipitation_velocities!(Y, p, + ::NonEquilMoistModel, ::Microphysics2Moment, turbconv_model::PrognosticEDMFX, ) - (; ᶜwₗ, ᶜwᵢ, ᶜwᵣ, ᶜwₛ, ᶜwₙₗ, ᶜwₙᵣ, ᶜwₜqₜ, ᶜwₕhₜ) = p.precomputed + (; ᶜwₜqₜ, ᶜwₕhₜ) = p.precomputed (; ᶜΦ) = p.core - (; ᶜwₗʲs, ᶜwᵢʲs, ᶜwᵣʲs, ᶜwₛʲs, ᶜwₙₗʲs, ᶜwₙᵣʲs) = p.precomputed - (; ᶜts⁰, ᶜtsʲs) = p.precomputed + (; ᶜts⁰, ᶜtsʲs, ᶜρʲs) = p.precomputed + ᶜρ = Y.c.ρ + ᶜsgsʲs = Y.c.sgsʲs cmc = CAP.microphysics_cloud_params(p.params) cm1p = CAP.microphysics_1m_params(p.params) cm2p = CAP.microphysics_2m_params(p.params) thp = CAP.thermodynamics_params(p.params) - FT = eltype(p.params) ᶜρ⁰ = @. lazy(TD.air_density(thp, ᶜts⁰)) ᶜρa⁰ = @. lazy(ρa⁰(Y.c.ρ, Y.c.sgsʲs, turbconv_model)) - n = n_mass_flux_subdomains(turbconv_model) - - # scratch to compute env velocities - ᶜw⁰ = p.scratch.ᶜtemp_scalar - # scratch to add positive masses of subdomains - # TODO use Y.c.ρq instead of ᶜρχ for computing gs velocities by averaging velocities - # over subdomains once negative subdomain mass issues are resolved - ᶜρχ = p.scratch.ᶜtemp_scalar_2 - # scratch for adding energy fluxes over subdomains - ᶜρwₕhₜ = p.scratch.ᶜtemp_scalar_3 - - # Compute gs sedimentation velocity based on subdomain velocities (assuming gs flux - # equals sum of sgs fluxes) - - ᶜq_liq⁰ = ᶜspecific_env_value(@name(q_liq), Y, p) - ᶜq_ice⁰ = ᶜspecific_env_value(@name(q_ice), Y, p) - ᶜq_rai⁰ = ᶜspecific_env_value(@name(q_rai), Y, p) - ᶜq_sno⁰ = ᶜspecific_env_value(@name(q_sno), Y, p) - ᶜn_liq⁰ = ᶜspecific_env_value(@name(n_liq), Y, p) - ᶜn_rai⁰ = ᶜspecific_env_value(@name(n_rai), Y, p) - - # Cloud liquid (number) - @. ᶜw⁰ = getindex( - CM2.rain_terminal_velocity( - cm2p.sb, - cm2p.rtv, - max(zero(Y.c.ρ), ᶜq_liq⁰), - ᶜρ⁰, - max(zero(Y.c.ρ), ᶜn_liq⁰), - ), - 1, - ) - @. ᶜwₙₗ = ᶜρa⁰ * ᶜn_liq⁰ * ᶜw⁰ - @. ᶜρχ = max(zero(Y.c.ρ), ᶜρa⁰ * ᶜn_liq⁰) - for j in 1:n - @. ᶜwₙₗ += Y.c.sgsʲs.:($$j).ρa * Y.c.sgsʲs.:($$j).n_liq * ᶜwₙₗʲs.:($$j) - @. ᶜρχ += max(zero(Y.c.ρ), Y.c.sgsʲs.:($$j).ρa * Y.c.sgsʲs.:($$j).n_liq) - end - @. ᶜwₙₗ = ifelse(ᶜρχ > FT(0), ᶜwₙₗ / ᶜρχ, FT(0)) - # Cloud liquid (mass) - @. ᶜw⁰ = getindex( - CM2.rain_terminal_velocity( - cm2p.sb, - cm2p.rtv, - max(zero(Y.c.ρ), ᶜq_liq⁰), - ᶜρ⁰, - max(zero(Y.c.ρ), ᶜn_liq⁰), - ), - 2, - ) - @. ᶜwₗ = ᶜρa⁰ * ᶜq_liq⁰ * ᶜw⁰ - @. ᶜρχ = max(zero(Y.c.ρ), ᶜρa⁰ * ᶜq_liq⁰) - for j in 1:n - @. ᶜwₗ += Y.c.sgsʲs.:($$j).ρa * Y.c.sgsʲs.:($$j).q_liq * ᶜwₗʲs.:($$j) - @. ᶜρχ += max(zero(Y.c.ρ), Y.c.sgsʲs.:($$j).ρa * Y.c.sgsʲs.:($$j).q_liq) - end - @. ᶜwₗ = ifelse(ᶜρχ > FT(0), ᶜwₗ / ᶜρχ, FT(0)) - - # contribution of env cloud liquid advection to htot advection - @. ᶜρwₕhₜ = ᶜρa⁰ * ᶜq_liq⁰ * ᶜw⁰ * (Iₗ(thp, ᶜts⁰) + ᶜΦ) - - # Cloud ice - # TODO sedimentation of ice is based on the 1M scheme - @. ᶜw⁰ = CMNe.terminal_velocity( - cmc.ice, - cmc.Ch2022.small_ice, - ᶜρ⁰, - max(zero(Y.c.ρ), ᶜq_ice⁰), - ) - @. ᶜwᵢ = ᶜρa⁰ * ᶜq_ice⁰ * ᶜw⁰ - @. ᶜρχ = max(zero(Y.c.ρ), ᶜρa⁰ * ᶜq_ice⁰) - for j in 1:n - @. ᶜwᵢ += Y.c.sgsʲs.:($$j).ρa * Y.c.sgsʲs.:($$j).q_ice * ᶜwᵢʲs.:($$j) - @. ᶜρχ += max(zero(Y.c.ρ), Y.c.sgsʲs.:($$j).ρa * Y.c.sgsʲs.:($$j).q_ice) - end - @. ᶜwᵢ = ifelse(ᶜρχ > FT(0), ᶜwᵢ / ᶜρχ, FT(0)) - - # contribution of env cloud ice advection to htot advection - @. ᶜρwₕhₜ += ᶜρa⁰ * ᶜq_ice⁰ * ᶜw⁰ * (Iᵢ(thp, ᶜts⁰) + ᶜΦ) - - # Rain (number) - @. ᶜw⁰ = getindex( - CM2.rain_terminal_velocity( - cm2p.sb, - cm2p.rtv, - max(zero(Y.c.ρ), ᶜq_rai⁰), - ᶜρ⁰, - max(zero(Y.c.ρ), ᶜn_rai⁰), - ), - 1, - ) - @. ᶜwₙᵣ = ᶜρa⁰ * ᶜn_rai⁰ * ᶜw⁰ - @. ᶜρχ = max(zero(Y.c.ρ), ᶜρa⁰ * ᶜn_rai⁰) - for j in 1:n - @. ᶜwₙᵣ += Y.c.sgsʲs.:($$j).ρa * Y.c.sgsʲs.:($$j).n_rai * ᶜwₙᵣʲs.:($$j) - @. ᶜρχ += max(zero(Y.c.ρ), Y.c.sgsʲs.:($$j).ρa * Y.c.sgsʲs.:($$j).n_rai) + qᵪ_names = (@name(q_liq), @name(q_ice), @name(q_rai), @name(q_sno)) + for qᵪ_name in qᵪ_names + nᵪ_name = get_nᵪ_name_from_qᵪ_name(qᵪ_name) + nᵪ_exists = nᵪ_name != @name() + # Functions that depend on `qᵪ_name` + wᵪ_func = terminal_velocity_mass_func_2M(cm2p, cmc, cm1p, qᵪ_name) + ρaqʲ(ᶜsgsʲ) = begin # Calculate `ρa ⋅ qᵪ` for subdomain `j` + qᵪʲ = get_field(ᶜsgsʲ, qᵪ_name) + ᶜsgsʲ.ρa * qᵪʲ + end + ρaqwʲ(ᶜsgsʲ, ᶜρʲ) = begin # Calculate `ρa ⋅ qᵪ ⋅ wᵪ` for subdomain `j` + qᵪʲ = get_field(ᶜsgsʲ, qᵪ_name) + nᵪʲ = nᵪ_exists ? get_field(ᶜsgsʲ, nᵪ_name) : zero(qᵪʲ) + wᵪʲ = wᵪ_func(ᶜρʲ, clip(qᵪʲ), clip(ᶜρʲ * nᵪʲ)) # TODO: Should consistently clip nᵪ⁰ and nᵪʲ + ᶜsgsʲ.ρa * qᵪʲ * wᵪʲ + end + wₙᵪ_func = terminal_velocity_number_func_2M(cm2p, cmc, cm1p, qᵪ_name) + ρanʲ(ᶜsgsʲ) = begin # Calculate `ρa ⋅ nᵪ` for subdomain `j` + nᵪʲ = get_field(ᶜsgsʲ, nᵪ_name) + ᶜsgsʲ.ρa * nᵪʲ + end + ρanwʲ(ᶜsgsʲ, ᶜρʲ) = begin # Calculate `ρa ⋅ nᵪ ⋅ wᵪ` for subdomain `j` + qᵪʲ = get_field(ᶜsgsʲ, qᵪ_name) + nᵪʲ = nᵪ_exists ? get_field(ᶜsgsʲ, nᵪ_name) : zero(qᵪʲ) + wₙᵪʲ = wₙᵪ_func(ᶜρʲ, clip(qᵪʲ), clip(ᶜρʲ * nᵪʲ)) # TODO: Should consistently clip nᵪ⁰ and nᵪʲ + ᶜsgsʲ.ρa * nᵪʲ * wₙᵪʲ + end + calc_hₜ = internal_plus_geopotential_energy_func(qᵪ_name) + # Get grid-scale quantities + ᶜqᵪ = ᶜspecific_gs_tracer(Y, qᵪ_name) + ᶜwᵪ = get_field(p.precomputed, get_ᶜwᵪ_name_from_qᵪ_name(qᵪ_name)) + ᶜwₙᵪ = get_field(p.precomputed, get_ᶜwₙᵪ_name_from_qᵪ_name(qᵪ_name)) + # Get environment quantities + ᶜqᵪ⁰ = ᶜspecific_env_value(qᵪ_name, Y, p) + ᶜnᵪ⁰ = nᵪ_exists ? ᶜspecific_env_value(nᵪ_name, Y, p) : (@. lazy(zero(ᶜqᵪ⁰))) + + # Mass + clip_ρaqʲ(ᶜsgsʲ) = clip(ρaqʲ(ᶜsgsʲ)) + ᶜρqᵪ = @. lazy(clip(ᶜρ⁰ * ᶜqᵪ⁰) + draft_sum(clip_ρaqʲ, ᶜsgsʲs)) + ᶜwᵪ⁰ = @. lazy(wᵪ_func(ᶜρ⁰, clip(ᶜqᵪ⁰), clip(ᶜρ⁰ * clip(ᶜnᵪ⁰)))) # TODO: Should consistently clip nᵪ⁰ and nᵪʲ + ᶜρaqwᵪ⁰ = @. lazy(ᶜρa⁰ * ᶜqᵪ⁰ * ᶜwᵪ⁰) + @. ᶜwᵪ = ρξw_div_ρξ_bounded(ᶜρaqwᵪ⁰ + draft_sum(ρaqwʲ, ᶜsgsʲs, ᶜρʲs), ᶜρqᵪ) + + # Number + if nᵪ_exists + clip_ρanʲ(ᶜsgsʲ) = clip(ρanʲ(ᶜsgsʲ)) + ᶜρnᵪ = @. lazy(clip(ᶜρ⁰ * ᶜnᵪ⁰) + draft_sum(clip_ρanʲ, ᶜsgsʲs)) + ᶜwₙᵪ⁰ = @. lazy(wₙᵪ_func(ᶜρ⁰, clip(ᶜqᵪ⁰), clip(ᶜρ⁰ * clip(ᶜnᵪ⁰)))) # TODO: Should consistently clip nᵪ⁰ and nᵪʲ + ᶜρanwᵪ⁰ = @. lazy(ᶜρa⁰ * ᶜnᵪ⁰ * ᶜwₙᵪ⁰) + @. ᶜwₙᵪ = ρξw_div_ρξ_bounded(ᶜρanwᵪ⁰ + draft_sum(ρanwʲ, ᶜsgsʲs, ᶜρʲs), ᶜρnᵪ) + end + + # Calculate contribution of tracer `q` to total water and energy advection + ρaqwₕhₜʲ(Φ) = (ᶜsgsʲ, ᶜρʲ, ᶜtsʲ) -> ρaqwʲ(ᶜsgsʲ, ᶜρʲ) * calc_hₜ(thp, ᶜtsʲ, Φ) + ᶜρaqwᵪhₜ⁰ = @. lazy(ᶜρaqwᵪ⁰ * calc_hₜ(thp, ᶜts⁰, ᶜΦ)) + @. ᶜwₕhₜ += WVec(ᶜρaqwᵪhₜ⁰ + draft_sum(ρaqwₕhₜʲ(ᶜΦ), ᶜsgsʲs, ᶜρʲs, ᶜtsʲs)) / ᶜρ + @. ᶜwₜqₜ += WVec(ᶜwᵪ * ᶜqᵪ) end - @. ᶜwₙᵣ = ifelse(ᶜρχ > FT(0), ᶜwₙᵣ / ᶜρχ, FT(0)) - - # Rain (mass) - @. ᶜw⁰ = getindex( - CM2.rain_terminal_velocity( - cm2p.sb, - cm2p.rtv, - max(zero(Y.c.ρ), ᶜq_rai⁰), - ᶜρ⁰, - max(zero(Y.c.ρ), ᶜn_rai⁰), - ), - 2, - ) - @. ᶜwᵣ = ᶜρa⁰ * ᶜq_rai⁰ * ᶜw⁰ - @. ᶜρχ = max(zero(Y.c.ρ), ᶜρa⁰ * ᶜq_rai⁰) - for j in 1:n - @. ᶜwᵣ += Y.c.sgsʲs.:($$j).ρa * Y.c.sgsʲs.:($$j).q_rai * ᶜwᵣʲs.:($$j) - @. ᶜρχ += max(zero(Y.c.ρ), Y.c.sgsʲs.:($$j).ρa * Y.c.sgsʲs.:($$j).q_rai) - end - @. ᶜwᵣ = ifelse(ᶜρχ > FT(0), ᶜwᵣ / ᶜρχ, FT(0)) - - # contribution of env rain advection to qtot advection - @. ᶜρwₕhₜ += ᶜρa⁰ * ᶜq_rai⁰ * ᶜw⁰ * (Iₗ(thp, ᶜts⁰) + ᶜΦ) - - # Snow - # TODO sedimentation of snow is based on the 1M scheme - @. ᶜw⁰ = CM1.terminal_velocity( - cm1p.ps, - cm1p.tv.snow, - ᶜρ⁰, - max(zero(Y.c.ρ), ᶜq_sno⁰), - ) - @. ᶜwₛ = ᶜρa⁰ * ᶜq_sno⁰ * ᶜw⁰ - @. ᶜρχ = max(zero(Y.c.ρ), ᶜρa⁰ * ᶜq_sno⁰) - for j in 1:n - @. ᶜwₛ += Y.c.sgsʲs.:($$j).ρa * Y.c.sgsʲs.:($$j).q_sno * ᶜwₛʲs.:($$j) - @. ᶜρχ += max(zero(Y.c.ρ), Y.c.sgsʲs.:($$j).ρa * Y.c.sgsʲs.:($$j).q_sno) - end - @. ᶜwₛ = ifelse(ᶜρχ > FT(0), ᶜwₛ / ᶜρχ, FT(0)) - - # contribution of env snow advection to htot advection - @. ᶜρwₕhₜ += ᶜρa⁰ * ᶜq_sno⁰ * ᶜw⁰ * (Iᵢ(thp, ᶜts⁰) + ᶜΦ) - - # Add contributions to energy and total water advection - # TODO do we need to add kinetic energy of subdomains? - for j in 1:n - @. ᶜρwₕhₜ += - Y.c.sgsʲs.:($$j).ρa * - ( - ᶜwₗʲs.:($$j) * Y.c.sgsʲs.:($$j).q_liq * (Iₗ(thp, ᶜtsʲs.:($$j)) + ᶜΦ) + - ᶜwᵢʲs.:($$j) * Y.c.sgsʲs.:($$j).q_ice * (Iᵢ(thp, ᶜtsʲs.:($$j)) + ᶜΦ) + - ᶜwᵣʲs.:($$j) * Y.c.sgsʲs.:($$j).q_rai * (Iₗ(thp, ᶜtsʲs.:($$j)) + ᶜΦ) + - ᶜwₛʲs.:($$j) * Y.c.sgsʲs.:($$j).q_sno * (Iᵢ(thp, ᶜtsʲs.:($$j)) + ᶜΦ) - ) - end - @. ᶜwₕhₜ = Geometry.WVector(ᶜρwₕhₜ) / Y.c.ρ - - @. ᶜwₜqₜ = - Geometry.WVector( - ᶜwₗ * Y.c.ρq_liq + - ᶜwᵢ * Y.c.ρq_ice + - ᶜwᵣ * Y.c.ρq_rai + - ᶜwₛ * Y.c.ρq_sno, - ) / Y.c.ρ - return nothing end @@ -559,12 +301,15 @@ function set_precipitation_velocities!( @. ᶜwᵢ = CMP3.ice_terminal_velocity_mass_weighted(args...; use_aspect_ratio) # compute their contributions to energy and total water advection - @. ᶜwₜqₜ = Geometry.WVector(ᶜwₗ * ρq_liq + ᶜwᵢ * ρq_ice + ᶜwᵣ * ρq_rai) / ρ + @. ᶜwₜqₜ = WVec(ᶜwₗ * ρq_liq + ᶜwᵢ * ρq_ice + ᶜwᵣ * ρq_rai) / ρ + calc_hₗ = total_energy_func(@name(q_liq)) + calc_hᵢ = total_energy_func(@name(q_ice)) + calc_hᵣ = total_energy_func(@name(q_rai)) @. ᶜwₕhₜ = - Geometry.WVector( - ᶜwₗ * ρq_liq * (Iₗ(thp, ᶜts) + ᶜΦ + $(Kin(ᶜwₗ, ᶜu))) + - ᶜwᵢ * ρq_ice * (Iᵢ(thp, ᶜts) + ᶜΦ + $(Kin(ᶜwᵢ, ᶜu))) + - ᶜwᵣ * ρq_rai * (Iₗ(thp, ᶜts) + ᶜΦ + $(Kin(ᶜwᵣ, ᶜu))), + WVec( + ᶜwₗ * ρq_liq * (calc_hₗ(thp, ᶜts, ᶜΦ, ᶜwₗ, UVW(ᶜu))) + + ᶜwᵢ * ρq_ice * (calc_hᵢ(thp, ᶜts, ᶜΦ, ᶜwᵢ, UVW(ᶜu))) + + ᶜwᵣ * ρq_rai * (calc_hᵣ(thp, ᶜts, ᶜΦ, ᶜwᵣ, UVW(ᶜu))), ) / ρ return nothing end @@ -589,88 +334,37 @@ function set_precipitation_cache!(Y, p, ::Microphysics0Moment, _) thermo_params = CAP.thermodynamics_params(params) @. ᶜS_ρq_tot = Y.c.ρ * q_tot_0M_precipitation_sources( - thermo_params, - cm_params, - dt, - Y.c.ρq_tot / Y.c.ρ, - ᶜts, + thermo_params, cm_params, dt, Y.c.ρq_tot / Y.c.ρ, ᶜts, ) - @. ᶜS_ρe_tot = - ᶜS_ρq_tot * - e_tot_0M_precipitation_sources_helper(thermo_params, ᶜts, ᶜΦ) - return nothing -end -function set_precipitation_cache!( - Y, - p, - ::Microphysics0Moment, - ::DiagnosticEDMFX, -) - # For environment we multiply by grid mean ρ and not byᶜρa⁰ - # assuming a⁰=1 - (; ᶜΦ) = p.core - (; ᶜSqₜᵖ⁰, ᶜSqₜᵖʲs, ᶜρaʲs) = p.precomputed - (; ᶜS_ρq_tot, ᶜS_ρe_tot) = p.precomputed - (; ᶜts, ᶜtsʲs) = p.precomputed - thermo_params = CAP.thermodynamics_params(p.params) - - n = n_mass_flux_subdomains(p.atmos.turbconv_model) - ρ = Y.c.ρ - - @. ᶜS_ρq_tot = ᶜSqₜᵖ⁰ * ρ - @. ᶜS_ρe_tot = - ᶜSqₜᵖ⁰ * - ρ * - e_tot_0M_precipitation_sources_helper(thermo_params, ᶜts, ᶜΦ) - for j in 1:n - @. ᶜS_ρq_tot += ᶜSqₜᵖʲs.:($$j) * ᶜρaʲs.:($$j) - @. ᶜS_ρe_tot += - ᶜSqₜᵖʲs.:($$j) * - ᶜρaʲs.:($$j) * - e_tot_0M_precipitation_sources_helper( - thermo_params, - ᶜtsʲs.:($$j), - ᶜΦ, - ) - end + @. ᶜS_ρe_tot = ᶜS_ρq_tot * e_tot_0M_precipitation_sources_helper(thermo_params, ᶜts, ᶜΦ) return nothing end -function set_precipitation_cache!( - Y, - p, - ::Microphysics0Moment, - ::PrognosticEDMFX, +function set_precipitation_cache!(Y, p, + ::Microphysics0Moment, turbconv::Union{DiagnosticEDMFX, PrognosticEDMFX}, ) (; ᶜΦ) = p.core (; ᶜSqₜᵖ⁰, ᶜSqₜᵖʲs) = p.precomputed (; ᶜS_ρq_tot, ᶜS_ρe_tot) = p.precomputed (; ᶜts⁰, ᶜtsʲs) = p.precomputed - thermo_params = CAP.thermodynamics_params(p.params) - ᶜρa⁰ = @. lazy(ρa⁰(Y.c.ρ, Y.c.sgsʲs, p.atmos.turbconv_model)) - - n = n_mass_flux_subdomains(p.atmos.turbconv_model) - - @. ᶜS_ρq_tot = ᶜSqₜᵖ⁰ * ᶜρa⁰ - @. ᶜS_ρe_tot = - ᶜSqₜᵖ⁰ * - ᶜρa⁰ * - e_tot_0M_precipitation_sources_helper(thermo_params, ᶜts⁰, ᶜΦ) - for j in 1:n - @. ᶜS_ρq_tot += ᶜSqₜᵖʲs.:($$j) * Y.c.sgsʲs.:($$j).ρa - @. ᶜS_ρe_tot += - ᶜSqₜᵖʲs.:($$j) * - Y.c.sgsʲs.:($$j).ρa * - e_tot_0M_precipitation_sources_helper( - thermo_params, - ᶜtsʲs.:($$j), - ᶜΦ, - ) - end + thp = CAP.thermodynamics_params(p.params) + ᶜρ = Y.c.ρ + + # Draft and environment area-weighted densities + ᶜsgs_ρaʲs = ᶜρaʲs_list(Y, p, turbconv) + ## For environment we multiply by grid mean `ρ` and not by `ᶜρa⁰`, assuming `a⁰ = 1` + ᶜρa⁰ = turbconv isa PrognosticEDMFX ? (@. lazy(ρa⁰(ᶜρ, Y.c.sgsʲs, turbconv))) : ᶜρ + ᶜts⁰ = turbconv isa PrognosticEDMFX ? p.precomputed.ᶜts⁰ : p.precomputed.ᶜts + + e_tot_src(Sqₜ, ρ, ts, Φ) = Sqₜ * ρ * e_tot_0M_precipitation_sources_helper(thp, ts, Φ) + e_tot_srcʲ(Φ) = (Sqₜ, ρ, ts) -> e_tot_src(Sqₜ, ρ, ts, Φ) + @. ᶜS_ρq_tot = ᶜSqₜᵖ⁰ * ᶜρa⁰ + draft_sum((S, ρ) -> S * ρ, ᶜSqₜᵖʲs, ᶜsgs_ρaʲs) + ᶜS_ρe_tot⁰ = @. lazy(e_tot_src(ᶜSqₜᵖ⁰, ᶜρa⁰, ᶜts⁰, ᶜΦ)) + @. ᶜS_ρe_tot = ᶜS_ρe_tot⁰ + draft_sum(e_tot_srcʲ(ᶜΦ), ᶜSqₜᵖʲs, ᶜsgs_ρaʲs, ᶜtsʲs) return nothing end function set_precipitation_cache!(Y, p, ::Microphysics1Moment, _) (; dt) = p - (; ᶜts, ᶜwᵣ, ᶜwₛ, ᶜu) = p.precomputed + (; ᶜts) = p.precomputed (; ᶜSqₗᵖ, ᶜSqᵢᵖ, ᶜSqᵣᵖ, ᶜSqₛᵖ) = p.precomputed ᶜq_tot = @. lazy(specific(Y.c.ρq_tot, Y.c.ρ)) @@ -681,7 +375,6 @@ function set_precipitation_cache!(Y, p, ::Microphysics1Moment, _) ᶜSᵖ = p.scratch.ᶜtemp_scalar ᶜSᵖ_snow = p.scratch.ᶜtemp_scalar_2 - ᶜ∇T = p.scratch.ᶜtemp_CT123 # get thermodynamics and 1-moment microphysics params (; params) = p @@ -690,58 +383,25 @@ function set_precipitation_cache!(Y, p, ::Microphysics1Moment, _) # compute precipitation source terms on the grid mean compute_precipitation_sources!( - ᶜSᵖ, - ᶜSᵖ_snow, - ᶜSqₗᵖ, - ᶜSqᵢᵖ, - ᶜSqᵣᵖ, - ᶜSqₛᵖ, - Y.c.ρ, - ᶜq_tot, - ᶜq_liq, - ᶜq_ice, - ᶜq_rai, - ᶜq_sno, - ᶜts, - dt, - cmp, - thp, + ᶜSᵖ, ᶜSᵖ_snow, ᶜSqₗᵖ, ᶜSqᵢᵖ, ᶜSqᵣᵖ, ᶜSqₛᵖ, + Y.c.ρ, ᶜq_tot, ᶜq_liq, ᶜq_ice, ᶜq_rai, ᶜq_sno, + ᶜts, dt, cmp, thp, ) # compute precipitation sinks on the grid mean compute_precipitation_sinks!( - ᶜSᵖ, - ᶜSqᵣᵖ, - ᶜSqₛᵖ, - Y.c.ρ, - ᶜq_tot, - ᶜq_liq, - ᶜq_ice, - ᶜq_rai, - ᶜq_sno, - ᶜts, - dt, - cmp, - thp, + ᶜSᵖ, ᶜSqᵣᵖ, ᶜSqₛᵖ, + Y.c.ρ, ᶜq_tot, ᶜq_liq, ᶜq_ice, ᶜq_rai, ᶜq_sno, + ᶜts, dt, cmp, thp, ) return nothing end -function set_precipitation_cache!( - Y, - p, - ::Microphysics1Moment, - ::DiagnosticEDMFX, -) +function set_precipitation_cache!(Y, p, ::Microphysics1Moment, ::DiagnosticEDMFX) # Nothing needs to be done on the grid mean. The Sources are computed # in edmf sub-domains. return nothing end -function set_precipitation_cache!( - Y, - p, - ::Microphysics1Moment, - ::PrognosticEDMFX, -) +function set_precipitation_cache!(Y, p, ::Microphysics1Moment, ::PrognosticEDMFX) # Nothing needs to be done on the grid mean. The Sources are computed # in edmf sub-domains. return nothing @@ -751,6 +411,7 @@ function set_precipitation_cache!(Y, p, ::Microphysics2Moment, _) (; ᶜts) = p.precomputed (; ᶜSqₗᵖ, ᶜSqᵢᵖ, ᶜSqᵣᵖ, ᶜSqₛᵖ) = p.precomputed (; ᶜSnₗᵖ, ᶜSnᵣᵖ) = p.precomputed + ᶜρ = Y.c.ρ ᶜSᵖ = p.scratch.ᶜtemp_scalar ᶜS₂ᵖ = p.scratch.ᶜtemp_scalar_2 @@ -760,26 +421,19 @@ function set_precipitation_cache!(Y, p, ::Microphysics2Moment, _) cmp = CAP.microphysics_2m_params(params) thp = CAP.thermodynamics_params(params) + ᶜn_liq = @. lazy(specific(Y.c.ρn_liq, ᶜρ)) + ᶜn_rai = @. lazy(specific(Y.c.ρn_rai, ᶜρ)) + ᶜq_tot = @. lazy(specific(Y.c.ρq_tot, ᶜρ)) + ᶜq_liq = @. lazy(specific(Y.c.ρq_liq, ᶜρ)) + ᶜq_ice = @. lazy(specific(Y.c.ρq_ice, ᶜρ)) + ᶜq_rai = @. lazy(specific(Y.c.ρq_rai, ᶜρ)) + ᶜq_sno = @. lazy(specific(Y.c.ρq_sno, ᶜρ)) + # compute warm precipitation sources on the grid mean (based on SB2006 2M scheme) compute_warm_precipitation_sources_2M!( - ᶜSᵖ, - ᶜS₂ᵖ, - ᶜSnₗᵖ, - ᶜSnᵣᵖ, - ᶜSqₗᵖ, - ᶜSqᵣᵖ, - Y.c.ρ, - lazy.(specific.(Y.c.ρn_liq, Y.c.ρ)), - lazy.(specific.(Y.c.ρn_rai, Y.c.ρ)), - lazy.(specific.(Y.c.ρq_tot, Y.c.ρ)), - lazy.(specific.(Y.c.ρq_liq, Y.c.ρ)), - lazy.(specific.(Y.c.ρq_ice, Y.c.ρ)), - lazy.(specific.(Y.c.ρq_rai, Y.c.ρ)), - lazy.(specific.(Y.c.ρq_sno, Y.c.ρ)), - ᶜts, - dt, - cmp, - thp, + ᶜSᵖ, ᶜS₂ᵖ, ᶜSnₗᵖ, ᶜSnᵣᵖ, ᶜSqₗᵖ, ᶜSqᵣᵖ, + ᶜρ, ᶜn_liq, ᶜn_rai, ᶜq_tot, ᶜq_liq, ᶜq_ice, ᶜq_rai, ᶜq_sno, + ᶜts, dt, cmp, thp, ) #TODO - implement 2M cold processes! @@ -788,21 +442,11 @@ function set_precipitation_cache!(Y, p, ::Microphysics2Moment, _) return nothing end -function set_precipitation_cache!( - Y, - p, - ::Microphysics2Moment, - ::DiagnosticEDMFX, -) +function set_precipitation_cache!(Y, p, ::Microphysics2Moment, ::DiagnosticEDMFX) error("Not implemented yet") return nothing end -function set_precipitation_cache!( - Y, - p, - ::Microphysics2Moment, - ::PrognosticEDMFX, -) +function set_precipitation_cache!(Y, p, ::Microphysics2Moment, ::PrognosticEDMFX) # Nothing needs to be done on the grid mean. The Sources are computed # in edmf sub-domains. return nothing @@ -842,17 +486,13 @@ end """ set_precipitation_surface_fluxes!(Y, p, precipitation model) -Computes the flux of rain and snow at the surface. For the 0-moment microphysics -it is an integral of the source terms in the column. For 1-moment microphysics -it is the flux through the bottom cell face. +Computes the flux of rain and snow at the surface. + +For the 0-moment microphysics it is an integral of the source terms in the column. +For 1-moment microphysics it is the flux through the bottom cell face. """ set_precipitation_surface_fluxes!(Y, p, _) = nothing -function set_precipitation_surface_fluxes!( - Y, - p, - microphysics_model::Microphysics0Moment, -) - ᶜT = p.scratch.ᶜtemp_scalar +function set_precipitation_surface_fluxes!(Y, p, ::Microphysics0Moment) (; ᶜts) = p.precomputed # assume ᶜts has been updated (; ᶜS_ρq_tot, ᶜS_ρe_tot) = p.precomputed (; surface_rain_flux, surface_snow_flux) = p.precomputed @@ -860,24 +500,21 @@ function set_precipitation_surface_fluxes!( # update total column energy source for surface energy balance Operators.column_integral_definite!( - col_integrated_precip_energy_tendency, - ᶜS_ρe_tot, + col_integrated_precip_energy_tendency, ᶜS_ρe_tot, ) # update surface precipitation fluxes in cache for coupler's use thermo_params = CAP.thermodynamics_params(p.params) T_freeze = TD.Parameters.T_freeze(thermo_params) FT = eltype(p.params) - @. ᶜT = TD.air_temperature(thermo_params, ᶜts) + ᶜT = @. lazy(TD.air_temperature(thermo_params, ᶜts)) ᶜ3d_rain = @. lazy(ifelse(ᶜT >= T_freeze, ᶜS_ρq_tot, FT(0))) ᶜ3d_snow = @. lazy(ifelse(ᶜT < T_freeze, ᶜS_ρq_tot, FT(0))) Operators.column_integral_definite!(surface_rain_flux, ᶜ3d_rain) Operators.column_integral_definite!(surface_snow_flux, ᶜ3d_snow) return nothing end -function set_precipitation_surface_fluxes!( - Y, - p, - microphysics_model::Union{Microphysics1Moment, Microphysics2Moment}, +function set_precipitation_surface_fluxes!(Y, p, + ::Union{Microphysics1Moment, Microphysics2Moment}, ) (; surface_rain_flux, surface_snow_flux) = p.precomputed (; col_integrated_precip_energy_tendency) = p.conservation_check @@ -889,8 +526,7 @@ function set_precipitation_surface_fluxes!( # Jacobian-weighted extrapolation from interior to surface, consistent with # the reconstruction of density on cell faces, ᶠρ = ᶠinterp(Y.c.ρ * ᶜJ) / ᶠJ - sfc_lev(x) = - Fields.Field(Fields.field_values(Fields.level(x, 1)), sfc_space) + sfc_lev(x) = Fields.Field(Fields.field_values(Fields.level(x, 1)), sfc_space) int_J = sfc_lev(ᶜJ) int_ρ = sfc_lev(Y.c.ρ) sfc_ρ = @. lazy(int_ρ * int_J / sfc_J) @@ -905,21 +541,16 @@ function set_precipitation_surface_fluxes!( @. ᶜq_sno = specific(Y.c.ρq_sno, Y.c.ρ) @. ᶜq_liq = specific(Y.c.ρq_liq, Y.c.ρ) @. ᶜq_ice = specific(Y.c.ρq_ice, Y.c.ρ) - sfc_qᵣ = - Fields.Field(Fields.field_values(Fields.level(ᶜq_rai, 1)), sfc_space) - sfc_qₛ = - Fields.Field(Fields.field_values(Fields.level(ᶜq_sno, 1)), sfc_space) - sfc_qₗ = - Fields.Field(Fields.field_values(Fields.level(ᶜq_liq, 1)), sfc_space) - sfc_qᵢ = - Fields.Field(Fields.field_values(Fields.level(ᶜq_ice, 1)), sfc_space) + sfc_qᵣ = Fields.Field(Fields.field_values(Fields.level(ᶜq_rai, 1)), sfc_space) + sfc_qₛ = Fields.Field(Fields.field_values(Fields.level(ᶜq_sno, 1)), sfc_space) + sfc_qₗ = Fields.Field(Fields.field_values(Fields.level(ᶜq_liq, 1)), sfc_space) + sfc_qᵢ = Fields.Field(Fields.field_values(Fields.level(ᶜq_ice, 1)), sfc_space) sfc_wᵣ = Fields.Field(Fields.field_values(Fields.level(ᶜwᵣ, 1)), sfc_space) sfc_wₛ = Fields.Field(Fields.field_values(Fields.level(ᶜwₛ, 1)), sfc_space) sfc_wₗ = Fields.Field(Fields.field_values(Fields.level(ᶜwₗ, 1)), sfc_space) sfc_wᵢ = Fields.Field(Fields.field_values(Fields.level(ᶜwᵢ, 1)), sfc_space) sfc_wₕhₜ = Fields.Field( - Fields.field_values(Fields.level(ᶜwₕhₜ.components.data.:1, 1)), - sfc_space, + Fields.field_values(Fields.level(ᶜwₕhₜ.components.data.:1, 1)), sfc_space, ) @. surface_rain_flux = sfc_ρ * (sfc_qᵣ * (-sfc_wᵣ) + sfc_qₗ * (-sfc_wₗ)) diff --git a/src/parameterized_tendencies/microphysics/microphysics_wrappers.jl b/src/parameterized_tendencies/microphysics/microphysics_wrappers.jl index ebb4c608c7..a6c1aac51a 100644 --- a/src/parameterized_tendencies/microphysics/microphysics_wrappers.jl +++ b/src/parameterized_tendencies/microphysics/microphysics_wrappers.jl @@ -22,6 +22,38 @@ function clip(q) return max(FT(0), q) end +import ClimaCore.MatrixFields: @name, FieldName + +function terminal_velocity_func_1M(cmc, cmp, q_name) + vₗ(ρ, q) = CMNe.terminal_velocity(cmc.liquid, cmc.Ch2022.rain, ρ, q) + vᵢ(ρ, q) = CMNe.terminal_velocity(cmc.ice, cmc.Ch2022.small_ice, ρ, q) + vᵣ(ρ, q) = CM1.terminal_velocity(cmp.pr, cmp.tv.rain, ρ, q) + vₛ(ρ, q) = CM1.terminal_velocity(cmp.ps, cmp.tv.snow, ρ, q) + q_name == @name(q_liq) && return vₗ + q_name == @name(q_ice) && return vᵢ + q_name == @name(q_rai) && return vᵣ + q_name == @name(q_sno) && return vₛ + throw(ArgumentError("Invalid q_name")) +end + +function _terminal_velocity_func_2M(cm2p, cmc, cm1p, q_name) + # TODO: sedimentation of ice is based on the 1M scheme + # TODO: sedimentation of snow is based on the 1M scheme + vₗ(ρ, q, n) = CM2.cloud_terminal_velocity(cm2p.sb.pdf_c, cm2p.ctv, q, ρ, n) + vᵢ(ρ, q, _) = terminal_velocity_func_1M(cmc, cm1p, q_name)(ρ, q) + vᵣ(ρ, q, n) = CM2.rain_terminal_velocity(cm2p.sb, cm2p.rtv, q, ρ, n) + vₛ(ρ, q, _) = terminal_velocity_func_1M(cmc, cm1p, q_name)(ρ, q) + q_name == @name(q_liq) && return vₗ + q_name == @name(q_ice) && return vᵢ + q_name == @name(q_rai) && return vᵣ + q_name == @name(q_sno) && return vₛ + throw(ArgumentError("Invalid q_name")) +end +terminal_velocity_mass_func_2M(cm2p, cmc, cm1p, q_name) = + last ∘ _terminal_velocity_func_2M(cm2p, cmc, cm1p, q_name) +terminal_velocity_number_func_2M(cm2p, cmc, cm1p, q_name) = + first ∘ _terminal_velocity_func_2M(cm2p, cmc, cm1p, q_name) + # Helper function to compute the limit of the tendency in the traingle limiter. # The limit is defined as the available amont / n times model time step. function limit(q, dt, n::Int) diff --git a/src/utils/abbreviations.jl b/src/utils/abbreviations.jl index bedfec0a32..54c098fb9f 100644 --- a/src/utils/abbreviations.jl +++ b/src/utils/abbreviations.jl @@ -13,6 +13,7 @@ const CT12 = Geometry.Contravariant12Vector const CT3 = Geometry.Contravariant3Vector const CT123 = Geometry.Contravariant123Vector const UVW = Geometry.UVWVector +const WVec = Geometry.WVector const divₕ = Operators.Divergence() const wdivₕ = Operators.WeakDivergence() diff --git a/src/utils/variable_manipulations.jl b/src/utils/variable_manipulations.jl index e34b61ddff..6a3ecb3757 100644 --- a/src/utils/variable_manipulations.jl +++ b/src/utils/variable_manipulations.jl @@ -85,8 +85,7 @@ gs_tracer_names(Y) = `Tuple` of the specific tracer names `@name(χ)` that correspond to the density-weighted tracer names `@name(ρχ)` in `gs_tracer_names(Y)`. """ -specific_gs_tracer_names(Y) = - unrolled_map(specific_tracer_name, gs_tracer_names(Y)) +specific_gs_tracer_names(Y) = unrolled_map(specific_tracer_name, gs_tracer_names(Y)) """ ᶜempty(Y) @@ -110,6 +109,17 @@ function ᶜgs_tracers(Y) return @. lazy(NamedTuple{ρχ_symbols}(tuple(ρχ_fields...))) end +""" + ᶜspecific_gs_tracer(Y, χ_name) + +Lazy center `Field` of the specific tracer `χ` that corresponds to the +grid-scale tracer `ρχ` given by `χ_name`. +""" +function ᶜspecific_gs_tracer(Y, χ_name) + ρχ = MatrixFields.get_field(Y.c, get_ρχ_name(χ_name)) + @. lazy(specific(ρχ, Y.c.ρ)) +end + """ ᶜspecific_gs_tracers(Y) @@ -118,15 +128,71 @@ grid-scale tracers given by `specific_gs_tracer_names(Y)`. """ function ᶜspecific_gs_tracers(Y) isempty(gs_tracer_names(Y)) && return ᶜempty(Y) - χ_symbols = - unrolled_map(MatrixFields.extract_first, specific_gs_tracer_names(Y)) + χ_symbols = unrolled_map(MatrixFields.extract_first, specific_gs_tracer_names(Y)) χ_fields = unrolled_map(gs_tracer_names(Y)) do ρχ_name - ρχ_field = MatrixFields.get_field(Y.c, ρχ_name) - @. lazy(specific(ρχ_field, Y.c.ρ)) + ᶜspecific_gs_tracer(Y, specific_tracer_name(ρχ_name)) end return @. lazy(NamedTuple{χ_symbols}(tuple(χ_fields...))) end +""" + get_ᶜwᵪ_name_from_qᵪ_name(qᵪ_name) + +Returns the name of the center `Field` `ᶜwᵪ` that corresponds to the +grid-scale tracer `qᵪ_name`: +- `q_liq` -> `ᶜwₗ` +- `q_ice` -> `ᶜwᵢ` +- `q_rai` -> `ᶜwᵣ` +- `q_sno` -> `ᶜwₛ` +- otherwise `@name()` + +# Examples + +```julia-repl +julia> get_ᶜwᵪ_name_from_qᵪ_name(@name(q_liq)) +@name(ᶜwₗ) +``` + +See also [`get_ᶜwₙᵪ_name_from_qᵪ_name`](@ref) +""" +get_ᶜwᵪ_name_from_qᵪ_name(qᵪ_name) = + qᵪ_name == @name(q_liq) ? @name(ᶜwₗ) : + qᵪ_name == @name(q_ice) ? @name(ᶜwᵢ) : + qᵪ_name == @name(q_rai) ? @name(ᶜwᵣ) : + qᵪ_name == @name(q_sno) ? @name(ᶜwₛ) : @name() + +""" + get_ᶜwₙᵪ_name_from_qᵪ_name(qᵪ_name) + +Returns the name of the center `Field` `ᶜwₙᵪ` that corresponds to the +grid-scale tracer `qᵪ_name`: +- `q_liq` -> `ᶜwₙₗ` +- `q_rai` -> `ᶜwₙᵣ` +- otherwise `@name()` + +See also [`get_ᶜwᵪ_name_from_qᵪ_name`](@ref) +""" +get_ᶜwₙᵪ_name_from_qᵪ_name(qᵪ_name) = + qᵪ_name == @name(q_liq) ? @name(ᶜwₙₗ) : + qᵪ_name == @name(q_rai) ? @name(ᶜwₙᵣ) : @name() + +""" + get_nᵪ_name_from_qᵪ_name(qᵪ_name) + +Returns the `nᵪ_name` that corresponds to the +grid-scale tracer `qᵪ_name`: +- `q_liq` -> `n_liq` +- `q_rai` -> `n_rai` +- otherwise `@name()` +""" +get_nᵪ_name_from_qᵪ_name(qᵪ_name) = + qᵪ_name == @name(q_liq) ? @name(n_liq) : + qᵪ_name == @name(q_rai) ? @name(n_rai) : @name() + +get_ρnᵪ_name_from_qᵪ_name(qᵪ_name) = + qᵪ_name == @name(q_liq) ? @name(ρn_liq) : + qᵪ_name == @name(q_rai) ? @name(ρn_rai) : @name() + """ foreach_gs_tracer(f, Y_or_similar_values...) @@ -241,7 +307,7 @@ Arguments: - `f`: A function to apply to each element of `sgsʲs`. - `sgsʲs`: An iterator over the draft subdomain states. """ -draft_sum(f, sgsʲs) = unrolled_sum(f, sgsʲs) +draft_sum(f, sgsʲs...) = unrolled_mapreduce(f, +, sgsʲs...) """ ᶜenv_value(grid_scale_value, f_draft, gs, turbconv_model) @@ -368,6 +434,12 @@ The function operates recursively on hierarchical field names: to its specific form. - If `ρχ_name` has internal structure (e.g. a composite name), the function recurses into the child names and applies the transformation at the lowest level. + +# Examples +```julia-repl +julia> get_χʲ_name_from_ρχ_name(@name(ρq_tot)) +@name(sgsʲs.:(1).q_tot) +``` """ function get_χʲ_name_from_ρχ_name(ρχ_name) parent_name = MatrixFields.FieldName(MatrixFields.extract_first(ρχ_name)) @@ -399,18 +471,36 @@ Arguments: Returns: - The area-weighted density of the environment (`ρa⁰`). """ +ρa⁰(ρ, sgsʲs, ::PrognosticEDMFX) = env_value(ρ, sgsʲ -> sgsʲ.ρa, sgsʲs) # ρ - Σ ρaʲ +ρa⁰(ρ, sgsʲs, ::DiagnosticEDMFX) = env_value(ρ, ᶜρaʲ -> ᶜρaʲ, sgsʲs) # ρ - Σ ρaʲ +ρa⁰(ρ, _, _) = ρ -function ρa⁰(ρ, sgsʲs, turbconv_model) - # ρ - Σ ρaʲ - if turbconv_model isa PrognosticEDMFX - return env_value(ρ, sgsʲ -> sgsʲ.ρa, sgsʲs) +""" + ᶜρaʲs_list(Y, p, turbconv_model) - elseif turbconv_model isa DiagnosticEDMFX - return env_value(ρ, ᶜρaʲ -> ᶜρaʲ, sgsʲs) - else - return ρ - end -end +Returns the list of subdomain density containers + +Pass each element of the list to [`ρaʲ`](@ref) to get the subdomain density. +""" +ᶜρaʲs_list(Y, p, ::PrognosticEDMFX) = Y.c.sgsʲs +ᶜρaʲs_list(Y, p, ::DiagnosticEDMFX) = p.precomputed.ᶜρaʲs + +""" + ρaʲ(sgsʲ, turbconv_model) + +Returns the subdomain density from the subdomain density container. + +Use with [`ᶜρaʲs_list`](@ref), which returns the list whose elements +are the first argument to this function. + +# Examples +```julia +ᶜsgs_ρaʲs = ᶜρaʲs_list(Y, p, turbconv_model) +draft_sum = draft_sum(ρaʲ, ᶜsgs_ρaʲs) +``` +""" +ρaʲ(sgsʲ, ::PrognosticEDMFX) = sgsʲ.ρa +ρaʲ(ᶜρaʲ, ::DiagnosticEDMFX) = ᶜρaʲ """ ᶜspecific_env_mse(Y, p)