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

Make tracer vert adv always explicit #2580

Merged
merged 1 commit into from
Feb 3, 2024
Merged

Make tracer vert adv always explicit #2580

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
32 changes: 20 additions & 12 deletions src/prognostic_equations/advection.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -67,14 +67,14 @@ NVTX.@annotate function horizontal_tracer_advection_tendency!(Yₜ, Y, p, t)
end

NVTX.@annotate function explicit_vertical_advection_tendency!(Yₜ, Y, p, t)
(; turbconv_model) = p.atmos
(; moisture_model, turbconv_model) = p.atmos
n = n_prognostic_mass_flux_subdomains(turbconv_model)
advect_tke = use_prognostic_tke(turbconv_model)
point_type = eltype(Fields.coordinate_field(Y.c))
(; dt) = p
ᶜJ = Fields.local_geometry_field(Y.c).J
(; ᶜf) = p.core
(; ᶜu, ᶠu³, ᶜK) = p.precomputed
(; ᶜh_tot, ᶜu, ᶠu³, ᶜK) = p.precomputed
(; edmfx_upwinding) = n > 0 || advect_tke ? p.atmos.numerics : all_nothing
(; ᶜuʲs, ᶜKʲs) = n > 0 ? p.precomputed : all_nothing
(; ᶠu³⁰) = advect_tke ? p.precomputed : all_nothing
Expand Down Expand Up @@ -108,10 +108,9 @@ NVTX.@annotate function explicit_vertical_advection_tendency!(Yₜ, Y, p, t)
# Without the CT12(), the right-hand side would be a CT1 or CT2 in 2D space.

Fields.bycolumn(axes(Y.c)) do colidx
if :ρe_tot in propertynames(Yₜ.c)
(; ᶜh_tot) = p.precomputed
# Upwinding corrections to active tracers (e_tot and q_tot)
if energy_upwinding != Val(:none)
for (coeff, upwinding) in ((1, energy_upwinding), (-1, Val(:none)))
energy_upwinding isa Val{:none} && continue
vertical_transport!(
coeff,
Yₜ.c.ρe_tot[colidx],
Expand All @@ -124,26 +123,35 @@ NVTX.@annotate function explicit_vertical_advection_tendency!(Yₜ, Y, p, t)
)
end
end
for (ᶜρχₜ, ᶜχ, χ_name) in matching_subfields(Yₜ.c, ᶜspecific)
χ_name == :e_tot && continue
if !(moisture_model isa DryModel) && tracer_upwinding != Val(:none)
for (coeff, upwinding) in ((1, tracer_upwinding), (-1, Val(:none)))
tracer_upwinding isa Val{:none} && continue
vertical_transport!(
coeff,
ᶜρχₜ[colidx],
Yₜ.c.ρq_tot[colidx],
ᶜJ[colidx],
Y.c.ρ[colidx],
ᶠu³[colidx],
ᶜχ[colidx],
ᶜspecific.q_tot[colidx],
dt,
upwinding,
)
end
end

end
# Full vertical advection of passive tracers (q_liq, q_rai, etc.)
for (ᶜρχₜ, ᶜχ, χ_name) in matching_subfields(Yₜ.c, ᶜspecific)
χ_name in (:e_tot, :q_tot) && continue
vertical_transport!(
ᶜρχₜ[colidx],
ᶜJ[colidx],
Y.c.ρ[colidx],
ᶠu³[colidx],
ᶜχ[colidx],
dt,
tracer_upwinding,
)
end

