From 46a213d9bf596c8064e7124ef20fad381a949347 Mon Sep 17 00:00:00 2001 From: Charles Kawczynski Date: Wed, 9 Oct 2024 11:21:11 -0400 Subject: [PATCH] Fuse some broadcast expressions Apply formatter --- src/ClimaAtmos.jl | 2 + .../diagnostic_edmf_precomputed_quantities.jl | 20 +- .../precipitation_precomputed_quantities.jl | 34 ++- src/cache/precomputed_quantities.jl | 11 +- .../prognostic_edmf_precomputed_quantities.jl | 54 ++-- .../microphysics/cloud_condensate.jl | 6 +- .../microphysics/microphysics_wrappers.jl | 274 +++++++++--------- .../microphysics/precipitation.jl | 66 +++-- src/utils/common_spaces.jl | 9 +- 9 files changed, 263 insertions(+), 213 deletions(-) diff --git a/src/ClimaAtmos.jl b/src/ClimaAtmos.jl index d09a10bddc..20699660e7 100644 --- a/src/ClimaAtmos.jl +++ b/src/ClimaAtmos.jl @@ -1,5 +1,7 @@ module ClimaAtmos +import ClimaCore.Fields: @fused_direct + using NVTX import Thermodynamics as TD diff --git a/src/cache/diagnostic_edmf_precomputed_quantities.jl b/src/cache/diagnostic_edmf_precomputed_quantities.jl index 81b10f935b..dc06ef366f 100644 --- a/src/cache/diagnostic_edmf_precomputed_quantities.jl +++ b/src/cache/diagnostic_edmf_precomputed_quantities.jl @@ -39,15 +39,17 @@ NVTX.@annotate function set_diagnostic_edmfx_draft_quantities_level!( Φ_level, ) FT = eltype(thermo_params) - @. ts_level = TD.PhaseEquil_phq( - thermo_params, - p_level, - mse_level - Φ_level, - q_tot_level, - 8, - FT(0.0003), - ) - @. ρ_level = TD.air_density(thermo_params, ts_level) + @fused_direct begin + @. ts_level = TD.PhaseEquil_phq( + thermo_params, + p_level, + mse_level - Φ_level, + q_tot_level, + 8, + FT(0.0003), + ) + @. ρ_level = TD.air_density(thermo_params, ts_level) + end return nothing end diff --git a/src/cache/precipitation_precomputed_quantities.jl b/src/cache/precipitation_precomputed_quantities.jl index a9c64c68b1..921212804d 100644 --- a/src/cache/precipitation_precomputed_quantities.jl +++ b/src/cache/precipitation_precomputed_quantities.jl @@ -21,22 +21,24 @@ function set_precipitation_precomputed_quantities!(Y, p, t) cmp = CAP.microphysics_precipitation_params(p.params) - # compute the precipitation specific humidities - @. ᶜqᵣ = qₚ(Y.c.ρq_rai, Y.c.ρ) - @. ᶜqₛ = qₚ(Y.c.ρq_sno, Y.c.ρ) + @fused_direct begin + # compute the precipitation specific humidities + @. ᶜqᵣ = qₚ(Y.c.ρq_rai, Y.c.ρ) + @. ᶜqₛ = qₚ(Y.c.ρq_sno, Y.c.ρ) - # compute the precipitation terminal velocity [m/s] - @. ᶜwᵣ = CM1.terminal_velocity( - cmp.pr, - cmp.tv.rain, - Y.c.ρ, - abs(Y.c.ρq_rai / Y.c.ρ), - ) - @. ᶜwₛ = CM1.terminal_velocity( - cmp.ps, - cmp.tv.snow, - Y.c.ρ, - abs(Y.c.ρq_sno / Y.c.ρ), - ) + # compute the precipitation terminal velocity [m/s] + @. ᶜwᵣ = CM1.terminal_velocity( + cmp.pr, + cmp.tv.rain, + Y.c.ρ, + abs(Y.c.ρq_rai / Y.c.ρ), + ) + @. ᶜwₛ = CM1.terminal_velocity( + cmp.ps, + cmp.tv.snow, + Y.c.ρ, + abs(Y.c.ρq_sno / Y.c.ρ), + ) + end return nothing end diff --git a/src/cache/precomputed_quantities.jl b/src/cache/precomputed_quantities.jl index 5d41912fd4..9e46d02973 100644 --- a/src/cache/precomputed_quantities.jl +++ b/src/cache/precomputed_quantities.jl @@ -477,8 +477,13 @@ NVTX.@annotate function set_precomputed_quantities!(Y, p, t) # @. ᶜK += Y.c.sgs⁰.ρatke / Y.c.ρ # TODO: We should think more about these increments before we use them. end - @. ᶜts = ts_gs(thermo_args..., ᶜspecific, ᶜK, ᶜΦ, Y.c.ρ) - @. ᶜp = TD.air_pressure(thermo_params, ᶜts) + (; ᶜh_tot) = p.precomputed + @fused_direct begin + @. ᶜts = ts_gs(thermo_args..., ᶜspecific, ᶜK, ᶜΦ, Y.c.ρ) + @. ᶜp = TD.air_pressure(thermo_params, ᶜts) + @. ᶜh_tot = + TD.total_specific_enthalpy(thermo_params, ᶜts, ᶜspecific.e_tot) + end if turbconv_model isa AbstractEDMF @. p.precomputed.ᶜgradᵥ_θ_virt = @@ -489,8 +494,6 @@ NVTX.@annotate function set_precomputed_quantities!(Y, p, t) ᶜgradᵥ(ᶠinterp(TD.liquid_ice_pottemp(thermo_params, ᶜts))) end - (; ᶜh_tot) = p.precomputed - @. ᶜh_tot = TD.total_specific_enthalpy(thermo_params, ᶜts, ᶜspecific.e_tot) if !isnothing(p.sfc_setup) SurfaceConditions.update_surface_conditions!(Y, p, t) diff --git a/src/cache/prognostic_edmf_precomputed_quantities.jl b/src/cache/prognostic_edmf_precomputed_quantities.jl index eb620888bb..8bd513a74c 100644 --- a/src/cache/prognostic_edmf_precomputed_quantities.jl +++ b/src/cache/prognostic_edmf_precomputed_quantities.jl @@ -25,27 +25,31 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_environment!( (; ᶜtke⁰, ᶜρa⁰, ᶠu₃⁰, ᶜu⁰, ᶠu³⁰, ᶜK⁰, ᶜts⁰, ᶜρ⁰, ᶜmse⁰, ᶜq_tot⁰) = p.precomputed - @. ᶜρa⁰ = ρa⁰(Y.c) - @. ᶜtke⁰ = divide_by_ρa(Y.c.sgs⁰.ρatke, ᶜρa⁰, 0, Y.c.ρ, turbconv_model) - @. ᶜmse⁰ = divide_by_ρa( - Y.c.ρ * (ᶜh_tot - ᶜK) - ρamse⁺(Y.c.sgsʲs), - ᶜρa⁰, - Y.c.ρ * (ᶜh_tot - ᶜK), - Y.c.ρ, - turbconv_model, - ) - @. ᶜq_tot⁰ = divide_by_ρa( - Y.c.ρq_tot - ρaq_tot⁺(Y.c.sgsʲs), - ᶜρa⁰, - Y.c.ρq_tot, - Y.c.ρ, - turbconv_model, - ) + @fused_direct begin + @. ᶜρa⁰ = ρa⁰(Y.c) + @. ᶜtke⁰ = divide_by_ρa(Y.c.sgs⁰.ρatke, ᶜρa⁰, 0, Y.c.ρ, turbconv_model) + @. ᶜmse⁰ = divide_by_ρa( + Y.c.ρ * (ᶜh_tot - ᶜK) - ρamse⁺(Y.c.sgsʲs), + ᶜρa⁰, + Y.c.ρ * (ᶜh_tot - ᶜK), + Y.c.ρ, + turbconv_model, + ) + @. ᶜq_tot⁰ = divide_by_ρa( + Y.c.ρq_tot - ρaq_tot⁺(Y.c.sgsʲs), + ᶜρa⁰, + Y.c.ρq_tot, + Y.c.ρ, + turbconv_model, + ) + end set_sgs_ᶠu₃!(u₃⁰, ᶠu₃⁰, Y, turbconv_model) set_velocity_quantities!(ᶜu⁰, ᶠu³⁰, ᶜK⁰, ᶠu₃⁰, Y.c.uₕ, ᶠuₕ³) # @. ᶜK⁰ += ᶜtke⁰ - @. ᶜts⁰ = TD.PhaseEquil_phq(thermo_params, ᶜp, ᶜmse⁰ - ᶜΦ, ᶜq_tot⁰) - @. ᶜρ⁰ = TD.air_density(thermo_params, ᶜts⁰) + @fused_direct begin + @. ᶜts⁰ = TD.PhaseEquil_phq(thermo_params, ᶜp, ᶜmse⁰ - ᶜΦ, ᶜq_tot⁰) + @. ᶜρ⁰ = TD.air_density(thermo_params, ᶜts⁰) + end return nothing end @@ -90,8 +94,10 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_draft_and_bc! set_velocity_quantities!(ᶜuʲ, ᶠu³ʲ, ᶜKʲ, ᶠu₃ʲ, Y.c.uₕ, ᶠuₕ³) @. ᶠKᵥʲ = (adjoint(CT3(ᶠu₃ʲ)) * ᶠu₃ʲ) / 2 - @. ᶜtsʲ = TD.PhaseEquil_phq(thermo_params, ᶜp, ᶜmseʲ - ᶜΦ, ᶜq_totʲ) - @. ᶜρʲ = TD.air_density(thermo_params, ᶜtsʲ) + @fused_direct begin + @. ᶜtsʲ = TD.PhaseEquil_phq(thermo_params, ᶜp, ᶜmseʲ - ᶜΦ, ᶜq_totʲ) + @. ᶜρʲ = TD.air_density(thermo_params, ᶜtsʲ) + end # EDMFX boundary condition: @@ -335,12 +341,14 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_closures!( ᶜtke_exch, ) - @. ᶜmixing_length = ᶜmixing_length_tuple.master turbconv_params = CAP.turbconv_params(params) c_m = CAP.tke_ed_coeff(turbconv_params) - @. ᶜK_u = c_m * ᶜmixing_length * sqrt(max(ᶜtke⁰, 0)) - @. ᶜK_h = ᶜK_u / ᶜprandtl_nvec + @fused_direct begin + @. ᶜmixing_length = ᶜmixing_length_tuple.master + @. ᶜK_u = c_m * ᶜmixing_length * sqrt(max(ᶜtke⁰, 0)) + @. ᶜK_h = ᶜK_u / ᶜprandtl_nvec + end ρatke_flux_values = Fields.field_values(ρatke_flux) ρ_int_values = Fields.field_values(Fields.level(ᶜρa⁰, 1)) diff --git a/src/parameterized_tendencies/microphysics/cloud_condensate.jl b/src/parameterized_tendencies/microphysics/cloud_condensate.jl index f4739fb0e6..424ad88c37 100644 --- a/src/parameterized_tendencies/microphysics/cloud_condensate.jl +++ b/src/parameterized_tendencies/microphysics/cloud_condensate.jl @@ -15,6 +15,8 @@ function cloud_condensate_tendency!(Yₜ, p, ::NonEquilMoistModel) thp = CAP.thermodynamics_params(params) cmc = CAP.microphysics_cloud_params(params) - @. Yₜ.c.ρq_liq += cloud_sources(cmc.liquid, thp, ᶜts, dt) - @. Yₜ.c.ρq_ice += cloud_sources(cmc.ice, thp, ᶜts, dt) + @fused_direct begin + @. Yₜ.c.ρq_liq += cloud_sources(cmc.liquid, thp, ᶜts, dt) + @. Yₜ.c.ρq_ice += cloud_sources(cmc.ice, thp, ᶜts, dt) + end end diff --git a/src/parameterized_tendencies/microphysics/microphysics_wrappers.jl b/src/parameterized_tendencies/microphysics/microphysics_wrappers.jl index bd499f3d82..700c394522 100644 --- a/src/parameterized_tendencies/microphysics/microphysics_wrappers.jl +++ b/src/parameterized_tendencies/microphysics/microphysics_wrappers.jl @@ -142,107 +142,112 @@ function compute_precipitation_sources!( thp, ) FT = eltype(thp) - # @. Sqₜᵖ = FT(0) should work after fixing - # https://github.com/CliMA/ClimaCore.jl/issues/1786 - @. Sqₜᵖ = ρ * FT(0) - @. Sqᵣᵖ = ρ * FT(0) - @. Sqₛᵖ = ρ * FT(0) - @. Seₜᵖ = ρ * FT(0) - - #! format: off - # rain autoconversion: q_liq -> q_rain - @. Sᵖ = ifelse( - mp.Ndp <= 0, - CM1.conv_q_liq_to_q_rai(mp.pr.acnv1M, qₗ(thp, ts), true), - CM2.conv_q_liq_to_q_rai(mp.var, qₗ(thp, ts), ρ, mp.Ndp), - ) - @. Sᵖ = min(limit(qₗ(thp, ts), dt, 5), Sᵖ) - @. Sqₜᵖ -= Sᵖ - @. Sqᵣᵖ += Sᵖ - @. Seₜᵖ -= Sᵖ * (Iₗ(thp, ts) + Φ) - - # snow autoconversion assuming no supersaturation: q_ice -> q_snow - @. Sᵖ = min( - limit(qᵢ(thp, ts), dt, 5), - CM1.conv_q_ice_to_q_sno_no_supersat(mp.ps.acnv1M, qᵢ(thp, ts), true), - ) - @. Sqₜᵖ -= Sᵖ - @. Sqₛᵖ += Sᵖ - @. Seₜᵖ -= Sᵖ * (Iᵢ(thp, ts) + Φ) - - # accretion: q_liq + q_rain -> q_rain - @. Sᵖ = min( - limit(qₗ(thp, ts), dt, 5), - CM1.accretion(mp.cl, mp.pr, mp.tv.rain, mp.ce, qₗ(thp, ts), qᵣ, ρ), - ) - @. Sqₜᵖ -= Sᵖ - @. Sqᵣᵖ += Sᵖ - @. Seₜᵖ -= Sᵖ * (Iₗ(thp, ts) + Φ) - - # accretion: q_ice + q_snow -> q_snow - @. Sᵖ = min( - limit(qᵢ(thp, ts), dt, 5), - CM1.accretion(mp.ci, mp.ps, mp.tv.snow, mp.ce, qᵢ(thp, ts), qₛ, ρ), - ) - @. Sqₜᵖ -= Sᵖ - @. Sqₛᵖ += Sᵖ - @. Seₜᵖ -= Sᵖ * (Iᵢ(thp, ts) + Φ) - - # accretion: q_liq + q_sno -> q_sno or q_rai - # sink of cloud water via accretion cloud water + snow - @. Sᵖ = min( - limit(qₗ(thp, ts), dt, 5), - CM1.accretion(mp.cl, mp.ps, mp.tv.snow, mp.ce, qₗ(thp, ts), qₛ, ρ), - ) - # if T < T_freeze cloud droplets freeze to become snow - # else the snow melts and both cloud water and snow become rain α(thp, ts) = cᵥₗ(thp) / Lf(thp, ts) * (Tₐ(thp, ts) - mp.ps.T_freeze) - @. Sᵖ_snow = ifelse( - Tₐ(thp, ts) < mp.ps.T_freeze, - Sᵖ, - FT(-1) * min(Sᵖ * α(thp, ts), limit(qₛ, dt, 5)), - ) - @. Sqₛᵖ += Sᵖ_snow - @. Sqₜᵖ -= Sᵖ - @. Sqᵣᵖ += ifelse(Tₐ(thp, ts) < mp.ps.T_freeze, FT(0), Sᵖ - Sᵖ_snow) - @. Seₜᵖ -= ifelse( - Tₐ(thp, ts) < mp.ps.T_freeze, - Sᵖ * (Iᵢ(thp, ts) + Φ), - Sᵖ * (Iₗ(thp, ts) + Φ) - Sᵖ_snow * (Iₗ(thp, ts) - Iᵢ(thp, ts)), - ) - - # accretion: q_ice + q_rai -> q_sno - @. Sᵖ = min( - limit(qᵢ(thp, ts), dt, 5), - CM1.accretion(mp.ci, mp.pr, mp.tv.rain, mp.ce, qᵢ(thp, ts), qᵣ, ρ), - ) - @. Sqₜᵖ -= Sᵖ - @. Sqₛᵖ += Sᵖ - @. Seₜᵖ -= Sᵖ * (Iᵢ(thp, ts) + Φ) - # sink of rain via accretion cloud ice - rain - @. Sᵖ = min( - limit(qᵣ, dt, 5), - CM1.accretion_rain_sink(mp.pr, mp.ci, mp.tv.rain, mp.ce, qᵢ(thp, ts), qᵣ, ρ), - ) - @. Sqᵣᵖ -= Sᵖ - @. Sqₛᵖ += Sᵖ - @. Seₜᵖ += Sᵖ * Lf(thp, ts) - - # accretion: q_rai + q_sno -> q_rai or q_sno - @. Sᵖ = ifelse( - Tₐ(thp, ts) < mp.ps.T_freeze, - min( + @fused_direct begin + @. Sqₜᵖ = FT(0) + @. Sqᵣᵖ = FT(0) + @. Sqₛᵖ = FT(0) + @. Seₜᵖ = FT(0) + + #! format: off + # rain autoconversion: q_liq -> q_rain + @. Sᵖ = ifelse( + mp.Ndp <= 0, + CM1.conv_q_liq_to_q_rai(mp.pr.acnv1M, qₗ(thp, ts), true), + CM2.conv_q_liq_to_q_rai(mp.var, qₗ(thp, ts), ρ, mp.Ndp), + ) + @. Sᵖ = min(limit(qₗ(thp, ts), dt, 5), Sᵖ) + @. Sqₜᵖ -= Sᵖ + @. Sqᵣᵖ += Sᵖ + @. Seₜᵖ -= Sᵖ * (Iₗ(thp, ts) + Φ) + + # snow autoconversion assuming no supersaturation: q_ice -> q_snow + @. Sᵖ = min( + limit(qᵢ(thp, ts), dt, 5), + CM1.conv_q_ice_to_q_sno_no_supersat(mp.ps.acnv1M, qᵢ(thp, ts), true), + ) + @. Sqₜᵖ -= Sᵖ + @. Sqₛᵖ += Sᵖ + @. Seₜᵖ -= Sᵖ * (Iᵢ(thp, ts) + Φ) + + # accretion: q_liq + q_rain -> q_rain + @. Sᵖ = min( + limit(qₗ(thp, ts), dt, 5), + CM1.accretion(mp.cl, mp.pr, mp.tv.rain, mp.ce, qₗ(thp, ts), qᵣ, ρ), + ) + + @. Sqₜᵖ -= Sᵖ + @. Sqᵣᵖ += Sᵖ + @. Seₜᵖ -= Sᵖ * (Iₗ(thp, ts) + Φ) + + # accretion: q_ice + q_snow -> q_snow + @. Sᵖ = min( + limit(qᵢ(thp, ts), dt, 5), + CM1.accretion(mp.ci, mp.ps, mp.tv.snow, mp.ce, qᵢ(thp, ts), qₛ, ρ), + ) + @. Sqₜᵖ -= Sᵖ + @. Sqₛᵖ += Sᵖ + + @. Seₜᵖ -= Sᵖ * (Iᵢ(thp, ts) + Φ) + # accretion: q_liq + q_sno -> q_sno or q_rai + # sink of cloud water via accretion cloud water + snow + @. Sᵖ = min( + limit(qₗ(thp, ts), dt, 5), + CM1.accretion(mp.cl, mp.ps, mp.tv.snow, mp.ce, qₗ(thp, ts), qₛ, ρ), + ) + # if T < T_freeze cloud droplets freeze to become snow + # else the snow melts and both cloud water and snow become rain + @. Sᵖ_snow = ifelse( + Tₐ(thp, ts) < mp.ps.T_freeze, + Sᵖ, + FT(-1) * min(Sᵖ * α(thp, ts), limit(qₛ, dt, 5)), + ) + + @. Sqₛᵖ += Sᵖ_snow + @. Sqₜᵖ -= Sᵖ + @. Sqᵣᵖ += ifelse(Tₐ(thp, ts) < mp.ps.T_freeze, FT(0), Sᵖ - Sᵖ_snow) + + @. Seₜᵖ -= ifelse( + Tₐ(thp, ts) < mp.ps.T_freeze, + Sᵖ * (Iᵢ(thp, ts) + Φ), + Sᵖ * (Iₗ(thp, ts) + Φ) - Sᵖ_snow * (Iₗ(thp, ts) - Iᵢ(thp, ts)), + ) + + # accretion: q_ice + q_rai -> q_sno + @. Sᵖ = min( + limit(qᵢ(thp, ts), dt, 5), + CM1.accretion(mp.ci, mp.pr, mp.tv.rain, mp.ce, qᵢ(thp, ts), qᵣ, ρ), + ) + @. Sqₜᵖ -= Sᵖ + @. Sqₛᵖ += Sᵖ + @. Seₜᵖ -= Sᵖ * (Iᵢ(thp, ts) + Φ) + # sink of rain via accretion cloud ice - rain + @. Sᵖ = min( limit(qᵣ, dt, 5), - CM1.accretion_snow_rain(mp.ps, mp.pr, mp.tv.rain, mp.tv.snow, mp.ce, qₛ, qᵣ, ρ), - ), - -min( - limit(qₛ, dt, 5), - CM1.accretion_snow_rain(mp.pr, mp.ps, mp.tv.snow, mp.tv.rain, mp.ce, qᵣ, qₛ, ρ), - ), - ) - @. Sqₛᵖ += Sᵖ - @. Sqᵣᵖ -= Sᵖ - @. Seₜᵖ += Sᵖ * Lf(thp, ts) + CM1.accretion_rain_sink(mp.pr, mp.ci, mp.tv.rain, mp.ce, qᵢ(thp, ts), qᵣ, ρ), + ) + + @. Sqᵣᵖ -= Sᵖ + @. Sqₛᵖ += Sᵖ + @. Seₜᵖ += Sᵖ * Lf(thp, ts) + + # accretion: q_rai + q_sno -> q_rai or q_sno + @. Sᵖ = ifelse( + Tₐ(thp, ts) < mp.ps.T_freeze, + min( + limit(qᵣ, dt, 5), + CM1.accretion_snow_rain(mp.ps, mp.pr, mp.tv.rain, mp.tv.snow, mp.ce, qₛ, qᵣ, ρ), + ), + -min( + limit(qₛ, dt, 5), + CM1.accretion_snow_rain(mp.pr, mp.ps, mp.tv.snow, mp.tv.rain, mp.ce, qᵣ, qₛ, ρ), + ), + ) + + @. Sqₛᵖ += Sᵖ + @. Sqᵣᵖ -= Sᵖ + @. Seₜᵖ += Sᵖ * Lf(thp, ts) + end #! format: on end @@ -281,8 +286,12 @@ function compute_precipitation_heating!( @. ᶜ∇T += CT123(gradₕ(Tₐ(thp, ᶜts))) # dot product with effective velocity of precipitation # (times q and specific heat) - @. ᶜSeₜᵖ -= dot(ᶜ∇T, (ᶜu - C123(Geometry.WVector(ᶜwᵣ)))) * cᵥₗ(thp) * ᶜqᵣ - @. ᶜSeₜᵖ -= dot(ᶜ∇T, (ᶜu - C123(Geometry.WVector(ᶜwₛ)))) * cᵥᵢ(thp) * ᶜqₛ + @fused_direct begin + @. ᶜSeₜᵖ -= + dot(ᶜ∇T, (ᶜu - C123(Geometry.WVector(ᶜwᵣ)))) * cᵥₗ(thp) * ᶜqᵣ + @. ᶜSeₜᵖ -= + dot(ᶜ∇T, (ᶜu - C123(Geometry.WVector(ᶜwₛ)))) * cᵥᵢ(thp) * ᶜqₛ + end end """ compute_precipitation_sinks!(Sᵖ, Sqₜᵖ, Sqᵣᵖ, Sqₛᵖ, Seₜᵖ, ρ, qᵣ, qₛ, ts, Φ, dt, mp, thp) @@ -320,34 +329,37 @@ function compute_precipitation_sinks!( sps = (mp.ps, mp.tv.snow, mp.aps, thp) rps = (mp.pr, mp.tv.rain, mp.aps, thp) - #! format: off - # evaporation: q_rai -> q_vap - @. Sᵖ = -min( - limit(qᵣ, dt, 5), - -CM1.evaporation_sublimation(rps..., PP(thp, ts), qᵣ, ρ, Tₐ(thp, ts)), - ) - @. Sqₜᵖ -= Sᵖ - @. Sqᵣᵖ += Sᵖ - @. Seₜᵖ -= Sᵖ * (Iₗ(thp, ts) + Φ) - - # melting: q_sno -> q_rai - @. Sᵖ = min( - limit(qₛ, dt, 5), - CM1.snow_melt(sps..., qₛ, ρ, Tₐ(thp, ts)), - ) - @. Sqᵣᵖ += Sᵖ - @. Sqₛᵖ -= Sᵖ - @. Seₜᵖ -= Sᵖ * Lf(thp, ts) - - # deposition/sublimation: q_vap <-> q_sno - @. Sᵖ = CM1.evaporation_sublimation(sps..., PP(thp, ts), qₛ, ρ, Tₐ(thp, ts)) - @. Sᵖ = ifelse( - Sᵖ > FT(0), - min(limit(qᵥ(thp, ts), dt, 5), Sᵖ), - -min(limit(qₛ, dt, 5), FT(-1) * Sᵖ), - ) - @. Sqₜᵖ -= Sᵖ - @. Sqₛᵖ += Sᵖ - @. Seₜᵖ -= Sᵖ * (Iᵢ(thp, ts) + Φ) - #! format: on + @fused_direct begin + #! format: off + # evaporation: q_rai -> q_vap + @. Sᵖ = -min( + limit(qᵣ, dt, 5), + -CM1.evaporation_sublimation(rps..., PP(thp, ts), qᵣ, ρ, Tₐ(thp, ts)), + ) + @. Sqₜᵖ -= Sᵖ + @. Sqᵣᵖ += Sᵖ + @. Seₜᵖ -= Sᵖ * (Iₗ(thp, ts) + Φ) + + # melting: q_sno -> q_rai + @. Sᵖ = min( + limit(qₛ, dt, 5), + CM1.snow_melt(sps..., qₛ, ρ, Tₐ(thp, ts)), + ) + + @. Sqᵣᵖ += Sᵖ + @. Sqₛᵖ -= Sᵖ + @. Seₜᵖ -= Sᵖ * Lf(thp, ts) + + # deposition/sublimation: q_vap <-> q_sno + @. Sᵖ = CM1.evaporation_sublimation(sps..., PP(thp, ts), qₛ, ρ, Tₐ(thp, ts)) + @. Sᵖ = ifelse( + Sᵖ > FT(0), + min(limit(qᵥ(thp, ts), dt, 5), Sᵖ), + -min(limit(qₛ, dt, 5), FT(-1) * Sᵖ), + ) + @. Sqₜᵖ -= Sᵖ + @. Sqₛᵖ += Sᵖ + @. Seₜᵖ -= Sᵖ * (Iᵢ(thp, ts) + Φ) + #! format: on + end end diff --git a/src/parameterized_tendencies/microphysics/precipitation.jl b/src/parameterized_tendencies/microphysics/precipitation.jl index 1f041ac2c8..4df2ba19bb 100644 --- a/src/parameterized_tendencies/microphysics/precipitation.jl +++ b/src/parameterized_tendencies/microphysics/precipitation.jl @@ -170,8 +170,10 @@ function precipitation_tendency!( compute_precipitation_surface_fluxes!(Y, p, precip_model) # Add the source terms to the tendencies - @. Yₜ.c.ρq_tot += ᶜS_ρq_tot - @. Yₜ.c.ρ += ᶜS_ρq_tot + @fused_direct begin + @. Yₜ.c.ρq_tot += ᶜS_ρq_tot + @. Yₜ.c.ρ += ᶜS_ρq_tot + end @. Yₜ.c.ρe_tot += ᶜS_ρe_tot return nothing @@ -344,11 +346,13 @@ function precipitation_tendency!( compute_precipitation_surface_fluxes!(Y, p, precip_model) # Update grid mean tendencies - @. Yₜ.c.ρ += Y.c.ρ * ᶜSqₜᵖ - @. Yₜ.c.ρq_tot += Y.c.ρ * ᶜSqₜᵖ - @. Yₜ.c.ρe_tot += Y.c.ρ * ᶜSeₜᵖ - @. Yₜ.c.ρq_rai += Y.c.ρ * ᶜSqᵣᵖ - @. Yₜ.c.ρq_sno += Y.c.ρ * ᶜSqₛᵖ + @fused_direct begin + @. Yₜ.c.ρ += Y.c.ρ * ᶜSqₜᵖ + @. Yₜ.c.ρq_tot += Y.c.ρ * ᶜSqₜᵖ + @. Yₜ.c.ρe_tot += Y.c.ρ * ᶜSeₜᵖ + @. Yₜ.c.ρq_rai += Y.c.ρ * ᶜSqᵣᵖ + @. Yₜ.c.ρq_sno += Y.c.ρ * ᶜSqₛᵖ + end return nothing end @@ -375,20 +379,24 @@ function precipitation_tendency!( # Update from environment precipitation sources # and the grid mean precipitation sinks - @. Yₜ.c.ρ += Y.c.ρ * (ᶜSqₜᵖ⁰ + ᶜSqₜᵖ) - @. Yₜ.c.ρq_tot += Y.c.ρ * (ᶜSqₜᵖ⁰ + ᶜSqₜᵖ) - @. Yₜ.c.ρe_tot += Y.c.ρ * (ᶜSeₜᵖ⁰ + ᶜSeₜᵖ) - @. Yₜ.c.ρq_rai += Y.c.ρ * (ᶜSqᵣᵖ⁰ + ᶜSqᵣᵖ) - @. Yₜ.c.ρq_sno += Y.c.ρ * (ᶜSqₛᵖ⁰ + ᶜSqₛᵖ) + @fused_direct begin + @. Yₜ.c.ρ += Y.c.ρ * (ᶜSqₜᵖ⁰ + ᶜSqₜᵖ) + @. Yₜ.c.ρq_tot += Y.c.ρ * (ᶜSqₜᵖ⁰ + ᶜSqₜᵖ) + @. Yₜ.c.ρe_tot += Y.c.ρ * (ᶜSeₜᵖ⁰ + ᶜSeₜᵖ) + @. Yₜ.c.ρq_rai += Y.c.ρ * (ᶜSqᵣᵖ⁰ + ᶜSqᵣᵖ) + @. Yₜ.c.ρq_sno += Y.c.ρ * (ᶜSqₛᵖ⁰ + ᶜSqₛᵖ) + end # Update from the updraft precipitation sources n = n_mass_flux_subdomains(p.atmos.turbconv_model) for j in 1:n - @. Yₜ.c.ρ += ᶜρaʲs.:($$j) * ᶜSqₜᵖʲs.:($$j) - @. Yₜ.c.ρq_tot += ᶜρaʲs.:($$j) * ᶜSqₜᵖʲs.:($$j) - @. Yₜ.c.ρe_tot += ᶜρaʲs.:($$j) * ᶜSeₜᵖʲs.:($$j) - @. Yₜ.c.ρq_rai += ᶜρaʲs.:($$j) * ᶜSqᵣᵖʲs.:($$j) - @. Yₜ.c.ρq_sno += ᶜρaʲs.:($$j) * ᶜSqₛᵖʲs.:($$j) + @fused_direct begin + @. Yₜ.c.ρ += ᶜρaʲs.:($$j) * ᶜSqₜᵖʲs.:($$j) + @. Yₜ.c.ρq_tot += ᶜρaʲs.:($$j) * ᶜSqₜᵖʲs.:($$j) + @. Yₜ.c.ρe_tot += ᶜρaʲs.:($$j) * ᶜSeₜᵖʲs.:($$j) + @. Yₜ.c.ρq_rai += ᶜρaʲs.:($$j) * ᶜSqᵣᵖʲs.:($$j) + @. Yₜ.c.ρq_sno += ᶜρaʲs.:($$j) * ᶜSqₛᵖʲs.:($$j) + end end end function precipitation_tendency!( @@ -412,19 +420,23 @@ function precipitation_tendency!( # Update from environment precipitation sources # and the grid mean precipitation sinks - @. Yₜ.c.ρ += ᶜρa⁰ * ᶜSqₜᵖ⁰ + Y.c.ρ * ᶜSqₜᵖ - @. Yₜ.c.ρq_tot += ᶜρa⁰ * ᶜSqₜᵖ⁰ + Y.c.ρ * ᶜSqₜᵖ - @. Yₜ.c.ρe_tot += ᶜρa⁰ * ᶜSeₜᵖ⁰ + Y.c.ρ * ᶜSeₜᵖ - @. Yₜ.c.ρq_rai += ᶜρa⁰ * ᶜSqᵣᵖ⁰ + Y.c.ρ * ᶜSqᵣᵖ - @. Yₜ.c.ρq_sno += ᶜρa⁰ * ᶜSqₛᵖ⁰ + Y.c.ρ * ᶜSqₛᵖ + @fused_direct begin + @. Yₜ.c.ρ += ᶜρa⁰ * ᶜSqₜᵖ⁰ + Y.c.ρ * ᶜSqₜᵖ + @. Yₜ.c.ρq_tot += ᶜρa⁰ * ᶜSqₜᵖ⁰ + Y.c.ρ * ᶜSqₜᵖ + @. Yₜ.c.ρe_tot += ᶜρa⁰ * ᶜSeₜᵖ⁰ + Y.c.ρ * ᶜSeₜᵖ + @. Yₜ.c.ρq_rai += ᶜρa⁰ * ᶜSqᵣᵖ⁰ + Y.c.ρ * ᶜSqᵣᵖ + @. Yₜ.c.ρq_sno += ᶜρa⁰ * ᶜSqₛᵖ⁰ + Y.c.ρ * ᶜSqₛᵖ + end # Update from the updraft precipitation sources n = n_mass_flux_subdomains(p.atmos.turbconv_model) for j in 1:n - @. Yₜ.c.ρ += Y.c.sgsʲs.:($$j).ρa * ᶜSqₜᵖʲs.:($$j) - @. Yₜ.c.ρq_tot += Y.c.sgsʲs.:($$j).ρa * ᶜSqₜᵖʲs.:($$j) - @. Yₜ.c.ρe_tot += Y.c.sgsʲs.:($$j).ρa * ᶜSeₜᵖʲs.:($$j) - @. Yₜ.c.ρq_rai += Y.c.sgsʲs.:($$j).ρa * ᶜSqᵣᵖʲs.:($$j) - @. Yₜ.c.ρq_sno += Y.c.sgsʲs.:($$j).ρa * ᶜSqₛᵖʲs.:($$j) + @fused_direct begin + @. Yₜ.c.ρ += Y.c.sgsʲs.:($$j).ρa * ᶜSqₜᵖʲs.:($$j) + @. Yₜ.c.ρq_tot += Y.c.sgsʲs.:($$j).ρa * ᶜSqₜᵖʲs.:($$j) + @. Yₜ.c.ρe_tot += Y.c.sgsʲs.:($$j).ρa * ᶜSeₜᵖʲs.:($$j) + @. Yₜ.c.ρq_rai += Y.c.sgsʲs.:($$j).ρa * ᶜSqᵣᵖʲs.:($$j) + @. Yₜ.c.ρq_sno += Y.c.sgsʲs.:($$j).ρa * ᶜSqₛᵖʲs.:($$j) + end end end diff --git a/src/utils/common_spaces.jl b/src/utils/common_spaces.jl index 8842d24fee..60be363678 100644 --- a/src/utils/common_spaces.jl +++ b/src/utils/common_spaces.jl @@ -1,6 +1,7 @@ using ClimaCore: Geometry, Domains, Meshes, Topologies, Spaces, Grids, Hypsography, Fields using ClimaComms +using ClimaCore: DataLayouts using ClimaUtilities: SpaceVaryingInputs.SpaceVaryingInput function periodic_line_mesh(; x_max, x_elem) @@ -41,7 +42,11 @@ function make_horizontal_space( ) if mesh isa Meshes.AbstractMesh1D topology = Topologies.IntervalTopology(comms_ctx, mesh) - space = Spaces.SpectralElementSpace1D(topology, quad) + space = Spaces.SpectralElementSpace1D( + topology, + quad; + horizontal_layout_type = DataLayouts.IJHF, + ) elseif mesh isa Meshes.AbstractMesh2D topology = Topologies.Topology2D( comms_ctx, @@ -52,6 +57,7 @@ function make_horizontal_space( topology, quad; enable_bubble = bubble, + horizontal_layout_type = DataLayouts.IJHF, ) end return space @@ -70,6 +76,7 @@ function make_horizontal_space(mesh, quad, comms_ctx, bubble) topology, quad; enable_bubble = bubble, + horizontal_layout_type = DataLayouts.IJHF, ) end return space