Skip to content

Commit

Permalink
separate advection and diffusion jacobian blocks
Browse files Browse the repository at this point in the history
  • Loading branch information
szy21 committed Mar 30, 2024
1 parent 7c623ea commit 9a02579
Show file tree
Hide file tree
Showing 13 changed files with 90 additions and 41 deletions.
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

0 comments on commit 9a02579

Please sign in to comment.