diff --git a/src/ClimaAtmos.jl b/src/ClimaAtmos.jl index 7ed4d5a5fc..f365631e80 100644 --- a/src/ClimaAtmos.jl +++ b/src/ClimaAtmos.jl @@ -8,6 +8,8 @@ import LazyBroadcast import LazyBroadcast: lazy import Thermodynamics as TD import Thermodynamics +import ClimaCore.MatrixFields: @name + include("compat.jl") include(joinpath("parameters", "Parameters.jl")) diff --git a/src/cache/cloud_fraction.jl b/src/cache/cloud_fraction.jl index bd3420aa85..92e301f58d 100644 --- a/src/cache/cloud_fraction.jl +++ b/src/cache/cloud_fraction.jl @@ -59,15 +59,10 @@ NVTX.@annotate function set_cloud_fraction!( TD.PhasePartition(thermo_params, ᶜts).ice, ) else - @. cloud_diagnostics_tuple = make_named_tuple( - ifelse( - specific(Y.c.ρq_liq, Y.c.ρ) + specific(Y.c.ρq_ice, Y.c.ρ) > 0, - 1, - 0, - ), - specific(Y.c.ρq_liq, Y.c.ρ), - specific(Y.c.ρq_ice, Y.c.ρ), - ) + q_liq = @. lazy(specific(Y.c.ρq_liq, Y.c.ρ)) + q_ice = @. lazy(specific(Y.c.ρq_ice, Y.c.ρ)) + @. cloud_diagnostics_tuple = + make_named_tuple(ifelse(q_liq + q_ice > 0, 1, 0), q_liq, q_ice) end end @@ -194,8 +189,9 @@ NVTX.@annotate function set_cloud_fraction!( FT = eltype(params) thermo_params = CAP.thermodynamics_params(params) (; ᶜts⁰, cloud_diagnostics_tuple) = p.precomputed - (; ᶜρʲs, ᶜtsʲs, ᶜρa⁰, ᶜρ⁰) = p.precomputed + (; ᶜρʲs, ᶜtsʲs) = p.precomputed (; turbconv_model) = p.atmos + ᶜρa⁰ = @. lazy(ρa⁰(Y.c.ρ, Y.c.sgsʲs, turbconv_model)) # TODO - we should make this default when using diagnostic edmf # environment @@ -217,9 +213,9 @@ NVTX.@annotate function set_cloud_fraction!( # weight cloud diagnostics by environmental area @. cloud_diagnostics_tuple *= NamedTuple{(:cf, :q_liq, :q_ice)}( tuple( - draft_area(ᶜρa⁰, ᶜρ⁰), - draft_area(ᶜρa⁰, ᶜρ⁰), - draft_area(ᶜρa⁰, ᶜρ⁰), + draft_area(ᶜρa⁰, TD.air_density(thermo_params, ᶜts⁰)), + draft_area(ᶜρa⁰, TD.air_density(thermo_params, ᶜts⁰)), + draft_area(ᶜρa⁰, TD.air_density(thermo_params, ᶜts⁰)), ), ) # updrafts diff --git a/src/cache/diagnostic_edmf_precomputed_quantities.jl b/src/cache/diagnostic_edmf_precomputed_quantities.jl index 51bafd3160..0a3d3f2cee 100644 --- a/src/cache/diagnostic_edmf_precomputed_quantities.jl +++ b/src/cache/diagnostic_edmf_precomputed_quantities.jl @@ -61,7 +61,7 @@ NVTX.@annotate function set_diagnostic_edmfx_env_quantities_level!( local_geometry_halflevel, turbconv_model, ) - @. u³⁰_halflevel = divide_by_ρa( + @. u³⁰_halflevel = specific( ρ_level * u³_halflevel - unrolled_dotproduct(ρaʲs_level, u³ʲs_halflevel), ρ_level, @@ -93,23 +93,28 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_bottom_bc!( FT = eltype(Y) n = n_mass_flux_subdomains(turbconv_model) (; ᶜΦ) = p.core - (; ᶜp, ᶠu³, ᶜh_tot, ᶜK) = p.precomputed - (; q_tot) = p.precomputed.ᶜspecific + (; ᶜp, ᶠu³, ᶜK) = p.precomputed (; ustar, obukhov_length, buoyancy_flux, ρ_flux_h_tot, ρ_flux_q_tot) = p.precomputed.sfc_conditions (; ᶜρaʲs, ᶠu³ʲs, ᶜKʲs, ᶜmseʲs, ᶜq_totʲs, ᶜtsʲs, ᶜρʲs) = p.precomputed (; ᶠu³⁰, ᶜK⁰) = p.precomputed (; params) = p + thermo_params = CAP.thermodynamics_params(params) turbconv_params = CAP.turbconv_params(params) + ᶜts = p.precomputed.ᶜts #TODO replace + + ᶜq_tot = @. lazy(specific(Y.c.ρq_tot, Y.c.ρ)) + ᶜe_tot = @. lazy(specific(Y.c.ρe_tot, Y.c.ρ)) + ᶜh_tot = @. lazy(TD.total_specific_enthalpy(thermo_params, ᶜts, ᶜe_tot)) ρ_int_level = Fields.field_values(Fields.level(Y.c.ρ, 1)) uₕ_int_level = Fields.field_values(Fields.level(Y.c.uₕ, 1)) u³_int_halflevel = Fields.field_values(Fields.level(ᶠu³, half)) h_tot_int_level = Fields.field_values(Fields.level(ᶜh_tot, 1)) K_int_level = Fields.field_values(Fields.level(ᶜK, 1)) - q_tot_int_level = Fields.field_values(Fields.level(q_tot, 1)) + q_tot_int_level = Fields.field_values(Fields.level(ᶜq_tot, 1)) p_int_level = Fields.field_values(Fields.level(ᶜp, 1)) Φ_int_level = Fields.field_values(Fields.level(ᶜΦ, 1)) @@ -305,8 +310,7 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!( (; dt) = p dt = float(dt) (; ᶜΦ, ᶜgradᵥ_ᶠΦ) = p.core - (; ᶜp, ᶠu³, ᶜts, ᶜh_tot, ᶜK) = p.precomputed - (; q_tot) = p.precomputed.ᶜspecific + (; ᶜp, ᶠu³, ᶜts, ᶜK) = p.precomputed (; ᶜρaʲs, ᶠu³ʲs, @@ -321,13 +325,13 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!( ᶠnh_pressure³_buoyʲs, ᶠnh_pressure³_dragʲs, ) = p.precomputed - (; ᶠu³⁰, ᶜK⁰, ᶜtke⁰) = p.precomputed + (; ᶠu³⁰, ᶜK⁰) = p.precomputed + + if microphysics_model isa Microphysics1Moment - ᶜq_liqʲs = p.precomputed.ᶜq_liqʲs - ᶜq_iceʲs = p.precomputed.ᶜq_iceʲs - q_rai = p.precomputed.ᶜqᵣ - q_sno = p.precomputed.ᶜqₛ + q_rai = @. lazy(specific(Y.c.ρq_rai, Y.c.ρ)) + q_sno = @. lazy(specific(Y.c.ρq_sno, Y.c.ρ)) end thermo_params = CAP.thermodynamics_params(params) @@ -347,13 +351,24 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!( Fields.field_values(Fields.level(Fields.coordinate_field(Y.f).z, half)) # integral + ᶜq_tot = @. lazy(specific(Y.c.ρq_tot, Y.c.ρ)) + ᶜh_tot = @. lazy( + TD.total_specific_enthalpy( + thermo_params, + ᶜts, + specific(Y.c.ρe_tot, Y.c.ρ), + ), + ) + + ᶜtke⁰ = @. lazy(specific_tke(Y.c.ρ, Y.c.sgs⁰.ρatke, Y.c.ρ, turbconv_model)) + for i in 2:Spaces.nlevels(axes(Y.c)) ρ_level = Fields.field_values(Fields.level(Y.c.ρ, i)) uₕ_level = Fields.field_values(Fields.level(Y.c.uₕ, i)) u³_halflevel = Fields.field_values(Fields.level(ᶠu³, i - half)) K_level = Fields.field_values(Fields.level(ᶜK, i)) h_tot_level = Fields.field_values(Fields.level(ᶜh_tot, i)) - q_tot_level = Fields.field_values(Fields.level(q_tot, i)) + q_tot_level = Fields.field_values(Fields.level(ᶜq_tot, i)) p_level = Fields.field_values(Fields.level(ᶜp, i)) Φ_level = Fields.field_values(Fields.level(ᶜΦ, i)) local_geometry_level = Fields.field_values( @@ -378,7 +393,7 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_do_integral!( u³⁰_data_prev_halflevel = u³⁰_prev_halflevel.components.data.:1 K_prev_level = Fields.field_values(Fields.level(ᶜK, i - 1)) h_tot_prev_level = Fields.field_values(Fields.level(ᶜh_tot, i - 1)) - q_tot_prev_level = Fields.field_values(Fields.level(q_tot, i - 1)) + q_tot_prev_level = Fields.field_values(Fields.level(ᶜq_tot, i - 1)) ts_prev_level = Fields.field_values(Fields.level(ᶜts, i - 1)) p_prev_level = Fields.field_values(Fields.level(ᶜp, i - 1)) z_prev_level = Fields.field_values(Fields.level(ᶜz, i - 1)) @@ -965,7 +980,6 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_env_closures! (; dt) = p (; ᶜp, ᶜu, ᶜts) = p.precomputed (; ustar, obukhov_length) = p.precomputed.sfc_conditions - (; ᶜtke⁰) = p.precomputed (; ᶜlinear_buoygrad, ᶜstrain_rate_norm) = p.precomputed (; ρatke_flux) = p.precomputed turbconv_params = CAP.turbconv_params(params) @@ -973,7 +987,7 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_env_closures! ᶜlg = Fields.local_geometry_field(Y.c) if p.atmos.turbconv_model isa DiagnosticEDMFX - (; ᶜρaʲs, ᶠu³ʲs, ᶜdetrʲs, ᶠu³⁰, ᶜu⁰) = p.precomputed + (; ᶠu³⁰, ᶜu⁰) = p.precomputed elseif p.atmos.turbconv_model isa EDOnlyEDMFX ᶠu³⁰ = p.precomputed.ᶠu³ ᶜu⁰ = ᶜu @@ -1042,11 +1056,12 @@ NVTX.@annotate function set_diagnostic_edmf_precomputed_quantities_env_precipita (; ᶜts, ᶜSqₜᵖ⁰) = p.precomputed # Environment precipitation sources (to be applied to grid mean) + ᶜq_tot = @. lazy(specific(Y.c.ρq_tot, Y.c.ρ)) @. ᶜSqₜᵖ⁰ = q_tot_0M_precipitation_sources( thermo_params, microphys_0m_params, dt, - specific(Y.c.ρq_tot, Y.c.ρ), + ᶜq_tot, ᶜts, ) return nothing diff --git a/src/cache/precipitation_precomputed_quantities.jl b/src/cache/precipitation_precomputed_quantities.jl index c7bb9c083b..57da69e7e0 100644 --- a/src/cache/precipitation_precomputed_quantities.jl +++ b/src/cache/precipitation_precomputed_quantities.jl @@ -253,10 +253,11 @@ function set_precipitation_cache!( ::PrognosticEDMFX, ) (; ᶜΦ) = p.core - (; ᶜSqₜᵖ⁰, ᶜSqₜᵖʲs, ᶜρa⁰) = p.precomputed + (; ᶜ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) @@ -283,7 +284,11 @@ function set_precipitation_cache!(Y, p, ::Microphysics1Moment, _) (; ᶜts, ᶜwᵣ, ᶜwₛ, ᶜu) = p.precomputed (; ᶜSqₗᵖ, ᶜSqᵢᵖ, ᶜSqᵣᵖ, ᶜSqₛᵖ) = p.precomputed - (; q_tot, q_liq, q_ice, q_rai, q_sno) = p.precomputed.ᶜspecific + ᶜq_tot = @. lazy(specific(Y.c.ρq_tot, Y.c.ρ)) + ᶜq_rai = @. lazy(specific(Y.c.ρq_rai, Y.c.ρ)) + ᶜq_sno = @. lazy(specific(Y.c.ρq_sno, Y.c.ρ)) + ᶜq_liq = @. lazy(specific(Y.c.ρq_liq, Y.c.ρ)) + ᶜq_ice = @. lazy(specific(Y.c.ρq_ice, Y.c.ρ)) ᶜSᵖ = p.scratch.ᶜtemp_scalar ᶜSᵖ_snow = p.scratch.ᶜtemp_scalar_2 @@ -303,11 +308,11 @@ function set_precipitation_cache!(Y, p, ::Microphysics1Moment, _) ᶜSqᵣᵖ, ᶜSqₛᵖ, Y.c.ρ, - q_tot, - q_liq, - q_ice, - q_rai, - q_sno, + ᶜq_tot, + ᶜq_liq, + ᶜq_ice, + ᶜq_rai, + ᶜq_sno, ᶜts, dt, cmp, @@ -320,11 +325,11 @@ function set_precipitation_cache!(Y, p, ::Microphysics1Moment, _) ᶜSqᵣᵖ, ᶜSqₛᵖ, Y.c.ρ, - q_tot, - q_liq, - q_ice, - q_rai, - q_sno, + ᶜq_tot, + ᶜq_liq, + ᶜq_ice, + ᶜq_rai, + ᶜq_sno, ᶜts, dt, cmp, @@ -469,14 +474,27 @@ function set_precipitation_surface_fluxes!( sfc_ρ = @. lazy(int_ρ * int_J / sfc_J) # Constant extrapolation to surface, consistent with simple downwinding - sfc_wᵣ = sfc_lev(ᶜwᵣ) - sfc_wₛ = sfc_lev(ᶜwₛ) - sfc_wₗ = sfc_lev(ᶜwₗ) - sfc_wᵢ = sfc_lev(ᶜwᵢ) - sfc_qᵣ = lazy.(specific.(sfc_lev(Y.c.ρq_rai), sfc_ρ)) - sfc_qₛ = lazy.(specific.(sfc_lev(Y.c.ρq_sno), sfc_ρ)) - sfc_qₗ = lazy.(specific.(sfc_lev(Y.c.ρq_liq), sfc_ρ)) - sfc_qᵢ = lazy.(specific.(sfc_lev(Y.c.ρq_ice), sfc_ρ)) + # Temporary scratch variables are used here until CC.field_values supports fields + ᶜq_rai = p.scratch.ᶜtemp_scalar + ᶜq_sno = p.scratch.ᶜtemp_scalar_2 + ᶜq_liq = p.scratch.ᶜtemp_scalar_3 + ᶜq_ice = p.scratch.ᶜtemp_scalar_4 + @. ᶜq_rai = specific(Y.c.ρq_rai, Y.c.ρ) + @. ᶜ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_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) @. surface_rain_flux = sfc_ρ * (sfc_qᵣ * (-sfc_wᵣ) + sfc_qₗ * (-sfc_wₗ)) @. surface_snow_flux = sfc_ρ * (sfc_qₛ * (-sfc_wₛ) + sfc_qᵢ * (-sfc_wᵢ)) diff --git a/src/cache/precomputed_quantities.jl b/src/cache/precomputed_quantities.jl index 357e845c28..f523f248f2 100644 --- a/src/cache/precomputed_quantities.jl +++ b/src/cache/precomputed_quantities.jl @@ -3,6 +3,7 @@ ##### import Thermodynamics as TD import ClimaCore: Spaces, Fields +using Base.Broadcast: materialize """ implicit_precomputed_quantities(Y, atmos) @@ -11,28 +12,20 @@ Allocates precomputed quantities that are treated implicitly (i.e., updated on each iteration of the implicit solver). This includes all quantities related to velocity and thermodynamics that are used in the implicit tendency. -The following grid-scale quantities are treated implicitly: - - `ᶜspecific`: specific quantities on cell centers (for every prognostic - quantity `ρχ`, there is a corresponding specific quantity `χ`) +The following grid-scale quantities are treated implicitly and are precomputed: - `ᶜu`: covariant velocity on cell centers - `ᶠu`: contravariant velocity on cell faces - `ᶜK`: kinetic energy on cell centers - `ᶜts`: thermodynamic state on cell centers - `ᶜp`: air pressure on cell centers - - `ᶜh_tot`: total enthalpy on cell centers If the `turbconv_model` is `PrognosticEDMFX`, there also two SGS versions of every quantity except for `ᶜp` (which is shared across all subdomains): - `_⁰`: value for the environment - `_ʲs`: a tuple of values for the mass-flux subdomains In addition, there are several other SGS quantities for `PrognosticEDMFX`: - - `ᶜtke⁰`: turbulent kinetic energy of the environment on cell centers - - `ᶜρa⁰`: area-weighted air density of the environment on cell centers - - `ᶜmse⁰`: moist static energy of the environment on cell centers - - `ᶜq_tot⁰`: total specific humidity of the environment on cell centers - - `ᶜρ⁰`: air density of the environment on cell centers - `ᶜρʲs`: a tuple of the air densities of the mass-flux subdomains on cell centers -For every other `AbstractEDMF`, only `ᶜtke⁰` is added as a precomputed quantity. + TODO: Rename `ᶜK` to `ᶜκ`. """ @@ -49,34 +42,16 @@ function implicit_precomputed_quantities(Y, atmos) ᶜK = similar(Y.c, FT), ᶜts = similar(Y.c, TST), ᶜp = similar(Y.c, FT), - ᶜh_tot = similar(Y.c, FT), ) - sgs_quantities = - turbconv_model isa AbstractEDMF ? (; ᶜtke⁰ = similar(Y.c, FT)) : (;) - moisture_sgs_quantities = - ( - turbconv_model isa PrognosticEDMFX && - moisture_model isa NonEquilMoistModel && - microphysics_model isa Microphysics1Moment - ) ? - (; - ᶜq_liq⁰ = similar(Y.c, FT), - ᶜq_ice⁰ = similar(Y.c, FT), - ᶜq_rai⁰ = similar(Y.c, FT), - ᶜq_sno⁰ = similar(Y.c, FT), - ) : (;) + sgs_quantities = (;) prognostic_sgs_quantities = turbconv_model isa PrognosticEDMFX ? (; - ᶜρa⁰ = similar(Y.c, FT), - ᶜmse⁰ = similar(Y.c, FT), - ᶜq_tot⁰ = similar(Y.c, FT), ᶠu₃⁰ = similar(Y.f, C3{FT}), ᶜu⁰ = similar(Y.c, C123{FT}), ᶠu³⁰ = similar(Y.f, CT3{FT}), ᶜK⁰ = similar(Y.c, FT), ᶜts⁰ = similar(Y.c, TST), - ᶜρ⁰ = similar(Y.c, FT), ᶜuʲs = similar(Y.c, NTuple{n, C123{FT}}), ᶠu³ʲs = similar(Y.f, NTuple{n, CT3{FT}}), ᶜKʲs = similar(Y.c, NTuple{n, FT}), @@ -84,7 +59,6 @@ function implicit_precomputed_quantities(Y, atmos) ᶜtsʲs = similar(Y.c, NTuple{n, TST}), ᶜρʲs = similar(Y.c, NTuple{n, FT}), ᶠnh_pressure₃_dragʲs = similar(Y.f, NTuple{n, C3{FT}}), - moisture_sgs_quantities..., ) : (;) return (; gs_quantities..., sgs_quantities..., prognostic_sgs_quantities...) end @@ -188,10 +162,7 @@ function precomputed_quantities(Y, atmos) edonly_quantities = atmos.turbconv_model isa EDOnlyEDMFX ? - (; - ᶜtke⁰ = similar(Y.c, FT), - ρatke_flux = similar(Fields.level(Y.f, half), C3{FT}), - ) : (;) + (; ρatke_flux = similar(Fields.level(Y.f, half), C3{FT}),) : (;) sgs_quantities = (; ᶜgradᵥ_θ_virt = Fields.Field(C3{FT}, cspace), @@ -328,11 +299,12 @@ function set_sgs_ᶠu₃!(w_function, ᶠu₃, Y, turbconv_model) end function add_sgs_ᶜK!(ᶜK, Y, ᶜρa⁰, ᶠu₃⁰, turbconv_model) - @. ᶜK += ᶜρa⁰ * ᶜinterp(dot(ᶠu₃⁰ - Yf.u₃, CT3(ᶠu₃⁰ - Yf.u₃))) / 2 / Yc.ρ + @. ᶜK += ᶜρa⁰ * ᶜinterp(dot(ᶠu₃⁰ - Y.f.u₃, CT3(ᶠu₃⁰ - Y.f.u₃))) / 2 / Y.c.ρ for j in 1:n_mass_flux_subdomains(turbconv_model) ᶜρaʲ = Y.c.sgsʲs.:($j).ρa ᶠu₃ʲ = Y.f.sgsʲs.:($j).u₃ - @. ᶜK += ᶜρaʲ * ᶜinterp(dot(ᶠu₃ʲ - Yf.u₃, CT3(ᶠu₃ʲ - Yf.u₃))) / 2 / Yc.ρ + @. ᶜK += + ᶜρaʲ * ᶜinterp(dot(ᶠu₃ʲ - Y.f.u₃, CT3(ᶠu₃ʲ - Y.f.u₃))) / 2 / Y.c.ρ end return nothing end @@ -439,13 +411,12 @@ quantities are updated. NVTX.@annotate function set_implicit_precomputed_quantities!(Y, p, t) (; turbconv_model, moisture_model, microphysics_model) = p.atmos (; ᶜΦ) = p.core - (; ᶜspecific, ᶜu, ᶠu³, ᶠu, ᶜK, ᶜts, ᶜp, ᶜh_tot) = p.precomputed + (; ᶜu, ᶠu³, ᶠu, ᶜK, ᶜts, ᶜp) = p.precomputed ᶠuₕ³ = p.scratch.ᶠtemp_CT3 n = n_mass_flux_subdomains(turbconv_model) thermo_params = CAP.thermodynamics_params(p.params) thermo_args = (thermo_params, moisture_model, microphysics_model) - ᶜspecific .= ᶜspecific_gs_tracers(Y) @. ᶠuₕ³ = $compute_ᶠuₕ³(Y.c.uₕ, Y.c.ρ) # TODO: We might want to move this to dss! (and rename dss! to something @@ -473,19 +444,13 @@ NVTX.@annotate function set_implicit_precomputed_quantities!(Y, p, t) end @. ᶜts = ts_gs(thermo_args..., Y.c, ᶜK, ᶜΦ, Y.c.ρ) @. ᶜp = TD.air_pressure(thermo_params, ᶜts) - @. ᶜh_tot = TD.total_specific_enthalpy( - thermo_params, - ᶜts, - specific(Y.c.ρe_tot, Y.c.ρ), - ) if turbconv_model isa PrognosticEDMFX set_prognostic_edmf_precomputed_quantities_draft!(Y, p, ᶠuₕ³, t) set_prognostic_edmf_precomputed_quantities_environment!(Y, p, ᶠuₕ³, t) set_prognostic_edmf_precomputed_quantities_implicit_closures!(Y, p, t) - elseif turbconv_model isa AbstractEDMF - (; ᶜtke⁰) = p.precomputed - @. ᶜtke⁰ = Y.c.sgs⁰.ρatke / Y.c.ρ + elseif !(isnothing(turbconv_model)) + # Do nothing for other turbconv models for now end end diff --git a/src/cache/prognostic_edmf_precomputed_quantities.jl b/src/cache/prognostic_edmf_precomputed_quantities.jl index 6f94fc5567..79ac32cad0 100644 --- a/src/cache/prognostic_edmf_precomputed_quantities.jl +++ b/src/cache/prognostic_edmf_precomputed_quantities.jl @@ -8,7 +8,7 @@ import ClimaCore: Spaces, Fields """ set_prognostic_edmf_precomputed_quantities!(Y, p, ᶠuₕ³, t) -Updates the edmf environment precomputed quantities stored in `p` for edmfx. +Updates the edmf environment precomputed quantities stored in `p` for prognostic edmfx. """ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_environment!( Y, @@ -21,66 +21,25 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_environment!( thermo_params = CAP.thermodynamics_params(p.params) (; turbconv_model) = p.atmos (; ᶜΦ,) = p.core - (; ᶜp, ᶜh_tot, ᶜK) = p.precomputed - (; ᶜtke⁰, ᶜρa⁰, ᶠu₃⁰, ᶜu⁰, ᶠu³⁰, ᶜK⁰, ᶜts⁰, ᶜρ⁰, ᶜmse⁰, ᶜq_tot⁰) = - p.precomputed - if p.atmos.moisture_model isa NonEquilMoistModel && - p.atmos.microphysics_model isa Microphysics1Moment - (; ᶜq_liq⁰, ᶜq_ice⁰, ᶜq_rai⁰, ᶜq_sno⁰) = p.precomputed - end + (; ᶜp, ᶜK) = p.precomputed + (; ᶠu₃⁰, ᶜu⁰, ᶠu³⁰, ᶜK⁰, ᶜts⁰) = p.precomputed + + ᶜρa⁰ = @. lazy(ρa⁰(Y.c.ρ, Y.c.sgsʲs, turbconv_model)) + ᶜtke⁰ = @. lazy(specific_tke(Y.c.ρ, Y.c.sgs⁰.ρatke, ᶜρa⁰, turbconv_model)) - @. ᶜρ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, - ) - if p.atmos.moisture_model isa NonEquilMoistModel && - p.atmos.microphysics_model isa Microphysics1Moment - @. ᶜq_liq⁰ = divide_by_ρa( - Y.c.ρq_liq - ρaq_liq⁺(Y.c.sgsʲs), - ᶜρa⁰, - Y.c.ρq_liq, - Y.c.ρ, - turbconv_model, - ) - @. ᶜq_ice⁰ = divide_by_ρa( - Y.c.ρq_ice - ρaq_ice⁺(Y.c.sgsʲs), - ᶜρa⁰, - Y.c.ρq_ice, - Y.c.ρ, - turbconv_model, - ) - @. ᶜq_rai⁰ = divide_by_ρa( - Y.c.ρq_rai - ρaq_rai⁺(Y.c.sgsʲs), - ᶜρa⁰, - Y.c.ρq_rai, - Y.c.ρ, - turbconv_model, - ) - @. ᶜq_sno⁰ = divide_by_ρa( - Y.c.ρq_sno - ρaq_sno⁺(Y.c.sgsʲs), - ᶜρa⁰, - Y.c.ρq_sno, - 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⁰ + ᶜq_tot⁰ = ᶜspecific_env_value(Val(:q_tot), Y, p) + + ᶜmse⁰ = ᶜspecific_env_mse(Y, p) + if p.atmos.moisture_model isa NonEquilMoistModel && p.atmos.microphysics_model isa Microphysics1Moment + ᶜq_liq⁰ = ᶜspecific_env_value(Val(:q_liq), Y, p) + ᶜq_ice⁰ = ᶜspecific_env_value(Val(:q_ice), Y, p) + ᶜq_rai⁰ = ᶜspecific_env_value(Val(:q_rai), Y, p) + ᶜq_sno⁰ = ᶜspecific_env_value(Val(:q_sno), Y, p) @. ᶜts⁰ = TD.PhaseNonEquil_phq( thermo_params, ᶜp, @@ -88,9 +47,9 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_environment!( TD.PhasePartition(ᶜq_tot⁰, ᶜq_liq⁰ + ᶜq_rai⁰, ᶜq_ice⁰ + ᶜq_sno⁰), ) else + @. ᶜts⁰ = TD.PhaseEquil_phq(thermo_params, ᶜp, ᶜmse⁰ - ᶜΦ, ᶜq_tot⁰) end - @. ᶜρ⁰ = TD.air_density(thermo_params, ᶜts⁰) return nothing end @@ -173,7 +132,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_bottom_bc!( turbconv_params = CAP.turbconv_params(p.params) (; ᶜΦ,) = p.core - (; ᶜspecific, ᶜp, ᶜh_tot, ᶜK, ᶜtsʲs, ᶜρʲs) = p.precomputed + (; ᶜp, ᶜK, ᶜtsʲs, ᶜρʲs, ᶜts) = p.precomputed (; ustar, obukhov_length, buoyancy_flux) = p.precomputed.sfc_conditions for j in 1:n @@ -217,6 +176,13 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_bottom_bc!( # TODO: replace this with the actual surface area fraction when # using prognostic surface area @. ᶜaʲ_int_val = FT(turbconv_params.surface_area) + ᶜh_tot = @. lazy( + TD.total_specific_enthalpy( + thermo_params, + ᶜts, + specific(Y.c.ρe_tot, Y.c.ρ), + ), + ) ᶜh_tot_int_val = Fields.field_values(Fields.level(ᶜh_tot, 1)) ᶜK_int_val = Fields.field_values(Fields.level(ᶜK, 1)) ᶜmseʲ_int_val = Fields.field_values(Fields.level(ᶜmseʲ, 1)) @@ -233,7 +199,9 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_bottom_bc!( ) # ... and the first interior point for EDMFX ᶜq_totʲ. - ᶜq_tot_int_val = Fields.field_values(Fields.level(ᶜspecific.q_tot, 1)) + + ᶜq_tot = @. lazy(specific(Y.c.ρq_tot, Y.c.ρ)) + ᶜq_tot_int_val = Fields.field_values(Fields.level(ᶜq_tot, 1)) ᶜq_totʲ_int_val = Fields.field_values(Fields.level(ᶜq_totʲ, 1)) @. ᶜq_totʲ_int_val = sgs_scalar_first_interior_bc( ᶜz_int_val - z_sfc_val, @@ -249,23 +217,24 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_bottom_bc!( if p.atmos.moisture_model isa NonEquilMoistModel && p.atmos.microphysics_model isa Microphysics1Moment # TODO - any better way to define the cloud and precip tracer flux? - ᶜq_liq_int_val = - Fields.field_values(Fields.level(ᶜspecific.q_liq, 1)) + + ᶜq_liq = @. lazy(specific(Y.c.ρq_liq, Y.c.ρ)) + ᶜq_ice = @. lazy(specific(Y.c.ρq_ice, Y.c.ρ)) + ᶜq_rai = @. lazy(specific(Y.c.ρq_rai, Y.c.ρ)) + ᶜq_sno = @. lazy(specific(Y.c.ρq_sno, Y.c.ρ)) + ᶜq_liq_int_val = Fields.field_values(Fields.level(ᶜq_liq, 1)) ᶜq_liqʲ_int_val = Fields.field_values(Fields.level(ᶜq_liqʲ, 1)) @. ᶜq_liqʲ_int_val = ᶜq_liq_int_val - ᶜq_ice_int_val = - Fields.field_values(Fields.level(ᶜspecific.q_ice, 1)) + ᶜq_ice_int_val = Fields.field_values(Fields.level(ᶜq_ice, 1)) ᶜq_iceʲ_int_val = Fields.field_values(Fields.level(ᶜq_iceʲ, 1)) @. ᶜq_iceʲ_int_val = ᶜq_ice_int_val - ᶜq_rai_int_val = - Fields.field_values(Fields.level(ᶜspecific.q_rai, 1)) + ᶜq_rai_int_val = Fields.field_values(Fields.level(ᶜq_rai, 1)) ᶜq_raiʲ_int_val = Fields.field_values(Fields.level(ᶜq_raiʲ, 1)) @. ᶜq_raiʲ_int_val = ᶜq_rai_int_val - ᶜq_sno_int_val = - Fields.field_values(Fields.level(ᶜspecific.q_sno, 1)) + ᶜq_sno_int_val = Fields.field_values(Fields.level(ᶜq_sno, 1)) ᶜq_snoʲ_int_val = Fields.field_values(Fields.level(ᶜq_snoʲ, 1)) @. ᶜq_snoʲ_int_val = ᶜq_sno_int_val end @@ -367,7 +336,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_explicit_clos FT = eltype(params) n = n_mass_flux_subdomains(turbconv_model) - (; ᶜtke⁰, ᶜu, ᶜp, ᶜρa⁰, ᶠu³⁰, ᶜts⁰, ᶜq_tot⁰) = p.precomputed + (; ᶜu, ᶜp, ᶠu³⁰, ᶜts⁰) = p.precomputed (; ᶜlinear_buoygrad, ᶜstrain_rate_norm, ρatke_flux) = p.precomputed (; ᶜuʲs, @@ -386,6 +355,8 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_explicit_clos ᶜdz = Fields.Δz_field(axes(Y.c)) ᶜlg = Fields.local_geometry_field(Y.c) ᶠlg = Fields.local_geometry_field(Y.f) + ᶜρa⁰ = @. lazy(ρa⁰(Y.c.ρ, Y.c.sgsʲs, turbconv_model)) + ᶜtke⁰ = @. lazy(specific_tke(Y.c.ρ, Y.c.sgs⁰.ρatke, ᶜρa⁰, turbconv_model)) ᶜvert_div = p.scratch.ᶜtemp_scalar ᶜmassflux_vert_div = p.scratch.ᶜtemp_scalar_2 @@ -470,6 +441,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_explicit_clos (; ᶜgradᵥ_θ_virt⁰, ᶜgradᵥ_q_tot⁰, ᶜgradᵥ_θ_liq_ice⁰) = p.precomputed # First order approximation: Use environmental mean fields. @. ᶜgradᵥ_θ_virt⁰ = ᶜgradᵥ(ᶠinterp(TD.virtual_pottemp(thermo_params, ᶜts⁰))) # ∂θv∂z_unsat + ᶜq_tot⁰ = ᶜspecific_env_value(Val(:q_tot), Y, p) @. ᶜgradᵥ_q_tot⁰ = ᶜgradᵥ(ᶠinterp(ᶜq_tot⁰)) # ∂qt∂z_sat @. ᶜgradᵥ_θ_liq_ice⁰ = ᶜgradᵥ(ᶠinterp(TD.liquid_ice_pottemp(thermo_params, ᶜts⁰))) # ∂θl∂z_sat @@ -531,7 +503,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_precipitation (; params, dt) = p thp = CAP.thermodynamics_params(params) cmp = CAP.microphysics_0m_params(params) - (; ᶜts⁰, ᶜq_tot⁰, ᶜtsʲs, ᶜSqₜᵖʲs, ᶜSqₜᵖ⁰) = p.precomputed + (; ᶜts⁰, ᶜtsʲs, ᶜSqₜᵖʲs, ᶜSqₜᵖ⁰) = p.precomputed # Sources from the updrafts n = n_mass_flux_subdomains(p.atmos.turbconv_model) @@ -545,6 +517,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_precipitation ) end # sources from the environment + ᶜq_tot⁰ = ᶜspecific_env_value(Val(:q_tot), Y, p) @. ᶜSqₜᵖ⁰ = q_tot_0M_precipitation_sources(thp, cmp, dt, ᶜq_tot⁰, ᶜts⁰) return nothing end @@ -560,10 +533,10 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_precipitation thp = CAP.thermodynamics_params(params) cmp = CAP.microphysics_1m_params(params) cmc = CAP.microphysics_cloud_params(params) + (; turbconv_model) = p.atmos (; ᶜSqₗᵖʲs, ᶜSqᵢᵖʲs, ᶜSqᵣᵖʲs, ᶜSqₛᵖʲs, ᶜρʲs, ᶜtsʲs) = p.precomputed - (; ᶜSqₗᵖ⁰, ᶜSqᵢᵖ⁰, ᶜSqᵣᵖ⁰, ᶜSqₛᵖ⁰, ᶜρ⁰, ᶜts⁰) = p.precomputed - (; ᶜq_tot⁰, ᶜq_liq⁰, ᶜq_ice⁰, ᶜq_rai⁰, ᶜq_sno⁰) = p.precomputed + (; ᶜSqₗᵖ⁰, ᶜSqᵢᵖ⁰, ᶜSqᵣᵖ⁰, ᶜSqₛᵖ⁰, ᶜts⁰) = p.precomputed # TODO - can I re-use them between js and env? ᶜSᵖ = p.scratch.ᶜtemp_scalar @@ -634,6 +607,12 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_precipitation end # Precipitation sources and sinks from the environment + ᶜq_tot⁰ = ᶜspecific_env_value(Val(:q_tot), Y, p) + ᶜq_liq⁰ = ᶜspecific_env_value(Val(:q_liq), Y, p) + ᶜq_ice⁰ = ᶜspecific_env_value(Val(:q_ice), Y, p) + ᶜq_rai⁰ = ᶜspecific_env_value(Val(:q_rai), Y, p) + ᶜq_sno⁰ = ᶜspecific_env_value(Val(:q_sno), Y, p) + ᶜρ⁰ = @. lazy(TD.air_density(thp, ᶜts⁰)) compute_precipitation_sources!( ᶜSᵖ, ᶜSᵖ_snow, diff --git a/src/diagnostics/Diagnostics.jl b/src/diagnostics/Diagnostics.jl index 7248e5a5d8..3ba86ad17a 100644 --- a/src/diagnostics/Diagnostics.jl +++ b/src/diagnostics/Diagnostics.jl @@ -64,6 +64,9 @@ import ..eddy_viscosity import ..turbulent_prandtl_number import ..smagorinsky_lilly_length import ..ᶜcompute_eddy_diffusivity_coefficient +import ..ρa⁰ +import ..specific_tke +import ..ᶜspecific_env_value # We need the abbreviations for symbols like curl, grad, and so on diff --git a/src/diagnostics/edmfx_diagnostics.jl b/src/diagnostics/edmfx_diagnostics.jl index 20655075fc..e0971124a1 100644 --- a/src/diagnostics/edmfx_diagnostics.jl +++ b/src/diagnostics/edmfx_diagnostics.jl @@ -633,15 +633,16 @@ compute_aren!(_, _, _, _, turbconv_model::T) where {T} = function compute_aren!(out, state, cache, time, turbconv_model::PrognosticEDMFX) thermo_params = CAP.thermodynamics_params(cache.params) + ᶜρa⁰ = @. lazy(ρa⁰(state.c.ρ, state.c.sgsʲs, turbconv_model)) if isnothing(out) return draft_area.( - cache.precomputed.ᶜρa⁰, + ᶜρa⁰, TD.air_density.(thermo_params, cache.precomputed.ᶜts⁰), ) else out .= draft_area.( - cache.precomputed.ᶜρa⁰, + ᶜρa⁰, TD.air_density.(thermo_params, cache.precomputed.ᶜts⁰), ) end @@ -943,6 +944,7 @@ function compute_clwen!( TD.liquid_specific_humidity.(thermo_params, cache.precomputed.ᶜts⁰) end end + function compute_clwen!( out, state, @@ -951,11 +953,11 @@ function compute_clwen!( moisture_model::NonEquilMoistModel, turbconv_model::PrognosticEDMFX, ) - thermo_params = CAP.thermodynamics_params(cache.params) + ᶜq_liq⁰ = ᶜspecific_env_value(Val(:q_liq), state, cache) if isnothing(out) - return cache.precomputed.ᶜq_liq⁰ + return Base.materialize(ᶜq_liq⁰) else - out .= cache.precomputed.ᶜq_liq⁰ + out .= ᶜq_liq⁰ end end @@ -1007,6 +1009,7 @@ function compute_clien!( out .= TD.ice_specific_humidity.(thermo_params, cache.precomputed.ᶜts⁰) end end + function compute_clien!( out, state, @@ -1015,11 +1018,11 @@ function compute_clien!( moisture_model::NonEquilMoistModel, turbconv_model::PrognosticEDMFX, ) - thermo_params = CAP.thermodynamics_params(cache.params) + ᶜq_ice⁰ = ᶜspecific_env_value(Val(:q_ice), state, cache) if isnothing(out) - return cache.precomputed.ᶜq_ice⁰ + return Base.materialize(ᶜq_ice⁰) else - out .= cache.precomputed.ᶜq_ice⁰ + out .= ᶜq_ice⁰ end end @@ -1064,11 +1067,11 @@ function compute_husraen!( microphysics_model_model::Microphysics1Moment, turbconv_model::PrognosticEDMFX, ) - thermo_params = CAP.thermodynamics_params(cache.params) + ᶜq_rai⁰ = ᶜspecific_env_value(Val(:q_rai), state, cache) if isnothing(out) - return cache.precomputed.ᶜq_rai⁰ + return Base.materialize(ᶜq_rai⁰) else - out .= cache.precomputed.ᶜq_rai⁰ + out .= ᶜq_rai⁰ end end @@ -1113,11 +1116,11 @@ function compute_hussnen!( microphysics_model_model::Microphysics1Moment, turbconv_model::PrognosticEDMFX, ) - thermo_params = CAP.thermodynamics_params(cache.params) + ᶜq_sno⁰ = ᶜspecific_env_value(Val(:q_sno), state, cache) if isnothing(out) - return cache.precomputed.ᶜq_sno⁰ + return Base.materialize(ᶜq_sno⁰) else - out .= cache.precomputed.ᶜq_sno⁰ + out .= ᶜq_sno⁰ end end @@ -1147,10 +1150,20 @@ function compute_tke!( time, turbconv_model::Union{EDOnlyEDMFX, PrognosticEDMFX, DiagnosticEDMFX}, ) + if turbconv_model isa PrognosticEDMFX + sgsʲs = state.c.sgsʲs + ᶜρa⁰ = @. lazy(ρa⁰(state.c.ρ, sgsʲs, turbconv_model)) + else + ᶜρa⁰ = state.c.ρ + end + + ᶜtke = @. lazy( + specific_tke(state.c.ρ, state.c.sgs⁰.ρatke, ᶜρa⁰, turbconv_model), + ) if isnothing(out) - return copy(cache.precomputed.ᶜtke⁰) + return Base.materialize(ᶜtke) else - out .= cache.precomputed.ᶜtke⁰ + out .= ᶜtke end end @@ -1314,9 +1327,15 @@ function compute_edt!( turbconv_model::Union{PrognosticEDMFX, DiagnosticEDMFX}, ) turbconv_params = CAP.turbconv_params(cache.params) - (; ᶜlinear_buoygrad, ᶜstrain_rate_norm, ᶜtke⁰) = cache.precomputed + (; ᶜlinear_buoygrad, ᶜstrain_rate_norm) = cache.precomputed (; params) = cache + ᶜρa⁰ = + turbconv_model isa PrognosticEDMFX ? + (@. lazy(ρa⁰(state.c.ρ, state.c.sgsʲs, turbconv_model))) : state.c.ρ + ᶜtke⁰ = @. lazy( + specific_tke(state.c.ρ, state.c.sgs⁰.ρatke, ᶜρa⁰, turbconv_model), + ) ᶜmixing_length_field = ᶜmixing_length(state, cache) ᶜK_u = @. lazy(eddy_viscosity(turbconv_params, ᶜtke⁰, ᶜmixing_length_field)) ᶜprandtl_nvec = @. lazy( @@ -1399,7 +1418,13 @@ function compute_evu!( turbconv_model::Union{PrognosticEDMFX, DiagnosticEDMFX}, ) turbconv_params = CAP.turbconv_params(cache.params) - (; ᶜtke⁰) = cache.precomputed + + ᶜρa⁰ = + turbconv_model isa PrognosticEDMFX ? + (@. lazy(ρa⁰(state.c.ρ, state.c.sgsʲs, turbconv_model))) : state.c.ρ + ᶜtke⁰ = @. lazy( + specific_tke(state.c.ρ, state.c.sgs⁰.ρatke, ᶜρa⁰, turbconv_model), + ) ᶜmixing_length_field = ᶜmixing_length(state, cache) ᶜK_u = @. lazy(eddy_viscosity(turbconv_params, ᶜtke⁰, ᶜmixing_length_field)) diff --git a/src/parameterized_tendencies/les_sgs_models/smagorinsky_lilly.jl b/src/parameterized_tendencies/les_sgs_models/smagorinsky_lilly.jl index 1ac289a006..962f4fb767 100644 --- a/src/parameterized_tendencies/les_sgs_models/smagorinsky_lilly.jl +++ b/src/parameterized_tendencies/les_sgs_models/smagorinsky_lilly.jl @@ -96,7 +96,8 @@ horizontal_smagorinsky_lilly_tendency!(Yₜ, Y, p, t, ::Nothing) = nothing vertical_smagorinsky_lilly_tendency!(Yₜ, Y, p, t, ::Nothing) = nothing function horizontal_smagorinsky_lilly_tendency!(Yₜ, Y, p, t, ::SmagorinskyLilly) - (; ᶜτ_smag, ᶠτ_smag, ᶜD_smag, ᶜspecific, ᶜh_tot) = p.precomputed + (; ᶜτ_smag, ᶠτ_smag, ᶜD_smag, ᶜts) = p.precomputed + thermo_params = CAP.thermodynamics_params(p.params) ## Momentum tendencies ᶠρ = @. p.scratch.ᶠtemp_scalar = ᶠinterp(Y.c.ρ) @@ -104,6 +105,13 @@ function horizontal_smagorinsky_lilly_tendency!(Yₜ, Y, p, t, ::SmagorinskyLill @. Yₜ.f.u₃ -= C3(wdivₕ(ᶠρ * ᶠτ_smag) / ᶠρ) ## Total energy tendency + ᶜh_tot = @. lazy( + TD.total_specific_enthalpy( + thermo_params, + ᶜts, + specific(Y.c.ρe_tot, Y.c.ρ), + ), + ) @. Yₜ.c.ρe_tot += wdivₕ(Y.c.ρ * ᶜD_smag * gradₕ(ᶜh_tot)) ## Tracer diffusion and associated mass changes @@ -124,9 +132,10 @@ import UnrolledUtilities as UU function vertical_smagorinsky_lilly_tendency!(Yₜ, Y, p, t, ::SmagorinskyLilly) FT = eltype(Y) (; sfc_temp_C3, ᶠtemp_scalar) = p.scratch - (; ᶜτ_smag, ᶠτ_smag, ᶠD_smag, ᶜspecific, ᶜh_tot, sfc_conditions) = - p.precomputed + (; ᶜτ_smag, ᶠτ_smag, ᶠD_smag, sfc_conditions) = p.precomputed (; ρ_flux_uₕ, ρ_flux_h_tot) = sfc_conditions + (; ᶜts) = p.precomputed + thermo_params = CAP.thermodynamics_params(p.params) # Define operators ᶠgradᵥ = Operators.GradientC2F() # apply BCs to ᶜdivᵥ, which wraps ᶠgradᵥ @@ -153,6 +162,13 @@ function vertical_smagorinsky_lilly_tendency!(Yₜ, Y, p, t, ::SmagorinskyLilly) @. Yₜ.f.u₃ -= C3(ᶠdivᵥ(Y.c.ρ * ᶜτ_smag) / ᶠρ) ## Total energy tendency + ᶜh_tot = @. lazy( + TD.total_specific_enthalpy( + thermo_params, + ᶜts, + specific(Y.c.ρe_tot, Y.c.ρ), + ), + ) @. Yₜ.c.ρe_tot -= ᶜdivᵥ_ρe_tot(-(ᶠρ * ᶠD_smag * ᶠgradᵥ(ᶜh_tot))) ## Tracer diffusion and associated mass changes diff --git a/src/parameterized_tendencies/microphysics/precipitation.jl b/src/parameterized_tendencies/microphysics/precipitation.jl index 3ed2eae58c..8647b52de8 100644 --- a/src/parameterized_tendencies/microphysics/precipitation.jl +++ b/src/parameterized_tendencies/microphysics/precipitation.jl @@ -148,7 +148,9 @@ function precipitation_tendency!( # Source terms from EDMFX updrafts (; ᶜSqₗᵖʲs, ᶜSqᵢᵖʲs, ᶜSqᵣᵖʲs, ᶜSqₛᵖʲs) = p.precomputed # Source terms from EDMFX environment - (; ᶜSqₗᵖ⁰, ᶜSqᵢᵖ⁰, ᶜSqᵣᵖ⁰, ᶜSqₛᵖ⁰, ᶜρa⁰) = p.precomputed + (; ᶜSqₗᵖ⁰, ᶜSqᵢᵖ⁰, ᶜSqᵣᵖ⁰, ᶜSqₛᵖ⁰) = p.precomputed + + ᶜρa⁰ = @. lazy(ρa⁰(Y.c.ρ, Y.c.sgsʲs, turbconv_model)) # Update from environment precipitation and cloud formation sources/sinks @. Yₜ.c.ρq_liq += ᶜρa⁰ * ᶜSqₗᵖ⁰ diff --git a/src/parameterized_tendencies/radiation/radiation.jl b/src/parameterized_tendencies/radiation/radiation.jl index 18cc7771ef..e1d6fd9035 100644 --- a/src/parameterized_tendencies/radiation/radiation.jl +++ b/src/parameterized_tendencies/radiation/radiation.jl @@ -475,8 +475,8 @@ function radiation_tendency!(Yₜ, Y, p, t, radiation_mode::RadiationDYCOMS) q_tot_isoline = FT(0.008) Operators.column_reduce!( (nt1, nt2) -> - abs(specific(nt1.ρq_tot, nt1.ρ) - q_tot_isoline) < - abs(specific(nt2.ρq_tot, nt2.ρ) - q_tot_isoline) ? nt1 : nt2, + abs(specific.(nt1.ρq_tot, nt1.ρ) - q_tot_isoline) < + abs(specific.(nt2.ρq_tot, nt2.ρ) - q_tot_isoline) ? nt1 : nt2, isoline_z_ρ_ρq, Base.broadcasted(NT ∘ tuple, ᶜz, Y.c.ρ, Y.c.ρq_tot), ) diff --git a/src/prognostic_equations/advection.jl b/src/prognostic_equations/advection.jl index fb31a95910..e4fc6d4754 100644 --- a/src/prognostic_equations/advection.jl +++ b/src/prognostic_equations/advection.jl @@ -36,7 +36,9 @@ Modifies `Yₜ.c.ρ`, `Yₜ.c.ρe_tot`, `Yₜ.c.uₕ`, and EDMFX-related fields NVTX.@annotate function horizontal_dynamics_tendency!(Yₜ, Y, p, t) n = n_mass_flux_subdomains(p.atmos.turbconv_model) (; ᶜΦ) = p.core - (; ᶜu, ᶜK, ᶜp) = p.precomputed + (; ᶜu, ᶜK, ᶜp, ᶜts) = p.precomputed + (; params) = p + thermo_params = CAP.thermodynamics_params(params) if p.atmos.turbconv_model isa PrognosticEDMFX (; ᶜuʲs) = p.precomputed @@ -49,7 +51,13 @@ NVTX.@annotate function horizontal_dynamics_tendency!(Yₜ, Y, p, t) end end - (; ᶜh_tot) = p.precomputed + ᶜh_tot = @. lazy( + TD.total_specific_enthalpy( + thermo_params, + ᶜts, + specific(Y.c.ρe_tot, Y.c.ρ), + ), + ) @. Yₜ.c.ρe_tot -= wdivₕ(Y.c.ρ * ᶜh_tot * ᶜu) if p.atmos.turbconv_model isa PrognosticEDMFX @@ -174,11 +182,11 @@ NVTX.@annotate function explicit_vertical_advection_tendency!(Yₜ, Y, p, t) (; dt) = p ᶜJ = Fields.local_geometry_field(Y.c).J (; ᶜf³, ᶠf¹², ᶜΦ) = p.core - (; ᶜu, ᶠu³, ᶜK) = p.precomputed + (; ᶜu, ᶠu³, ᶜK, ᶜts) = p.precomputed (; edmfx_upwinding) = n > 0 || advect_tke ? p.atmos.numerics : all_nothing (; ᶜuʲs, ᶜKʲs, ᶠKᵥʲs) = n > 0 ? p.precomputed : all_nothing (; energy_upwinding, tracer_upwinding) = p.atmos.numerics - (; ᶜspecific) = p.precomputed + thermo_params = CAP.thermodynamics_params(p.params) ᶠu³⁰ = advect_tke ? @@ -186,9 +194,26 @@ NVTX.@annotate function explicit_vertical_advection_tendency!(Yₜ, Y, p, t) turbconv_model isa EDOnlyEDMFX ? p.precomputed.ᶠu³ : p.precomputed.ᶠu³⁰ ) : nothing - ᶜρa⁰ = advect_tke ? (n > 0 ? p.precomputed.ᶜρa⁰ : Y.c.ρ) : nothing - ᶜρ⁰ = advect_tke ? (n > 0 ? p.precomputed.ᶜρ⁰ : Y.c.ρ) : nothing - ᶜtke⁰ = advect_tke ? p.precomputed.ᶜtke⁰ : nothing + ᶜρa⁰ = + advect_tke ? + ( + turbconv_model isa PrognosticEDMFX ? + (@. lazy(ρa⁰(Y.c.ρ, Y.c.sgsʲs, turbconv_model))) : Y.c.ρ + ) : nothing + ᶜρ⁰ = if advect_tke + if n > 0 + (; ᶜts⁰) = p.precomputed + @. lazy(TD.air_density(thermo_params, ᶜts⁰)) + else + Y.c.ρ + end + else + nothing + end + ᶜtke⁰ = + advect_tke ? + (@. lazy(specific_tke(Y.c.ρ, Y.c.sgs⁰.ρatke, ᶜρa⁰, turbconv_model))) : + nothing ᶜa_scalar = p.scratch.ᶜtemp_scalar ᶜω³ = p.scratch.ᶜtemp_CT3 ᶠω¹² = p.scratch.ᶠtemp_CT12 @@ -223,7 +248,13 @@ NVTX.@annotate function explicit_vertical_advection_tendency!(Yₜ, Y, p, t) # ... and upwinding correction of energy and total water. # (The central advection of energy and total water is done implicitly.) if energy_upwinding != Val(:none) - (; ᶜh_tot) = p.precomputed + ᶜh_tot = @. lazy( + TD.total_specific_enthalpy( + thermo_params, + ᶜts, + specific(Y.c.ρe_tot, Y.c.ρ), + ), + ) vtt = vertical_transport(ᶜρ, ᶠu³, ᶜh_tot, float(dt), energy_upwinding) vtt_central = vertical_transport(ᶜρ, ᶠu³, ᶜh_tot, float(dt), Val(:none)) @. Yₜ.c.ρe_tot += vtt - vtt_central diff --git a/src/prognostic_equations/eddy_diffusion_closures.jl b/src/prognostic_equations/eddy_diffusion_closures.jl index abe325e589..1a57135b12 100644 --- a/src/prognostic_equations/eddy_diffusion_closures.jl +++ b/src/prognostic_equations/eddy_diffusion_closures.jl @@ -541,11 +541,16 @@ end function ᶜmixing_length(Y, p, property::Val{P} = Val{:master}()) where {P} (; params) = p (; ustar, obukhov_length) = p.precomputed.sfc_conditions - (; ᶜtke⁰) = p.precomputed (; ᶜlinear_buoygrad, ᶜstrain_rate_norm) = p.precomputed ᶜz = Fields.coordinate_field(Y.c).z z_sfc = Fields.level(Fields.coordinate_field(Y.f).z, Fields.half) ᶜdz = Fields.Δz_field(axes(Y.c)) + + turbconv_model = p.atmos.turbconv_model + ᶜρa⁰ = + turbconv_model isa PrognosticEDMFX ? + (@. lazy(ρa⁰(Y.c.ρ, Y.c.sgsʲs, turbconv_model))) : Y.c.ρ + ᶜtke⁰ = @. lazy(specific_tke(Y.c.ρ, Y.c.sgs⁰.ρatke, ᶜρa⁰, turbconv_model)) sfc_tke = Fields.level(ᶜtke⁰, 1) ᶜprandtl_nvec = p.scratch.ᶜtemp_scalar_5 @@ -677,11 +682,12 @@ function ᶜtke_exchange(Y, p) (; turbconv_model) = p.atmos n = n_mass_flux_subdomains(turbconv_model) ᶜρa⁰ = - p.atmos.turbconv_model isa PrognosticEDMFX ? p.precomputed.ᶜρa⁰ : Y.c.ρ - + p.atmos.turbconv_model isa PrognosticEDMFX ? + (@. lazy(ρa⁰(Y.c.ρ, Y.c.sgsʲs, turbconv_model))) : Y.c.ρ + ᶜtke⁰ = @. lazy(specific_tke(Y.c.ρ, Y.c.sgs⁰.ρatke, ᶜρa⁰, turbconv_model)) if p.atmos.turbconv_model isa PrognosticEDMFX - (; ᶜdetrʲs, ᶜtke⁰, ᶠu³⁰, ᶠu³ʲs) = p.precomputed + (; ᶜdetrʲs, ᶠu³⁰, ᶠu³ʲs) = p.precomputed ᶜtke_exch = p.scratch.ᶜtemp_scalar_2 @. ᶜtke_exch = 0 for j in 1:n @@ -694,7 +700,7 @@ function ᶜtke_exchange(Y, p) return ᶜtke_exch elseif p.atmos.turbconv_model isa DiagnosticEDMFX - (; ᶜdetrʲs, ᶜtke⁰, ᶠu³⁰, ᶠu³ʲs, ᶜρaʲs) = p.precomputed + (; ᶜdetrʲs, ᶠu³⁰, ᶠu³ʲs, ᶜρaʲs) = p.precomputed ᶜtke_exch = p.scratch.ᶜtemp_scalar_2 @. ᶜtke_exch = 0 for j in 1:n diff --git a/src/prognostic_equations/edmfx_entr_detr.jl b/src/prognostic_equations/edmfx_entr_detr.jl index c3eeef3734..1cd8663775 100644 --- a/src/prognostic_equations/edmfx_entr_detr.jl +++ b/src/prognostic_equations/edmfx_entr_detr.jl @@ -528,44 +528,47 @@ function edmfx_entr_detr_tendency!(Yₜ, Y, p, t, turbconv_model::PrognosticEDMF n = n_mass_flux_subdomains(turbconv_model) (; ᶜturb_entrʲs, ᶜentrʲs, ᶜdetrʲs) = p.precomputed - (; ᶜq_tot⁰, ᶜmse⁰, ᶠu₃⁰) = p.precomputed + (; ᶠu₃⁰) = p.precomputed + ᶜmse⁰ = ᶜspecific_env_mse(Y, p) if p.atmos.moisture_model isa NonEquilMoistModel && p.atmos.microphysics_model isa Microphysics1Moment - (; ᶜq_liq⁰, ᶜq_ice⁰, ᶜq_rai⁰, ᶜq_sno⁰) = p.precomputed + ᶜq_liq⁰ = ᶜspecific_env_value(Val(:q_liq), Y, p) + ᶜq_ice⁰ = ᶜspecific_env_value(Val(:q_ice), Y, p) + ᶜq_rai⁰ = ᶜspecific_env_value(Val(:q_rai), Y, p) + ᶜq_sno⁰ = ᶜspecific_env_value(Val(:q_sno), Y, p) end for j in 1:n + ᶜentrʲ = ᶜentrʲs.:($j) + ᶜdetrʲ = ᶜdetrʲs.:($j) + ᶜturb_entrʲ = ᶜturb_entrʲs.:($j) + ᶜmseʲ = Y.c.sgsʲs.:($j).mse + ᶜq_totʲ = Y.c.sgsʲs.:($j).q_tot - @. Yₜ.c.sgsʲs.:($$j).ρa += - Y.c.sgsʲs.:($$j).ρa * (ᶜentrʲs.:($$j) - ᶜdetrʲs.:($$j)) + ᶜq_tot⁰ = ᶜspecific_env_value(Val(:q_tot), Y, p) - @. Yₜ.c.sgsʲs.:($$j).mse += - (ᶜentrʲs.:($$j) .+ ᶜturb_entrʲs.:($$j)) * - (ᶜmse⁰ - Y.c.sgsʲs.:($$j).mse) + @. Yₜ.c.sgsʲs.:($$j).ρa += Y.c.sgsʲs.:($$j).ρa * (ᶜentrʲ - ᶜdetrʲ) + + @. Yₜ.c.sgsʲs.:($$j).mse += (ᶜentrʲ .+ ᶜturb_entrʲ) * (ᶜmse⁰ - ᶜmseʲ) @. Yₜ.c.sgsʲs.:($$j).q_tot += - (ᶜentrʲs.:($$j) .+ ᶜturb_entrʲs.:($$j)) * - (ᶜq_tot⁰ - Y.c.sgsʲs.:($$j).q_tot) + (ᶜentrʲ .+ ᶜturb_entrʲ) * (ᶜq_tot⁰ - ᶜq_totʲ) if p.atmos.moisture_model isa NonEquilMoistModel && p.atmos.microphysics_model isa Microphysics1Moment @. Yₜ.c.sgsʲs.:($$j).q_liq += - (ᶜentrʲs.:($$j) .+ ᶜturb_entrʲs.:($$j)) * - (ᶜq_liq⁰ - Y.c.sgsʲs.:($$j).q_liq) + (ᶜentrʲ .+ ᶜturb_entrʲ) * (ᶜq_liq⁰ - Y.c.sgsʲs.:($$j).q_liq) @. Yₜ.c.sgsʲs.:($$j).q_ice += - (ᶜentrʲs.:($$j) .+ ᶜturb_entrʲs.:($$j)) * - (ᶜq_ice⁰ - Y.c.sgsʲs.:($$j).q_ice) + (ᶜentrʲ .+ ᶜturb_entrʲ) * (ᶜq_ice⁰ - Y.c.sgsʲs.:($$j).q_ice) @. Yₜ.c.sgsʲs.:($$j).q_rai += - (ᶜentrʲs.:($$j) .+ ᶜturb_entrʲs.:($$j)) * - (ᶜq_rai⁰ - Y.c.sgsʲs.:($$j).q_rai) + (ᶜentrʲ .+ ᶜturb_entrʲ) * (ᶜq_rai⁰ - Y.c.sgsʲs.:($$j).q_rai) @. Yₜ.c.sgsʲs.:($$j).q_sno += - (ᶜentrʲs.:($$j) .+ ᶜturb_entrʲs.:($$j)) * - (ᶜq_sno⁰ - Y.c.sgsʲs.:($$j).q_sno) + (ᶜentrʲ .+ ᶜturb_entrʲ) * (ᶜq_sno⁰ - Y.c.sgsʲs.:($$j).q_sno) end @. Yₜ.f.sgsʲs.:($$j).u₃ += - (ᶠinterp(ᶜentrʲs.:($$j)) .+ ᶠinterp(ᶜturb_entrʲs.:($$j))) * + (ᶠinterp(ᶜentrʲ) .+ ᶠinterp(ᶜturb_entrʲ)) * (ᶠu₃⁰ - Y.f.sgsʲs.:($$j).u₃) end return nothing diff --git a/src/prognostic_equations/edmfx_sgs_flux.jl b/src/prognostic_equations/edmfx_sgs_flux.jl index 01abfe1f4b..5fe28f2d69 100644 --- a/src/prognostic_equations/edmfx_sgs_flux.jl +++ b/src/prognostic_equations/edmfx_sgs_flux.jl @@ -40,15 +40,12 @@ function edmfx_sgs_mass_flux_tendency!( n = n_mass_flux_subdomains(turbconv_model) (; edmfx_sgsflux_upwinding) = p.atmos.numerics - (; ᶠu³, ᶜh_tot) = p.precomputed + (; ᶠu³) = p.precomputed (; ᶠu³ʲs, ᶜKʲs, ᶜρʲs) = p.precomputed - (; ᶜρa⁰, ᶜρ⁰, ᶠu³⁰, ᶜK⁰, ᶜmse⁰, ᶜq_tot⁰) = p.precomputed - if ( - p.atmos.moisture_model isa NonEquilMoistModel && - p.atmos.microphysics_model isa Microphysics1Moment - ) - (; ᶜq_liq⁰, ᶜq_ice⁰, ᶜq_rai⁰, ᶜq_sno⁰) = p.precomputed - end + (; ᶠu³⁰, ᶜK⁰, ᶜts⁰, ᶜts) = p.precomputed + thermo_params = CAP.thermodynamics_params(p.params) + ᶜρ⁰ = @. lazy(TD.air_density(thermo_params, ᶜts⁰)) + ᶜρa⁰ = @. lazy(ρa⁰(Y.c.ρ, Y.c.sgsʲs, turbconv_model)) (; dt) = p ᶜJ = Fields.local_geometry_field(Y.c).J @@ -59,6 +56,13 @@ function edmfx_sgs_mass_flux_tendency!( # [best after removal of precomputed quantities] ᶠu³_diff = p.scratch.ᶠtemp_CT3 ᶜa_scalar = p.scratch.ᶜtemp_scalar + ᶜh_tot = @. lazy( + TD.total_specific_enthalpy( + thermo_params, + ᶜts, + specific(Y.c.ρe_tot, Y.c.ρ), + ), + ) for j in 1:n @. ᶠu³_diff = ᶠu³ʲs.:($$j) - ᶠu³ @. ᶜa_scalar = @@ -75,6 +79,8 @@ function edmfx_sgs_mass_flux_tendency!( end # Add the environment fluxes @. ᶠu³_diff = ᶠu³⁰ - ᶠu³ + + ᶜmse⁰ = ᶜspecific_env_mse(Y, p) @. ᶜa_scalar = (ᶜmse⁰ + ᶜK⁰ - ᶜh_tot) * draft_area(ᶜρa⁰, ᶜρ⁰) vtt = vertical_transport( ᶜρ⁰, @@ -102,6 +108,7 @@ function edmfx_sgs_mass_flux_tendency!( @. Yₜ.c.ρq_tot += vtt end # Add the environment fluxes + ᶜq_tot⁰ = ᶜspecific_env_value(Val(:q_tot), Y, p) @. ᶠu³_diff = ᶠu³⁰ - ᶠu³ @. ᶜa_scalar = (ᶜq_tot⁰ - specific(Y.c.ρq_tot, Y.c.ρ)) * draft_area(ᶜρa⁰, ᶜρ⁰) @@ -119,6 +126,10 @@ function edmfx_sgs_mass_flux_tendency!( p.atmos.moisture_model isa NonEquilMoistModel && p.atmos.microphysics_model isa Microphysics1Moment ) + ᶜq_liq⁰ = ᶜspecific_env_value(Val(:q_liq), Y, p) + ᶜq_ice⁰ = ᶜspecific_env_value(Val(:q_ice), Y, p) + ᶜq_rai⁰ = ᶜspecific_env_value(Val(:q_rai), Y, p) + ᶜq_sno⁰ = ᶜspecific_env_value(Val(:q_sno), Y, p) # Liquid, ice, rain and snow specific humidity fluxes for j in 1:n @. ᶠu³_diff = ᶠu³ʲs.:($$j) - ᶠu³ @@ -236,14 +247,22 @@ function edmfx_sgs_mass_flux_tendency!( a_max = CAP.max_area(turbconv_params) n = n_mass_flux_subdomains(turbconv_model) (; edmfx_sgsflux_upwinding) = p.atmos.numerics - (; ᶠu³, ᶜh_tot) = p.precomputed - (; ᶜρaʲs, ᶜρʲs, ᶠu³ʲs, ᶜKʲs, ᶜmseʲs, ᶜq_totʲs) = p.precomputed + (; ᶠu³) = p.precomputed + (; ᶜρaʲs, ᶜρʲs, ᶠu³ʲs, ᶜKʲs, ᶜmseʲs, ᶜq_totʲs, ᶜts) = p.precomputed (; dt) = p ᶜJ = Fields.local_geometry_field(Y.c).J FT = eltype(Y) if p.atmos.edmfx_model.sgs_mass_flux isa Val{true} - # Enthalpy fluxes + thermo_params = CAP.thermodynamics_params(p.params) + # energy + ᶜh_tot = @. lazy( + TD.total_specific_enthalpy( + thermo_params, + ᶜts, + specific(Y.c.ρe_tot, Y.c.ρ), + ), + ) ᶠu³_diff = p.scratch.ᶠtemp_CT3 ᶜa_scalar = p.scratch.ᶜtemp_scalar for j in 1:n @@ -272,7 +291,7 @@ function edmfx_sgs_mass_flux_tendency!( for j in 1:n @. ᶠu³_diff = ᶠu³ʲs.:($$j) - ᶠu³ # @. ᶜa_scalar = - # (ᶜq_totʲs.:($$j) - specific(Y.c.ρq_tot, Y.c.ρ)) * + # (ᶜq_totʲs.:($$j) - specific(Y.c.ρq_tot, Y.c.ρ) * # draft_area(ᶜρaʲs.:($$j), ᶜρʲs.:($$j)) # TODO: remove this filter when mass flux is treated implicitly @. ᶜa_scalar = @@ -296,19 +315,51 @@ function edmfx_sgs_mass_flux_tendency!( end # TODO: add environment flux? end + # TODO: the following adds the environment flux to the tendency + # Make active and test later + # @. ᶠu³_diff = p.precomputed.ᶠu³⁰ - ᶠu³ + # ρa⁰(Y.c.ρ, Y.c.sgsʲs, turbconv_model) + # ᶜρ⁰ = p.scratch.ᶜtemp_scalar_2 + # @. ᶜρ⁰ = TD.air_density( + # CAP.thermodynamics_params(p.params), + # p.precomputed.ᶜts⁰, + # ) + # ᶜmse⁰ = @.lazy(ᶜspecific_env_mse(Y, p)) + # @. ᶜa_scalar = + # (ᶜmse⁰ + p.precomputed.ᶜK⁰ - ᶜh_tot) * draft_area(ᶜρa⁰, ᶜρ⁰) + # vtt = vertical_transport( + # ᶜρ⁰, + # ᶠu³_diff, + # ᶜa_scalar, + # dt, + # edmfx_sgsflux_upwinding, + # ) + # @. Yₜ.c.ρe_tot += vtt + # if !(p.atmos.moisture_model isa DryModel) + # ᶜq_tot⁰ = @specific_env_value(:q_tot, Y.c, turbconv_model)) + # @. ᶜa_scalar = + # (ᶜq_tot⁰ - specific(Y.c.ρq_tot, Y.c.ρ)) * + # draft_area(ᶜρa⁰, ᶜρ⁰) + # vtt = vertical_transport( + # ᶜρ⁰, + # ᶠu³_diff, + # ᶜa_scalar, + # dt, + # edmfx_sgsflux_upwinding, + # ) + # @. Yₜ.c.ρq_tot += vtt + # end end - # TODO: add vertical momentum fluxes - - return nothing end """ edmfx_sgs_diffusive_flux_tendency!(Yₜ, Y, p, t, turbconv_model) -Computes and applies tendencies to the grid-mean prognostic variables due to the -divergence of subgrid-scale (SGS) diffusive fluxes, representing turbulent mixing -by the EDMFX environment component. +Computes and applies the tendency to the grid-mean state `Y` due to SGS +diffusive fluxes from the EDMFX environment. This involves calculating the +divergence of turbulent fluxes, which are parameterized using eddy diffusivity +and viscosity closures. This function parameterizes these fluxes using an eddy-diffusivity/viscosity approach (K-theory) for the environment (sgs⁰). Tendencies are calculated for @@ -322,9 +373,9 @@ in place. Arguments: - `Yₜ`: The tendency state vector for grid-mean variables. -- `Y`: The current state vector. -- `p`: Cache containing parameters, precomputed fields, atmospheric - model settings, and scratch space. +- `Y`: The current state vector (used for grid-mean and SGS properties). +- `p`: Cache containing parameters, precomputed fields, atmospheric model settings, + and scratch space. - `t`: Current simulation time. - `turbconv_model`: The turbulence convection model instance. """ @@ -341,15 +392,11 @@ function edmfx_sgs_diffusive_flux_tendency!( (; dt, params) = p turbconv_params = CAP.turbconv_params(params) c_d = CAP.tke_diss_coeff(turbconv_params) - (; ᶜρa⁰, ᶜu⁰, ᶜK⁰, ᶜmse⁰, ᶜq_tot⁰, ᶜtke⁰) = p.precomputed - if ( - p.atmos.moisture_model isa NonEquilMoistModel && - p.atmos.microphysics_model isa Microphysics1Moment - ) - (; ᶜq_liq⁰, ᶜq_ice⁰, ᶜq_rai⁰, ᶜq_sno⁰) = p.precomputed - end + (; ᶜu⁰, ᶜK⁰, ᶜlinear_buoygrad, ᶜstrain_rate_norm) = p.precomputed (; ρatke_flux) = p.precomputed ᶠgradᵥ = Operators.GradientC2F() + ᶜρa⁰ = @. lazy(ρa⁰(Y.c.ρ, Y.c.sgsʲs, turbconv_model)) + ᶜtke⁰ = @. lazy(specific_tke(Y.c.ρ, Y.c.sgs⁰.ρatke, ᶜρa⁰, turbconv_model)) if p.atmos.edmfx_model.sgs_diffusive_flux isa Val{true} @@ -378,6 +425,8 @@ function edmfx_sgs_diffusive_flux_tendency!( top = Operators.SetValue(C3(FT(0))), bottom = Operators.SetValue(C3(FT(0))), ) + + ᶜmse⁰ = ᶜspecific_env_mse(Y, p) @. Yₜ.c.ρe_tot -= ᶜdivᵥ_ρe_tot(-(ᶠρaK_h * ᶠgradᵥ(ᶜmse⁰ + ᶜK⁰))) if use_prognostic_tke(turbconv_model) # Turbulent TKE transport (diffusion) @@ -406,6 +455,7 @@ function edmfx_sgs_diffusive_flux_tendency!( top = Operators.SetValue(C3(FT(0))), bottom = Operators.SetValue(C3(FT(0))), ) + ᶜq_tot⁰ = ᶜspecific_env_value(Val(:q_tot), Y, p) @. ᶜρχₜ_diffusion = ᶜdivᵥ_ρq_tot(-(ᶠρaK_h * ᶠgradᵥ(ᶜq_tot⁰))) @. Yₜ.c.ρq_tot -= ᶜρχₜ_diffusion @. Yₜ.c.ρ -= ᶜρχₜ_diffusion # Effect of moisture diffusion on (moist) air mass @@ -414,6 +464,10 @@ function edmfx_sgs_diffusive_flux_tendency!( p.atmos.moisture_model isa NonEquilMoistModel && p.atmos.microphysics_model isa Microphysics1Moment ) + ᶜq_liq⁰ = ᶜspecific_env_value(Val(:q_liq), Y, p) + ᶜq_ice⁰ = ᶜspecific_env_value(Val(:q_ice), Y, p) + ᶜq_rai⁰ = ᶜspecific_env_value(Val(:q_rai), Y, p) + ᶜq_sno⁰ = ᶜspecific_env_value(Val(:q_sno), Y, p) # Liquid, ice, rain and snow specific humidity diffusion α_vert_diff_tracer = CAP.α_vert_diff_tracer(params) @@ -461,10 +515,12 @@ function edmfx_sgs_diffusive_flux_tendency!( 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, ᶜh_tot, ᶜtke⁰) = p.precomputed + (; ᶜu, ᶜts) = p.precomputed (; ρatke_flux) = p.precomputed ᶠgradᵥ = Operators.GradientC2F() + ᶜtke⁰ = @. lazy(specific_tke(Y.c.ρ, Y.c.sgs⁰.ρatke, Y.c.ρ, turbconv_model)) if p.atmos.edmfx_model.sgs_diffusive_flux isa Val{true} @@ -494,6 +550,13 @@ function edmfx_sgs_diffusive_flux_tendency!( 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) diff --git a/src/prognostic_equations/edmfx_tke.jl b/src/prognostic_equations/edmfx_tke.jl index e0b7b566e2..f72563946f 100644 --- a/src/prognostic_equations/edmfx_tke.jl +++ b/src/prognostic_equations/edmfx_tke.jl @@ -27,9 +27,10 @@ edmfx_tke_tendency!(Yₜ, Y, p, t, turbconv_model) = nothing function edmfx_tke_tendency!(Yₜ, Y, p, t, turbconv_model::EDOnlyEDMFX) (; params) = p - (; ᶜstrain_rate_norm, ᶜlinear_buoygrad, ᶜtke⁰) = p.precomputed + (; ᶜstrain_rate_norm, ᶜlinear_buoygrad) = p.precomputed turbconv_params = CAP.turbconv_params(p.params) - + ᶜρa⁰ = Y.c.ρ # EDOnly + ᶜtke⁰ = @. lazy(specific_tke(Y.c.ρ, Y.c.sgs⁰.ρatke, ᶜρa⁰, turbconv_model)) ᶜmixing_length_field = p.scratch.ᶜtemp_scalar ᶜmixing_length_field .= ᶜmixing_length(Y, p) ᶜK_u = @. lazy(eddy_viscosity(turbconv_params, ᶜtke⁰, ᶜmixing_length_field)) @@ -53,11 +54,14 @@ function edmfx_tke_tendency!( ) n = n_mass_flux_subdomains(turbconv_model) (; ᶜturb_entrʲs, ᶜentrʲs, ᶜdetrʲs, ᶠu³ʲs) = p.precomputed - (; ᶠu³⁰, ᶠu³, ᶜstrain_rate_norm, ᶜlinear_buoygrad, ᶜtke⁰) = p.precomputed + (; ᶠu³⁰, ᶠu³, ᶜstrain_rate_norm, ᶜlinear_buoygrad) = p.precomputed turbconv_params = CAP.turbconv_params(p.params) FT = eltype(p.params) - ᶜρa⁰ = turbconv_model isa PrognosticEDMFX ? p.precomputed.ᶜρa⁰ : Y.c.ρ + + ᶜρa⁰ = + turbconv_model isa PrognosticEDMFX ? + (@. lazy(ρa⁰(Y.c.ρ, Y.c.sgsʲs, turbconv_model))) : Y.c.ρ nh_pressure3_buoyʲs = turbconv_model isa PrognosticEDMFX ? p.precomputed.ᶠnh_pressure₃_buoyʲs : p.precomputed.ᶠnh_pressure³_buoyʲs @@ -82,6 +86,11 @@ function edmfx_tke_tendency!( ᶜmixing_length_field = p.scratch.ᶜtemp_scalar_2 ᶜmixing_length_field .= ᶜmixing_length(Y, p) + ᶜρa⁰ = + turbconv_model isa PrognosticEDMFX ? + (@. lazy(ρa⁰(Y.c.ρ, Y.c.sgsʲs, turbconv_model))) : Y.c.ρ + ᶜtke⁰ = + @. lazy(specific_tke(Y.c.ρ, Y.c.sgs⁰.ρatke, ᶜρa⁰, turbconv_model)) ᶜK_u = @. lazy( eddy_viscosity(turbconv_params, ᶜtke⁰, ᶜmixing_length_field), ) @@ -99,6 +108,9 @@ function edmfx_tke_tendency!( # buoyancy production @. Yₜ.c.sgs⁰.ρatke -= ᶜρa⁰ * ᶜK_h * ᶜlinear_buoygrad + ᶜtke⁰ = + @. lazy(specific_tke(Y.c.ρ, Y.c.sgs⁰.ρatke, ᶜρa⁰, turbconv_model)) + # entrainment and detraiment # using ᶜu⁰ and local geometry results in allocation for j in 1:n diff --git a/src/prognostic_equations/forcing/external_forcing.jl b/src/prognostic_equations/forcing/external_forcing.jl index 0d1f023fde..36385ec154 100644 --- a/src/prognostic_equations/forcing/external_forcing.jl +++ b/src/prognostic_equations/forcing/external_forcing.jl @@ -337,7 +337,7 @@ function external_forcing_tendency!( # horizontal advection, vertical fluctuation, nudging, subsidence (need to add), (; params) = p thermo_params = CAP.thermodynamics_params(params) - (; ᶜts, ᶜh_tot) = p.precomputed + (; ᶜts) = p.precomputed (; ᶜdTdt_fluc, ᶜdqtdt_fluc, @@ -358,6 +358,13 @@ function external_forcing_tendency!( @. ᶜuₕ_nudge = C12(Geometry.UVVector(ᶜu_nudge, ᶜv_nudge), ᶜlg) @. Yₜ.c.uₕ -= (Y.c.uₕ - ᶜuₕ_nudge) * ᶜinv_τ_wind + ᶜh_tot = @. lazy( + TD.total_specific_enthalpy( + thermo_params, + ᶜts, + specific(Y.c.ρe_tot, Y.c.ρ), + ), + ) # nudging tendency ᶜdTdt_nudging = p.scratch.ᶜtemp_scalar ᶜdqtdt_nudging = p.scratch.ᶜtemp_scalar_2 @@ -580,7 +587,7 @@ function external_forcing_tendency!(Yₜ, Y, p, t, ::ISDACForcing) FT = Spaces.undertype(axes(Y.c)) (; params) = p thermo_params = CAP.thermodynamics_params(params) - (; ᶜts, ᶜh_tot, ᶜp) = p.precomputed + (; ᶜts, ᶜp) = p.precomputed ᶜinv_τ_scalar = APL.ISDAC_inv_τ_scalar(FT) # s⁻¹ ᶜinv_τ_wind = APL.ISDAC_inv_τ_wind(FT) # s⁻¹ diff --git a/src/prognostic_equations/forcing/subsidence.jl b/src/prognostic_equations/forcing/subsidence.jl index c722a50e03..cb8538359f 100644 --- a/src/prognostic_equations/forcing/subsidence.jl +++ b/src/prognostic_equations/forcing/subsidence.jl @@ -65,10 +65,32 @@ If `subsidence_model` is `Nothing`, no subsidence tendency is applied. """ subsidence_tendency!(Yₜ, Y, p, t, ::Nothing) = nothing # No subsidence + +""" + subsidence_tendency!(Yₜ, Y, p, t, subsidence_model::Subsidence) + +Applies subsidence tendencies to total energy (`ρe_tot`), total specific humidity +(`ρq_tot`), and other moisture species (`ρq_liq`, `ρq_ice`) if a `NonEquilMoistModel` +is used. + +The subsidence velocity profile `w_sub(z)` is obtained from `subsidence_model.prof`. +This profile is used to construct a face-valued vertical velocity field `ᶠsubsidence³`. +The `subsidence!` helper function is then called (currently with a first-order +upwind scheme) to compute and apply the vertical advective tendency for each relevant +scalar quantity `χ`. + +Arguments: +- `Yₜ`: The tendency state vector, modified in place. +- `Y`: The current state vector, used for density (`ρ`). +- `p`: Cache containing parameters, and the subsidence model object. +- `t`: Current simulation time. +- `subsidence`: The subsidence model object. +""" function subsidence_tendency!(Yₜ, Y, p, t, ::Subsidence) (; moisture_model) = p.atmos subsidence_profile = p.atmos.subsidence.prof - (; ᶜh_tot) = p.precomputed + thermo_params = CAP.thermodynamics_params(p.params) + (; ᶜts) = p.precomputed ᶠz = Fields.coordinate_field(axes(Y.f)).z ᶠlg = Fields.local_geometry_field(Y.f) @@ -76,9 +98,16 @@ function subsidence_tendency!(Yₜ, Y, p, t, ::Subsidence) @. ᶠsubsidence³ = subsidence_profile(ᶠz) * CT3(unit_basis_vector_data(CT3, ᶠlg)) - # Large-scale subsidence - subsidence!(Yₜ.c.ρe_tot, Y.c.ρ, ᶠsubsidence³, ᶜh_tot, Val{:first_order}()) + # LS Subsidence + ᶜh_tot = @. lazy( + TD.total_specific_enthalpy( + thermo_params, + ᶜts, + specific(Y.c.ρe_tot, Y.c.ρ), + ), + ) ᶜq_tot = @. lazy(specific(Y.c.ρq_tot, Y.c.ρ)) + subsidence!(Yₜ.c.ρe_tot, Y.c.ρ, ᶠsubsidence³, ᶜh_tot, Val{:first_order}()) subsidence!(Yₜ.c.ρq_tot, Y.c.ρ, ᶠsubsidence³, ᶜq_tot, Val{:first_order}()) if moisture_model isa NonEquilMoistModel ᶜq_liq = @. lazy(specific(Y.c.ρq_liq, Y.c.ρ)) diff --git a/src/prognostic_equations/hyperdiffusion.jl b/src/prognostic_equations/hyperdiffusion.jl index c93adb5023..284876ca70 100644 --- a/src/prognostic_equations/hyperdiffusion.jl +++ b/src/prognostic_equations/hyperdiffusion.jl @@ -117,7 +117,11 @@ NVTX.@annotate function prep_hyperdiffusion_tendency!(Yₜ, Y, p, t) wdivₕ(gradₕ(specific(Y.c.ρe_tot, Y.c.ρ) + ᶜp / Y.c.ρ - ᶜh_ref)) if diffuse_tke - (; ᶜtke⁰) = p.precomputed + ᶜρa⁰ = + turbconv_model isa PrognosticEDMFX ? + (@. lazy(ρa⁰(Y.c.ρ, Y.c.sgsʲs, turbconv_model))) : Y.c.ρ + ᶜtke⁰ = + @. lazy(specific_tke(Y.c.ρ, Y.c.sgs⁰.ρatke, ᶜρa⁰, turbconv_model)) (; ᶜ∇²tke⁰) = p.hyperdiff @. ᶜ∇²tke⁰ = wdivₕ(gradₕ(ᶜtke⁰)) end @@ -148,14 +152,12 @@ NVTX.@annotate function apply_hyperdiffusion_tendency!(Yₜ, Y, p, t) diffuse_tke = use_prognostic_tke(turbconv_model) ᶜJ = Fields.local_geometry_field(Y.c).J point_type = eltype(Fields.coordinate_field(Y.c)) - (; ᶜp) = p.precomputed (; ᶜ∇²u, ᶜ∇²specific_energy) = p.hyperdiff if turbconv_model isa PrognosticEDMFX - (; ᶜρa⁰) = p.precomputed + ᶜρa⁰ = @. lazy(ρa⁰(Y.c.ρ, Y.c.sgsʲs, turbconv_model)) (; ᶜ∇²uₕʲs, ᶜ∇²uᵥʲs, ᶜ∇²uʲs, ᶜ∇²mseʲs) = p.hyperdiff end if use_prognostic_tke(turbconv_model) - (; ᶜtke⁰) = p.precomputed (; ᶜ∇²tke⁰) = p.hyperdiff end diff --git a/src/prognostic_equations/implicit/implicit_tendency.jl b/src/prognostic_equations/implicit/implicit_tendency.jl index 1468746c67..60f4f90b68 100644 --- a/src/prognostic_equations/implicit/implicit_tendency.jl +++ b/src/prognostic_equations/implicit/implicit_tendency.jl @@ -140,7 +140,15 @@ function implicit_vertical_advection_tendency!(Yₜ, Y, p, t) ᶜJ = Fields.local_geometry_field(Y.c).J ᶠJ = Fields.local_geometry_field(Y.f).J (; ᶠgradᵥ_ᶜΦ) = p.core - (; ᶜh_tot, ᶠu³, ᶜp) = p.precomputed + (; ᶠu³, ᶜp, ᶜts) = p.precomputed + thermo_params = CAP.thermodynamics_params(p.params) + ᶜh_tot = @. lazy( + TD.total_specific_enthalpy( + thermo_params, + ᶜts, + specific(Y.c.ρe_tot, Y.c.ρ), + ), + ) @. Yₜ.c.ρ -= ᶜdivᵥ(ᶠinterp(Y.c.ρ * ᶜJ) / ᶠJ * ᶠu³) diff --git a/src/prognostic_equations/implicit/manual_sparse_jacobian.jl b/src/prognostic_equations/implicit/manual_sparse_jacobian.jl index 2a899ddf1f..ce905eaba1 100644 --- a/src/prognostic_equations/implicit/manual_sparse_jacobian.jl +++ b/src/prognostic_equations/implicit/manual_sparse_jacobian.jl @@ -338,7 +338,7 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t) (; matrix) = cache (; params) = p (; ᶜΦ, ᶠgradᵥ_ᶜΦ) = p.core - (; ᶠu³, ᶜK, ᶜts, ᶜp, ᶜh_tot) = p.precomputed + (; ᶠu³, ᶜK, ᶜts, ᶜp) = p.precomputed (; ∂ᶜK_∂ᶜuₕ, ∂ᶜK_∂ᶠu₃, @@ -366,8 +366,15 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t) Δcp_v = FT(CAP.cp_v(params)) - cp_d # This term appears a few times in the Jacobian, and is technically # minus ∂e_int_∂q_tot - ∂e_int_∂q_tot = T_0 * (Δcv_v - R_d) - FT(CAP.e_int_v0(params)) thermo_params = CAP.thermodynamics_params(params) + ∂e_int_∂q_tot = T_0 * (Δcv_v - R_d) - FT(CAP.e_int_v0(params)) + ᶜh_tot = @. lazy( + TD.total_specific_enthalpy( + thermo_params, + ᶜts, + specific(Y.c.ρe_tot, Y.c.ρ), + ), + ) ᶜρ = Y.c.ρ ᶜuₕ = Y.c.uₕ @@ -417,13 +424,21 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t) @. ∂ᶜρ_err_∂ᶠu₃ = dtγ * ᶜadvection_matrix ⋅ DiagonalMatrixRow(g³³(ᶠgⁱʲ)) tracer_info = (@name(c.ρe_tot), @name(c.ρq_tot)) + MatrixFields.unrolled_foreach(tracer_info) do ρχ_name MatrixFields.has_field(Y, ρχ_name) || return ᶜχ = if ρχ_name === @name(c.ρe_tot) - p.precomputed.ᶜh_tot + @. lazy( + TD.total_specific_enthalpy( + thermo_params, + ᶜts, + specific(Y.c.ρe_tot, Y.c.ρ), + ), + ) else @. lazy(specific(Y.c.ρq_tot, Y.c.ρ)) end + if use_derivative(topography_flag) ∂ᶜρχ_err_∂ᶜuₕ = matrix[ρχ_name, @name(c.uₕ)] @. ∂ᶜρχ_err_∂ᶜuₕ = @@ -446,15 +461,15 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t) ᶠinterp_matrix() ) @. ∂ᶠu₃_err_∂ᶜρe_tot = dtγ * ᶠp_grad_matrix ⋅ DiagonalMatrixRow(ᶜkappa_m) + ᶜe_tot = @. lazy(specific(Y.c.ρe_tot, Y.c.ρ)) if MatrixFields.has_field(Y, @name(c.ρq_tot)) + ᶜq_tot = @. lazy(specific(Y.c.ρq_tot, Y.c.ρ)) ∂ᶠu₃_err_∂ᶜρq_tot = matrix[@name(f.u₃), @name(c.ρq_tot)] @. ∂ᶠu₃_err_∂ᶜρq_tot = dtγ * ᶠp_grad_matrix ⋅ DiagonalMatrixRow(( ᶜkappa_m * ∂e_int_∂q_tot + - ᶜ∂kappa_m∂q_tot * ( - cp_d * T_0 + specific(Y.c.ρe_tot, Y.c.ρ) - ᶜK - ᶜΦ + - ∂e_int_∂q_tot * specific(Y.c.ρq_tot, Y.c.ρ) - ) + ᶜ∂kappa_m∂q_tot * + (cp_d * T_0 + ᶜe_tot - ᶜK - ᶜΦ + ∂e_int_∂q_tot * ᶜq_tot) )) end @@ -518,15 +533,15 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t) ∂ᶜρq_tot_err_∂ᶜρq_tot = matrix[@name(c.ρq_tot), @name(c.ρq_tot)] @. ∂ᶜρq_tot_err_∂ᶜρq_tot = zero(typeof(∂ᶜρq_tot_err_∂ᶜρq_tot)) - (I,) - #TODO: tetsing explicit vs implicit + #TODO: testing explicit vs implicit #@. ∂ᶜρq_tot_err_∂ᶜρq_tot = # dtγ * -(ᶜprecipdivᵥ_matrix()) ⋅ # DiagonalMatrixRow(ᶠinterp(ᶜρ * ᶜJ) / ᶠJ) ⋅ ᶠright_bias_matrix() ⋅ # DiagonalMatrixRow( # -1 / ᶜρ * ifelse( - # specific(Y.c.ρq_tot, Y.c.ρ) == 0, + # ᶜq_tot == 0, # (Geometry.WVector(FT(0)),), - # p.precomputed.ᶜwₜqₜ / specific(Y.c.ρq_tot, Y.c.ρ), + # p.precomputed.ᶜwₜqₜ / _tot, # ), # ) - (I,) @@ -544,6 +559,7 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t) end if use_derivative(diffusion_flag) + (; turbconv_model) = p.atmos turbconv_params = CAP.turbconv_params(params) FT = eltype(params) (; vertical_diffusion) = p.atmos @@ -560,7 +576,13 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t) ) ᶜK_u = ᶜK_h else - (; ᶜtke⁰, ᶜlinear_buoygrad, ᶜstrain_rate_norm) = p.precomputed + (; ᶜlinear_buoygrad, ᶜstrain_rate_norm) = p.precomputed + ᶜρa⁰ = + p.atmos.turbconv_model isa PrognosticEDMFX ? + (@. lazy(ρa⁰(Y.c.ρ, Y.c.sgsʲs, turbconv_model))) : Y.c.ρ + ᶜtke⁰ = @. lazy( + specific_tke(Y.c.ρ, Y.c.sgs⁰.ρatke, ᶜρa⁰, turbconv_model), + ) ᶜmixing_length_field = p.scratch.ᶜtemp_scalar_3 ᶜmixing_length_field .= ᶜmixing_length(Y, p) ᶜK_u = @. lazy( @@ -597,28 +619,24 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t) ∂ᶜρe_tot_err_∂ᶜρ = matrix[@name(c.ρe_tot), @name(c.ρ)] @. ∂ᶜρe_tot_err_∂ᶜρ = dtγ * ᶜdiffusion_h_matrix ⋅ DiagonalMatrixRow( - ( - -(1 + ᶜkappa_m) * specific(Y.c.ρe_tot, Y.c.ρ) - - ᶜkappa_m * ∂e_int_∂q_tot * specific(Y.c.ρq_tot, Y.c.ρ) - ) / ᶜρ, + (-(1 + ᶜkappa_m) * ᶜe_tot - ᶜkappa_m * ∂e_int_∂q_tot * ᶜq_tot) / + ᶜρ, ) @. ∂ᶜρe_tot_err_∂ᶜρe_tot += dtγ * ᶜdiffusion_h_matrix ⋅ DiagonalMatrixRow((1 + ᶜkappa_m) / ᶜρ) if MatrixFields.has_field(Y, @name(c.ρq_tot)) + ᶜq_tot = @. lazy(specific(Y.c.ρq_tot, Y.c.ρ)) ∂ᶜρe_tot_err_∂ᶜρq_tot = matrix[@name(c.ρe_tot), @name(c.ρq_tot)] ∂ᶜρq_tot_err_∂ᶜρ = matrix[@name(c.ρq_tot), @name(c.ρ)] @. ∂ᶜρe_tot_err_∂ᶜρq_tot += dtγ * ᶜdiffusion_h_matrix ⋅ DiagonalMatrixRow(( ᶜkappa_m * ∂e_int_∂q_tot / ᶜρ + - ᶜ∂kappa_m∂q_tot * ( - cp_d * T_0 + specific(Y.c.ρe_tot, Y.c.ρ) - ᶜK - ᶜΦ + - ∂e_int_∂q_tot * specific(Y.c.ρq_tot, Y.c.ρ) - ) + ᶜ∂kappa_m∂q_tot * + (cp_d * T_0 + ᶜe_tot - ᶜK - ᶜΦ + ∂e_int_∂q_tot * ᶜq_tot) )) @. ∂ᶜρq_tot_err_∂ᶜρ = - dtγ * ᶜdiffusion_h_matrix ⋅ - DiagonalMatrixRow(-(specific(Y.c.ρq_tot, Y.c.ρ)) / ᶜρ) + dtγ * ᶜdiffusion_h_matrix ⋅ DiagonalMatrixRow(-(ᶜq_tot) / ᶜρ) @. ∂ᶜρq_tot_err_∂ᶜρq_tot += dtγ * ᶜdiffusion_h_matrix ⋅ DiagonalMatrixRow(1 / ᶜρ) end @@ -644,13 +662,18 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t) turbconv_params = CAP.turbconv_params(params) c_d = CAP.tke_diss_coeff(turbconv_params) (; dt) = p - (; ᶜtke⁰) = p.precomputed + turbconv_model = p.atmos.turbconv_model ᶜρa⁰ = p.atmos.turbconv_model isa PrognosticEDMFX ? - p.precomputed.ᶜρa⁰ : ᶜρ + (@. lazy(ρa⁰(Y.c.ρ, Y.c.sgsʲs, turbconv_model))) : Y.c.ρ + ᶜtke⁰ = @. lazy( + specific_tke(Y.c.ρ, Y.c.sgs⁰.ρatke, ᶜρa⁰, turbconv_model), + ) ᶜρatke⁰ = Y.c.sgs⁰.ρatke - ᶜmixing_length_field = ᶜmixing_length(Y, p) + # scratch to prevent GPU Kernel parameter memory error + ᶜmixing_length_field = p.scratch.ᶜtemp_scalar_3 + ᶜmixing_length_field .= ᶜmixing_length(Y, p) @inline tke_dissipation_rate_tendency(tke⁰, mixing_length) = tke⁰ >= 0 ? c_d * sqrt(tke⁰) / mixing_length : 1 / float(dt) @@ -724,7 +747,7 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t) TD.gas_constant_air(thermo_params, ᶜtsʲs.:(1)) / TD.cv_m(thermo_params, ᶜtsʲs.:(1)) - # Note this is the derivative of R_m / cp_m with respect to q_tot + # Note this is the derivative of R_m / cp_m with respect to ᶜq_tot # but we call it ∂kappa_m∂q_totʲ ᶜ∂kappa_m∂q_totʲ = p.scratch.ᶜtemp_scalar_2 @. ᶜ∂kappa_m∂q_totʲ = @@ -980,14 +1003,14 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t) Δcv_v * TD.gas_constant_air(thermo_params, ᶜts) ) / abs2(TD.cv_m(thermo_params, ᶜts)) + ᶜq_tot = @. lazy(specific(Y.c.ρq_tot, Y.c.ρ)) + @. ∂ᶜρe_tot_err_∂ᶜρ += dtγ * ᶜadvdivᵥ_matrix() ⋅ ∂ᶜupdraft_mass_flux_∂ᶜscalar ⋅ DiagonalMatrixRow( ( - -(1 + ᶜkappa_m) * specific(Y.c.ρe_tot, Y.c.ρ) - - ᶜkappa_m * - ∂e_int_∂q_tot * - specific(Y.c.ρq_tot, Y.c.ρ) + -(1 + ᶜkappa_m) * ᶜe_tot - + ᶜkappa_m * ∂e_int_∂q_tot * ᶜq_tot ) / ᶜρ, ) @@ -996,8 +1019,8 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t) DiagonalMatrixRow(( ᶜkappa_m * ∂e_int_∂q_tot / ᶜρ + ᶜ∂kappa_m∂q_tot * ( - cp_d * T_0 + specific(Y.c.ρe_tot, Y.c.ρ) - ᶜK - ᶜΦ + - ∂e_int_∂q_tot * specific(Y.c.ρq_tot, Y.c.ρ) + cp_d * T_0 + ᶜe_tot - ᶜK - ᶜΦ + + ∂e_int_∂q_tot * ᶜq_tot ) )) @@ -1013,7 +1036,7 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t) ## grid-mean ρq_tot @. ∂ᶜρq_tot_err_∂ᶜρ += dtγ * ᶜadvdivᵥ_matrix() ⋅ ∂ᶜupdraft_mass_flux_∂ᶜscalar ⋅ - DiagonalMatrixRow(-(specific(Y.c.ρq_tot, Y.c.ρ)) / ᶜρ) + DiagonalMatrixRow(-(ᶜq_tot) / ᶜρ) @. ∂ᶜρq_tot_err_∂ᶜρq_tot += dtγ * ᶜadvdivᵥ_matrix() ⋅ ∂ᶜupdraft_mass_flux_∂ᶜscalar ⋅ @@ -1052,10 +1075,7 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t) @. ∂ᶜρq_tot_err_∂ᶠu₃ += dtγ * ᶜadvdivᵥ_matrix() ⋅ DiagonalMatrixRow( ᶠinterp( - ( - Y.c.sgsʲs.:(1).q_tot - - specific(Y.c.ρq_tot, Y.c.ρ) - ) * + (Y.c.sgsʲs.:(1).q_tot - ᶜq_tot) * ᶜρʲs.:(1) * ᶜJ * draft_area(Y.c.sgsʲs.:(1).ρa, ᶜρʲs.:(1)), @@ -1067,10 +1087,7 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t) @. ∂ᶜρq_tot_err_∂ᶠu₃ʲ = dtγ * -(ᶜadvdivᵥ_matrix()) ⋅ DiagonalMatrixRow( ᶠinterp( - ( - Y.c.sgsʲs.:(1).q_tot - - specific(Y.c.ρq_tot, Y.c.ρ) - ) * + (Y.c.sgsʲs.:(1).q_tot - ᶜq_tot) * ᶜρʲs.:(1) * ᶜJ * draft_area(Y.c.sgsʲs.:(1).ρa, ᶜρʲs.:(1)), @@ -1090,9 +1107,8 @@ function update_jacobian!(alg::ManualSparseJacobian, cache, Y, p, dtγ, t) matrix[@name(c.ρq_tot), @name(c.sgsʲs.:(1).ρa)] @. ∂ᶜρq_tot_err_∂ᶜρa = dtγ * -(ᶜadvdivᵥ_matrix()) ⋅ DiagonalMatrixRow( - (ᶠu³ʲs.:(1) - ᶠu³) * ᶠinterp(( - Y.c.sgsʲs.:(1).q_tot - specific(Y.c.ρq_tot, Y.c.ρ) - )) / ᶠJ, + (ᶠu³ʲs.:(1) - ᶠu³) * + ᶠinterp((Y.c.sgsʲs.:(1).q_tot - ᶜq_tot)) / ᶠJ, ) ⋅ ᶠinterp_matrix() ⋅ DiagonalMatrixRow(ᶜJ) end elseif rs isa RayleighSponge diff --git a/src/prognostic_equations/remaining_tendency.jl b/src/prognostic_equations/remaining_tendency.jl index ae94a52563..3fe38dcdd5 100644 --- a/src/prognostic_equations/remaining_tendency.jl +++ b/src/prognostic_equations/remaining_tendency.jl @@ -1,4 +1,3 @@ - """ hyperdiffusion_tendency!(Yₜ, Yₜ_lim, Y, p, t) @@ -137,7 +136,6 @@ dependencies or operator splitting assumptions. """ NVTX.@annotate function additional_tendency!(Yₜ, Y, p, t) - (; ᶜh_tot, ᶜspecific) = p.precomputed ᶜuₕ = Y.c.uₕ ᶠu₃ = Y.f.u₃ ᶜρ = Y.c.ρ @@ -148,6 +146,13 @@ NVTX.@annotate function additional_tendency!(Yₜ, Y, p, t) thermo_params = CAP.thermodynamics_params(params) (; ᶜp, sfc_conditions, ᶜts) = p.precomputed + ᶜh_tot = @. lazy( + TD.total_specific_enthalpy( + thermo_params, + ᶜts, + specific(Y.c.ρe_tot, Y.c.ρ), + ), + ) vst_uₕ = viscous_sponge_tendency_uₕ(ᶜuₕ, viscous_sponge) vst_u₃ = viscous_sponge_tendency_u₃(ᶠu₃, viscous_sponge) vst_ρe_tot = viscous_sponge_tendency_ρe_tot(ᶜρ, ᶜh_tot, viscous_sponge) @@ -168,7 +173,6 @@ NVTX.@annotate function additional_tendency!(Yₜ, Y, p, t) @. Yₜ.f.u₃.components.data.:1 += vst_u₃ @. Yₜ.c.ρe_tot += vst_ρe_tot - # TODO: can we write this out explicitly? foreach_gs_tracer(Yₜ, Y) do ᶜρχₜ, ᶜρχ, ρχ_name ᶜχ = @. lazy(specific(ᶜρχ, Y.c.ρ)) vst_tracer = viscous_sponge_tendency_tracer(ᶜρ, ᶜχ, viscous_sponge) diff --git a/src/prognostic_equations/surface_flux.jl b/src/prognostic_equations/surface_flux.jl index 10b6b1ef90..e76e083040 100644 --- a/src/prognostic_equations/surface_flux.jl +++ b/src/prognostic_equations/surface_flux.jl @@ -104,7 +104,9 @@ function surface_flux_tendency!(Yₜ, Y, p, t) p.atmos.disable_surface_flux_tendency && return FT = eltype(Y) - (; ᶜh_tot, ᶜspecific, sfc_conditions) = p.precomputed + (; params) = p + (; sfc_conditions, ᶜts) = p.precomputed + thermo_params = CAP.thermodynamics_params(params) if !disable_momentum_vertical_diffusion(p.atmos.vertical_diffusion) btt = @@ -112,6 +114,13 @@ function surface_flux_tendency!(Yₜ, Y, p, t) @. Yₜ.c.uₕ -= btt end + ᶜh_tot = @. lazy( + TD.total_specific_enthalpy( + thermo_params, + ᶜts, + specific(Y.c.ρe_tot, Y.c.ρ), + ), + ) btt = boundary_tendency_scalar(ᶜh_tot, sfc_conditions.ρ_flux_h_tot) @. Yₜ.c.ρe_tot -= btt ρ_flux_χ = p.scratch.sfc_temp_C3 diff --git a/src/prognostic_equations/vertical_diffusion_boundary_layer.jl b/src/prognostic_equations/vertical_diffusion_boundary_layer.jl index d25a889824..1cf106919e 100644 --- a/src/prognostic_equations/vertical_diffusion_boundary_layer.jl +++ b/src/prognostic_equations/vertical_diffusion_boundary_layer.jl @@ -46,14 +46,20 @@ This function is dispatched based on the type of the vertical diffusion model for scalars. Zero-flux boundary conditions are explicitly applied. - **Note on mass conservation for `q_tot` diffusion**: The current implementation also modifies the tendency of total moist air density `Yₜ.c.ρ` based on the - diffusion tendency of total specific humidity `ρq_tot`: + diffusion tendency of total specific humidity `ρq_tot`: `Yₜ.c.ρ -= ᶜρχₜ_diffusion_for_q_tot`. -Arguments for all methods: -- `Yₜ`: The tendency state vector, modified in place. +This function is acting as a wrapper around the specific implementations +for different turbulence and convection models. + +The primary role of this function is to dispatch to the correct turbulence model's +tendency function. It operates on the state `Y` and its tendency `Yₜ`, using +the model-specific cache `p`. + +Arguments: +- `Yₜ`: The tendency state vector. - `Y`: The current state vector. - `p`: Cache containing parameters (e.g., `p.params` for `CAP.α_vert_diff_tracer`), - precomputed fields (e.g., `ᶜK_u`, `ᶜK_h`, `ᶜh_tot`, `ᶜspecific` tracer values), atmospheric model configurations (like `p.atmos.vertical_diffusion`), and scratch space. - `t`: Current simulation time (not directly used in diffusion calculations). - `vert_diff_model` (for dispatched methods): The specific vertical diffusion model instance. @@ -83,7 +89,8 @@ function vertical_diffusion_boundary_layer_tendency!( FT = eltype(Y) (; vertical_diffusion) = p.atmos α_vert_diff_tracer = CAP.α_vert_diff_tracer(p.params) - (; ᶜu, ᶜh_tot, ᶜspecific, ᶜp) = p.precomputed + thermo_params = CAP.thermodynamics_params(p.params) + (; ᶜu, ᶜp, ᶜts) = p.precomputed ᶠgradᵥ = Operators.GradientC2F() # apply BCs to ᶜdivᵥ, which wraps ᶠgradᵥ ᶜK_h = p.scratch.ᶜtemp_scalar if vertical_diffusion isa DecayWithHeightDiffusion @@ -108,6 +115,13 @@ function vertical_diffusion_boundary_layer_tendency!( top = Operators.SetValue(C3(0)), bottom = Operators.SetValue(C3(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(-(ᶠinterp(Y.c.ρ) * ᶠinterp(ᶜK_h) * ᶠgradᵥ(ᶜh_tot))) diff --git a/src/solver/types.jl b/src/solver/types.jl index eb38db31a5..07991c67df 100644 --- a/src/solver/types.jl +++ b/src/solver/types.jl @@ -363,13 +363,13 @@ This allows running simulations with TKE-based vertical diffusion. struct EDOnlyEDMFX <: AbstractEDMF end struct PrognosticEDMFX{N, TKE, FT} <: AbstractEDMF - a_half::FT # WARNING: this should never be used outside of divide_by_ρa + a_half::FT # WARNING: this should never be used outside of `specific` end PrognosticEDMFX{N, TKE}(a_half::FT) where {N, TKE, FT} = PrognosticEDMFX{N, TKE, FT}(a_half) struct DiagnosticEDMFX{N, TKE, FT} <: AbstractEDMF - a_half::FT # WARNING: this should never be used outside of divide_by_ρa + a_half::FT # WARNING: this should never be used outside of `specific` end DiagnosticEDMFX{N, TKE}(a_half::FT) where {N, TKE, FT} = DiagnosticEDMFX{N, TKE, FT}(a_half) diff --git a/src/utils/variable_manipulations.jl b/src/utils/variable_manipulations.jl index 443c5eb359..04479f33c1 100644 --- a/src/utils/variable_manipulations.jl +++ b/src/utils/variable_manipulations.jl @@ -1,4 +1,5 @@ import ClimaCore.MatrixFields: @name +import ClimaCore.RecursiveApply: ⊞, ⊠, rzero, rpromote_type """ specific(ρχ, ρ) @@ -42,9 +43,9 @@ Arguments: - `ρ_fallback`: The grid-mean density used for the fallback value. - `turbconv_model`: The turbulence convection model, containing parameters for regularization (e.g., `a_half`). """ -@inline specific(ρχ, ρ) = ρχ / ρ +specific(ρχ, ρ) = ρχ / ρ -@inline function specific(ρaχ, ρa, ρχ, ρ, turbconv_model) +function specific(ρaχ, ρa, ρχ, ρ, turbconv_model) # TODO: Replace turbconv_model struct by parameters, and include a_half in # parameters, not in config weight = sgs_weight_function(ρa / ρ, turbconv_model.a_half) @@ -177,199 +178,314 @@ foreach_gs_tracer(f::F, Y_or_similar_values...) where {F} = f(ρχ_or_χ_fields..., ρχ_name) end + """ sgs_weight_function(a, a_half) -Computes the weight of the SGS variables in the linear interpolation used in -`divide_by_ρa`. This is a continuously differentiable and monotonically -increasing function of `a` that is equal to 0 when `a ≤ 0`, is equal to 1 when -`a ≥ 1`, is equal to `1 / 2` when `a = a_half`, grows very rapidly near -`a = a_half`, and grows very slowly at all other values of `a`. If `a_half` is -sufficiently small, this function is essentially equal to 1 for all `a` more -than a few times larger than `a_half` (up to floating-point precision). - -We will now provide a description of how this function was constructed. We need -the function to be equal to 0 when `a ≤ 0` and equal to 1 when `a ≥ 1`. Since -the function must also be continuously differentiable, its derivative at these -values of `a` has to be 0. To obtain a function with these properties, we use a -piecewise definition: - - For all `a < 0`, the function is equal to 0. - - For all `a > 1`, the function is equal to 1. - - For all `0 ≤ a ≤ 1`, the function is a sigmoid that connects the point - `(0, 0)` to the point `(1, 1)`, with a derivative of 0 at these points. -Most well-known sigmoid functions connect the "points" `(-Inf, 0)` and -`(Inf, 1)`, not `(0, 0)` and `(1, 1)`. To obtain the desired sigmoid curve, we -begin with two simple sigmoid functions that go from `(-Inf, 0)` to `(Inf, 1)` -at different rates. In this case, we use two `tanh` functions, scaled and -translated so that they lie between 0 and 1: - - `fast_sigmoid(a) = (1 + tanh(a)) / 2` and - - `slow_sigmoid(a) = (1 + tanh(a / 2)) / 2`. -Note that the second sigmoid is commonly called the "logistic" function. We then -take the inverse of the sigmoid that grows more slowly, and we make that the -input of the sigmoid that grows more quickly: - - `sigmoid(a) = fast_sigmoid(slow_sigmoid⁻¹(a)) = - (1 + tanh(2 * atanh(2 * a - 1))) / 2`. -The resulting function goes from `(0, 0)` to `(1, 1)`, and, since the outer -sigmoid grows more quickly, it has the same asymptotic behavior as the outer -sigmoid, which means that its derivative at the boundary points is 0. If we had -instead put the sigmoid that grows more slowly on the outside, the asymptotic -behavior would come from the inverted inner sigmoid, which means that the -derivative at the boundary points would be `Inf`. - -The sigmoid function we have constructed reaches `1 / 2` when `a = 1 / 2`. More -generally, we need the weight function to reach `1 / 2` when `a` is some small -value `a_half`. To achieve this, we replace the input to the sigmoid function -with a smooth, monotonically increasing function that goes through `(0, 0)`, -`(a_half, 1 / 2)`, and `(1, 1)`. The simplest option is the power function - - `power(a) = a^(-1 / log2(a_half))`. -However, this function does not work well because, when `a_half < 1 / 2`, its -derivative at `a = 0` is `Inf`, and its derivative at `a = 1` is some positive -number. Making this power function the input to the sigmoid function causes the -derivative of the sigmoid to become `Inf` at `a = 0` when `a_half < 1 / 4`, and -it causes the sigmoid to grow too slowly from `1 / 2` to 1, only reaching 1 when -`a` is significantly larger than `a_half`. In order to fix this, we transform -the power function by replacing `a` with `1 - a`, `a_half` with `1 - a_half`, -and `power(a)` with `1 - power(a)`, which gives us - - `power(a) = 1 - (1 - a)^(-1 / log2(1 - a_half))`. -This transformed function works better because, when `a_half < 1 / 2`, its -derivative at `a = 0` is some positive number, and its derivative at `a = 1` is -0. When we make this the input to the sigmoid function, the result has a -continuous derivative and is essentially equal to 1 for all `a` more than a few -times larger than `a_half`. So, for all `0 ≤ a ≤ 1`, we define the weight -function as - - `weight(a) = sigmoid(power(a)) = - (1 + tanh(2 * atanh(1 - 2 * (1 - a)^(-1 / log2(1 - a_half))))) / 2`. -""" -sgs_weight_function(a, a_half) = - if a <= 0 # autodiff generates NaNs when a is 0 +Computes a smooth, monotonic weight function `w(a)` that ranges from 0 to 1. + +This function is used as the interpolation weight in the regularized `specific` +function. It ensures a numerically stable and smooth transition between a subgrid-scale +(SGS) quantity and its grid-mean counterpart, especially when the SGS area fraction `a` +is small. + +**Key Properties:** +- `w(a) = 0` for `a ≤ 0`. +- `w(a) = 1` for `a ≥ 1`. +- `w(a_half) = 0.5`. +- The function is continuously differentiable, with derivatives equal to zero at + `a = 0` and `a = 1`, which ensures smooth blending. +- The functions grows very rapidly near `a = a_half`, and grows very slowly at all other + values of `a`. +- For small `a_half`, the weight rapidly approaches 1 for values of `a` that are + a few times larger than `a_half`. + +**Construction Method:** +The function is piecewise. For `a` between 0 and 1, it is a custom sigmoid curve +constructed in two main steps to satisfy the key properties: +1. **Bounded Sigmoid Creation**: A base sigmoid is created that maps the interval + `(0, 1)` to `(0, 1)` with zero derivatives at the endpoints. This is achieved + by composing a standard `tanh` function with the inverse of a slower-growing + `tanh` function. +2. **Midpoint Control**: To ensure the function passes through the control point + `(a_half, 0.5)`, the input `a` is first transformed by a specially designed + power function (`1 - (1 - a)^k`) before being passed to the bounded sigmoid. + This transformation maps `a_half` to `0.5` while preserving differentiability + at the boundaries. + +Arguments: +- `a`: The input SGS area fraction (often approximated as `ρa / ρ`). +- `a_half`: The value of `a` at which the weight function should be 0.5, controlling + the transition point of the sigmoid curve. + +Returns: +- The computed weight, a value between 0 and 1. +""" +function sgs_weight_function(a, a_half) + if a < 0 zero(a) elseif a > min(1, 42 * a_half) # autodiff generates NaNs when a is large one(a) else (1 + tanh(2 * atanh(1 - 2 * (1 - a)^(-1 / log2(1 - a_half))))) / 2 end +end """ - divide_by_ρa(ρaχ, ρa, ρχ, ρ, turbconv_model) + draft_sum(f, sgsʲs) -Computes `ρaχ / ρa`, regularizing the result to avoid issues when `a` is small. -This is done by performing a linear interpolation from `ρaχ / ρa` to `ρχ / ρ`, -using `sgs_weight_function(ρa / ρ, turbconv_model.a_half)` as the weight of -`ρaχ / ρa` in the interpolation. Note that `ρa / ρ` is the "anelastic -approximation" of `a`; we cannot directly use `a` to compute the weight because -this function needs to be called before `a` has been computed. Also, note that -using this function instead of directly computing `ρaχ / ρa` breaks the -assumption of domain decomposition when the approximated `a` is small. -""" -function divide_by_ρa(ρaχ, ρa, ρχ, ρ, turbconv_model) - weight = sgs_weight_function(ρa / ρ, turbconv_model.a_half) - # If ρa = 0, we know that ρa / ρ = 0, which means that weight = 0. However, - # 0 * ρaχ / 0 = NaN, regardless of what ρaχ is, so the linear interpolation - # will always return NaN when ρa = 0. To avoid this problem, we need to add - # a special case for ρa = 0. - return ρa == 0 ? ρχ / ρ : weight * ρaχ / ρa + (1 - weight) * ρχ / ρ -end +Computes the sum of a function `f` applied to each draft subdomain +state `sgsʲ` in the iterator `sgsʲs`. +Arguments: +- `f`: A function to apply to each element of `sgsʲs`. +- `sgsʲs`: An iterator over the draft subdomain states. """ - ρa⁺(gs) +draft_sum(f, sgsʲs) = unrolled_sum(f, sgsʲs) -Computes the total mass-flux subdomain area-weighted density, assuming that the -mass-flux subdomain states are stored in `gs.sgsʲs`. """ -ρa⁺(gs) = mapreduce_with_init(sgsʲ -> sgsʲ.ρa, +, gs.sgsʲs) + ᶜenv_value(grid_scale_value, f_draft, gs, turbconv_model) -""" - ρah_tot⁺(sgsʲs) +Computes the value of a quantity `ρaχ` in the environment subdomain by subtracting +the sum of its values in all draft subdomains from the grid-scale value. Available +for general variables in PrognosticEDMFX and environmental area (ᶜρa⁰) in DiagnosticEDMFX. -Computes the total mass-flux subdomain area-weighted ρh_tot, assuming that the -mass-flux subdomain states are stored in `sgsʲs`. -""" -ρah_tot⁺(sgsʲs) = mapreduce_with_init(sgsʲ -> sgsʲ.ρa * sgsʲ.h_tot, +, sgsʲs) +This is based on the domain decomposition principle for density-area weighted +quantities: `GridMean(ρχ) = Env(ρaχ) + Sum(Drafts(ρaχ))`. -""" - ρamse⁺(sgsʲs) +The function handles both PrognosticEDMFX and DiagnosticEDMFX models: +- For PrognosticEDMFX: Uses gs.sgsʲs to access draft subdomain states +- For DiagnosticEDMFX: Uses p.precomputed.ᶜρaʲs for draft area-weighted densities -Computes the total mass-flux subdomain area-weighted ρmse, assuming that the -mass-flux subdomain states are stored in `sgsʲs`. +Arguments: +- `grid_scale_value`: The `ρa`-weighted grid-scale value of the quantity. +- `f_draft`: A function that extracts the corresponding value from a draft subdomain state. +- `gs`: The grid-scale iteration object, which contains the draft subdomain states `gs.sgsʲs` (for PrognosticEDMFX) from the state `Y.c`, or `ᶜρaʲs` in the cache for DiagnosticEDMFX. +- `turbconv_model`: The turbulence convection model, used to determine how to access draft data. """ -ρamse⁺(sgsʲs) = mapreduce_with_init(sgsʲ -> sgsʲ.ρa * sgsʲ.mse, +, sgsʲs) +function ᶜenv_value(grid_scale_value, f_draft, gs) + return @. lazy(grid_scale_value - draft_sum(f_draft, gs)) +end -""" - ρaq_tot⁺(sgsʲs) -Computes the total mass-flux subdomain area-weighted ρq_tot, assuming that the -mass-flux subdomain states are stored in `sgsʲs`. -""" -ρaq_tot⁺(sgsʲs) = mapreduce_with_init(sgsʲ -> sgsʲ.ρa * sgsʲ.q_tot, +, sgsʲs) +function env_value(grid_scale_value, f_draft, gs) + return grid_scale_value - draft_sum(f_draft, gs) +end -""" - ρaq_liq⁺(sgsʲs) -Computes the liquid water mass-flux subdomain area-weighted ρq_liq, assuming that the -mass-flux subdomain states are stored in `sgsʲs`. -""" -ρaq_liq⁺(sgsʲs) = mapreduce_with_init(sgsʲ -> sgsʲ.ρa * sgsʲ.q_liq, +, sgsʲs) """ - ρaq_ice⁺(sgsʲs) + ᶜspecific_env_value(::Val{χ_name}, Y, p) -Computes the ice water mass-flux subdomain area-weighted ρq_ice, assuming that the -mass-flux subdomain states are stored in `sgsʲs`. -""" -ρaq_ice⁺(sgsʲs) = mapreduce_with_init(sgsʲ -> sgsʲ.ρa * sgsʲ.q_ice, +, sgsʲs) +Calculates the specific value of a quantity `χ` in the environment (`χ⁰`). -""" - ρaq_rai⁺(sgsʲs) +This function uses the domain decomposition principle to first find the +density-area-weighted environment value (`ρa⁰χ⁰`) and the environment +density-weighted environmental area (`ρa⁰`). It then computes the specific value using the +regularized `specific` function, which provides a stable result even when the +environment area fraction is very small. -Computes the rain mass-flux subdomain area-weighted ρq_rai, assuming that the -mass-flux subdomain states are stored in `sgsʲs`. -""" -ρaq_rai⁺(sgsʲs) = mapreduce_with_init(sgsʲ -> sgsʲ.ρa * sgsʲ.q_rai, +, sgsʲs) +Arguments: +- `::Val{χ_name}`: A `Val` type containing the symbol for the specific quantity `χ` (e.g., `Val(:h_tot)`, `Val(:q_tot)`). +- `Y`: The state, containing grid-mean and draft subdomain states. +- `p`: The cache, containing precomputed quantities and turbconv_model. + +Returns: +- The specific value of the quantity `χ` in the environment. +""" +function ᶜspecific_env_value(::Val{χ_name}, Y, p) where {χ_name} + turbconv_model = p.atmos.turbconv_model + + # Grid-scale density-weighted variable name, e.g., :ρq_tot + ρχ_name = Symbol(:ρ, χ_name) + + ᶜρχ = getproperty(Y.c, ρχ_name) + + # environment density-area-weighted mse (`ρa⁰χ⁰`). + # Numerator: ρa⁰χ⁰ = ρχ - (Σ ρaʲ * χʲ) + if turbconv_model isa PrognosticEDMFX + #Numerator: ρa⁰χ⁰ = ρχ - (Σ sgsʲ.ρa * sgsʲ.χ) + + ᶜρaχ⁰ = ᶜenv_value( + ᶜρχ, + sgsʲ -> getfield(sgsʲ, :ρa) * getfield(sgsʲ, χ_name), + Y.c.sgsʲs, + ) + # Denominator: ρa⁰ = ρ - Σ ρaʲ + ᶜρa⁰ = @. lazy(ρa⁰(Y.c.ρ, Y.c.sgsʲs, turbconv_model)) + + elseif turbconv_model isa DiagnosticEDMFX || turbconv_model isa EDOnlyEDMFX + ᶜχʲs = getproperty(p.precomputed, Symbol(:ᶜ, χ_name, :ʲs)) + n = n_mass_flux_subdomains(turbconv_model) + + # Σ ρaʲ * χʲ + ᶜρaχʲs_sum = p.scratch.ᶜtemp_scalar_3 + @. ᶜρaχʲs_sum = 0 + for j in 1:n + ᶜρaʲ = p.precomputed.ᶜρaʲs.:($j) + ᶜχʲ = ᶜχʲs.:($j) + @. ᶜρaχʲs_sum += ᶜρaʲ * ᶜχʲ + end -""" - ρaq_sno⁺(sgsʲs) + ᶜρaχ⁰ = @. lazy(ᶜρχ - ᶜρaχʲs_sum) + # Denominator: ρa⁰ = ρ - Σ ρaʲ, assume ᶜρa⁰ = ρ + ᶜρa⁰ = Y.c.ρ + end -Computes the snow mass-flux subdomain area-weighted ρq_sno, assuming that the -mass-flux subdomain states are stored in `sgsʲs`. -""" -ρaq_sno⁺(sgsʲs) = mapreduce_with_init(sgsʲ -> sgsʲ.ρa * sgsʲ.q_sno, +, sgsʲs) + return @. lazy(specific( + ᶜρaχ⁰, # ρaχ for environment + ᶜρa⁰, # ρa for environment + ᶜρχ, # Fallback ρχ is the grid-mean value + Y.c.ρ, # Fallback ρ is the grid-mean value + turbconv_model, + )) +end """ - ρa⁰(gs) + ρa⁰(ρ, sgsʲs, turbconv_model) + +Computes the environment area-weighted density (`ρa⁰`). + +This function calculates the environment area-weighted density by subtracting the sum of all draft subdomain area-weighted densities (`ρaʲ`) from the grid-mean density (`ρ`), following the domain decomposition principle (`GridMean = Environment + Sum(Drafts)`). + +Arguments: +- `ρ`: Grid-mean density. +- `sgsʲs`: Iterable of draft subdomain quantities. + - For `PrognosticEDMFX`: typically `Y.c.sgsʲs` + - For `DiagnosticEDMFX`: typically `p.precomputed.ᶜρaʲs` +- `turbconv_model`: The turbulence convection model (e.g., `PrognosticEDMFX`, `DiagnosticEDMFX`, or others). -Computes the environment area-weighted density, assuming that the mass-flux -subdomain states are stored in `gs.sgsʲs`. +Returns: +- The area-weighted density of the environment (`ρa⁰`). """ -ρa⁰(gs) = gs.ρ - mapreduce_with_init(sgsʲ -> sgsʲ.ρa, +, gs.sgsʲs) + +function ρa⁰(ρ, sgsʲs, turbconv_model) + # ρ - Σ ρaʲ + if turbconv_model isa PrognosticEDMFX + return env_value(ρ, sgsʲ -> sgsʲ.ρa, sgsʲs) + + elseif turbconv_model isa DiagnosticEDMFX + return env_value(ρ, ᶜρaʲ -> ᶜρaʲ, sgsʲs) + else + return ρ + end +end + """ - u₃⁺(ρaʲs, u₃ʲs, ρ, u₃, turbconv_model) + specific_tke(ρ, ρatke, ρa⁰, turbconv_model) + +Computes the specific turbulent kinetic energy (TKE) in the environment. + +This function returns the specific TKE of the environment by regularizing the +area-weighted TKE density (`ρatke`) with the environment area-weighted density +(`ρa⁰`). In the limit of vanishing environment area fraction, the fallback value +for TKE is zero. + +Arguments: +- `ρ`: Grid-mean density. +- `ρatke`: Area-weighted TKE density in the environment. +- `ρa⁰`: Area-weighted density of the environment. +- `turbconv_model`: The turbulence convection model (e.g., `PrognosticEDMFX`, `DiagnosticEDMFX`, or others). + +Returns: +- The specific TKE of the environment (`tke⁰`). +""" +function specific_tke(ρ, ρatke, ρa⁰, turbconv_model) + + if turbconv_model isa PrognosticEDMFX || turbconv_model isa DiagnosticEDMFX + return specific( + ρatke, # ρaχ for environment TKE + ρa⁰, # ρa for environment, now computed internally + 0, # Fallback ρχ is zero for TKE + ρ, # Fallback ρ + turbconv_model, + ) + else + return specific(ρatke, ρa⁰) + end +end + -Computes the average mass-flux subdomain vertical velocity `u₃⁺` by dividing the -total momentum `ρaw⁺` by the total area-weighted density `ρa⁺`, both of which -are computed from the tuples of subdomain densities and velocities `ρaʲs` and -`u₃ʲs`. The division is computed using `divide_by_ρa` to avoid issues when `a⁺` -is small. """ -u₃⁺(ρaʲs, u₃ʲs, ρ, u₃, turbconv_model) = divide_by_ρa( - unrolled_dotproduct(ρaʲs, u₃ʲs), - reduce(+, ρaʲs), - ρ * u₃, - ρ, - turbconv_model, -) + ᶜspecific_env_mse(Y, p) + +Computes the specific moist static energy (`mse`) in the environment (`mse⁰`). + +This is a specialized helper function because `mse` is not a grid-scale prognostic +variable. It first computes the grid-scale moist static energy density (`ρmse`) +from other grid-scale quantities (`ρ`, total specific enthalpy `h_tot`, specific +kinetic energy `K`). It then uses the `ᶜenv_value` helper to compute the environment's +portion of `ρmse` and `ρa` via domain decomposition, and finally calculates the specific +value using the regularized `specific` function. + +Arguments: +- `Y`: The state containing `Y.c.ρ` and `Y.c.sgsʲs` (for PrognosticEDMFX). +- `p`: The cache, containing the turbconv_model and precomputed quantities. + +Returns: +- A `ClimaCore.Fields.Field` containing the specific moist static energy of the + environment (`mse⁰`). +""" +function ᶜspecific_env_mse(Y, p) + turbconv_model = p.atmos.turbconv_model + (; ᶜK, ᶜts) = p.precomputed + thermo_params = CAP.thermodynamics_params(p.params) + ᶜh_tot = @. lazy( + TD.total_specific_enthalpy( + thermo_params, + ᶜts, + specific(Y.c.ρe_tot, Y.c.ρ), + ), + ) + + # grid-scale moist static energy density `ρ * mse`. + ᶜρmse = @. lazy(Y.c.ρ * (ᶜh_tot - ᶜK)) + + # environment density-area-weighted mse (`ρa⁰mse⁰`). + # Numerator: ρa⁰mse⁰ = ρmse - (Σ ρaʲ * mseʲ) + + if turbconv_model isa PrognosticEDMFX + ρa⁰mse⁰ = ᶜenv_value(ᶜρmse, sgsʲ -> sgsʲ.ρa * sgsʲ.mse, Y.c.sgsʲs) + ᶜρa⁰ = @. lazy(ρa⁰(Y.c.ρ, Y.c.sgsʲs, turbconv_model)) + elseif turbconv_model isa DiagnosticEDMFX || turbconv_model isa EDOnlyEDMFX + + n = n_mass_flux_subdomains(turbconv_model) + (; ᶜρaʲs) = p.precomputed + ᶜρamseʲ_sum = p.scratch.ᶜtemp_scalar_2 + @. ᶜρamseʲ_sum = 0 + for j in 1:n + ᶜρaʲ = ᶜρaʲs.:($j) + ᶜmseʲ = p.precomputed.ᶜmseʲs.:($j) + @. ᶜρamseʲ_sum += ᶜρaʲ * ᶜmseʲ + end + ρa⁰mse⁰ = @. lazy(ᶜρmse - ᶜρamseʲ_sum) + # Denominator: ρa⁰ = ρ - Σ ρaʲ, assume ᶜρa⁰ = ρ + ᶜρa⁰ = Y.c.ρ + end + + return @. lazy(specific(ρa⁰mse⁰, ᶜρa⁰, ᶜρmse, Y.c.ρ, turbconv_model)) +end """ u₃⁰(ρaʲs, u₃ʲs, ρ, u₃, turbconv_model) -Computes the environment vertical velocity `u₃⁰` by dividing the environment -momentum `ρaw⁰` by the environment area-weighted density `ρa⁰`, both of which -are computed from the domain decomposition of the grid-scale quantities `ρw` and -`ρ` into the mass-flux subdomain quantities `ρawʲs` and `ρaʲs` and the -environment quantities. The division is computed using `divide_by_ρa` to avoid -issues when `a⁰` is small. +Computes the environment vertical velocity `u₃⁰`. + +This function calculates the environment's total vertical momentum (`ρa⁰u₃⁰`) and +its total area-weighted density (`ρa⁰`) using the domain decomposition principle +(GridMean = Env + Sum(Drafts)). It then computes the final specific velocity `u₃⁰` +using the regularized `specific` function to ensure numerical stability when the +environment area fraction `a⁰` is small. + +Arguments: +- `ρaʲs`: A tuple of area-weighted densities for each draft subdomain. +- `u₃ʲs`: A tuple of vertical velocities for each draft subdomain. +- `ρ`: The grid-mean air density. +- `u₃`: The grid-mean vertical velocity. +- `turbconv_model`: The turbulence convection model, containing regularization parameters. """ -u₃⁰(ρaʲs, u₃ʲs, ρ, u₃, turbconv_model) = divide_by_ρa( +u₃⁰(ρaʲs, u₃ʲs, ρ, u₃, turbconv_model) = specific( ρ * u₃ - unrolled_dotproduct(ρaʲs, u₃ʲs), ρ - reduce(+, ρaʲs), ρ * u₃, @@ -377,14 +493,50 @@ u₃⁰(ρaʲs, u₃ʲs, ρ, u₃, turbconv_model) = divide_by_ρa( turbconv_model, ) -import ClimaCore.RecursiveApply: ⊞, ⊠, rzero, rpromote_type +""" + mapreduce_with_init(f, op, iter...) + +A wrapper for Julia's `mapreduce` function that automatically determines +the initial value (`init`) for the reduction. + +This is useful for iterators whose elements are custom structs or +`ClimaCore.Geometry.AxisTensor`s, where the zero element cannot be inferred +as a simple scalar. It uses `ClimaCore.RecursiveApply` tools (`rzero`, +`rpromote_type`) to create a type-stable, correctly-structured zero element +based on the output of the function `f` applied to the first elements of the +iterators. + +Arguments: +- `f`: The function to apply to each element. +- `op`: The reduction operator (e.g., `+`, `*`). +- `iter...`: One or more iterators. +""" function mapreduce_with_init(f, op, iter...) r₀ = rzero(rpromote_type(typeof(f(map(first, iter)...)))) mapreduce(f, op, iter...; init = r₀) end -# Inference fails for certain mapreduce calls inside cuda -# kernels, so let's define a recursive unrolled dot product: +""" + unrolled_dotproduct(a::Tuple, b::Tuple) + +Computes the dot product of two `Tuple`s (`a` and `b`) using a recursive, +manually unrolled implementation. + +This function is designed to be type-stable and efficient for CUDA kernels, +where standard `mapreduce` implementations can otherwise suffer from type-inference +failures. + +It uses `ClimaCore.RecursiveApply` operators (`⊞` for addition, `⊠` for +multiplication), which allows it to handle dot products of tuples containing +complex, nested types such as `ClimaCore.Geometry.AxisTensor`s. + +Arguments: +- `a`: The first `Tuple`. +- `b`: The second `Tuple`, which must have the same length as `a`. + +Returns: +- The result of the dot product, `Σᵢ a[i] * b[i]`. +""" promote_type_mul(n::Number, x::Geometry.AxisTensor) = typeof(x) promote_type_mul(x::Geometry.AxisTensor, n::Number) = typeof(x) @inline function unrolled_dotproduct(a::Tuple, b::Tuple)