diff --git a/src/prognostic_equations/edmfx_sgs_flux.jl b/src/prognostic_equations/edmfx_sgs_flux.jl index 3ede2bfd86..97ae9d4576 100644 --- a/src/prognostic_equations/edmfx_sgs_flux.jl +++ b/src/prognostic_equations/edmfx_sgs_flux.jl @@ -382,160 +382,13 @@ function edmfx_sgs_diffusive_flux_tendency!( Y, p, t, - turbconv_model::PrognosticEDMFX, + turbconv_model::Union{EDOnlyEDMFX, DiagnosticEDMFX, PrognosticEDMFX}, ) - FT = Spaces.undertype(axes(Y.c)) - (; dt, params) = p - thermo_params = CAP.thermodynamics_params(params) - turbconv_params = CAP.turbconv_params(params) - c_d = CAP.tke_diss_coeff(turbconv_params) - (; ᶜK, ᶜu⁰, ᶜK⁰, ᶜlinear_buoygrad, ᶜstrain_rate_norm, ᶜρʲs, ᶜts, ᶜts⁰) = p.precomputed - (; ρatke_flux) = p.precomputed - ᶠgradᵥ = Operators.GradientC2F() - ᶜρa⁰ = @. lazy(ρa⁰(Y.c.ρ, Y.c.sgsʲs, turbconv_model)) - ᶜtke⁰ = @. lazy(specific(Y.c.sgs⁰.ρatke, Y.c.ρ)) - - if p.atmos.edmfx_model.sgs_diffusive_flux isa Val{true} - - (; ᶜlinear_buoygrad, ᶜstrain_rate_norm) = p.precomputed - # scratch to prevent GPU Kernel parameter memory error - ᶜmixing_length_field = p.scratch.ᶜtemp_scalar_2 - ᶜmixing_length_field .= ᶜmixing_length(Y, p) - ᶜK_u = @. lazy( - eddy_viscosity(turbconv_params, ᶜtke⁰, ᶜmixing_length_field), - ) - ᶜprandtl_nvec = @. lazy( - turbulent_prandtl_number( - params, - ᶜlinear_buoygrad, - ᶜstrain_rate_norm, - ), - ) - ᶜK_h = @. lazy(eddy_diffusivity(ᶜK_u, ᶜprandtl_nvec)) - ᶠρaK_h = p.scratch.ᶠtemp_scalar - @. ᶠρaK_h = ᶠinterp(Y.c.ρ) * ᶠinterp(ᶜK_h) - ᶠρaK_u = p.scratch.ᶠtemp_scalar_2 - @. ᶠρaK_u = ᶠinterp(Y.c.ρ) * ᶠinterp(ᶜK_u) - # Total enthalpy diffusion - ᶜdivᵥ_ρe_tot = Operators.DivergenceF2C( - top = Operators.SetValue(C3(FT(0))), - bottom = Operators.SetValue(C3(FT(0))), - ) - - ᶜh_tot = @. lazy( - TD.total_specific_enthalpy( - thermo_params, - ᶜts, - specific(Y.c.ρe_tot, Y.c.ρ), - ), - ) - @. Yₜ.c.ρe_tot -= ᶜdivᵥ_ρe_tot(-(ᶠρaK_h * ᶠgradᵥ(ᶜh_tot))) - if use_prognostic_tke(turbconv_model) - # Turbulent TKE transport (diffusion) - ᶜdivᵥ_ρatke = Operators.DivergenceF2C( - top = Operators.SetValue(C3(FT(0))), - bottom = Operators.SetValue(ρatke_flux), - ) - # Add flux divergence and dissipation term, relaxing TKE to zero - # in one time step if tke < 0 - @. Yₜ.c.sgs⁰.ρatke -= - ᶜdivᵥ_ρatke(-(ᶠρaK_u * ᶠgradᵥ(ᶜtke⁰))) + ifelse( - ᶜtke⁰ >= FT(0), - tke_dissipation( - turbconv_params, - Y.c.sgs⁰.ρatke, - ᶜtke⁰, - ᶜmixing_length_field, - ), - Y.c.sgs⁰.ρatke / dt, - ) - end - if !(p.atmos.moisture_model isa DryModel) - # Specific humidity diffusion - ᶜρχₜ_diffusion = p.scratch.ᶜtemp_scalar - ᶜdivᵥ_ρq_tot = Operators.DivergenceF2C( - top = Operators.SetValue(C3(FT(0))), - bottom = Operators.SetValue(C3(FT(0))), - ) - @. ᶜρχₜ_diffusion = - ᶜdivᵥ_ρq_tot(-(ᶠρaK_h * ᶠgradᵥ(specific(Y.c.ρq_tot, Y.c.ρ)))) - @. Yₜ.c.ρq_tot -= ᶜρχₜ_diffusion - @. Yₜ.c.ρ -= ᶜρχₜ_diffusion # Effect of moisture diffusion on (moist) air mass - end - - if p.atmos.moisture_model isa NonEquilMoistModel && ( - p.atmos.microphysics_model isa Microphysics1Moment || - p.atmos.microphysics_model isa Microphysics2Moment - ) - α_precip = CAP.α_vert_diff_tracer(params) - ᶜρχₜ_diffusion = p.scratch.ᶜtemp_scalar - ᶜdivᵥ_ρq = Operators.DivergenceF2C( - top = Operators.SetValue(C3(FT(0))), - bottom = Operators.SetValue(C3(FT(0))), - ) - ᶜρʲ = ᶜρʲs.:(1) - ᶜρ⁰ = @. lazy(TD.air_density(thermo_params, ᶜts⁰)) - - microphysics_tracers = ( - (@name(c.ρq_liq), @name(c.sgsʲs.:(1).q_liq), @name(q_liq), FT(1)), - (@name(c.ρq_ice), @name(c.sgsʲs.:(1).q_ice), @name(q_ice), FT(1)), - (@name(c.ρq_rai), @name(c.sgsʲs.:(1).q_rai), @name(q_rai), α_precip), - (@name(c.ρq_sno), @name(c.sgsʲs.:(1).q_sno), @name(q_sno), α_precip), - (@name(c.ρn_liq), @name(c.sgsʲs.:(1).n_liq), @name(n_liq), FT(1)), - (@name(c.ρn_rai), @name(c.sgsʲs.:(1).n_rai), @name(n_rai), α_precip), - ) - - MatrixFields.unrolled_foreach( - microphysics_tracers, - ) do (ρχ_name, χʲ_name, χ_name, α) - MatrixFields.has_field(Y, ρχ_name) || return - - ᶜρχₜ = MatrixFields.get_field(Yₜ, ρχ_name) - ᶜχʲ = MatrixFields.get_field(Y, χʲ_name) - ᶜχ⁰ = ᶜspecific_env_value(χ_name, Y, p) - - # add updraft diffusion - @. ᶜρχₜ_diffusion = - draft_area(Y.c.sgsʲs.:(1).ρa, ᶜρʲ) * - ᶜdivᵥ_ρq( - -(ᶠinterp(ᶜρʲ) * ᶠinterp(ᶜK_h) * α * ᶠgradᵥ(ᶜχʲ)), - ) - @. ᶜρχₜ -= ᶜρχₜ_diffusion - - # add environment diffusion - @. ᶜρχₜ_diffusion = - draft_area(ᶜρa⁰, ᶜρ⁰) * - ᶜdivᵥ_ρq( - -(ᶠinterp(ᶜρ⁰) * ᶠinterp(ᶜK_h) * α * ᶠgradᵥ(ᶜχ⁰)), - ) - @. ᶜρχₜ -= ᶜρχₜ_diffusion - end - end - - # Momentum diffusion - ᶠstrain_rate = compute_strain_rate_face_vertical(ᶜu⁰) - @. Yₜ.c.uₕ -= C12(ᶜdivᵥ(-(2 * ᶠρaK_u * ᶠstrain_rate)) / Y.c.ρ) - end - return nothing -end - -function edmfx_sgs_diffusive_flux_tendency!( - Yₜ, - Y, - p, - t, - turbconv_model::Union{EDOnlyEDMFX, DiagnosticEDMFX}, -) - - # Assumes envinronmental area fraction is 1 (so draft area fraction is negligible) - # TODO: Relax this assumption and construct diagnostic EDMF fluxes in parallel to - # prognostic fluxes FT = Spaces.undertype(axes(Y.c)) (; dt, params) = p turbconv_params = CAP.turbconv_params(params) thermo_params = CAP.thermodynamics_params(params) - c_d = CAP.tke_diss_coeff(turbconv_params) (; ᶜu, ᶜts) = p.precomputed (; ρatke_flux) = p.precomputed ᶠgradᵥ = Operators.GradientC2F() @@ -609,27 +462,24 @@ function edmfx_sgs_diffusive_flux_tendency!( @. ᶜρχₜ_diffusion = ᶜdivᵥ_ρq_tot(-(ᶠρaK_h * ᶠgradᵥ(specific(Y.c.ρq_tot, Y.c.ρ)))) @. Yₜ.c.ρq_tot -= ᶜρχₜ_diffusion - @. Yₜ.c.ρ -= ᶜρχₜ_diffusion + @. Yₜ.c.ρ -= ᶜρχₜ_diffusion # Effect of moisture diffusion on (moist) air mass end - cloud_tracers = (@name(c.ρq_liq), @name(c.ρq_ice), @name(c.ρn_liq)) - precip_tracers = (@name(c.ρq_rai), @name(c.ρq_sno), @name(c.ρn_rai)) - - α = CAP.α_vert_diff_tracer(params) + α_precip = CAP.α_vert_diff_tracer(params) ᶜρχₜ_diffusion = p.scratch.ᶜtemp_scalar ᶜdivᵥ_ρq = Operators.DivergenceF2C( top = Operators.SetValue(C3(FT(0))), bottom = Operators.SetValue(C3(FT(0))), ) - MatrixFields.unrolled_foreach(cloud_tracers) do ρχ_name - MatrixFields.has_field(Y, ρχ_name) || return - ᶜρχ = MatrixFields.get_field(Y, ρχ_name) - ᶜχ = (@. lazy(specific(ᶜρχ, Y.c.ρ))) - @. ᶜρχₜ_diffusion = ᶜdivᵥ_ρq(-(ᶠρaK_h * ᶠgradᵥ(ᶜχ))) - ᶜρχₜ = MatrixFields.get_field(Yₜ, ρχ_name) - @. ᶜρχₜ -= ᶜρχₜ_diffusion - end - MatrixFields.unrolled_foreach(precip_tracers) do ρχ_name + microphysics_tracers = ( + (@name(c.ρq_liq), FT(1)), + (@name(c.ρq_ice), FT(1)), + (@name(c.ρq_rai), α_precip), + (@name(c.ρq_sno), α_precip), + (@name(c.ρn_liq), FT(1)), + (@name(c.ρn_rai), α_precip), + ) + MatrixFields.unrolled_foreach(microphysics_tracers) do (ρχ_name, α) MatrixFields.has_field(Y, ρχ_name) || return ᶜρχ = MatrixFields.get_field(Y, ρχ_name) ᶜχ = (@. lazy(specific(ᶜρχ, Y.c.ρ)))