diff --git a/src/cache/precomputed_quantities.jl b/src/cache/precomputed_quantities.jl index 9835b817ad..fe74dd3d04 100644 --- a/src/cache/precomputed_quantities.jl +++ b/src/cache/precomputed_quantities.jl @@ -42,7 +42,7 @@ function implicit_precomputed_quantities(Y, atmos) TST = thermo_state_type(moisture_model, FT) n = n_mass_flux_subdomains(turbconv_model) gs_quantities = (; - ᶜspecific = specific_gs.(Y.c), + ᶜspecific = Base.materialize(ᶜspecific_gs_tracers(Y)), ᶜu = similar(Y.c, C123{FT}), ᶠu³ = similar(Y.f, CT3{FT}), ᶠu = similar(Y.f, CT123{FT}), @@ -461,7 +461,7 @@ NVTX.@annotate function set_implicit_precomputed_quantities!(Y, p, t) thermo_params = CAP.thermodynamics_params(p.params) thermo_args = (thermo_params, moisture_model, precip_model) - @. ᶜspecific = specific_gs(Y.c) + ᶜ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 diff --git a/src/prognostic_equations/hyperdiffusion.jl b/src/prognostic_equations/hyperdiffusion.jl index e6006808e9..662691b662 100644 --- a/src/prognostic_equations/hyperdiffusion.jl +++ b/src/prognostic_equations/hyperdiffusion.jl @@ -6,6 +6,22 @@ import ClimaCore.Geometry as Geometry import ClimaCore.Fields as Fields import ClimaCore.Spaces as Spaces +""" + ν₄(hyperdiff, Y) + +A `NamedTuple` of the hyperdiffusivity `ν₄_scalar` and the hyperviscosity +`ν₄_vorticity`. These quantities are assumed to scale with `h^3`, where `h` is +the mean nodal distance, following the empirical results of Lauritzen et al. +(2018, https://doi.org/10.1029/2017MS001257). When `h == 1`, these quantities +are equal to `hyperdiff.ν₄_scalar_coeff` and `hyperdiff.ν₄_vorticity_coeff`. +""" +function ν₄(hyperdiff, Y) + h = Spaces.node_horizontal_length_scale(Spaces.horizontal_space(axes(Y.c))) + ν₄_scalar = hyperdiff.ν₄_scalar_coeff * h^3 + ν₄_vorticity = hyperdiff.ν₄_vorticity_coeff * h^3 + return (; ν₄_scalar, ν₄_vorticity) +end + hyperdiffusion_cache(Y, atmos) = hyperdiffusion_cache( Y, atmos.hyperdiff, @@ -34,7 +50,7 @@ function hyperdiffusion_cache( gs_quantities = (; ᶜ∇²u = similar(Y.c, C123{FT}), ᶜ∇²specific_energy = similar(Y.c, FT), - ᶜ∇²specific_tracers = remove_energy_var.(specific_gs.(Y.c)), + ᶜ∇²specific_tracers = Base.materialize(ᶜspecific_gs_tracers(Y)), ) # Sub-grid scale quantities @@ -85,7 +101,7 @@ NVTX.@annotate function prep_hyperdiffusion_tendency!(Yₜ, Y, p, t) n = n_mass_flux_subdomains(turbconv_model) diffuse_tke = use_prognostic_tke(turbconv_model) - (; ᶜp, ᶜspecific) = p.precomputed + (; ᶜp) = p.precomputed (; ᶜh_ref) = p.core (; ᶜ∇²u, ᶜ∇²specific_energy) = p.hyperdiff if turbconv_model isa PrognosticEDMFX @@ -125,20 +141,14 @@ NVTX.@annotate function apply_hyperdiffusion_tendency!(Yₜ, Y, p, t) (; hyperdiff, turbconv_model) = p.atmos isnothing(hyperdiff) && return nothing - (; ν₄_vorticity_coeff, ν₄_scalar_coeff, divergence_damping_factor) = - hyperdiff - - h_space = Spaces.horizontal_space(axes(Y.c)) - h_length_scale = Spaces.node_horizontal_length_scale(h_space) # mean nodal distance - - ν₄_scalar = ν₄_scalar_coeff * h_length_scale^3 - ν₄_vorticity = ν₄_vorticity_coeff * h_length_scale^3 + (; divergence_damping_factor) = hyperdiff + (; ν₄_scalar, ν₄_vorticity) = ν₄(hyperdiff, Y) n = n_mass_flux_subdomains(turbconv_model) diffuse_tke = use_prognostic_tke(turbconv_model) ᶜJ = Fields.local_geometry_field(Y.c).J point_type = eltype(Fields.coordinate_field(Y.c)) - (; ᶜp, ᶜspecific) = p.precomputed + (; ᶜp) = p.precomputed (; ᶜ∇²u, ᶜ∇²specific_energy) = p.hyperdiff if turbconv_model isa PrognosticEDMFX (; ᶜρa⁰) = p.precomputed @@ -241,11 +251,12 @@ NVTX.@annotate function prep_tracer_hyperdiffusion_tendency!(Yₜ, Y, p, t) (; hyperdiff, turbconv_model) = p.atmos isnothing(hyperdiff) && return nothing - (; ᶜspecific) = p.precomputed (; ᶜ∇²specific_tracers) = p.hyperdiff - for χ_name in propertynames(ᶜ∇²specific_tracers) - @. ᶜ∇²specific_tracers.:($$χ_name) = wdivₕ(gradₕ(ᶜspecific.:($$χ_name))) + # TODO: Fix RecursiveApply bug in gradₕ to fuse this operation. + # ᶜ∇²specific_tracers .= wdivₕ.(gradₕ.(ᶜspecific_gs_tracers(Y))) + foreach_gs_tracer(Y, ᶜ∇²specific_tracers) do ᶜρχ, ᶜ∇²χ, _ + @. ᶜ∇²χ = wdivₕ(gradₕ(specific(ᶜρχ, Y.c.ρ))) end if turbconv_model isa PrognosticEDMFX @@ -276,31 +287,24 @@ NVTX.@annotate function apply_tracer_hyperdiffusion_tendency!(Yₜ, Y, p, t) (; hyperdiff, turbconv_model) = p.atmos isnothing(hyperdiff) && return nothing - α_hyperdiff_tracer = CAP.α_hyperdiff_tracer(p.params) - (; ν₄_scalar_coeff) = hyperdiff - h_space = Spaces.horizontal_space(axes(Y.c)) - h_length_scale = Spaces.node_horizontal_length_scale(h_space) # mean nodal distance - ν₄_scalar = ν₄_scalar_coeff * h_length_scale^3 - n = n_mass_flux_subdomains(turbconv_model) + # Rescale the hyperdiffusivity for precipitating species. + (; ν₄_scalar) = ν₄(hyperdiff, Y) + ν₄_scalar_for_precip = CAP.α_hyperdiff_tracer(p.params) * ν₄_scalar + n = n_mass_flux_subdomains(turbconv_model) (; ᶜ∇²specific_tracers) = p.hyperdiff # TODO: Since we are not applying the limiter to density (or area-weighted # density), the mass redistributed by hyperdiffusion will not be conserved # by the limiter. Is this a significant problem? - # TODO: Figure out why caching the duplicated tendencies in ᶜtemp_scalar - # triggers allocations. - for (ᶜρχₜ, ᶜ∇²χ, χ_name) in matching_subfields(Yₜ.c, ᶜ∇²specific_tracers) - ν₄_scalar = ifelse( - χ_name in (:q_rai, :q_sno, :n_rai), - α_hyperdiff_tracer * ν₄_scalar, - ν₄_scalar, - ) - @. ᶜρχₜ -= ν₄_scalar * wdivₕ(Y.c.ρ * gradₕ(ᶜ∇²χ)) - - # Exclude contributions from hyperdiffusion of condensate, - # precipitating species from mass tendency. - if χ_name == :q_tot + foreach_gs_tracer(Yₜ, ᶜ∇²specific_tracers) do ᶜρχₜ, ᶜ∇²χ, ρχ_name + ν₄_scalar_for_χ = + ρχ_name in (@name(ρq_rai), @name(ρq_sno), @name(ρn_rai)) ? + ν₄_scalar_for_precip : ν₄_scalar + @. ᶜρχₜ -= ν₄_scalar_for_χ * wdivₕ(Y.c.ρ * gradₕ(ᶜ∇²χ)) + + # Take into account the effect of total water diffusion on density. + if ρχ_name == @name(ρq_tot) @. Yₜ.c.ρ -= ν₄_scalar * wdivₕ(Y.c.ρ * gradₕ(ᶜ∇²χ)) end end @@ -323,13 +327,9 @@ NVTX.@annotate function apply_tracer_hyperdiffusion_tendency!(Yₜ, Y, p, t) @. Yₜ.c.sgsʲs.:($$j).q_ice -= ν₄_scalar * wdivₕ(gradₕ(ᶜ∇²q_iceʲs.:($$j))) @. Yₜ.c.sgsʲs.:($$j).q_rai -= - α_hyperdiff_tracer * - ν₄_scalar * - wdivₕ(gradₕ(ᶜ∇²q_raiʲs.:($$j))) + ν₄_scalar_for_precip * wdivₕ(gradₕ(ᶜ∇²q_raiʲs.:($$j))) @. Yₜ.c.sgsʲs.:($$j).q_sno -= - α_hyperdiff_tracer * - ν₄_scalar * - wdivₕ(gradₕ(ᶜ∇²q_snoʲs.:($$j))) + ν₄_scalar_for_precip * wdivₕ(gradₕ(ᶜ∇²q_snoʲs.:($$j))) end end end diff --git a/src/utils/variable_manipulations.jl b/src/utils/variable_manipulations.jl index 021c0c79dc..443c5eb359 100644 --- a/src/utils/variable_manipulations.jl +++ b/src/utils/variable_manipulations.jl @@ -53,64 +53,128 @@ Arguments: return ρa == 0 ? ρχ / ρ : weight * ρaχ / ρa + (1 - weight) * ρχ / ρ end +# Internal method that checks if its input is @name(ρχ) for some variable χ. +@generated is_ρ_weighted_name( + ::MatrixFields.FieldName{name_chain}, +) where {name_chain} = + length(name_chain) == 1 && startswith(string(name_chain[1]), "ρ") + +# Internal method that converts @name(ρχ) to @name(χ) for some variable χ. +@generated function specific_tracer_name( + ::MatrixFields.FieldName{ρχ_name_chain}, +) where {ρχ_name_chain} + χ_symbol = Symbol(string(ρχ_name_chain[1])[(ncodeunits("ρ") + 1):end]) + return :(@name($χ_symbol)) +end + """ - tracer_names(field) + gs_tracer_names(Y) -Filters and returns the names of the variables from a given state -vector component, excluding `ρ`, `ρe_tot`, and `uₕ` and SGS fields. +`Tuple` of `@name`s for the grid-scale tracers in the center field `Y.c` +(excluding `ρ`, `ρe_tot`, velocities, and SGS fields). +""" +gs_tracer_names(Y) = + unrolled_filter(MatrixFields.top_level_names(Y.c)) do name + is_ρ_weighted_name(name) && !(name in (@name(ρ), @name(ρe_tot))) + end -Arguments: +""" + specific_gs_tracer_names(Y) + +`Tuple` of the specific tracer names `@name(χ)` that correspond to the +density-weighted tracer names `@name(ρχ)` in `gs_tracer_names(Y)`. +""" +specific_gs_tracer_names(Y) = + unrolled_map(specific_tracer_name, gs_tracer_names(Y)) + +""" + ᶜempty(Y) + +Lazy center `Field` of empty `NamedTuple`s. +""" +ᶜempty(Y) = lazy.(Returns((;)).(Y.c)) -- `field`: A component of the state vector `Y.c`. +""" + ᶜgs_tracers(Y) -Returns: +Lazy center `Field` of `NamedTuple`s that contain the values of all grid-scale +tracers given by `gs_tracer_names(Y)`. +""" +function ᶜgs_tracers(Y) + isempty(gs_tracer_names(Y)) && return ᶜempty(Y) + ρχ_symbols = unrolled_map(MatrixFields.extract_first, gs_tracer_names(Y)) + ρχ_fields = unrolled_map(gs_tracer_names(Y)) do ρχ_name + MatrixFields.get_field(Y.c, ρχ_name) + end + return @. lazy(NamedTuple{ρχ_symbols}(tuple(ρχ_fields...))) +end -- A `Tuple` of `ClimaCore.MatrixFields.FieldName`s corresponding to the tracers. """ -tracer_names(field) = - unrolled_filter(MatrixFields.top_level_names(field)) do name - !( - name in - (@name(ρ), @name(ρe_tot), @name(uₕ), @name(sgs⁰), @name(sgsʲs)) - ) + ᶜspecific_gs_tracers(Y) + +Lazy center `Field` of `NamedTuple`s that contain the values of all specific +grid-scale tracers given by `specific_gs_tracer_names(Y)`. +""" +function ᶜspecific_gs_tracers(Y) + isempty(gs_tracer_names(Y)) && return ᶜempty(Y) + χ_symbols = + unrolled_map(MatrixFields.extract_first, specific_gs_tracer_names(Y)) + χ_fields = unrolled_map(gs_tracer_names(Y)) do ρχ_name + ρχ_field = MatrixFields.get_field(Y.c, ρχ_name) + @. lazy(specific(ρχ_field, Y.c.ρ)) end + return @. lazy(NamedTuple{χ_symbols}(tuple(χ_fields...))) +end """ - foreach_gs_tracer(f::F, Yₜ, Y) where {F} + foreach_gs_tracer(f, Y_or_similar_values...) + +Applies a function `f` to each grid-scale tracer in the state `Y` or any similar +value like the tendency `Yₜ`. This is used to implement performant loops over +all tracers given by `gs_tracer_names(Y)`. -Applies a given function `f` to each grid-scale scalar variable (except `ρ` and `ρe_tot`) -in the state `Y` and its corresponding tendency `Yₜ`. -This utility abstracts the process of iterating over all scalars. It uses -`tracer_names` to identify the relevant variables and `unrolled_foreach` to -ensure a performant loop. For each tracer, it calls the provided function `f` -with the tendency field, the state field, and a boolean flag indicating if -the current tracer is `ρq_tot` (to allow for special handling). +Although the first input value needs to be similar to `Y`, the remaining values +can also be center `Field`s similar to `Y.c`, and they can use specific tracers +given by `specific_gs_tracer_names(Y)` instead of density-weighted tracers. Arguments: -- `f`: A function to apply to each grid-scale scalar. It must have the signature `f - (ᶜρχₜ, ᶜρχ, ρχ_name)`, where `ᶜρχₜ` is the tendency field, `ᶜρχ` - is the state field, and `ρχ_name` is a `MatrixFields.@name` object. -- `Yₜ`: The tendency state vector. -- `Y`: The current state vector. +- `f`: The function applied to each grid-scale tracer, which must have the + signature `f(ρχ_or_χ_fields..., ρχ_name)`, where `ρχ_or_χ_fields` are + grid-scale tracer subfields (either density-weighted or specific) and + `ρχ_name` is the `MatrixFields.FieldName` of the tracer. +- `Y_or_similar_values`: The state `Y` or similar values like the tendency `Yₜ`. -# Example +# Examples ```julia foreach_gs_tracer(Yₜ, Y) do ᶜρχₜ, ᶜρχ, ρχ_name - # Apply some operation, e.g., a sponge layer - @. ᶜρχₜ += some_sponge_function(ᶜρχ) + ᶜρχₜ .+= tendency_of_ρχ(ᶜρχ) if ρχ_name == @name(ρq_tot) - # Perform an additional operation only for ρq_tot + ᶜρχₜ .+= additional_tendency_of_ρq_tot(ᶜρχ) + end +end +``` + +```julia +foreach_gs_tracer(Yₜ, Base.materialize(ᶜspecific_gs_tracers(Y))) do ᶜρχₜ, ᶜχ, ρχ_name + ᶜρχₜ .+= Y.c.ρ .* tendency_of_χ(ᶜχ) + if ρχ_name == @name(ρq_tot) + ᶜρχₜ .+= Y.c.ρ .* additional_tendency_of_q_tot(ᶜχ) end end ``` """ -foreach_gs_tracer(f::F, Yₜ, Y) where {F} = - unrolled_foreach(tracer_names(Y.c)) do scalar_name - ᶜρχₜ = MatrixFields.get_field(Yₜ.c, scalar_name) - ᶜρχ = MatrixFields.get_field(Y.c, scalar_name) - f(ᶜρχₜ, ᶜρχ, scalar_name) +foreach_gs_tracer(f::F, Y_or_similar_values...) where {F} = + unrolled_foreach(gs_tracer_names(Y_or_similar_values[1])) do ρχ_name + ρχ_or_χ_fields = unrolled_map(Y_or_similar_values) do value + field = value isa Fields.Field ? value : value.c + ρχ_or_χ_name = + MatrixFields.has_field(field, ρχ_name) ? ρχ_name : + specific_tracer_name(ρχ_name) + MatrixFields.get_field(field, ρχ_or_χ_name) + end + f(ρχ_or_χ_fields..., ρχ_name) end """ @@ -206,94 +270,6 @@ function divide_by_ρa(ρaχ, ρa, ρχ, ρ, turbconv_model) return ρa == 0 ? ρχ / ρ : weight * ρaχ / ρa + (1 - weight) * ρχ / ρ end -# Helper functions for manipulating symbols in the generated functions: -has_prefix(symbol, prefix_symbol) = - startswith(string(symbol), string(prefix_symbol)) -remove_prefix(symbol, prefix_symbol) = - Symbol(string(symbol)[(ncodeunits(string(prefix_symbol)) + 1):end]) -# Note that we need to use ncodeunits instead of length because prefix_symbol -# can contain non-ASCII characters like 'ρ'. - -""" - specific_gs(gs) - -Converts every variable of the form `ρχ` in the grid-scale state `gs` into the -specific variable `χ` by dividing it by `ρ`. All other variables in `gs` are -omitted from the result. -""" -@generated function specific_gs(gs) - gs_names = Base._nt_names(gs) - relevant_gs_names = - filter(name -> has_prefix(name, :ρ) && name != :ρ, gs_names) - specific_gs_names = map(name -> remove_prefix(name, :ρ), relevant_gs_names) - specific_gs_values = map(name -> :(gs.$name / gs.ρ), relevant_gs_names) - return :(NamedTuple{$specific_gs_names}(($(specific_gs_values...),))) -end - -""" - specific_sgs(sgs, gs, turbconv_model) - -Converts every variable of the form `ρaχ` in the sub-grid-scale state `sgs` into -the specific variable `χ` by dividing it by `ρa`. All other variables in `sgs` -are omitted from the result. The division is computed as -`divide_by_ρa(ρaχ, ρa, ρχ, ρ, turbconv_model)`, which is preferable to simply -calling `ρaχ / ρa` because it avoids numerical issues that arise when `a` is -small. The values of `ρ` and `ρχ` are taken from `gs`, but, when `ρχ` is not -available in `gs` (e.g., when `χ` is a second moment variable like `tke`), its -value is assumed to be equal to the value of `ρaχ` in `sgs`. -""" -@generated function specific_sgs(sgs, gs, turbconv_model) - sgs_names = Base._nt_names(sgs) - gs_names = Base._nt_names(gs) - relevant_sgs_names = - filter(name -> has_prefix(name, :ρa) && name != :ρa, sgs_names) - specific_sgs_names = - map(name -> remove_prefix(name, :ρa), relevant_sgs_names) - relevant_gs_names = map(name -> Symbol(:ρ, name), specific_sgs_names) - specific_sgs_values = map( - (sgs_name, gs_name) -> :(divide_by_ρa( - sgs.$sgs_name, - sgs.ρa, - $(gs_name in gs_names ? :(gs.$gs_name) : :(sgs.$sgs_name)), - gs.ρ, - turbconv_model, - )), - relevant_sgs_names, - relevant_gs_names, - ) - return :(NamedTuple{$specific_sgs_names}(($(specific_sgs_values...),))) -end - -""" - matching_subfields(tendency_field, specific_field) - -Given a field that contains the tendencies of variables of the form `ρχ` or -`ρaχ` and another field that contains the values of specific variables `χ`, -returns all tuples `(tendency_field.<ρχ or ρaχ>, specific_field.<χ>, :<χ>)`. -Variables in `tendency_field` that do not have matching variables in -`specific_field` are omitted, as are variables in `specific_field` that do not -have matching variables in `tendency_field`. This function is needed to avoid -allocations due to failures in type inference, which are triggered when the -`propertynames` of these fields are manipulated during runtime in order to pick -out the matching subfields (as of Julia 1.8). -""" -@generated function matching_subfields(tendency_field, specific_field) - tendency_names = Base._nt_names(eltype(tendency_field)) - specific_names = Base._nt_names(eltype(specific_field)) - prefix = :ρa in tendency_names ? :ρa : :ρ - relevant_specific_names = - filter(name -> Symbol(prefix, name) in tendency_names, specific_names) - subfield_tuples = map( - name -> :(( - tendency_field.$(Symbol(prefix, name)), - specific_field.$name, - $(QuoteNode(name)), - )), - relevant_specific_names, - ) - return :(($(subfield_tuples...),)) -end - """ ρa⁺(gs) @@ -401,19 +377,6 @@ u₃⁰(ρaʲs, u₃ʲs, ρ, u₃, turbconv_model) = divide_by_ρa( turbconv_model, ) -""" - remove_energy_var(specific_state) - -Creates a copy of `specific_state` with the energy variable -removed, where `specific_state` is the result of calling, e.g., `specific_gs`, -`specific_sgsʲs`, or `specific_sgs⁰`. -""" -remove_energy_var(specific_state::NamedTuple) = - Base.structdiff(specific_state, NamedTuple{(:e_tot,)}) -remove_energy_var(specific_state::Tuple) = - map(remove_energy_var, specific_state) - - import ClimaCore.RecursiveApply: ⊞, ⊠, rzero, rpromote_type function mapreduce_with_init(f, op, iter...) r₀ = rzero(rpromote_type(typeof(f(map(first, iter)...))))