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.

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

separate advection and diffusion jacobian blocks #2856

Merged
merged 1 commit into from
Mar 30, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
3 changes: 3 additions & 0 deletions config/default_configs/default_edmf_config.yml
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,9 @@ advection_test:
gs_tendency:
help: "Turns on all grid-scale tendencies [`true` (default), `false`]"
value: true
implicit_sgs_advection:
help: "Whether to treat the subgrid-scale vertical advection tendency implicitly [`false` (default), `true`]"
value: false
edmf_coriolis:
help: "EDMF coriolis [`nothing` (default), `Bomex`,`LifeCycleTan2018`,`Rico`,`ARM_SGP`,`DYCOMS_RF01`,`DYCOMS_RF02`,`GABLS`]"
value: ~
Expand Down
1 change: 1 addition & 0 deletions config/model_configs/prognostic_edmfx_bomex_box.yml
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ ls_adv: "Bomex"
surface_setup: "Bomex"
turbconv: "prognostic_edmfx"
implicit_diffusion: true
implicit_sgs_advection: true
approximate_linear_solve_iters: 2
max_newton_iters_ode: 3
edmfx_upwinding: first_order
Expand Down
1 change: 1 addition & 0 deletions config/model_configs/prognostic_edmfx_bomex_column.yml
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ ls_adv: "Bomex"
surface_setup: "Bomex"
turbconv: "prognostic_edmfx"
implicit_diffusion: true
implicit_sgs_advection: true
approximate_linear_solve_iters: 2
max_newton_iters_ode: 3
edmfx_upwinding: first_order
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ ls_adv: "Bomex"
surface_setup: "Bomex"
turbconv: "prognostic_edmfx"
implicit_diffusion: true
implicit_sgs_advection: true
approximate_linear_solve_iters: 2
max_newton_iters_ode: 3
edmfx_upwinding: first_order
Expand Down
1 change: 1 addition & 0 deletions config/model_configs/prognostic_edmfx_dycoms_rf01_box.yml
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ rad: DYCOMS_RF01
surface_setup: DYCOMS_RF01
turbconv: prognostic_edmfx
implicit_diffusion: true
implicit_sgs_advection: true
approximate_linear_solve_iters: 2
max_newton_iters_ode: 3
edmfx_upwinding: first_order
Expand Down
1 change: 1 addition & 0 deletions config/model_configs/prognostic_edmfx_gabls_box.yml
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ edmf_coriolis: GABLS
surface_setup: GABLS
turbconv: "prognostic_edmfx"
implicit_diffusion: true
implicit_sgs_advection: true
approximate_linear_solve_iters: 2
max_newton_iters_ode: 3
edmfx_upwinding: "first_order"
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ initial_condition: "SimplePlume"
surface_setup: "SimplePlume"
turbconv: "prognostic_edmfx"
implicit_diffusion: true
implicit_sgs_advection: true
approximate_linear_solve_iters: 2
max_newton_iters_ode: 3
gs_tendency: false
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,8 @@ surface_setup: DefaultExchangeCoefficients
rad: gray
vert_diff: "false"
turbconv: prognostic_edmfx
implicit_diffusion: true
implicit_sgs_advection: true
prognostic_tke: true
edmfx_upwinding: first_order
edmfx_entr_model: "Generalized"
Expand Down
104 changes: 65 additions & 39 deletions src/prognostic_equations/implicit/implicit_solver.jl
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,10 @@ use_derivative(::UseDerivative) = true
use_derivative(::IgnoreDerivative) = false

