Inline set_sgs_ᶠu₃ to improve inference

CliMA · Oct 18, 2023 · 70cce53 · 70cce53
1 parent 6b1e657
commit 70cce53
Showing 1 changed file with 11 additions and 1 deletion.
diff --git a/src/cache/edmf_precomputed_quantities.jl b/src/cache/edmf_precomputed_quantities.jl
@@ -43,7 +43,17 @@ function set_edmf_precomputed_quantities!(Y, p, ᶠuₕ³, t)
 
     @. ᶜspecific⁰ = specific_full_sgs⁰(Y.c, turbconv_model)
     @. ᶜρa⁰ = ρa⁰(Y.c)
-    set_sgs_ᶠu₃!(u₃⁰, ᶠu₃⁰, Y, turbconv_model)
+
+    ρaʲs(sgsʲs) = map(sgsʲ -> sgsʲ.ρa, sgsʲs)
+    u₃ʲs(sgsʲs) = map(sgsʲ -> sgsʲ.u₃, sgsʲs)
+    @. ᶠu₃⁰ = u₃⁰(
+        ᶠinterp(ρaʲs(Y.c.sgsʲs)),
+        u₃ʲs(Y.f.sgsʲs),
+        ᶠinterp(Y.c.ρ),
+        Y.f.u₃,
+        turbconv_model,
+    )
+
     set_velocity_quantities!(ᶜu⁰, ᶠu³⁰, ᶜK⁰, ᶠu₃⁰, Y.c.uₕ, ᶠuₕ³)
     @. ᶜK⁰ += ᶜspecific⁰.tke
     @. ᶜts⁰ = ts_sgs(thermo_args..., ᶜspecific⁰, ᶜK⁰, ᶜΦ, ᶜp)