Use dispatch for pressure_work_tendency

CliMA · Oct 18, 2023 · 3209d1a · 3209d1a
1 parent 065e6a5
commit 3209d1a
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 8 deletions.
diff --git a/src/prognostic_equations/implicit/implicit_tendency.jl b/src/prognostic_equations/implicit/implicit_tendency.jl
@@ -20,9 +20,9 @@ NVTX.@annotate function implicit_tendency!(Yₜ, Y, p, t)
             )
         end
         # NOTE: All ρa tendencies should be applied before calling this function
-        pressure_work_tendency!(Yₜ, Y, p, t, colidx, p.atmos.turbconv_model)
+        pressure_work_tendency!(Yₜ, Y, p, t, colidx)
 
-        # NOTE: This will zero out all monmentum tendencies in the edmfx advection test
+        # NOTE: This will zero out all momentum tendencies in the edmfx advection test
         # please DO NOT add additional velocity tendencies after this function
         zero_velocity_tendency!(Yₜ, Y, p, t, colidx)
     end

diff --git a/src/prognostic_equations/pressure_work.jl b/src/prognostic_equations/pressure_work.jl
@@ -2,16 +2,25 @@
 ##### Pressure work
 #####
 
-function pressure_work_tendency!(Yₜ, Y, p, t, colidx, ::EDMFX)
+function pressure_work_tendency!(Yₜ, Y, p, t, colidx, ::EDMFX, ::TotalEnergy)
     n = n_mass_flux_subdomains(p.atmos.turbconv_model)
     (; ᶜp) = p.precomputed
     for j in 1:n
-        if :ρe_tot in propertynames(Y.c)
-            @. Yₜ.c.sgsʲs.:($$j).ρae_tot[colidx] -=
-                ᶜp[colidx] / Y.c.ρ[colidx] * Yₜ.c.sgsʲs.:($$j).ρa[colidx]
-        end
+        @. Yₜ.c.sgsʲs.:($$j).ρae_tot[colidx] -=
+            ᶜp[colidx] / Y.c.ρ[colidx] * Yₜ.c.sgsʲs.:($$j).ρa[colidx]
     end
     return nothing
 end
 
-pressure_work_tendency!(Yₜ, Y, p, t, colidx, ::Any) = nothing
+pressure_work_tendency!(Yₜ, Y, p, t, colidx) = pressure_work_tendency!(
+    Yₜ,
+    Y,
+    p,
+    t,
+    colidx,
+    p.atmos.turbconv_model,
+    p.atmos.energy_form,
+)
+
+pressure_work_tendency!(Yₜ, Y, p, t, colidx, turbconv_model, energy_form) =
+    nothing