From 2a7090f7eede435472f07c505887907f4a34cfa0 Mon Sep 17 00:00:00 2001 From: Charles Kawczynski Date: Fri, 26 Jan 2024 10:34:56 -0500 Subject: [PATCH] Move the advection of passive tracers to the explicit tendency --- src/prognostic_equations/advection.jl | 32 +++++---- .../implicit/implicit_solver.jl | 65 +++++-------------- .../implicit/implicit_tendency.jl | 42 +++++------- 3 files changed, 55 insertions(+), 84 deletions(-) diff --git a/src/prognostic_equations/advection.jl b/src/prognostic_equations/advection.jl index 4d70c2e733..6469f9b66b 100644 --- a/src/prognostic_equations/advection.jl +++ b/src/prognostic_equations/advection.jl @@ -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 @@ -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], @@ -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] × diff --git a/src/prognostic_equations/implicit/implicit_solver.jl b/src/prognostic_equations/implicit/implicit_solver.jl index 0ab8370533..2ec81de086 100644 --- a/src/prognostic_equations/implicit/implicit_solver.jl +++ b/src/prognostic_equations/implicit/implicit_solver.jl @@ -158,6 +158,7 @@ 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) ? @@ -165,16 +166,16 @@ function ImplicitEquationJacobian( 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), @@ -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, @@ -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)) @@ -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.ρ)] diff --git a/src/prognostic_equations/implicit/implicit_tendency.jl b/src/prognostic_equations/implicit/implicit_tendency.jl index eb355c6d09..582c3b3c69 100644 --- a/src/prognostic_equations/implicit/implicit_tendency.jl +++ b/src/prognostic_equations/implicit/implicit_tendency.jl @@ -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