From df3139fb410450052cba445d7f73d9fe88588cd9 Mon Sep 17 00:00:00 2001
From: Zhaoyi Shen <11598433+szy21@users.noreply.github.com>
Date: Thu, 19 Oct 2023 13:47:52 -0700
Subject: [PATCH] change buoyancy

---
 src/cache/diagnostic_edmf_precomputed_quantities.jl | 13 +++++++------
 src/cache/prognostic_edmf_precomputed_quantities.jl | 12 ++++++------
 src/prognostic_equations/advection.jl               |  5 +++--
 3 files changed, 16 insertions(+), 14 deletions(-)

diff --git a/src/cache/diagnostic_edmf_precomputed_quantities.jl b/src/cache/diagnostic_edmf_precomputed_quantities.jl
index 61d40de08f..a375409221 100644
--- a/src/cache/diagnostic_edmf_precomputed_quantities.jl
+++ b/src/cache/diagnostic_edmf_precomputed_quantities.jl
@@ -289,6 +289,7 @@ function set_diagnostic_edmf_precomputed_quantities_do_integral!(Y, p, t)
         ∇Φ³_prev_level = Fields.field_values(Fields.level(ᶜ∇Φ³, i - 1))
         ∇Φ³_data_prev_level = ∇Φ³_prev_level.components.data.:1
         ρ_prev_level = Fields.field_values(Fields.level(Y.c.ρ, i - 1))
+        u³_prev_halflevel = Fields.field_values(Fields.level(ᶠu³, i - 1 - half))
         u³⁰_prev_halflevel =
             Fields.field_values(Fields.level(ᶠu³⁰, i - 1 - half))
         u³⁰_data_prev_halflevel = u³⁰_prev_halflevel.components.data.:1
@@ -359,13 +360,13 @@ function set_diagnostic_edmf_precomputed_quantities_do_integral!(Y, p, t)
                     local_geometry_prev_halflevel,
                 ),
                 TD.relative_humidity(thermo_params, tsʲ_prev_level),
-                ᶜphysical_buoyancy(params, ρ_ref_prev_level, ρʲ_prev_level),
+                ᶜphysical_buoyancy(params, ρ_prev_level, ρʲ_prev_level),
                 get_physical_w(
-                    u³⁰_prev_halflevel,
+                    u³_prev_halflevel,
                     local_geometry_prev_halflevel,
                 ),
                 TD.relative_humidity(thermo_params, ts_prev_level),
-                ᶜphysical_buoyancy(params, ρ_ref_prev_level, ρ_prev_level),
+                FT(0),
                 dt,
                 p.atmos.edmfx_entr_model,
             )
@@ -383,13 +384,13 @@ function set_diagnostic_edmf_precomputed_quantities_do_integral!(Y, p, t)
                     local_geometry_prev_halflevel,
                 ),
                 TD.relative_humidity(thermo_params, tsʲ_prev_level),
-                ᶜphysical_buoyancy(params, ρ_ref_prev_level, ρʲ_prev_level),
+                ᶜphysical_buoyancy(params, ρ_prev_level, ρʲ_prev_level),
                 get_physical_w(
-                    u³⁰_prev_halflevel,
+                    u³_prev_halflevel,
                     local_geometry_prev_halflevel,
                 ),
                 TD.relative_humidity(thermo_params, ts_prev_level),
-                ᶜphysical_buoyancy(params, ρ_ref_prev_level, ρ_prev_level),
+                FT(0),
                 dt,
                 p.atmos.edmfx_detr_model,
             )
diff --git a/src/cache/prognostic_edmf_precomputed_quantities.jl b/src/cache/prognostic_edmf_precomputed_quantities.jl
index 537e6c91d1..d17af47771 100644
--- a/src/cache/prognostic_edmf_precomputed_quantities.jl
+++ b/src/cache/prognostic_edmf_precomputed_quantities.jl
@@ -197,10 +197,10 @@ function set_prognostic_edmf_precomputed_quantities_closures!(Y, p, t)
             draft_area(Y.c.sgsʲs.:($$j).ρa, ᶜρʲs.:($$j)),
             get_physical_w(ᶜuʲs.:($$j), ᶜlg),
             TD.relative_humidity(thermo_params, ᶜtsʲs.:($$j)),
-            ᶜphysical_buoyancy(params, ᶜρ_ref, ᶜρʲs.:($$j)),
-            get_physical_w(ᶜu⁰, ᶜlg),
+            ᶜphysical_buoyancy(params, Y.c.ρ, ᶜρʲs.:($$j)),
+            get_physical_w(ᶜu, ᶜlg),
             TD.relative_humidity(thermo_params, ᶜts⁰),
-            ᶜphysical_buoyancy(params, ᶜρ_ref, ᶜρ⁰),
+            FT(0),
             dt,
             p.atmos.edmfx_entr_model,
         )
@@ -214,10 +214,10 @@ function set_prognostic_edmf_precomputed_quantities_closures!(Y, p, t)
             draft_area(Y.c.sgsʲs.:($$j).ρa, ᶜρʲs.:($$j)),
             get_physical_w(ᶜuʲs.:($$j), ᶜlg),
             TD.relative_humidity(thermo_params, ᶜtsʲs.:($$j)),
-            ᶜphysical_buoyancy(params, ᶜρ_ref, ᶜρʲs.:($$j)),
-            get_physical_w(ᶜu⁰, ᶜlg),
+            ᶜphysical_buoyancy(params, Y.c.ρ, ᶜρʲs.:($$j)),
+            get_physical_w(ᶜu, ᶜlg),
             TD.relative_humidity(thermo_params, ᶜts⁰),
-            ᶜphysical_buoyancy(params, ᶜρ_ref, ᶜρ⁰),
+            FT(0),
             dt,
             p.atmos.edmfx_detr_model,
         )
diff --git a/src/prognostic_equations/advection.jl b/src/prognostic_equations/advection.jl
index c3a500edcc..ef96f2659b 100644
--- a/src/prognostic_equations/advection.jl
+++ b/src/prognostic_equations/advection.jl
@@ -123,10 +123,11 @@ NVTX.@annotate function explicit_vertical_advection_tendency!(Yₜ, Y, p, t)
                 ᶠgradᵥ(ᶜKʲs.:($$j)[colidx])
 
             # TODO: Move this to implicit_vertical_advection_tendency!.
+            # For the updraft u_3 equation, we assume the grid-mean to be hydrostatic
+            # and calcuate the buoyancy term relative to the grid-mean density.
             @. Yₜ.f.sgsʲs.:($$j).u₃[colidx] -=
                 (
-                    ᶠgradᵥ(ᶜp[colidx] - ᶜp_ref[colidx]) +
-                    ᶠinterp(ᶜρʲs.:($$j)[colidx] - ᶜρ_ref[colidx]) *
+                    ᶠinterp(ᶜρʲs.:($$j)[colidx] - Y.c.ρ[colidx]) *
                     ᶠgradᵥ_ᶜΦ[colidx]
                 ) / ᶠinterp(ᶜρʲs.:($$j)[colidx])
         end