"""
ImplicitEquationJacobian(Y, atmos; approximate_solve_iters, diffusion_flag, topography_flag, transform_flag)
ImplicitEquationJacobian(
Y, atmos;
approximate_solve_iters, diffusion_flag, topography_flag, sgs_advection_flag, transform_flag
)

A wrapper for the matrix ``∂E/∂Y``, where ``E(Y)`` is the "error" of the
implicit step with the state ``Y``.
Expand Down Expand Up @@ -76,11 +79,15 @@ is reached.
- `diffusion_flag::DerivativeFlag`: whether the derivative of the
diffusion tendency with respect to the quantities being diffused should be
computed or approximated as 0; must be either `UseDerivative()` or
`Ignoreerivative()` instead of a `Bool` to ensure type-stability
`IgnoreDerivative()` instead of a `Bool` to ensure type-stability
- `topography_flag::DerivativeFlag`: whether the derivative of vertical
contravariant velocity with respect to horizontal covariant velocity should be
computed or approximated as 0; must be either `UseDerivative()` or
`IgnoreDerivative()` instead of a `Bool` to ensure type-stability
- `sgs_advection_flag::DerivativeFlag`: whether the derivative of the
subgrid-scale advection tendency with respect to the updraft quantities should be
computed or approximated as 0; must be either `UseDerivative()` or
`IgnoreDerivative()` instead of a `Bool` to ensure type-stability
- `transform_flag::Bool`: whether the error should be transformed from ``E(Y)``
to ``E(Y)/Δt``, which is required for non-Rosenbrock timestepping schemes from
OrdinaryDiffEq.jl
Expand All @@ -90,6 +97,7 @@ struct ImplicitEquationJacobian{
S <: MatrixFields.FieldMatrixSolver,
F1 <: DerivativeFlag,
F2 <: DerivativeFlag,
F3 <: DerivativeFlag,
T <: Fields.FieldVector,
R <: Base.RefValue,
}
Expand All @@ -102,6 +110,7 @@ struct ImplicitEquationJacobian{
# flags that determine how E'(Y) is approximated
diffusion_flag::F1
topography_flag::F2
sgs_advection_flag::F3

# required by Krylov.jl to evaluate ldiv! with AbstractVector inputs
temp_b::T
Expand All @@ -120,6 +129,8 @@ function ImplicitEquationJacobian(
IgnoreDerivative(),
topography_flag = has_topography(axes(Y.c)) ? UseDerivative() :
IgnoreDerivative(),
sgs_advection_flag = atmos.sgs_adv_mode == Implicit() ? UseDerivative() :
IgnoreDerivative(),
transform_flag = false,
)
FT = Spaces.undertype(axes(Y.c))
Expand Down Expand Up @@ -214,14 +225,14 @@ function ImplicitEquationJacobian(
)
end

sgs_blocks = if atmos.turbconv_model isa PrognosticEDMFX
sgs_advection_blocks = if atmos.turbconv_model isa PrognosticEDMFX
@assert n_prognostic_mass_flux_subdomains(atmos.turbconv_model) == 1
sgs_scalar_names = (
@name(c.sgsʲs.:(1).q_tot),
@name(c.sgsʲs.:(1).mse),
@name(c.sgsʲs.:(1).ρa)
)
if use_derivative(diffusion_flag)
if use_derivative(sgs_advection_flag)
(
MatrixFields.unrolled_map(
name -> (name, name) => similar(Y.c, TridiagonalRow),
Expand Down Expand Up @@ -261,7 +272,7 @@ function ImplicitEquationJacobian(

matrix = MatrixFields.FieldMatrix(
identity_blocks...,
sgs_blocks...,
sgs_advection_blocks...,
advection_blocks...,
diffusion_blocks...,
)
Expand All @@ -287,49 +298,60 @@ function ImplicitEquationJacobian(
)

alg₂ = MatrixFields.BlockLowerTriangularSolve(@name(c.uₕ))
alg = if use_derivative(diffusion_flag)
alg₁_subalg₂ = if atmos.turbconv_model isa PrognosticEDMFX
(;
alg₂ = MatrixFields.BlockLowerTriangularSolve(
@name(c.sgsʲs.:(1).q_tot);
alg₂ = MatrixFields.BlockLowerTriangularSolve(
@name(c.sgsʲs.:(1).mse);
alg₂ = MatrixFields.BlockLowerTriangularSolve(
@name(c.sgsʲs.:(1).ρa),
@name(f.sgsʲs.:(1).u₃);
alg =
if use_derivative(diffusion_flag) || use_derivative(sgs_advection_flag)
alg₁_subalg₂ =
if atmos.turbconv_model isa PrognosticEDMFX &&
use_derivative(sgs_advection_flag)
diff_subalg =
use_derivative(diffusion_flag) ?
(;
alg₂ = MatrixFields.BlockLowerTriangularSolve(
names₁_group₂...,
)
) : (;)
(;
alg₂ = MatrixFields.BlockLowerTriangularSolve(
@name(c.sgsʲs.:(1).q_tot);
alg₂ = MatrixFields.BlockLowerTriangularSolve(
@name(c.sgsʲs.:(1).mse);
alg₂ = MatrixFields.BlockLowerTriangularSolve(
@name(c.sgsʲs.:(1).ρa),
@name(f.sgsʲs.:(1).u₃);
diff_subalg...,
),
),
),
),
)
)
)
else
is_in_Y(@name(c.ρq_tot)) ?
(;
alg₂ = MatrixFields.BlockLowerTriangularSolve(
names₁_group₂...,
)
) : (;)
end
alg₁ = MatrixFields.BlockLowerTriangularSolve(
names₁_group₁...;
alg₁_subalg₂...,
)
MatrixFields.ApproximateBlockArrowheadIterativeSolve(
names₁...;
alg₁,
alg₂,
P_alg₁ = MatrixFields.MainDiagonalPreconditioner(),
n_iters = approximate_solve_iters,
)
else
is_in_Y(@name(c.ρq_tot)) ?
(;
alg₂ = MatrixFields.BlockLowerTriangularSolve(names₁_group₂...)
) : (;)
MatrixFields.BlockArrowheadSolve(names₁...; alg₂)
end
alg₁ = MatrixFields.BlockLowerTriangularSolve(
names₁_group₁...;
alg₁_subalg₂...,
)
MatrixFields.ApproximateBlockArrowheadIterativeSolve(
names₁...;
alg₁,
alg₂,
P_alg₁ = MatrixFields.MainDiagonalPreconditioner(),
n_iters = approximate_solve_iters,
)
else
MatrixFields.BlockArrowheadSolve(names₁...; alg₂)
end

return ImplicitEquationJacobian(
matrix,
MatrixFields.FieldMatrixSolver(alg, matrix, Y),
diffusion_flag,
topography_flag,
sgs_advection_flag,
similar(Y),
similar(Y),
transform_flag,
Expand Down Expand Up @@ -397,9 +419,13 @@ NVTX.@annotate function Wfact!(A, Y, p, dtγ, t)
(
use_derivative(A.diffusion_flag) &&
p.atmos.turbconv_model isa PrognosticEDMFX ?
(; p.precomputed.ᶜρa⁰) : (;)
)...,
(
use_derivative(A.sgs_advection_flag) &&
p.atmos.turbconv_model isa PrognosticEDMFX ?
(;
p.core.ᶜgradᵥ_ᶠΦ,
p.precomputed.ᶜρa⁰,
p.precomputed.ᶜρʲs,
p.precomputed.ᶠu³ʲs,
p.precomputed.ᶜtsʲs,
Expand Down Expand Up @@ -441,7 +467,7 @@ NVTX.@annotate function Wfact!(A, Y, p, dtγ, t)
end

function update_implicit_equation_jacobian!(A, Y, p, dtγ, colidx)
(; matrix, diffusion_flag, topography_flag) = A
(; matrix, diffusion_flag, sgs_advection_flag, topography_flag) = A
(; ᶜspecific, ᶜK, ᶜts, ᶜp, ᶜΦ, ᶠgradᵥ_ᶜΦ, ᶜρ_ref, ᶜp_ref) = p
(;
ᶜtemp_C3,
Expand Down Expand Up @@ -695,7 +721,7 @@ function update_implicit_equation_jacobian!(A, Y, p, dtγ, colidx)
end

if p.atmos.turbconv_model isa PrognosticEDMFX
if use_derivative(diffusion_flag)
if use_derivative(sgs_advection_flag)
(; ᶜgradᵥ_ᶠΦ, ᶜρʲs, ᶠu³ʲs, ᶜtsʲs) = p
is_third_order = edmfx_upwinding == Val(:third_order)
ᶠupwind = is_third_order ? ᶠupwind3 : ᶠupwind1
Expand Down
5 changes: 4 additions & 1 deletion src/prognostic_equations/implicit/implicit_tendency.jl
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ NVTX.@annotate function implicit_tendency!(Yₜ, Y, p, t)
Yₜ .= zero(eltype(Yₜ))
Fields.bycolumn(axes(Y.c)) do colidx
implicit_vertical_advection_tendency!(Yₜ, Y, p, t, colidx)
if p.atmos.diff_mode == Implicit()
if p.atmos.sgs_adv_mode == Implicit()
edmfx_sgs_vertical_advection_tendency!(
Yₜ,
Y,
Expand All @@ -18,6 +18,9 @@ NVTX.@annotate function implicit_tendency!(Yₜ, Y, p, t)
colidx,
p.atmos.turbconv_model,
)
end

if p.atmos.diff_mode == Implicit()
vertical_diffusion_boundary_layer_tendency!(
Yₜ,
Y,
Expand Down
5 changes: 4 additions & 1 deletion src/prognostic_equations/remaining_tendency.jl
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@ NVTX.@annotate function additional_tendency!(Yₜ, Y, p, t)
edmf_coriolis_tendency!(Yₜ, Y, p, t, colidx, p.atmos.edmf_coriolis)
large_scale_advection_tendency!(Yₜ, Y, p, t, colidx, p.atmos.ls_adv)

if p.atmos.diff_mode == Explicit()
if p.atmos.sgs_adv_mode == Explicit()
edmfx_sgs_vertical_advection_tendency!(
Yₜ,
Y,
Expand All @@ -43,6 +43,9 @@ NVTX.@annotate function additional_tendency!(Yₜ, Y, p, t)
colidx,
p.atmos.turbconv_model,
)
end

if p.atmos.diff_mode == Explicit()
vertical_diffusion_boundary_layer_tendency!(
Yₜ,
Y,
Expand Down
4 changes: 4 additions & 0 deletions src/solver/type_getters.jl
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,9 @@ function get_atmos(config::AtmosConfig, params)
implicit_diffusion = parsed_args["implicit_diffusion"]
@assert implicit_diffusion in (true, false)

implicit_sgs_advection = parsed_args["implicit_sgs_advection"]
@assert implicit_sgs_advection in (true, false)

model_config = get_model_config(parsed_args)
vert_diff =
get_vertical_diffusion_model(diffuse_momentum, parsed_args, params, FT)
Expand Down Expand Up @@ -82,6 +85,7 @@ function get_atmos(config::AtmosConfig, params)
hyperdiff = get_hyperdiffusion_model(parsed_args, FT),
vert_diff,
diff_mode = implicit_diffusion ? Implicit() : Explicit(),
sgs_adv_mode = implicit_sgs_advection ? Implicit() : Explicit(),
viscous_sponge = get_viscous_sponge_model(parsed_args, params, FT),
rayleigh_sponge = get_rayleigh_sponge_model(parsed_args, params, FT),
sfc_temperature = get_sfc_temperature_form(parsed_args),
Expand Down
2 changes: 2 additions & 0 deletions src/solver/types.jl
Original file line number Diff line number Diff line change
Expand Up @@ -340,6 +340,7 @@ Base.@kwdef struct AtmosModel{
HD,
VD,
DM,
SAM,
VS,
RS,
ST,
Expand Down Expand Up @@ -371,6 +372,7 @@ Base.@kwdef struct AtmosModel{
hyperdiff::HD = nothing
vert_diff::VD = nothing
diff_mode::DM = nothing
sgs_adv_mode::SAM = nothing
viscous_sponge::VS = nothing
rayleigh_sponge::RS = nothing
sfc_temperature::ST = nothing
Expand Down
Loading