From b774c6b6c3228af48aba9b221ba2823dc496ff4e Mon Sep 17 00:00:00 2001 From: Dennis Yatunin Date: Thu, 30 Nov 2023 17:35:56 -0800 Subject: [PATCH] Add implict solver for diagnostic EDMF --- .buildkite/pipeline.yml | 8 +++ config/default_configs/default_config.yml | 3 + .../diagnostic_edmfx_bomex_box.yml | 5 +- .../diagnostic_edmfx_bomex_stretched_box.yml | 3 +- .../diagnostic_edmfx_dycoms_rf01_box.yml | 6 +- ...gnostic_edmfx_dycoms_rf01_explicit_box.yml | 29 ++++++++++ .../diagnostic_edmfx_gabls_box.yml | 2 + .../diagnostic_edmfx_rico_box.yml | 1 - .../diagnostic_edmfx_trmm_box.yml | 1 - src/prognostic_equations/edmfx_sgs_flux.jl | 42 ++++++++++---- src/prognostic_equations/edmfx_tke.jl | 13 +---- .../implicit/implicit_solver.jl | 56 +++++++++++-------- .../implicit/implicit_tendency.jl | 8 +++ .../remaining_tendency.jl | 16 +++--- src/solver/type_getters.jl | 5 +- 15 files changed, 136 insertions(+), 62 deletions(-) create mode 100644 config/model_configs/diagnostic_edmfx_dycoms_rf01_explicit_box.yml diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index ce75093009..43f49a8f6c 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -582,6 +582,14 @@ steps: agents: slurm_mem: 20GB + - label: ":genie: Diagnostic EDMFX DYCOMS_RF01 in a box (explicit)" + command: > + julia --color=yes --project=examples examples/hybrid/driver.jl + --config_file $CONFIG_PATH/diagnostic_edmfx_dycoms_rf01_explicit_box.yml + artifact_paths: "diagnostic_edmfx_dycoms_rf01_explicitbox/*" + agents: + slurm_mem: 20GB + - label: ":genie: Diagnostic EDMFX DYCOMS_RF01 in a box" command: > julia --color=yes --project=examples examples/hybrid/driver.jl diff --git a/config/default_configs/default_config.yml b/config/default_configs/default_config.yml index 79e3826583..0915e344de 100644 --- a/config/default_configs/default_config.yml +++ b/config/default_configs/default_config.yml @@ -246,6 +246,9 @@ prognostic_surface: implicit_diffusion: help: "Whether to treat the vertical diffusion tendency implicitly [`false` (default), `true`]" value: false +approximate_linear_solve_iters: + help: "Number of iterations for the approximate linear solve (used when `implicit_diffusion` is true)" + value: 1 override_τ_precip: help: "If true, sets τ_precip to dt. Otherwise, τ_precip is set to the value in the toml dictionary" value: true diff --git a/config/model_configs/diagnostic_edmfx_bomex_box.yml b/config/model_configs/diagnostic_edmfx_bomex_box.yml index 8753b82091..3d4a9dab01 100644 --- a/config/model_configs/diagnostic_edmfx_bomex_box.yml +++ b/config/model_configs/diagnostic_edmfx_bomex_box.yml @@ -1,4 +1,3 @@ - job_id: "diagnostic_edmfx_bomex_box" initial_condition: "Bomex" subsidence: "Bomex" @@ -6,6 +5,8 @@ edmf_coriolis: "Bomex" ls_adv: "Bomex" surface_setup: "Bomex" turbconv: "diagnostic_edmfx" +implicit_diffusion: true +approximate_linear_solve_iters: 2 prognostic_tke: true edmfx_upwinding: "first_order" edmfx_entr_model: "Generalized" @@ -24,7 +25,7 @@ y_elem: 2 z_elem: 60 z_max: 3e3 z_stretch: false -dt: "50secs" +dt: "200secs" t_end: "6hours" dt_save_to_disk: "10mins" toml: [toml/diagnostic_edmfx_box.toml] diff --git a/config/model_configs/diagnostic_edmfx_bomex_stretched_box.yml b/config/model_configs/diagnostic_edmfx_bomex_stretched_box.yml index a469e26666..a49bb7a691 100644 --- a/config/model_configs/diagnostic_edmfx_bomex_stretched_box.yml +++ b/config/model_configs/diagnostic_edmfx_bomex_stretched_box.yml @@ -1,4 +1,3 @@ - job_id: "diagnostic_edmfx_bomex_stretched_box" initial_condition: "Bomex" subsidence: "Bomex" @@ -6,6 +5,8 @@ edmf_coriolis: "Bomex" ls_adv: "Bomex" surface_setup: "Bomex" turbconv: "diagnostic_edmfx" +implicit_diffusion: true +approximate_linear_solve_iters: 2 prognostic_tke: true edmfx_upwinding: "first_order" edmfx_entr_model: "Generalized" diff --git a/config/model_configs/diagnostic_edmfx_dycoms_rf01_box.yml b/config/model_configs/diagnostic_edmfx_dycoms_rf01_box.yml index c875dfaa7a..30385a58af 100644 --- a/config/model_configs/diagnostic_edmfx_dycoms_rf01_box.yml +++ b/config/model_configs/diagnostic_edmfx_dycoms_rf01_box.yml @@ -5,6 +5,8 @@ edmf_coriolis: DYCOMS_RF01 rad: DYCOMS_RF01 surface_setup: DYCOMS_RF01 turbconv: diagnostic_edmfx +implicit_diffusion: true +approximate_linear_solve_iters: 2 prognostic_tke: true edmfx_upwinding: first_order edmfx_entr_model: "Generalized" @@ -23,7 +25,7 @@ y_elem: 2 z_elem: 30 z_max: 1500 z_stretch: false -dt: 40secs +dt: 200secs t_end: 4hours -dt_save_to_disk: 2mins +dt_save_to_disk: 10mins toml: [toml/diagnostic_edmfx_box.toml] diff --git a/config/model_configs/diagnostic_edmfx_dycoms_rf01_explicit_box.yml b/config/model_configs/diagnostic_edmfx_dycoms_rf01_explicit_box.yml new file mode 100644 index 0000000000..9c28d29a08 --- /dev/null +++ b/config/model_configs/diagnostic_edmfx_dycoms_rf01_explicit_box.yml @@ -0,0 +1,29 @@ +job_id: diagnostic_edmfx_dycoms_rf01_explicit_box +initial_condition: DYCOMS_RF01 +subsidence: DYCOMS +edmf_coriolis: DYCOMS_RF01 +rad: DYCOMS_RF01 +surface_setup: DYCOMS_RF01 +turbconv: diagnostic_edmfx +prognostic_tke: true +edmfx_upwinding: first_order +edmfx_entr_model: "Generalized" +edmfx_detr_model: "Generalized" +edmfx_nh_pressure: true +edmfx_sgs_mass_flux: true +edmfx_sgs_diffusive_flux: true +moist: equil +config: box +hyperdiff: "true" +kappa_4: 1e21 +x_max: 1e8 +y_max: 1e8 +x_elem: 2 +y_elem: 2 +z_elem: 30 +z_max: 1500 +z_stretch: false +dt: 40secs +t_end: 4hours +dt_save_to_disk: 10mins +toml: [toml/diagnostic_edmfx_box.toml] diff --git a/config/model_configs/diagnostic_edmfx_gabls_box.yml b/config/model_configs/diagnostic_edmfx_gabls_box.yml index 864cff562b..2424d9ee34 100644 --- a/config/model_configs/diagnostic_edmfx_gabls_box.yml +++ b/config/model_configs/diagnostic_edmfx_gabls_box.yml @@ -3,6 +3,8 @@ initial_condition: GABLS edmf_coriolis: GABLS surface_setup: GABLS turbconv: diagnostic_edmfx +implicit_diffusion: true +approximate_linear_solve_iters: 2 prognostic_tke: true edmfx_upwinding: first_order edmfx_entr_model: "Generalized" diff --git a/config/model_configs/diagnostic_edmfx_rico_box.yml b/config/model_configs/diagnostic_edmfx_rico_box.yml index 74d1338bbd..27f8412e47 100644 --- a/config/model_configs/diagnostic_edmfx_rico_box.yml +++ b/config/model_configs/diagnostic_edmfx_rico_box.yml @@ -1,4 +1,3 @@ - job_id: diagnostic_edmfx_rico_box initial_condition: Rico subsidence: Rico diff --git a/config/model_configs/diagnostic_edmfx_trmm_box.yml b/config/model_configs/diagnostic_edmfx_trmm_box.yml index 9f5518960b..22a0cfd221 100644 --- a/config/model_configs/diagnostic_edmfx_trmm_box.yml +++ b/config/model_configs/diagnostic_edmfx_trmm_box.yml @@ -1,4 +1,3 @@ - job_id: diagnostic_edmfx_trmm_box initial_condition: TRMM_LBA rad: TRMM_LBA diff --git a/src/prognostic_equations/edmfx_sgs_flux.jl b/src/prognostic_equations/edmfx_sgs_flux.jl index 48a585c937..bcd107b3c5 100644 --- a/src/prognostic_equations/edmfx_sgs_flux.jl +++ b/src/prognostic_equations/edmfx_sgs_flux.jl @@ -169,15 +169,17 @@ function edmfx_sgs_diffusive_flux_tendency!( FT = Spaces.undertype(axes(Y.c)) (; sfc_conditions) = p.precomputed - (; ᶜρa⁰, ᶜu⁰, ᶜK⁰, ᶜmse⁰, ᶜq_tot⁰) = p.precomputed - (; ᶜK_u, ᶜK_h) = p.precomputed + (; ᶜρa⁰, ᶜu⁰, ᶜK⁰, ᶜmse⁰, ᶜq_tot⁰, ᶜtke⁰) = p.precomputed + (; ᶜK_u, ᶜK_h, ρatke_flux) = p.precomputed ᶠgradᵥ = Operators.GradientC2F() if p.atmos.edmfx_sgs_diffusive_flux - # energy ᶠρaK_h = p.scratch.ᶠtemp_scalar @. ᶠρaK_h[colidx] = ᶠinterp(ᶜρa⁰[colidx]) * ᶠinterp(ᶜK_h[colidx]) + ᶠρaK_u = p.scratch.ᶠtemp_scalar + @. ᶠρaK_u[colidx] = ᶠinterp(ᶜρa⁰[colidx]) * ᶠinterp(ᶜK_u[colidx]) + # energy ᶜdivᵥ_ρe_tot = Operators.DivergenceF2C( top = Operators.SetValue(C3(FT(0))), bottom = Operators.SetValue(sfc_conditions.ρ_flux_h_tot[colidx]), @@ -185,7 +187,16 @@ function edmfx_sgs_diffusive_flux_tendency!( @. Yₜ.c.ρe_tot[colidx] -= ᶜdivᵥ_ρe_tot( -(ᶠρaK_h[colidx] * ᶠgradᵥ(ᶜmse⁰[colidx] + ᶜK⁰[colidx])), ) - + if use_prognostic_tke(turbconv_model) + # turbulent transport (diffusive flux) + # boundary condition for the diffusive flux + ᶜdivᵥ_ρatke = Operators.DivergenceF2C( + top = Operators.SetValue(C3(FT(0))), + bottom = Operators.SetValue(ρatke_flux[colidx]), + ) + @. Yₜ.c.sgs⁰.ρatke[colidx] -= + ᶜdivᵥ_ρatke(-(ᶠρaK_u[colidx] * ᶠgradᵥ(ᶜtke⁰[colidx]))) + end if !(p.atmos.moisture_model isa DryModel) # specific humidity ᶜρχₜ_diffusion = p.scratch.ᶜtemp_scalar @@ -202,8 +213,6 @@ function edmfx_sgs_diffusive_flux_tendency!( end # momentum - ᶠρaK_u = p.scratch.ᶠtemp_scalar - @. ᶠρaK_u[colidx] = ᶠinterp(ᶜρa⁰[colidx]) * ᶠinterp(ᶜK_u[colidx]) ᶠstrain_rate = p.scratch.ᶠtemp_UVWxUVW compute_strain_rate_face!(ᶠstrain_rate[colidx], ᶜu⁰[colidx]) @. Yₜ.c.uₕ[colidx] -= C12( @@ -234,15 +243,17 @@ function edmfx_sgs_diffusive_flux_tendency!( FT = Spaces.undertype(axes(Y.c)) (; sfc_conditions) = p.precomputed - (; ᶜu, ᶜh_tot, ᶜspecific) = p.precomputed - (; ᶜK_u, ᶜK_h) = p.precomputed + (; ᶜu, ᶜh_tot, ᶜspecific, ᶜtke⁰) = p.precomputed + (; ᶜK_u, ᶜK_h, ρatke_flux) = p.precomputed ᶠgradᵥ = Operators.GradientC2F() if p.atmos.edmfx_sgs_diffusive_flux - # energy ᶠρaK_h = p.scratch.ᶠtemp_scalar @. ᶠρaK_h[colidx] = ᶠinterp(Y.c.ρ[colidx]) * ᶠinterp(ᶜK_h[colidx]) + ᶠρaK_u = p.scratch.ᶠtemp_scalar + @. ᶠρaK_u[colidx] = ᶠinterp(Y.c.ρ[colidx]) * ᶠinterp(ᶜK_u[colidx]) + # energy ᶜdivᵥ_ρe_tot = Operators.DivergenceF2C( top = Operators.SetValue(C3(FT(0))), bottom = Operators.SetValue(sfc_conditions.ρ_flux_h_tot[colidx]), @@ -250,6 +261,17 @@ function edmfx_sgs_diffusive_flux_tendency!( @. Yₜ.c.ρe_tot[colidx] -= ᶜdivᵥ_ρe_tot(-(ᶠρaK_h[colidx] * ᶠgradᵥ(ᶜh_tot[colidx]))) + if use_prognostic_tke(turbconv_model) + # turbulent transport (diffusive flux) + # boundary condition for the diffusive flux + ᶜdivᵥ_ρatke = Operators.DivergenceF2C( + top = Operators.SetValue(C3(FT(0))), + bottom = Operators.SetValue(ρatke_flux[colidx]), + ) + @. Yₜ.c.sgs⁰.ρatke[colidx] -= + ᶜdivᵥ_ρatke(-(ᶠρaK_u[colidx] * ᶠgradᵥ(ᶜtke⁰[colidx]))) + end + if !(p.atmos.moisture_model isa DryModel) # specific humidity ᶜρχₜ_diffusion = p.scratch.ᶜtemp_scalar @@ -267,8 +289,6 @@ function edmfx_sgs_diffusive_flux_tendency!( end # momentum - ᶠρaK_u = p.scratch.ᶠtemp_scalar - @. ᶠρaK_u[colidx] = ᶠinterp(Y.c.ρ[colidx]) * ᶠinterp(ᶜK_u[colidx]) ᶠstrain_rate = p.scratch.ᶠtemp_UVWxUVW compute_strain_rate_face!(ᶠstrain_rate[colidx], ᶜu[colidx]) @. Yₜ.c.uₕ[colidx] -= C12( diff --git a/src/prognostic_equations/edmfx_tke.jl b/src/prognostic_equations/edmfx_tke.jl index 3d5ae4e86d..211aa7b557 100644 --- a/src/prognostic_equations/edmfx_tke.jl +++ b/src/prognostic_equations/edmfx_tke.jl @@ -21,22 +21,11 @@ function edmfx_tke_tendency!( (; ᶜentrʲs, ᶜdetrʲs, ᶠu³ʲs) = p.precomputed (; ᶠu³⁰, ᶜstrain_rate_norm, ᶜlinear_buoygrad, ᶜtke⁰, ᶜmixing_length) = p.precomputed - (; ᶜK_u, ᶜK_h, ρatke_flux) = p.precomputed + (; ᶜK_u, ᶜK_h) = p.precomputed (; dt) = p.simulation ᶜρa⁰ = turbconv_model isa PrognosticEDMFX ? p.precomputed.ᶜρa⁰ : Y.c.ρ - ᶠgradᵥ = Operators.GradientC2F() - ᶠρaK_u = p.scratch.ᶠtemp_scalar if use_prognostic_tke(turbconv_model) - # turbulent transport (diffusive flux) - @. ᶠρaK_u[colidx] = ᶠinterp(ᶜρa⁰[colidx]) * ᶠinterp(ᶜK_u[colidx]) - # boundary condition for the diffusive flux - ᶜdivᵥ_ρatke = Operators.DivergenceF2C( - top = Operators.SetValue(C3(FT(0))), - bottom = Operators.SetValue(ρatke_flux[colidx]), - ) - @. Yₜ.c.sgs⁰.ρatke[colidx] -= - ᶜdivᵥ_ρatke(-(ᶠρaK_u[colidx] * ᶠgradᵥ(ᶜtke⁰[colidx]))) # shear production @. Yₜ.c.sgs⁰.ρatke[colidx] += 2 * ᶜρa⁰[colidx] * ᶜK_u[colidx] * ᶜstrain_rate_norm[colidx] diff --git a/src/prognostic_equations/implicit/implicit_solver.jl b/src/prognostic_equations/implicit/implicit_solver.jl index f7212ed603..e09f338d2d 100644 --- a/src/prognostic_equations/implicit/implicit_solver.jl +++ b/src/prognostic_equations/implicit/implicit_solver.jl @@ -13,7 +13,7 @@ struct IgnoreEnthalpyDerivative <: EnthalpyDerivativeFlag end use_derivative(flag::EnthalpyDerivativeFlag) = flag == UseEnthalpyDerivative() """ - ImplicitEquationJacobian(Y, atmos, diffusion_flag, enthalpy_flag, transform_flag) + ImplicitEquationJacobian(Y, atmos, diffusion_flag, enthalpy_flag, transform_flag, approximate_solve_iters) A wrapper for the matrix ``∂E/∂Y``, where ``E(Y)`` is the "error" of the implicit step with the state ``Y``. @@ -86,6 +86,8 @@ is reached. - `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 +- `approximate_solve_iters::Int`: number of iterations to take for the + approximate linear solve required by `UseDiffusionDerivative()` """ struct ImplicitEquationJacobian{ M <: MatrixFields.FieldMatrix, @@ -120,6 +122,7 @@ function ImplicitEquationJacobian( diffusion_flag, enthalpy_flag, transform_flag, + approximate_solve_iters, ) FT = Spaces.undertype(axes(Y.c)) one_C3xACT3 = C3(FT(1)) * CT3(FT(1))' @@ -137,10 +140,10 @@ function ImplicitEquationJacobian( @name(c.ρq_ice), @name(c.ρq_rai), @name(c.ρq_sno), + @name(c.sgs⁰.ρatke), ) available_tracer_names = MatrixFields.unrolled_filter(is_in_Y, tracer_names) other_names = ( - @name(c.sgs⁰.ρatke), @name(c.sgsʲs), @name(f.sgsʲs), @name(c.turbconv), @@ -185,8 +188,10 @@ function ImplicitEquationJacobian( alg = use_derivative(diffusion_flag) ? - MatrixFields.ApproximateBlockArrowheadIterativeSolve(@name(c)) : - MatrixFields.BlockArrowheadSolve(@name(c)) + MatrixFields.ApproximateBlockArrowheadIterativeSolve( + @name(c); + n_iters = approximate_solve_iters, + ) : MatrixFields.BlockArrowheadSolve(@name(c)) # By default, the ApproximateBlockArrowheadIterativeSolve takes 1 iteration # and approximates the A[c, c] block with a MainDiagonalPreconditioner. @@ -242,7 +247,6 @@ _linsolve!(x, A, b, update_matrix = false; kwargs...) = ldiv!(x, A, b) function Wfact!(A, Y, p, dtγ, t) NVTX.@range "Wfact!" color = colorant"green" begin # Remove unnecessary values from p to avoid allocations in bycolumn. - (; rayleigh_sponge) = p.atmos p′ = (; p.precomputed.ᶜspecific, p.precomputed.ᶠu³, @@ -253,6 +257,11 @@ function Wfact!(A, Y, p, dtγ, t) use_derivative(A.diffusion_flag) ? (; p.precomputed.ᶜK_u, p.precomputed.ᶜK_h) : (;) )..., + ( + use_derivative(A.diffusion_flag) && + p.atmos.turbconv_model isa PrognosticEDMFX ? + (; p.precomputed.ᶜρa⁰) : (;) + )..., p.core.∂ᶜK_∂ᶠu₃, p.core.ᶜΦ, p.core.ᶠgradᵥ_ᶜΦ, @@ -263,7 +272,7 @@ function Wfact!(A, Y, p, dtγ, t) p.params, p.atmos, ( - rayleigh_sponge isa RayleighSponge ? + p.atmos.rayleigh_sponge isa RayleighSponge ? (; p.rayleigh_sponge.ᶠβ_rayleigh_w) : (;) )..., ) @@ -444,6 +453,12 @@ function update_implicit_equation_jacobian!(A, Y, p, dtγ, colidx) end end + # TODO: Move the vertical advection of ρatke into the implicit tendency. + if MatrixFields.has_field(Y, @name(c.sgs⁰.ρatke)) + ∂ᶜρatke_err_∂ᶠu₃ = matrix[@name(c.sgs⁰.ρatke), @name(f.u₃)] + ∂ᶜρatke_err_∂ᶠu₃[colidx] .= (zero(eltype(∂ᶜρatke_err_∂ᶠu₃)),) + end + # We can express the implicit tendency for vertical velocity as # ᶠu₃ₜ = # -(ᶠgradᵥ(ᶜp - ᶜp_ref) + ᶠinterp(ᶜρ - ᶜρ_ref) * ᶠgradᵥ_ᶜΦ) / @@ -549,27 +564,24 @@ function update_implicit_equation_jacobian!(A, Y, p, dtγ, colidx) ) ⋅ ᶠgradᵥ_matrix() - (I,) end - ∂ᶜρe_tot_err_∂ᶜρe_tot = matrix[@name(c.ρe_tot), @name(c.ρe_tot)] - @. ∂ᶜρe_tot_err_∂ᶜρe_tot[colidx] = - dtγ * ᶜadvdivᵥ_matrix() ⋅ - DiagonalMatrixRow(ᶠinterp(ᶜρ[colidx]) * ᶠinterp(p.ᶜK_h[colidx])) ⋅ - ᶠgradᵥ_matrix() ⋅ DiagonalMatrixRow((1 + R_d / cv_d) / ᶜρ[colidx]) - - (I,) - - tracer_names = ( - @name(c.ρq_tot), - @name(c.ρq_liq), - @name(c.ρq_ice), - @name(c.ρq_rai), - @name(c.ρq_sno), + ᶜρ⁰ = p.atmos.turbconv_model isa PrognosticEDMFX ? p.ᶜρa⁰ : ᶜρ + scalar_info = ( + (@name(c.ρe_tot), ᶜρ, p.ᶜK_h, 1 + R_d / cv_d), + (@name(c.ρq_tot), ᶜρ, p.ᶜK_h, 1), + (@name(c.ρq_liq), ᶜρ, p.ᶜK_h, 1), + (@name(c.ρq_ice), ᶜρ, p.ᶜK_h, 1), + (@name(c.ρq_rai), ᶜρ, p.ᶜK_h, 1), + (@name(c.ρq_sno), ᶜρ, p.ᶜK_h, 1), + (@name(c.sgs⁰.ρatke), ᶜρ⁰, p.ᶜK_u, 1), ) - MatrixFields.unrolled_foreach(tracer_names) do ρχ_name + MatrixFields.unrolled_foreach(scalar_info) do (ρχ_name, ᶜρ_χ, ᶜK_χ, s_χ) MatrixFields.has_field(Y, ρχ_name) || return ∂ᶜρχ_err_∂ᶜρχ = matrix[ρχ_name, ρχ_name] @. ∂ᶜρχ_err_∂ᶜρχ[colidx] = dtγ * ᶜadvdivᵥ_matrix() ⋅ DiagonalMatrixRow( - ᶠinterp(ᶜρ[colidx]) * ᶠinterp(p.ᶜK_h[colidx]), - ) ⋅ ᶠgradᵥ_matrix() ⋅ DiagonalMatrixRow(1 / ᶜρ[colidx]) - (I,) + ᶠinterp(ᶜρ_χ[colidx]) * ᶠinterp(ᶜK_χ[colidx]), + ) ⋅ ᶠgradᵥ_matrix() ⋅ DiagonalMatrixRow(s_χ / ᶜρ_χ[colidx]) - + (I,) end end end diff --git a/src/prognostic_equations/implicit/implicit_tendency.jl b/src/prognostic_equations/implicit/implicit_tendency.jl index d3f1617525..177bf6f2a3 100644 --- a/src/prognostic_equations/implicit/implicit_tendency.jl +++ b/src/prognostic_equations/implicit/implicit_tendency.jl @@ -18,6 +18,14 @@ NVTX.@annotate function implicit_tendency!(Yₜ, Y, p, t) colidx, p.atmos.vert_diff, ) + edmfx_sgs_diffusive_flux_tendency!( + Yₜ, + Y, + p, + t, + colidx, + p.atmos.turbconv_model, + ) end # NOTE: All ρa tendencies should be applied before calling this function pressure_work_tendency!(Yₜ, Y, p, t, colidx, p.atmos.turbconv_model) diff --git a/src/prognostic_equations/remaining_tendency.jl b/src/prognostic_equations/remaining_tendency.jl index 5dc0bd581a..021d42cf33 100644 --- a/src/prognostic_equations/remaining_tendency.jl +++ b/src/prognostic_equations/remaining_tendency.jl @@ -32,6 +32,14 @@ NVTX.@annotate function additional_tendency!(Yₜ, Y, p, t) colidx, p.atmos.vert_diff, ) + edmfx_sgs_diffusive_flux_tendency!( + Yₜ, + Y, + p, + t, + colidx, + p.atmos.turbconv_model, + ) end radiation_tendency!(Yₜ, Y, p, t, colidx, p.atmos.radiation_mode) @@ -44,14 +52,6 @@ NVTX.@annotate function additional_tendency!(Yₜ, Y, p, t) colidx, p.atmos.turbconv_model, ) - edmfx_sgs_diffusive_flux_tendency!( - Yₜ, - Y, - p, - t, - colidx, - p.atmos.turbconv_model, - ) edmfx_nh_pressure_tendency!(Yₜ, Y, p, t, colidx, p.atmos.turbconv_model) edmfx_tke_tendency!(Yₜ, Y, p, t, colidx, p.atmos.turbconv_model) precipitation_tendency!(Yₜ, Y, p, t, colidx, p.atmos.precip_model) diff --git a/src/solver/type_getters.jl b/src/solver/type_getters.jl index 5c2e19265d..042de67064 100644 --- a/src/solver/type_getters.jl +++ b/src/solver/type_getters.jl @@ -375,7 +375,7 @@ additional_integrator_kwargs(::CTS.DistributedODEAlgorithm) = (; is_cts_algo(::SciMLBase.AbstractODEAlgorithm) = false is_cts_algo(::CTS.DistributedODEAlgorithm) = true -function jac_kwargs(ode_algo, Y, atmos) +function jac_kwargs(ode_algo, Y, atmos, parsed_args) if is_implicit(ode_algo) A = ImplicitEquationJacobian( Y, @@ -384,6 +384,7 @@ function jac_kwargs(ode_algo, Y, atmos) IgnoreDiffusionDerivative(), IgnoreEnthalpyDerivative(), use_transform(ode_algo), + parsed_args["approximate_linear_solve_iters"], ) if use_transform(ode_algo) return (; jac_prototype = A, Wfact_t = Wfact!) @@ -646,7 +647,7 @@ function args_integrator(parsed_args, Y, p, tspan, ode_algo, callback) func = if parsed_args["split_ode"] implicit_func = SciMLBase.ODEFunction( implicit_tendency!; - jac_kwargs(ode_algo, Y, atmos)..., + jac_kwargs(ode_algo, Y, atmos, parsed_args)..., tgrad = (∂Y∂t, Y, p, t) -> (∂Y∂t .= 0), ) if is_cts_algo(ode_algo)