Merge pull request #2580 from CliMA/ck/2510

Make tracer vert adv always explicit

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] ×

src/prognostic_equations/implicit/implicit_solver.jl

@@ -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),
@@ -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.ρ)]

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