use mse for prognostic edmf

src/cache/precomputed_quantities.jl

@@ -71,6 +71,7 @@ function precomputed_quantities(Y, atmos)
             ᶜu⁰ = similar(Y.c, C123{FT}),
             ᶠu³⁰ = similar(Y.f, CT3{FT}),
             ᶜK⁰ = similar(Y.c, FT),
+            ᶜmse⁰ = similar(Y.c, FT),
             ᶜh_tot⁰ = similar(Y.c, FT),
             ᶜq_tot⁰ = similar(Y.c, FT),
             ᶜts⁰ = similar(Y.c, TST),
@@ -86,6 +87,7 @@ function precomputed_quantities(Y, atmos)
             ᶜKʲs = similar(Y.c, NTuple{n, FT}),
             ᶜtsʲs = similar(Y.c, NTuple{n, TST}),
             ᶜρʲs = similar(Y.c, NTuple{n, FT}),
+            ᶜh_totʲs = similar(Y.c, NTuple{n, FT}),
             ᶜentrʲs = similar(Y.c, NTuple{n, FT}),
             ᶜdetrʲs = similar(Y.c, NTuple{n, FT}),
         ) : (;)

src/cache/prognostic_edmf_precomputed_quantities.jl

@@ -15,16 +15,16 @@ function set_prognostic_edmf_precomputed_quantities_environment!(Y, p, ᶠuₕ³
     thermo_params = CAP.thermodynamics_params(p.params)
     (; turbconv_model) = p.atmos
     (; ᶜΦ,) = p.core
-    (; ᶜp, ᶜh_tot) = p.precomputed
-    (; ᶜtke⁰, ᶜρa⁰, ᶠu₃⁰, ᶜu⁰, ᶠu³⁰, ᶜK⁰, ᶜts⁰, ᶜρ⁰, ᶜh_tot⁰, ᶜq_tot⁰) =
+    (; ᶜp, ᶜh_tot, ᶜK) = p.precomputed
+    (; ᶜtke⁰, ᶜρa⁰, ᶠu₃⁰, ᶜu⁰, ᶠu³⁰, ᶜK⁰, ᶜts⁰, ᶜρ⁰, ᶜmse⁰, ᶜh_tot⁰, ᶜq_tot⁰) =
         p.precomputed
 
     @. ᶜρa⁰ = ρa⁰(Y.c)
     @. ᶜtke⁰ = divide_by_ρa(Y.c.sgs⁰.ρatke, ᶜρa⁰, 0, Y.c.ρ, turbconv_model)
-    @. ᶜh_tot⁰ = divide_by_ρa(
-        Y.c.ρ * ᶜh_tot - ρah_tot⁺(Y.c.sgsʲs),
+    @. ᶜmse⁰ = divide_by_ρa(
+        Y.c.ρ * (ᶜh_tot - ᶜK) - ρamse⁺(Y.c.sgsʲs),
         ᶜρa⁰,
-        Y.c.ρ * ᶜh_tot,
+        Y.c.ρ * (ᶜh_tot - ᶜK),
         Y.c.ρ,
         turbconv_model,
     )
@@ -38,6 +38,7 @@ function set_prognostic_edmf_precomputed_quantities_environment!(Y, p, ᶠuₕ³
     set_sgs_ᶠu₃!(u₃⁰, ᶠu₃⁰, Y, turbconv_model)
     set_velocity_quantities!(ᶜu⁰, ᶠu³⁰, ᶜK⁰, ᶠu₃⁰, Y.c.uₕ, ᶠuₕ³)
     @. ᶜK⁰ += ᶜtke⁰
+    @. ᶜh_tot⁰ = ᶜmse⁰ + ᶜK⁰
     @. ᶜts⁰ = TD.PhaseEquil_phq(thermo_params, ᶜp, ᶜh_tot⁰ - ᶜK⁰ - ᶜΦ, ᶜq_tot⁰)
     @. ᶜρ⁰ = TD.air_density(thermo_params, ᶜts⁰)
     return nothing
@@ -63,7 +64,7 @@ function set_prognostic_edmf_precomputed_quantities_draft_and_bc!(Y, p, ᶠuₕ
 
     (; ᶜΦ,) = p.core
     (; ᶜspecific, ᶜp, ᶜh_tot) = p.precomputed
-    (; ᶜuʲs, ᶠu³ʲs, ᶜKʲs, ᶜtsʲs, ᶜρʲs) = p.precomputed
+    (; ᶜuʲs, ᶠu³ʲs, ᶜh_totʲs, ᶜKʲs, ᶜtsʲs, ᶜρʲs) = p.precomputed
     (; ustar, obukhov_length, buoyancy_flux) = p.precomputed.sfc_conditions
 
     for j in 1:n
@@ -73,12 +74,13 @@ function set_prognostic_edmf_precomputed_quantities_draft_and_bc!(Y, p, ᶠuₕ
         ᶠu₃ʲ = Y.f.sgsʲs.:($j).u₃
         ᶜtsʲ = ᶜtsʲs.:($j)
         ᶜρʲ = ᶜρʲs.:($j)
-        ᶜh_totʲ = Y.c.sgsʲs.:($j).h_tot
+        ᶜmseʲ = Y.c.sgsʲs.:($j).mse
+        ᶜh_totʲ = ᶜh_totʲs.:($j)
         ᶜq_totʲ = Y.c.sgsʲs.:($j).q_tot
 
         set_velocity_quantities!(ᶜuʲ, ᶠu³ʲ, ᶜKʲ, ᶠu₃ʲ, Y.c.uₕ, ᶠuₕ³)
-        @. ᶜtsʲ =
-            TD.PhaseEquil_phq(thermo_params, ᶜp, ᶜh_totʲ - ᶜKʲ - ᶜΦ, ᶜq_totʲ)
+        @. ᶜh_totʲ = ᶜmseʲ + ᶜKʲ
+        @. ᶜtsʲ = TD.PhaseEquil_phq(thermo_params, ᶜp, ᶜmseʲ - ᶜΦ, ᶜq_totʲ)
         @. ᶜρʲ = TD.air_density(thermo_params, ᶜtsʲ)
 
         # EDMFX boundary condition:
@@ -134,8 +136,10 @@ function set_prognostic_edmf_precomputed_quantities_draft_and_bc!(Y, p, ᶠuₕ
         )
 
         # Then overwrite the prognostic variables at first inetrior point.
+        ᶜmseʲ_int_val = Fields.field_values(Fields.level(ᶜmseʲ, 1))
         ᶜKʲ_int_val = Fields.field_values(Fields.level(ᶜKʲ, 1))
         ᶜΦ_int_val = Fields.field_values(Fields.level(ᶜΦ, 1))
+        @. ᶜmseʲ_int_val = ᶜh_totʲ_int_val - ᶜKʲ_int_val
         ᶜtsʲ_int_val = Fields.field_values(Fields.level(ᶜtsʲ, 1))
         @. ᶜtsʲ_int_val = TD.PhaseEquil_phq(
             thermo_params,
@@ -173,9 +177,7 @@ function set_prognostic_edmf_precomputed_quantities_closures!(Y, p, t)
     FT = eltype(params)
     n = n_mass_flux_subdomains(turbconv_model)
 
-    (; ᶜρ_ref) = p.core
-    (; ᶜspecific, ᶜtke⁰, ᶜu, ᶜp, ᶜρa⁰, ᶜu⁰, ᶠu³⁰, ᶜts⁰, ᶜρ⁰, ᶜq_tot⁰) =
-        p.precomputed
+    (; ᶜtke⁰, ᶜu, ᶜp, ᶜρa⁰, ᶠu³⁰, ᶜts⁰, ᶜρ⁰, ᶜq_tot⁰) = p.precomputed
     (;
         ᶜmixing_length,
         ᶜlinear_buoygrad,

src/initial_conditions/atmos_state.jl

@@ -120,12 +120,11 @@ function turbconv_center_variables(ls, turbconv_model::PrognosticEDMFX, gs_vars)
     a_draft = ls.turbconv_state.draft_area
     sgs⁰ = (; ρatke = ls.ρ * (1 - a_draft) * ls.turbconv_state.tke)
     ρa = ls.ρ * a_draft / n
-    h_tot =
+    mse =
         TD.specific_enthalpy(ls.thermo_params, ls.thermo_state) +
-        norm_sqr(ls.velocity) / 2 +
         CAP.grav(ls.params) * ls.geometry.coordinates.z
     q_tot = TD.total_specific_humidity(ls.thermo_params, ls.thermo_state)
-    sgsʲs = ntuple(_ -> (; ρa = ρa, h_tot = h_tot, q_tot = q_tot), Val(n))
+    sgsʲs = ntuple(_ -> (; ρa = ρa, mse = mse, q_tot = q_tot), Val(n))
     return (; sgs⁰, sgsʲs)
 end
 

src/prognostic_equations/advection.jl

@@ -33,9 +33,9 @@ NVTX.@annotate function horizontal_advection_tendency!(Yₜ, Y, p, t)
 
     if p.atmos.turbconv_model isa PrognosticEDMFX
         for j in 1:n
-            @. Yₜ.c.sgsʲs.:($$j).h_tot -=
-                wdivₕ(Y.c.sgsʲs.:($$j).h_tot * ᶜuʲs.:($$j)) -
-                Y.c.sgsʲs.:($$j).h_tot * wdivₕ(ᶜuʲs.:($$j))
+            @. Yₜ.c.sgsʲs.:($$j).mse -=
+                wdivₕ(Y.c.sgsʲs.:($$j).mse * ᶜuʲs.:($$j)) -
+                Y.c.sgsʲs.:($$j).mse * wdivₕ(ᶜuʲs.:($$j))
         end
     end
 
@@ -75,9 +75,10 @@ NVTX.@annotate function explicit_vertical_advection_tendency!(Yₜ, Y, p, t)
     n = n_prognostic_mass_flux_subdomains(turbconv_model)
     advect_tke = use_prognostic_tke(turbconv_model)
     point_type = eltype(Fields.coordinate_field(Y.c))
+    (; params) = p
     (; dt) = p.simulation
     ᶜJ = Fields.local_geometry_field(Y.c).J
-    (; ᶜf) = p.core
+    (; ᶜf, ᶜΦ) = p.core
     (; ᶜu, ᶠu³, ᶜK) = p.precomputed
     (; edmfx_upwinding) = n > 0 || advect_tke ? p.atmos.numerics : all_nothing
     (; ᶜp, ᶜuʲs, ᶠu³ʲs, ᶜKʲs, ᶜρʲs) = n > 0 ? p.precomputed : all_nothing
@@ -90,6 +91,7 @@ NVTX.@annotate function explicit_vertical_advection_tendency!(Yₜ, Y, p, t)
     ᶜω³ = p.scratch.ᶜtemp_CT3
     ᶠω¹² = p.scratch.ᶠtemp_CT12
     ᶠω¹²ʲs = p.scratch.ᶠtemp_CT12ʲs
+    FT = Spaces.undertype(axes(Y.c))
 
     if point_type <: Geometry.Abstract3DPoint
         @. ᶜω³ = curlₕ(Y.c.uₕ)
@@ -109,6 +111,10 @@ NVTX.@annotate function explicit_vertical_advection_tendency!(Yₜ, Y, p, t)
     end
     # Without the CT12(), the right-hand side would be a CT1 or CT2 in 2D space.
 
+    ᶠz = Fields.coordinate_field(Y.f).z
+    ᶠΦ = p.scratch.ᶠtemp_scalar
+    @. ᶠΦ = CAP.grav(params) * ᶠz
+
     Fields.bycolumn(axes(Y.c)) do colidx
         @. Yₜ.c.uₕ[colidx] -=
             ᶜinterp(
@@ -132,37 +138,41 @@ NVTX.@annotate function explicit_vertical_advection_tendency!(Yₜ, Y, p, t)
                     ᶠinterp(ᶜρʲs.:($$j)[colidx] - Y.c.ρ[colidx]) *
                     ᶠgradᵥ_ᶜΦ[colidx]
                 ) / ᶠinterp(ᶜρʲs.:($$j)[colidx])
+
+            # buoyancy term in mse equation
+            @. Yₜ.c.sgsʲs.:($$j).mse[colidx] +=
+                adjoint(CT3(ᶜinterp(Y.f.sgsʲs.:($$j).u₃[colidx]))) *
+                (ᶜρʲs.:($$j)[colidx] - Y.c.ρ[colidx]) *
+                ᶜgradᵥ(ᶠΦ[colidx]) / ᶜρʲs.:($$j)[colidx]
         end
 
         # TODO: Move this to implicit_vertical_advection_tendency!.
-        if p.atmos.turbconv_model isa PrognosticEDMFX
-            for j in 1:n
-                @. ᶜa_scalar[colidx] =
-                    draft_area(Y.c.sgsʲs.:($$j).ρa[colidx], ᶜρʲs.:($$j)[colidx])
-                vertical_transport!(
-                    Yₜ.c.sgsʲs.:($j).ρa[colidx],
-                    ᶜJ[colidx],
-                    ᶜρʲs.:($j)[colidx],
-                    ᶠu³ʲs.:($j)[colidx],
-                    ᶜa_scalar[colidx],
-                    dt,
-                    edmfx_upwinding,
-                )
-
-                vertical_advection!(
-                    Yₜ.c.sgsʲs.:($j).h_tot[colidx],
-                    ᶠu³ʲs.:($j)[colidx],
-                    Y.c.sgsʲs.:($j).h_tot[colidx],
-                    edmfx_upwinding,
-                )
-
-                vertical_advection!(
-                    Yₜ.c.sgsʲs.:($j).q_tot[colidx],
-                    ᶠu³ʲs.:($j)[colidx],
-                    Y.c.sgsʲs.:($j).q_tot[colidx],
-                    edmfx_upwinding,
-                )
-            end
+        for j in 1:n
+            @. ᶜa_scalar[colidx] =
+                draft_area(Y.c.sgsʲs.:($$j).ρa[colidx], ᶜρʲs.:($$j)[colidx])
+            vertical_transport!(
+                Yₜ.c.sgsʲs.:($j).ρa[colidx],
+                ᶜJ[colidx],
+                ᶜρʲs.:($j)[colidx],
+                ᶠu³ʲs.:($j)[colidx],
+                ᶜa_scalar[colidx],
+                dt,
+                edmfx_upwinding,
+            )
+
+            vertical_advection!(
+                Yₜ.c.sgsʲs.:($j).mse[colidx],
+                ᶠu³ʲs.:($j)[colidx],
+                Y.c.sgsʲs.:($j).mse[colidx],
+                edmfx_upwinding,
+            )
+
+            vertical_advection!(
+                Yₜ.c.sgsʲs.:($j).q_tot[colidx],
+                ᶠu³ʲs.:($j)[colidx],
+                Y.c.sgsʲs.:($j).q_tot[colidx],
+                edmfx_upwinding,
+            )
         end
 
         # TODO: Move this to implicit_vertical_advection_tendency!.

src/prognostic_equations/edmfx_entr_detr.jl

@@ -363,17 +363,17 @@ function edmfx_entr_detr_tendency!(
 
     n = n_mass_flux_subdomains(turbconv_model)
     (; ᶜentrʲs, ᶜdetrʲs) = p.precomputed
-    (; ᶜq_tot⁰, ᶜh_tot⁰, ᶠu₃⁰) = p.precomputed
+    (; ᶜq_tot⁰, ᶜmse⁰, ᶠu₃⁰) = p.precomputed
 
     for j in 1:n
 
         @. Yₜ.c.sgsʲs.:($$j).ρa[colidx] +=
             Y.c.sgsʲs.:($$j).ρa[colidx] *
             (ᶜentrʲs.:($$j)[colidx] - ᶜdetrʲs.:($$j)[colidx])
 
-        @. Yₜ.c.sgsʲs.:($$j).h_tot[colidx] +=
+        @. Yₜ.c.sgsʲs.:($$j).mse[colidx] +=
             ᶜentrʲs.:($$j)[colidx] *
-            (ᶜh_tot⁰[colidx] - Y.c.sgsʲs.:($$j).h_tot[colidx])
+            (ᶜmse⁰[colidx] - Y.c.sgsʲs.:($$j).mse[colidx])
 
         @. Yₜ.c.sgsʲs.:($$j).q_tot[colidx] +=
             ᶜentrʲs.:($$j)[colidx] *

src/prognostic_equations/edmfx_sgs_flux.jl

@@ -16,7 +16,7 @@ function edmfx_sgs_mass_flux_tendency!(
     n = n_mass_flux_subdomains(turbconv_model)
     (; edmfx_sgsflux_upwinding) = p.atmos.numerics
     (; ᶠu³, ᶜh_tot, ᶜspecific) = p.precomputed
-    (; ᶠu³ʲs, ᶜρʲs) = p.precomputed
+    (; ᶠu³ʲs, ᶜρʲs, ᶜh_totʲs) = p.precomputed
     (; ᶜρa⁰, ᶜρ⁰, ᶠu³⁰, ᶜh_tot⁰, ᶜq_tot⁰) = p.precomputed
     (; dt) = p.simulation
     ᶜJ = Fields.local_geometry_field(Y.c).J
@@ -28,7 +28,7 @@ function edmfx_sgs_mass_flux_tendency!(
         for j in 1:n
             @. ᶠu³_diff_colidx = ᶠu³ʲs.:($$j)[colidx] - ᶠu³[colidx]
             @. ᶜa_scalar_colidx =
-                (Y.c.sgsʲs.:($$j).h_tot[colidx] - ᶜh_tot[colidx]) *
+                (ᶜh_totʲs.:($$j)[colidx] - ᶜh_tot[colidx]) *
                 draft_area(Y.c.sgsʲs.:($$j).ρa[colidx], ᶜρʲs.:($$j)[colidx])
             vertical_transport!(
                 Yₜ.c.ρe_tot[colidx],

src/prognostic_equations/hyperdiffusion.jl

@@ -37,7 +37,7 @@ function hyperdiffusion_cache(Y, hyperdiff::ClimaHyperdiffusion, turbconv_model)
             ᶜ∇²tke⁰ = similar(Y.c, FT),
             ᶜ∇²uₕʲs = similar(Y.c, NTuple{n, C12{FT}}),
             ᶜ∇²uᵥʲs = similar(Y.c, NTuple{n, C3{FT}}),
-            ᶜ∇²h_totʲs = similar(Y.c, NTuple{n, FT}),
+            ᶜ∇²mseʲs = similar(Y.c, NTuple{n, FT}),
             ᶜ∇²q_totʲs = similar(Y.c, NTuple{n, FT}),
         ) :
         turbconv_model isa DiagnosticEDMFX ? (; ᶜ∇²tke⁰ = similar(Y.c, FT)) :
@@ -69,7 +69,7 @@ NVTX.@annotate function hyperdiffusion_tendency!(Yₜ, Y, p, t)
     (; ᶜ∇²u, ᶜ∇²specific_energy) = p.hyperdiff
     if turbconv_model isa PrognosticEDMFX
         (; ᶜρa⁰, ᶜtke⁰) = p.precomputed
-        (; ᶜ∇²tke⁰, ᶜ∇²uₕʲs, ᶜ∇²uᵥʲs, ᶜ∇²uʲs, ᶜ∇²h_totʲs) = p.hyperdiff
+        (; ᶜ∇²tke⁰, ᶜ∇²uₕʲs, ᶜ∇²uᵥʲs, ᶜ∇²uʲs, ᶜ∇²mseʲs) = p.hyperdiff
     end
     if turbconv_model isa DiagnosticEDMFX
         (; ᶜtke⁰) = p.precomputed
@@ -99,7 +99,7 @@ NVTX.@annotate function hyperdiffusion_tendency!(Yₜ, Y, p, t)
             @. ᶜ∇²uʲs.:($$j) =
                 C123(wgradₕ(divₕ(p.precomputed.ᶜuʲs.:($$j)))) -
                 C123(wcurlₕ(C123(curlₕ(p.precomputed.ᶜuʲs.:($$j)))))
-            @. ᶜ∇²h_totʲs.:($$j) = wdivₕ(gradₕ(Y.c.sgsʲs.:($$j).h_tot))
+            @. ᶜ∇²mseʲs.:($$j) = wdivₕ(gradₕ(Y.c.sgsʲs.:($$j).mse))
         end
     end
 
@@ -131,7 +131,7 @@ NVTX.@annotate function hyperdiffusion_tendency!(Yₜ, Y, p, t)
                         @. ᶜ∇²uʲs.:($$j) =
                             C123(ᶜ∇²uₕʲs.:($$j)) + C123(ᶜ∇²uᵥʲs.:($$j))
                     end
-                    dss_op!(ᶜ∇²h_totʲs, buffer.ᶜ∇²h_totʲs)
+                    dss_op!(ᶜ∇²mseʲs, buffer.ᶜ∇²mseʲs)
                 end
             end
         end
@@ -160,7 +160,7 @@ NVTX.@annotate function hyperdiffusion_tendency!(Yₜ, Y, p, t)
                     κ₄ * ᶠwinterp(ᶜJ * Y.c.ρ, ᶜ∇²uᵥʲs.:($$j))
             end
             # Note: It is more correct to have ρa inside and outside the divergence
-            @. Yₜ.c.sgsʲs.:($$j).h_tot -= κ₄ * wdivₕ(gradₕ(ᶜ∇²h_totʲs.:($$j)))
+            @. Yₜ.c.sgsʲs.:($$j).mse -= κ₄ * wdivₕ(gradₕ(ᶜ∇²mseʲs.:($$j)))
         end
     end
 

src/utils/variable_manipulations.jl

@@ -188,15 +188,23 @@ mass-flux subdomain states are stored in `gs.sgsʲs`.
 ρa⁺(gs) = mapreduce(sgsʲ -> sgsʲ.ρa, +, gs.sgsʲs)
 
 """
-    ρah_tot⁺(gs)
+    ρah_tot⁺(sgsʲs)
 
 Computes the total mass-flux subdomain area-weighted ρh_tot, assuming that the
 mass-flux subdomain states are stored in `sgsʲs`.
 """
 ρah_tot⁺(sgsʲs) = mapreduce(sgsʲ -> sgsʲ.ρa * sgsʲ.h_tot, +, sgsʲs)
 
 """
-    ρaq_tot⁺(gs)
+    ρamse⁺(sgsʲs)
+
+Computes the total mass-flux subdomain area-weighted ρmse, assuming that the
+mass-flux subdomain states are stored in `sgsʲs`.
+"""
+ρamse⁺(sgsʲs) = mapreduce(sgsʲ -> sgsʲ.ρa * sgsʲ.mse, +, sgsʲs)
+
+"""
+    ρaq_tot⁺(sgsʲs)
 
 Computes the total mass-flux subdomain area-weighted ρq_tot, assuming that the
 mass-flux subdomain states are stored in `sgsʲs`.