Move more tracers to explicit

CliMA · Jan 26, 2024 · 960224e · 960224e
1 parent f18f43d
commit 960224e
Show file tree

Hide file tree

Showing 3 changed files with 50 additions and 54 deletions.
diff --git a/src/prognostic_equations/advection.jl b/src/prognostic_equations/advection.jl
@@ -88,6 +88,7 @@ NVTX.@annotate function explicit_vertical_advection_tendency!(Yₜ, Y, p, t)
     ᶜρ⁰ = advect_tke ? (n > 0 ? p.precomputed.ᶜρ⁰ : Y.c.ρ) : nothing
     ᶜtke⁰ = advect_tke ? p.precomputed.ᶜtke⁰ : nothing
     ᶜa_scalar = p.scratch.ᶜtemp_scalar
+    ᶜqₚ = p.scratch.ᶜtemp_scalar
     ᶜω³ = p.scratch.ᶜtemp_CT3
     ᶠω¹² = p.scratch.ᶠtemp_CT12
     ᶠω¹²ʲs = p.scratch.ᶠtemp_CT12ʲs
@@ -113,7 +114,15 @@ NVTX.@annotate function explicit_vertical_advection_tendency!(Yₜ, Y, p, t)
 
     ᶠz = Fields.coordinate_field(Y.f).z
     ᶠΦ = p.scratch.ᶠtemp_scalar
+    ᶠu³ₚ = p.scratch.ᶠtemp_CT3
     @. ᶠΦ = CAP.grav(params) * ᶠz
+    lgf = Fields.local_geometry_field(Y.f)
+
+    # TODO: Add support for SetDivergence to DivergenceF2C.
+    ᶜdivᵥ_ρqₚ = Operators.DivergenceF2C(
+        top = Operators.SetValue(C3(FT(0))),
+        # bottom = Operators.SetDivergence(FT(0)),
+    )
 
     Fields.bycolumn(axes(Y.c)) do colidx
         if :ρe_tot in propertynames(Yₜ.c)
@@ -132,6 +141,45 @@ NVTX.@annotate function explicit_vertical_advection_tendency!(Yₜ, Y, p, t)
                 )
             end
         end
+
+        if precip_model isa Microphysics1Moment
+            # Advection of precipitation with the mean flow
+            # is done with other tracers above.
+            # Here we add the advection with precipitation terminal velocity
+            # using first order upwind and free outflow bottom boundary condition
+            @. ᶠu³ₚ[colidx] =
+                FT(-1) *
+                ᶠinterp(p.precomputed.ᶜwᵣ[colidx]) *
+                CT3(unit_basis_vector_data(CT3, lgf[colidx]))
+            @. ᶜqₚ[colidx] = Y.c.ρq_rai[colidx] / Y.c.ρ[colidx]
+            vertical_transport!(
+                Yₜ.c.ρq_rai[colidx],
+                ᶜJ[colidx],
+                Y.c.ρ[colidx],
+                ᶠu³ₚ[colidx],
+                ᶜqₚ[colidx],
+                dt,
+                precip_upwinding,
+                ᶜdivᵥ_ρqₚ,
+            )
+
+            @. ᶠu³ₚ[colidx] =
+                FT(-1) *
+                ᶠinterp(p.precomputed.ᶜwₛ[colidx]) *
+                CT3(unit_basis_vector_data(CT3, lgf[colidx]))
+            @. ᶜqₚ[colidx] = Y.c.ρq_sno[colidx] / Y.c.ρ[colidx]
+            vertical_transport!(
+                Yₜ.c.ρq_sno[colidx],
+                ᶜJ[colidx],
+                Y.c.ρ[colidx],
+                ᶠu³ₚ[colidx],
+                ᶜqₚ[colidx],
+                dt,
+                precip_upwinding,
+                ᶜdivᵥ_ρqₚ,
+            )
+        end
+
         for (ᶜρχₜ, ᶜχ, χ_name) in matching_subfields(Yₜ.c, ᶜspecific)
             χ_name == :e_tot && continue
             vertical_transport!(

diff --git a/src/prognostic_equations/implicit/implicit_solver.jl b/src/prognostic_equations/implicit/implicit_solver.jl
@@ -165,12 +165,12 @@ function ImplicitEquationJacobian(
                 name ->
                     (name, @name(c.uₕ)) =>
                         similar(Y.c, TridiagonalRow_ACTh),
-                (@name(c.ρ), @name(c.ρe_tot), available_tracer_names...),
+                (@name(c.ρ), @name(c.ρe_tot)),
             ) : ()
         )...,
         MatrixFields.unrolled_map(
             name -> (name, @name(f.u₃)) => similar(Y.c, BidiagonalRow_ACT3),
-            (@name(c.ρ), @name(c.ρe_tot), available_tracer_names...),
+            (@name(c.ρ), @name(c.ρe_tot)),
         )...,
         MatrixFields.unrolled_map(
             name -> (@name(f.u₃), name) => similar(Y.f, BidiagonalRow_C3),

diff --git a/src/prognostic_equations/implicit/implicit_tendency.jl b/src/prognostic_equations/implicit/implicit_tendency.jl
@@ -127,7 +127,6 @@ function implicit_vertical_advection_tendency!(Yₜ, Y, p, t, colidx)
     (; ᶜ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)
 
     @. Yₜ.c.ρ[colidx] -=
@@ -146,57 +145,6 @@ function implicit_vertical_advection_tendency!(Yₜ, Y, p, t, colidx)
         )
     end
 
