Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Use LazyBroadcast for subsidence tendency #3579

Draft
wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

Use LazyBroadcast for subsidence tendency #3579

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
76 changes: 26 additions & 50 deletions src/prognostic_equations/forcing/subsidence.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,59 +7,35 @@ import ClimaCore.Spaces as Spaces
import ClimaCore.Fields as Fields
import ClimaCore.Operators as Operators

#####
##### No subsidence
#####

subsidence_tendency!(Yₜ, Y, p, t, ::Nothing) = nothing

#####
##### Subsidence
#####

subsidence!(ᶜρχₜ, ᶜρ, ᶠu³, ᶜχ, ::Val{:none}) =
@. ᶜρχₜ -= ᶜρ * (ᶜsubdivᵥ(ᶠu³ * ᶠinterp(ᶜχ)) - ᶜχ * ᶜsubdivᵥ(ᶠu³))
subsidence!(ᶜρχₜ, ᶜρ, ᶠu³, ᶜχ, ::Val{:first_order}) =
@. ᶜρχₜ -= ᶜρ * (ᶜsubdivᵥ(ᶠupwind1(ᶠu³, ᶜχ)) - ᶜχ * ᶜsubdivᵥ(ᶠu³))
subsidence!(ᶜρχₜ, ᶜρ, ᶠu³, ᶜχ, ::Val{:third_order}) =
@. ᶜρχₜ -= ᶜρ * (ᶜsubdivᵥ(ᶠupwind3(ᶠu³, ᶜχ)) - ᶜχ * ᶜsubdivᵥ(ᶠu³))

function subsidence_tendency!(Yₜ, Y, p, t, ::Subsidence)
(; moisture_model) = p.atmos
subsidence_profile = p.atmos.subsidence.prof
(; ᶜh_tot, ᶜspecific) = p.precomputed

ᶠz = Fields.coordinate_field(axes(Y.f)).z
ᶠlg = Fields.local_geometry_field(Y.f)
ᶠsubsidence³ = p.scratch.ᶠtemp_CT3
@. ᶠsubsidence³ =
subsidence_profile(ᶠz) * CT3(unit_basis_vector_data(CT3, ᶠlg))

# LS Subsidence
subsidence!(Yₜ.c.ρe_tot, Y.c.ρ, ᶠsubsidence³, ᶜh_tot, Val{:first_order}())
subsidence!(
Yₜ.c.ρq_tot,
Y.c.ρ,
ᶠsubsidence³,
ᶜspecific.q_tot,
Val{:first_order}(),
)
if moisture_model isa NonEquilMoistModel
subsidence!(
Yₜ.c.ρq_liq,
Y.c.ρ,
ᶠsubsidence³,
ᶜspecific.q_liq,
Val{:first_order}(),
)
subsidence!(
Yₜ.c.ρq_ice,
Y.c.ρ,
ᶠsubsidence³,
ᶜspecific.q_ice,
Val{:first_order}(),
)
end
subsidence_tracer(::Val{:h_tot}, thermo_params, ts, ρ, ρe_tot) =
TD.total_specific_enthalpy(thermo_params, ts, ρe_tot / ρ)
subsidence_tracer(::Val{:q_tot}, thermo_params, ts, ρ, ρe_tot) =
TD.total_specific_humidity(thermo_params, ts)
subsidence_tracer(::Val{:q_liq}, thermo_params, ts, ρ, ρe_tot) =
TD.liquid_specific_humidity(thermo_params, ts)
subsidence_tracer(::Val{:q_ice}, thermo_params, ts, ρ, ρe_tot) =
TD.ice_specific_humidity(thermo_params, ts)

function subsidence_tendency(χname, subsidence, thermo_params, ᶜts, ᶜρ, ᶜρe_tot)
subsidence isa Nothing && return NullBroadcasted()
ᶠu³ = ᶠsubsidence³(axes(ᶜρ), subsidence)
ᶜχ = @lazy @. subsidence_tracer(χname, thermo_params, ᶜts, ᶜρ, ᶜρe_tot)
return subsidence_tendency(ᶜρ, ᶜχ, ᶠu³, Val{:first_order}())
end

