Put set_velocity_quantities back in, leverage funcs

Author: charleskawczynski

src/cache/precomputed_quantities.jl

@@ -237,12 +237,33 @@ function set_velocity_at_top!(Y, turbconv_model)
     return nothing
 end
 
+"""
+    compute_ᶜu(ᶜuₕ, ᶠu₃)
+
+Computes the cell-centered velocity, given
+
+ - `ᶜuₕ` the cell-centered horizontal velocity
+ - `ᶠu₃` the cell-face `Covariant3` velocity
+"""
+compute_ᶜu(ᶜuₕ, ᶠu₃) = @lazy @. C123(ᶜuₕ) + ᶜinterp(C123(ᶠu₃))
+
+"""
+    compute_ᶠu³(ᶠuₕ³, ᶠu₃)
+
+Computes the cell-face contravariant velocity, given
+
+ - `ᶠuₕ³` the cell-face horizontal `Contravariant3` velocity
+ - `ᶠu₃` the cell-face `Covariant3` velocity
+"""
+compute_ᶠu³(ᶠuₕ³, ᶠu₃) = ᶠuₕ³ + CT3(ᶠu₃)
+
 # This is used to set the grid-scale velocity quantities ᶜu, ᶠu³, ᶜK based on
 # ᶠu₃, and it is also used to set the SGS quantities based on ᶠu₃⁰ and ᶠu₃ʲ.
-compute_ᶜu(ᶜuₕ, ᶠu₃) = @lazy @. C123(ᶜuₕ) + ᶜinterp(C123(ᶠu₃))
-function compute_ᶠu³(ᶜuₕ, ᶜρ, ᶠu₃)
-    ᶠuₕ³ = compute_ᶠuₕ³(ᶜuₕ, ᶜρ)
-    return @lazy @. ᶠuₕ³ + CT3(ᶠu₃)
+function set_velocity_quantities!(ᶜu, ᶠu³, ᶜK, ᶠu₃, ᶜuₕ, ᶜρ)
+    ᶜu .= compute_ᶜu(ᶜuₕ, ᶠu₃)
+    ᶠu³ .= compute_ᶠu³.(compute_ᶠuₕ³(ᶜuₕ, ᶜρ), ᶠu₃)
+    ᶜK .= compute_kinetic(ᶜuₕ, ᶠu₃)
+    return nothing
 end
 
 function set_sgs_ᶠu₃!(w_function, ᶠu₃, Y, turbconv_model)
@@ -390,9 +411,7 @@ NVTX.@annotate function set_precomputed_quantities!(Y, p, t)
     set_velocity_at_surface!(Y, ᶠuₕ³, turbconv_model)
     set_velocity_at_top!(Y, turbconv_model)
 
-    ᶠu³ .= compute_ᶠu³(ᶜuₕ, ᶜρ, ᶠu₃)
-    ᶜu .= compute_ᶜu(ᶜuₕ, ᶠu₃)
-    ᶜK .= compute_kinetic(ᶜuₕ, ᶠu₃)
+    set_velocity_quantities!(ᶜu, ᶠu³, ᶜK, Y.f.u₃, Y.c.uₕ, ᶜρ)
 
     if n > 0
         # TODO: In the following increments to ᶜK, we actually need to add

src/cache/prognostic_edmf_precomputed_quantities.jl

@@ -44,9 +44,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_environment!(
         turbconv_model,
     )
     set_sgs_ᶠu₃!(u₃⁰, ᶠu₃⁰, Y, turbconv_model)
-    ᶠu³⁰ .= compute_ᶠu³(ᶜuₕ, ᶜρ, ᶠu₃)
-    ᶜu⁰ .= compute_ᶜu(ᶜuₕ, ᶠu₃⁰)
-    ᶜK⁰ .= compute_kinetic(ᶜuₕ, ᶠu₃⁰)
+    set_velocity_quantities!(ᶜu⁰, ᶠu³⁰, ᶜK⁰, ᶠu₃⁰, Y.c.uₕ, ᶜρ)
 
     # @. ᶜK⁰ += ᶜtke⁰
     @. ᶜts⁰ = TD.PhaseEquil_phq(thermo_params, ᶜp, ᶜmse⁰ - ᶜΦ, ᶜq_tot⁰)
@@ -95,9 +93,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_draft_and_bc!
         ᶜmseʲ = Y.c.sgsʲs.:($j).mse
         ᶜq_totʲ = Y.c.sgsʲs.:($j).q_tot
 
-        ᶠu³ʲ .= compute_ᶠu³(ᶜuₕ, ᶜρ, ᶠu₃)
-        ᶜuʲ .= compute_ᶜu(ᶜuₕ, ᶠu₃ʲ)
-        ᶜKʲ .= compute_kinetic(ᶜuₕ, ᶠu₃ʲ)
+        set_velocity_quantities!(ᶜuʲ, ᶠu³ʲ, ᶜKʲ, ᶠu₃ʲ, Y.c.uₕ, ᶜρ)
 
         @. ᶠKᵥʲ = (adjoint(CT3(ᶠu₃ʲ)) * ᶠu₃ʲ) / 2
         @. ᶜtsʲ = TD.PhaseEquil_phq(thermo_params, ᶜp, ᶜmseʲ - ᶜΦ, ᶜq_totʲ)