-    if precip_model isa Microphysics1Moment
-        # Advection of precipitation with the mean flow
-        # is done with other tracers above.
-        # Here we add the advection with precipitation terminal velocity
-        # using first order upwind and free outflow bottom boundary condition
-
-        ᶠu³ₚ = p.scratch.ᶠtemp_CT3
-        ᶜqₚ = p.scratch.ᶜtemp_scalar
-        lgf = Fields.local_geometry_field(Y.f)
-        FT = Spaces.undertype(axes(Y.c))
-
-        @. ᶠu³ₚ[colidx] =
-            FT(-1) *
-            ᶠinterp(p.precomputed.ᶜwᵣ[colidx]) *
-            CT3(unit_basis_vector_data(CT3, lgf[colidx]))
-        @. ᶜqₚ[colidx] = Y.c.ρq_rai[colidx] / Y.c.ρ[colidx]
-
-        # TODO: Add support for SetDivergence to DivergenceF2C.
-        ᶜdivᵥ_ρqₚ = Operators.DivergenceF2C(
-            top = Operators.SetValue(C3(FT(0))),
-            # bottom = Operators.SetDivergence(FT(0)),
-        )
-
-        vertical_transport!(
-            Yₜ.c.ρq_rai[colidx],
-            ᶜJ[colidx],
-            Y.c.ρ[colidx],
-            ᶠu³ₚ[colidx],
-            ᶜqₚ[colidx],
-            dt,
-            precip_upwinding,
-            ᶜdivᵥ_ρqₚ,
-        )
-
-        @. ᶠu³ₚ[colidx] =
-            FT(-1) *
-            ᶠinterp(p.precomputed.ᶜwₛ[colidx]) *
-            CT3(unit_basis_vector_data(CT3, lgf[colidx]))
-        @. ᶜqₚ[colidx] = Y.c.ρq_sno[colidx] / Y.c.ρ[colidx]
-        vertical_transport!(
-            Yₜ.c.ρq_sno[colidx],
-            ᶜJ[colidx],
-            Y.c.ρ[colidx],
-            ᶠu³ₚ[colidx],
-            ᶜqₚ[colidx],
-            dt,
-            precip_upwinding,
-            ᶜdivᵥ_ρqₚ,
-        )
-    end
-
     @. Yₜ.f.u₃[colidx] +=
         -(
             ᶠgradᵥ(ᶜp[colidx] - ᶜp_ref[colidx]) +