Fields.bycolumn(axes(Y.c)) do colidx
@. Yₜ.c.uₕ[colidx] -=
ᶜinterp(
ᶠω¹²[colidx] ×
Expand Down
65 changes: 17 additions & 48 deletions src/prognostic_equations/implicit/implicit_solver.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -158,23 +158,24 @@ function ImplicitEquationJacobian(
(@name(c.ρ), other_available_names...),
)

active_scalar_names = (@name(c.ρ), @name(c.ρe_tot), ρq_tot_if_available...)
advection_blocks = (
(
use_derivative(topography_flag) ?
MatrixFields.unrolled_map(
name ->
(name, @name(c.uₕ)) =>
similar(Y.c, TridiagonalRow_ACTh),
(@name(c.ρ), @name(c.ρe_tot), available_tracer_names...),
active_scalar_names,
) : ()
)...,
MatrixFields.unrolled_map(
name -> (name, @name(f.u₃)) => similar(Y.c, BidiagonalRow_ACT3),
(@name(c.ρ), @name(c.ρe_tot), available_tracer_names...),
active_scalar_names,
)...,
MatrixFields.unrolled_map(
name -> (@name(f.u₃), name) => similar(Y.f, BidiagonalRow_C3),
(@name(c.ρ), @name(c.ρe_tot), ρq_tot_if_available...),
active_scalar_names,
)...,
(@name(f.u₃), @name(c.uₕ)) => similar(Y.f, BidiagonalRow_C3xACTh),
(@name(f.u₃), @name(f.u₃)) => similar(Y.f, TridiagonalRow_C3xACT3),
Expand Down Expand Up @@ -297,7 +298,6 @@ function Wfact!(A, Y, p, dtγ, t)
# Remove unnecessary values from p to avoid allocations in bycolumn.
p′ = (;
p.precomputed.ᶜspecific,
p.precomputed.ᶠu³,
p.precomputed.ᶜK,
p.precomputed.ᶜts,
p.precomputed.ᶜp,
Expand Down Expand Up @@ -353,13 +353,10 @@ end

function update_implicit_equation_jacobian!(A, Y, p, dtγ, colidx)
(; matrix, diffusion_flag, topography_flag) = A
(; ᶜspecific, ᶠu³, ᶜK, ᶜts, ᶜp, ᶜΦ, ᶠgradᵥ_ᶜΦ, ᶜρ_ref, ᶜp_ref) = p
(; ᶜspecific, ᶜK, ᶜts, ᶜp, ᶜΦ, ᶠgradᵥ_ᶜΦ, ᶜρ_ref, ᶜp_ref) = p
(; ∂ᶜK_∂ᶜuₕ, ∂ᶜK_∂ᶠu₃, ᶠp_grad_matrix, ᶜadvection_matrix) = p
(; ᶜdiffusion_h_matrix, ᶜdiffusion_u_matrix, params) = p
# use CD2 for implicit, upwinding correction goes in explicit part
energy_upwinding = Val(:none)
tracer_upwinding = Val(:none)
(; precip_upwinding) = p.atmos.numerics # precipitation is always upwinded (rain always falls)
(; precip_upwinding) = p.atmos.numerics

FT = Spaces.undertype(axes(Y.c))
CTh = CTh_vector_type(axes(Y.c))
Expand Down Expand Up @@ -413,53 +410,25 @@ function update_implicit_equation_jacobian!(A, Y, p, dtγ, colidx)
@. ∂ᶜρ_err_∂ᶠu₃[colidx] =
dtγ * ᶜadvection_matrix[colidx] ⋅ DiagonalMatrixRow(g³³(ᶠgⁱʲ[colidx]))

ᶠu³_data = ᶠu³.components.data.:1
tracer_info = (
(@name(c.ρe_tot), @name(ᶜh_tot), energy_upwinding),
(@name(c.ρq_tot), @name(ᶜspecific.q_tot), tracer_upwinding),
(@name(c.ρq_liq), @name(ᶜspecific.q_liq), tracer_upwinding),
(@name(c.ρq_ice), @name(ᶜspecific.q_ice), tracer_upwinding),
(@name(c.ρq_rai), @name(ᶜspecific.q_rai), tracer_upwinding),
(@name(c.ρq_sno), @name(ᶜspecific.q_sno), tracer_upwinding),
(@name(c.ρe_tot), @name(ᶜh_tot)),
(@name(c.ρq_tot), @name(ᶜspecific.q_tot)),
)
MatrixFields.unrolled_foreach(tracer_info) do (ρχ_name, χ_name, upwinding)
MatrixFields.unrolled_foreach(tracer_info) do (ρχ_name, χ_name)
MatrixFields.has_field(Y, ρχ_name) || return
ᶜχ = MatrixFields.get_field(p, χ_name)
if use_derivative(topography_flag)
∂ᶜρχ_err_∂ᶜuₕ = matrix[ρχ_name, @name(c.uₕ)]
end
∂ᶜρχ_err_∂ᶠu₃ = matrix[ρχ_name, @name(f.u₃)]
if upwinding == Val(:none)
use_derivative(topography_flag) && @. ∂ᶜρχ_err_∂ᶜuₕ[colidx] =
dtγ * ᶜadvection_matrix[colidx] ⋅
DiagonalMatrixRow(ᶠinterp(ᶜχ[colidx])) ⋅
ᶠwinterp_matrix(ᶜJ[colidx] * ᶜρ[colidx]) ⋅
DiagonalMatrixRow(g³ʰ(ᶜgⁱʲ[colidx]))
@. ∂ᶜρχ_err_∂ᶠu₃[colidx] =
dtγ * ᶜadvection_matrix[colidx] ⋅
DiagonalMatrixRow(ᶠinterp(ᶜχ[colidx]) * g³³(ᶠgⁱʲ[colidx]))
else
ᶠupwind = upwinding == Val(:third_order) ? ᶠupwind3 : ᶠupwind1
ᶠset_upwind_bcs = Operators.SetBoundaryOperator(;
top = Operators.SetValue(zero(CT3{FT})),
bottom = Operators.SetValue(zero(CT3{FT})),
) # Need to wrap ᶠupwind in this for well-defined boundaries.
use_derivative(topography_flag) && @. ∂ᶜρχ_err_∂ᶜuₕ[colidx] =
dtγ * ᶜadvection_matrix[colidx] ⋅ DiagonalMatrixRow(
ᶠset_upwind_bcs(
ᶠupwind(CT3(sign(ᶠu³_data[colidx])), ᶜχ[colidx]),
) * adjoint(C3(sign(ᶠu³_data[colidx]))),
) ⋅ ᶠwinterp_matrix(ᶜJ[colidx] * ᶜρ[colidx]) ⋅
DiagonalMatrixRow(g³ʰ(ᶜgⁱʲ[colidx]))
@. ∂ᶜρχ_err_∂ᶠu₃[colidx] =
dtγ * ᶜadvection_matrix[colidx] ⋅ DiagonalMatrixRow(
ᶠset_upwind_bcs(
ᶠupwind(CT3(sign(ᶠu³_data[colidx])), ᶜχ[colidx]),
) *
adjoint(C3(sign(ᶠu³_data[colidx]))) *
g³³(ᶠgⁱʲ[colidx]),
)
end
use_derivative(topography_flag) && @. ∂ᶜρχ_err_∂ᶜuₕ[colidx] =
dtγ * ᶜadvection_matrix[colidx] ⋅
DiagonalMatrixRow(ᶠinterp(ᶜχ[colidx])) ⋅
ᶠwinterp_matrix(ᶜJ[colidx] * ᶜρ[colidx]) ⋅
DiagonalMatrixRow(g³ʰ(ᶜgⁱʲ[colidx]))
@. ∂ᶜρχ_err_∂ᶠu₃[colidx] =
dtγ * ᶜadvection_matrix[colidx] ⋅
DiagonalMatrixRow(ᶠinterp(ᶜχ[colidx]) * g³³(ᶠgⁱʲ[colidx]))
end

∂ᶠu₃_err_∂ᶜρ = matrix[@name(f.u₃), @name(c.ρ)]
Expand Down
42 changes: 18 additions & 24 deletions src/prognostic_equations/implicit/implicit_tendency.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -119,48 +119,42 @@ vertical_advection!(ᶜρχₜ, ᶠu³, ᶜχ, ::Val{:third_order}) =
@. ᶜρχₜ -= ᶜadvdivᵥ(ᶠupwind3(ᶠu³, ᶜχ)) - ᶜχ * ᶜadvdivᵥ(ᶠu³)

function implicit_vertical_advection_tendency!(Yₜ, Y, p, t, colidx)
(; turbconv_model, rayleigh_sponge, precip_model) = p.atmos
(; moisture_model, turbconv_model, rayleigh_sponge, precip_model) = p.atmos
(; dt) = p
n = n_mass_flux_subdomains(turbconv_model)
ᶜJ = Fields.local_geometry_field(Y.c).J
(; ᶠgradᵥ_ᶜΦ, ᶜρ_ref, ᶜp_ref) = p.core
(; ᶜspecific, ᶠu³, ᶜp) = p.precomputed
# use CD2 for implicit, upwinding correction goes in explicit part
energy_upwinding = Val(:none)
tracer_upwinding = Val(:none)
(; precip_upwinding) = p.atmos.numerics # precipitation is always upwinded (rain always falls)
(; ᶜh_tot, ᶜspecific, ᶠu³, ᶜp) = p.precomputed
(; precip_upwinding) = p.atmos.numerics

@. Yₜ.c.ρ[colidx] -=
ᶜdivᵥ(ᶠwinterp(ᶜJ[colidx], Y.c.ρ[colidx]) * ᶠu³[colidx])

if :ρe_tot in propertynames(Yₜ.c)
(; ᶜh_tot) = p.precomputed
# Central advection of active tracers (e_tot and q_tot)
vertical_transport!(
Yₜ.c.ρe_tot[colidx],
ᶜJ[colidx],
Y.c.ρ[colidx],
ᶠu³[colidx],
ᶜh_tot[colidx],
dt,
Val(:none),
)
if !(moisture_model isa DryModel)
vertical_transport!(
Yₜ.c.ρe_tot[colidx],
Yₜ.c.ρq_tot[colidx],
ᶜJ[colidx],
Y.c.ρ[colidx],
ᶠu³[colidx],
ᶜh_tot[colidx],
ᶜspecific.q_tot[colidx],
dt,
energy_upwinding,
)
end
for (ᶜρχₜ, ᶜχ, χ_name) in matching_subfields(Yₜ.c, ᶜspecific)
χ_name == :e_tot && continue
vertical_transport!(
ᶜρχₜ[colidx],
ᶜJ[colidx],
Y.c.ρ[colidx],
ᶠu³[colidx],
ᶜχ[colidx],
dt,
tracer_upwinding,
Val(:none),
)
end

if precip_model isa Microphysics1Moment
# Advection of precipitation with the mean flow
# is done with other tracers above.
# is done with other passive tracers in the explicit tendency.
# Here we add the advection with precipitation terminal velocity
# using first order upwind and free outflow bottom boundary condition

Expand Down
Loading