return nothing
function ᶠsubsidence³(ᶠspace, subsidence::Subsidence)
ᶠz = Fields.coordinate_field(ᶠspace).z
ᶠlg = Fields.local_geometry_field(ᶠspace)
(; prof) = subsidence
return @lazy @. prof(ᶠz) * CT3(unit_basis_vector_data(CT3, ᶠlg))
end
subsidence_tendency(ᶜρ, ᶜχ, ᶠu³, ::Val{:none}) =
@lazy @. - ᶜρ * (ᶜsubdivᵥ(ᶠu³ * ᶠinterp(ᶜχ)) - ᶜχ * ᶜsubdivᵥ(ᶠu³))
subsidence_tendency(ᶜρ, ᶜχ, ᶠu³, ::Val{:first_order}) =
@lazy @. - ᶜρ * (ᶜsubdivᵥ(ᶠupwind1(ᶠu³, ᶜχ)) - ᶜχ * ᶜsubdivᵥ(ᶠu³))
subsidence_tendency(ᶜρ, ᶜχ, ᶠu³, ::Val{:third_order}) =
@lazy @. - ᶜρ * (ᶜsubdivᵥ(ᶠupwind3(ᶠu³, ᶜχ)) - ᶜχ * ᶜsubdivᵥ(ᶠu³))
33 changes: 23 additions & 10 deletions src/prognostic_equations/remaining_tendency.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,10 +40,15 @@ NVTX.@annotate function additional_tendency!(Yₜ, Y, p, t)
ᶜuₕ = Y.c.uₕ
ᶠu₃ = Yₜ.f.u₃
ᶜρ = Y.c.ρ
ᶜρe_tot = Y.c.ρe_tot
(; forcing_type, moisture_model, rayleigh_sponge, viscous_sponge) = p.atmos
(; edmf_coriolis) = p.atmos
(; subsidence, edmf_coriolis) = p.atmos
(; params) = p
(; ᶜp, sfc_conditions) = p.precomputed
thermo_params = CAP.thermodynamics_params(params)
(; ᶜp, sfc_conditions, ᶜts) = p.precomputed
ᶠspace = axes(Y.f)
ᶠz = Fields.coordinate_field(ᶠspace).z
ᶠlg = Fields.local_geometry_field(ᶠspace)

vst_uₕ = viscous_sponge_tendency_uₕ(ᶜuₕ, viscous_sponge)
vst_u₃ = viscous_sponge_tendency_u₃(ᶠu₃, viscous_sponge)
Expand All @@ -52,6 +57,7 @@ NVTX.@annotate function additional_tendency!(Yₜ, Y, p, t)
hs_args = (ᶜuₕ, ᶜp, params, sfc_conditions.ts, moisture_model, forcing_type)
hs_tendency_uₕ = held_suarez_forcing_tendency_uₕ(hs_args...)
hs_tendency_ρe_tot = held_suarez_forcing_tendency_ρe_tot(ᶜρ, hs_args...)
subs_args = (subsidence, thermo_params, ᶜts, ᶜρ, ᶜρe_tot)
edmf_cor_tend_uₕ = edmf_coriolis_tendency_uₕ(ᶜuₕ, edmf_coriolis)

# TODO: fuse, once we fix
Expand All @@ -62,20 +68,27 @@ NVTX.@annotate function additional_tendency!(Yₜ, Y, p, t)
@. Yₜ.c.ρe_tot += vst_ρe_tot

# TODO: can we write this out explicitly?
if viscous_sponge isa ViscousSponge
for (ᶜρχₜ, ᶜχ, χ_name) in matching_subfields(Yₜ.c, ᶜspecific)
χ_name == :e_tot && continue
vst_tracer = viscous_sponge_tendency_tracer(ᶜρ, ᶜχ, viscous_sponge)
@. ᶜρχₜ += vst_tracer
@. Yₜ.c.ρ += vst_tracer
end
for (ᶜρχₜ, ᶜχ, χ_name) in matching_subfields(Yₜ.c, ᶜspecific)
χ_name == :e_tot && continue
vst_tracer = viscous_sponge_tendency_tracer(ᶜρ, ᶜχ, viscous_sponge)
@. ᶜρχₜ += vst_tracer
@. Yₜ.c.ρ += vst_tracer
end

# Held Suarez tendencies
@. Yₜ.c.uₕ += hs_tendency_uₕ
@. Yₜ.c.ρe_tot += hs_tendency_ρe_tot

subsidence_tendency!(Yₜ, Y, p, t, p.atmos.subsidence)
if moisture_model isa AbstractMoistModel
bc_subsidence_ρq_tot = subsidence_tendency(Val(:q_tot), subs_args...)
@. Yₜ.c.ρq_tot += bc_subsidence_ρq_tot
end
if moisture_model isa NonEquilMoistModel
bc_subsidence_ρq_liq = subsidence_tendency(Val(:q_liq), subs_args...)
bc_subsidence_ρq_ice = subsidence_tendency(Val(:q_ice), subs_args...)
@. Yₜ.c.ρq_liq += bc_subsidence_ρq_liq
@. Yₜ.c.ρq_ice += bc_subsidence_ρq_ice
end

@. Yₜ.c.uₕ += edmf_cor_tend_uₕ

Expand Down
5 changes: 3 additions & 2 deletions src/solver/types.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,9 +7,10 @@ import ClimaUtilities.ClimaArtifacts: @clima_artifact
import LazyArtifacts

abstract type AbstractMoistureModel end
abstract type AbstractMoistModel <: AbstractMoistureModel end
struct DryModel <: AbstractMoistureModel end
struct EquilMoistModel <: AbstractMoistureModel end
struct NonEquilMoistModel <: AbstractMoistureModel end
struct EquilMoistModel <: AbstractMoistModel end
struct NonEquilMoistModel <: AbstractMoistModel end

abstract type AbstractPrecipitationModel end
struct NoPrecipitation <: AbstractPrecipitationModel end
Expand Down
Loading