Add 1M cloud condensate and precip sedimentation in prognostic EDMF

trontrytel · trontrytel · commit cc50357e1a57 · 2025-07-18T14:05:16.000-07:00
diff --git a/.buildkite/Manifest-v1.11.toml b/.buildkite/Manifest-v1.11.toml
@@ -381,9 +381,9 @@ weakdeps = ["CUDA", "MPI"]
 
 [[deps.ClimaCore]]
 deps = ["Adapt", "BandedMatrices", "BlockArrays", "ClimaComms", "CubedSphere", "DataStructures", "ForwardDiff", "GaussQuadrature", "GilbertCurves", "HDF5", "InteractiveUtils", "IntervalSets", "KrylovKit", "LazyBroadcast", "LinearAlgebra", "MultiBroadcastFusion", "NVTX", "PkgVersion", "RecursiveArrayTools", "RootSolvers", "SparseArrays", "StaticArrays", "Statistics", "UnrolledUtilities"]
-git-tree-sha1 = "14d3d5810ce1e3c990450a2ce7abc6a1e162855f"
+git-tree-sha1 = "a273452127dbb052f2963e3c1095730a996d49a6"
 uuid = "d414da3d-4745-48bb-8d80-42e94e092884"
-version = "0.14.35"
+version = "0.14.36"
 weakdeps = ["CUDA", "Krylov"]
 
     [deps.ClimaCore.extensions]
diff --git a/Project.toml b/Project.toml
@@ -43,7 +43,7 @@ ArgParse = "1"
 Artifacts = "1"
 AtmosphericProfilesLibrary = "0.1.7"
 ClimaComms = "0.6.8"
-ClimaCore = "0.14.35"
+ClimaCore = "0.14.36"
 ClimaDiagnostics = "0.2.12"
 ClimaInterpolations = "0.1.0"
 ClimaParams = "0.10.35"
diff --git a/post_processing/ci_plots.jl b/post_processing/ci_plots.jl
@@ -1400,6 +1400,7 @@ EDMFBoxPlots = Union{
 EDMFBoxPlotsWithPrecip = Union{
     Val{:prognostic_edmfx_rico_column},
     Val{:prognostic_edmfx_trmm_column},
+    Val{:prognostic_edmfx_dycoms_rf02_column},
 }
 """
     plot_edmf_vert_profile!(grid_loc, var_group)
diff --git a/src/cache/precomputed_quantities.jl b/src/cache/precomputed_quantities.jl
@@ -165,6 +165,12 @@ function precomputed_quantities(Y, atmos)
             ᶜSqᵢᵖʲs = similar(Y.c, NTuple{n, FT}),
             ᶜSqᵣᵖʲs = similar(Y.c, NTuple{n, FT}),
             ᶜSqₛᵖʲs = similar(Y.c, NTuple{n, FT}),
+            ᶜwₗʲs = similar(Y.c, NTuple{n, FT}),
+            ᶜwᵢʲs = similar(Y.c, NTuple{n, FT}),
+            ᶜwᵣʲs = similar(Y.c, NTuple{n, FT}),
+            ᶜwₛʲs = similar(Y.c, NTuple{n, FT}),
+            ᶜwₜʲs = similar(Y.c, NTuple{n, FT}),
+            ᶜwₕʲs = similar(Y.c, NTuple{n, FT}),
             ᶜSqₗᵖ⁰ = similar(Y.c, FT),
             ᶜSqᵢᵖ⁰ = similar(Y.c, FT),
             ᶜSqᵣᵖ⁰ = similar(Y.c, FT),
diff --git a/src/cache/prognostic_edmf_precomputed_quantities.jl b/src/cache/prognostic_edmf_precomputed_quantities.jl
@@ -608,13 +608,65 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_precipitation
     (; ᶜSqₗᵖ⁰, ᶜSqᵢᵖ⁰, ᶜSqᵣᵖ⁰, ᶜSqₛᵖ⁰, ᶜρ⁰, ᶜts⁰) = p.precomputed
     (; ᶜq_tot⁰, ᶜq_liq⁰, ᶜq_ice⁰, ᶜq_rai⁰, ᶜq_sno⁰) = p.precomputed
 
+    (; ᶜwₗʲs, ᶜwᵢʲs, ᶜwᵣʲs, ᶜwₛʲs, ᶜwₜʲs, ᶜwₕʲs) = p.precomputed
+
     # TODO - can I re-use them between js and env?
     ᶜSᵖ = p.scratch.ᶜtemp_scalar
     ᶜSᵖ_snow = p.scratch.ᶜtemp_scalar_2
 
     n = n_mass_flux_subdomains(p.atmos.turbconv_model)
+    FT = eltype(params)
 
     for j in 1:n
+
+        # compute terminal velocity for precipitation
+        @. ᶜwᵣʲs.:($$j) = CM1.terminal_velocity(
+            cmp.pr,
+            cmp.tv.rain,
+            ᶜρʲs.:($$j),
+            max(zero(Y.c.ρ), Y.c.sgsʲs.:($$j).q_rai),
+        )
+        @. ᶜwₛʲs.:($$j) = CM1.terminal_velocity(
+            cmp.ps,
+            cmp.tv.snow,
+            ᶜρʲs.:($$j),
+            max(zero(Y.c.ρ), Y.c.sgsʲs.:($$j).q_sno),
+        )
+        # compute sedimentation velocity for cloud condensate [m/s]
+        @. ᶜwₗʲs.:($$j) = CMNe.terminal_velocity(
+            cmc.liquid,
+            cmc.Ch2022.rain,
+            ᶜρʲs.:($$j),
+            max(zero(Y.c.ρ), Y.c.sgsʲs.:($$j).q_liq),
+        )
+        @. ᶜwᵢʲs.:($$j) = CMNe.terminal_velocity(
+            cmc.ice,
+            cmc.Ch2022.small_ice,
+            ᶜρʲs.:($$j),
+            max(zero(Y.c.ρ), Y.c.sgsʲs.:($$j).q_ice),
+        )
+        # compute their contirbutions to energy and total water advection
+        @. ᶜwₜʲs.:($$j) = ifelse(
+            Y.c.sgsʲs.:($$j).ρa * Y.c.sgsʲs.:($$j).q_tot > FT(0),
+            (
+                ᶜwₗʲs.:($$j) * Y.c.sgsʲs.:($$j).q_liq +
+                ᶜwᵢʲs.:($$j) * Y.c.sgsʲs.:($$j).q_ice +
+                ᶜwᵣʲs.:($$j) * Y.c.sgsʲs.:($$j).q_rai +
+                ᶜwₛʲs.:($$j) * Y.c.sgsʲs.:($$j).q_sno
+            ) / Y.c.sgsʲs.:($$j).ρa / Y.c.sgsʲs.:($$j).q_tot,
+            FT(0)
+        )
+        @. ᶜwₕʲs.:($$j) = ifelse(
+            Y.c.sgsʲs.:($$j).ρa * abs(Y.c.sgsʲs.:($$j).mse) > FT(0),
+            (
+                ᶜwₗʲs.:($$j) * Y.c.sgsʲs.:($$j).q_liq * (Iₗ(thp, ᶜtsʲs.:($$j)) + ᶜΦ) +
+                ᶜwᵢʲs.:($$j) * Y.c.sgsʲs.:($$j).q_ice * (Iᵢ(thp, ᶜtsʲs.:($$j)) + ᶜΦ) +
+                ᶜwᵣʲs.:($$j) * Y.c.sgsʲs.:($$j).q_rai * (Iₗ(thp, ᶜtsʲs.:($$j)) + ᶜΦ) +
+                ᶜwₛʲs.:($$j) * Y.c.sgsʲs.:($$j).q_sno * (Iᵢ(thp, ᶜtsʲs.:($$j)) + ᶜΦ)
+            ) / Y.c.sgsʲs.:($$j).ρa / abs(Y.c.sgsʲs.:($$j).mse),
+            FT(0)
+        )
+
         # Precipitation sources and sinks from the updrafts
         compute_precipitation_sources!(
             ᶜSᵖ,
diff --git a/src/prognostic_equations/advection.jl b/src/prognostic_equations/advection.jl
@@ -10,7 +10,7 @@ import ClimaCore.Geometry as Geometry
     horizontal_dynamics_tendency!(Yₜ, Y, p, t)
 
 Computes tendencies due to horizontal advection for prognostic variables of the
-grid mean and EDMFX subdomains, and also applies horizontal pressure gradient and 
+grid mean and EDMFX subdomains, and also applies horizontal pressure gradient and
 gravitational acceleration terms for horizontal momentum.
 
 Specifically, this function calculates:
@@ -339,6 +339,7 @@ function edmfx_sgs_vertical_advection_tendency!(
     end
 
     for j in 1:n
+        # Flux form vertical advection of area farction with the grid mean velocity
         @. ᶜa_scalar = draft_area(Y.c.sgsʲs.:($$j).ρa, ᶜρʲs.:($$j))
         vtt = vertical_transport(
             ᶜρʲs.:($j),
@@ -349,48 +350,138 @@ function edmfx_sgs_vertical_advection_tendency!(
         )
         @. Yₜ.c.sgsʲs.:($$j).ρa += vtt
 
+        # Advective form advection of mse and q_tot with the grid mean velocity
         va = vertical_advection(
             ᶠu³ʲs.:($j),
             Y.c.sgsʲs.:($j).mse,
             edmfx_upwinding,
         )
         @. Yₜ.c.sgsʲs.:($$j).mse += va
-
         va = vertical_advection(
             ᶠu³ʲs.:($j),
             Y.c.sgsʲs.:($j).q_tot,
             edmfx_upwinding,
         )
         @. Yₜ.c.sgsʲs.:($$j).q_tot += va
+
         if p.atmos.moisture_model isa NonEquilMoistModel &&
            p.atmos.microphysics_model isa Microphysics1Moment
-            # TODO - add precipitation terminal velocity
-            # TODO - add cloud sedimentation velocity
-            # TODO - add their contributions to mean energy and mass
-            va = vertical_advection(
-                ᶠu³ʲs.:($j),
-                Y.c.sgsʲs.:($j).q_liq,
-                edmfx_upwinding,
-            )
+           # TODO - add precipitation terminal velocity in implicit solver/tendency with if/else
+           # TODO - add cloud sedimentation velocity in implicit solver/tendency with if/else
+
+            (; ᶜwₗʲs, ᶜwᵢʲs, ᶜwᵣʲs, ᶜwₛʲs, ᶜwₜʲs, ᶜwₕʲs) = p.precomputed
+
+            # Advective form advection of moisture tracers with the grid mean velocity
+            va = vertical_advection(ᶠu³ʲs.:($j), Y.c.sgsʲs.:($j).q_liq, edmfx_upwinding)
             @. Yₜ.c.sgsʲs.:($$j).q_liq += va
-            va = vertical_advection(
-                ᶠu³ʲs.:($j),
-                Y.c.sgsʲs.:($j).q_ice,
-                edmfx_upwinding,
-            )
+            va = vertical_advection(ᶠu³ʲs.:($j), Y.c.sgsʲs.:($j).q_ice, edmfx_upwinding)
             @. Yₜ.c.sgsʲs.:($$j).q_ice += va
-            va = vertical_advection(
-                ᶠu³ʲs.:($j),
-                Y.c.sgsʲs.:($j).q_rai,
-                edmfx_upwinding,
-            )
+            va = vertical_advection(ᶠu³ʲs.:($j), Y.c.sgsʲs.:($j).q_rai, edmfx_upwinding)
             @. Yₜ.c.sgsʲs.:($$j).q_rai += va
-            va = vertical_advection(
-                ᶠu³ʲs.:($j),
-                Y.c.sgsʲs.:($j).q_sno,
-                edmfx_upwinding,
-            )
+            va = vertical_advection(ᶠu³ʲs.:($j), Y.c.sgsʲs.:($j).q_sno, edmfx_upwinding)
             @. Yₜ.c.sgsʲs.:($$j).q_sno += va
+
+            FT = eltype(params)
+            (; ᶜwₗʲs, ᶜwᵢʲs, ᶜwᵣʲs, ᶜwₛʲs, ᶜtsʲs) = p.precomputed
+
+            ᶠwₗ³ʲs = (@. lazy( CT3(ᶠinterp(Geometry.WVector(-1 * ᶜwₗʲs.:($$j))))))
+            ᶠwᵢ³ʲs = (@. lazy( CT3(ᶠinterp(Geometry.WVector(-1 * ᶜwᵢʲs.:($$j))))))
+            ᶠwᵣ³ʲs = (@. lazy( CT3(ᶠinterp(Geometry.WVector(-1 * ᶜwᵣʲs.:($$j))))))
+            ᶠwₛ³ʲs = (@. lazy( CT3(ᶠinterp(Geometry.WVector(-1 * ᶜwₛʲs.:($$j))))))
+
+            ᶜa = (@. lazy(draft_area(Y.c.sgsʲs.:($$j).ρa, ᶜρʲs.:($$j))))
+
+            # Flux form vertical advection of rho * area with sedimentation contributions
+            vtt = vertical_transport(ᶜρʲs.:($j), ᶠwₗ³ʲs, (@. lazy(ᶜa * Y.c.sgsʲs.:($$j).q_liq)), dt, edmfx_upwinding)
+            @. Yₜ.c.sgsʲs.:($$j).ρa += vtt
+            vtt = vertical_transport(ᶜρʲs.:($j), ᶠwᵢ³ʲs, (@. lazy(ᶜa * Y.c.sgsʲs.:($$j).q_ice)), dt, edmfx_upwinding)
+            @. Yₜ.c.sgsʲs.:($$j).ρa += vtt
+            vtt = vertical_transport(ᶜρʲs.:($j), ᶠwᵣ³ʲs, (@. lazy(ᶜa * Y.c.sgsʲs.:($$j).q_rai)), dt, edmfx_upwinding)
+            @. Yₜ.c.sgsʲs.:($$j).ρa += vtt
+            vtt = vertical_transport(ᶜρʲs.:($j), ᶠwₛ³ʲs, (@. lazy(ᶜa * Y.c.sgsʲs.:($$j).q_sno)), dt, edmfx_upwinding)
+            @. Yₜ.c.sgsʲs.:($$j).ρa += vtt
+
+            # q_tot and moisture tracers advective form advection with sedimentation velocity
+            va = vertical_advection(ᶠwₗ³ʲs, Y.c.sgsʲs.:($j).q_liq, edmfx_upwinding)
+            @. Yₜ.c.sgsʲs.:($$j).q_tot += (1 - Y.c.sgsʲs.:($$j).q_tot) * va
+            @. Yₜ.c.sgsʲs.:($$j).q_liq += va
+            va = vertical_advection(ᶠwᵢ³ʲs, Y.c.sgsʲs.:($j).q_ice, edmfx_upwinding)
+            @. Yₜ.c.sgsʲs.:($$j).q_tot += (1 - Y.c.sgsʲs.:($$j).q_tot) * va
+            @. Yₜ.c.sgsʲs.:($$j).q_ice += va
+            va = vertical_advection(ᶠwᵣ³ʲs, Y.c.sgsʲs.:($j).q_rai, edmfx_upwinding)
+            @. Yₜ.c.sgsʲs.:($$j).q_tot += (1 - Y.c.sgsʲs.:($$j).q_tot) * va
+            @. Yₜ.c.sgsʲs.:($$j).q_rai += va
+            va = vertical_advection(ᶠwₛ³ʲs, Y.c.sgsʲs.:($j).q_sno, edmfx_upwinding)
+            @. Yₜ.c.sgsʲs.:($$j).q_tot += (1 - Y.c.sgsʲs.:($$j).q_tot) * va
+            @. Yₜ.c.sgsʲs.:($$j).q_sno += va
+
+            # mse advective form advection with sedimentation velocity
+            (; ᶜΦ) = p.core
+            thp = CAP.thermodynamics_params(params)
+            ᶜmseₗ = (@. lazy(Y.c.sgsʲs.:($$j).q_liq * (TD.internal_energy_liquid(thp, ᶜtsʲs.:($$j)) + TD.gas_constant_air(thp, ᶜtsʲs.:($$j)) * TD.air_temperature(thp, ᶜtsʲs.:($$j)) + ᶜΦ)))
+            ᶜmseᵢ = (@. lazy(Y.c.sgsʲs.:($$j).q_ice * (TD.internal_energy_ice(thp,    ᶜtsʲs.:($$j)) + TD.gas_constant_air(thp, ᶜtsʲs.:($$j)) * TD.air_temperature(thp, ᶜtsʲs.:($$j)) + ᶜΦ)))
+            ᶜmseᵣ = (@. lazy(Y.c.sgsʲs.:($$j).q_rai * (TD.internal_energy_liquid(thp, ᶜtsʲs.:($$j)) + TD.gas_constant_air(thp, ᶜtsʲs.:($$j)) * TD.air_temperature(thp, ᶜtsʲs.:($$j)) + ᶜΦ)))
+            ᶜmseₛ = (@. lazy(Y.c.sgsʲs.:($$j).q_sno * (TD.internal_energy_ice(thp,    ᶜtsʲs.:($$j)) + TD.gas_constant_air(thp, ᶜtsʲs.:($$j)) * TD.air_temperature(thp, ᶜtsʲs.:($$j)) + ᶜΦ)))
+
+            # RICO works well but we are missing mse * q_ advection term
+            va = vertical_advection(ᶠwₗ³ʲs, ᶜmseₗ, edmfx_upwinding)
+            @. Yₜ.c.sgsʲs.:($$j).mse += va
+            va = vertical_advection(ᶠwᵢ³ʲs, ᶜmseᵢ, edmfx_upwinding)
+            @. Yₜ.c.sgsʲs.:($$j).mse += va
+            va = vertical_advection(ᶠwᵣ³ʲs, ᶜmseᵣ, edmfx_upwinding)
+            @. Yₜ.c.sgsʲs.:($$j).mse += va
+            va = vertical_advection(ᶠwₛ³ʲs, ᶜmseₛ, edmfx_upwinding)
+            @. Yₜ.c.sgsʲs.:($$j).mse += va
+#=
+            va = vertical_advection(ᶠwₗ³ʲs, Y.c.sgsʲs.:($j).q_liq, edmfx_upwinding)
+            @. Yₜ.c.sgsʲs.:($$j).mse -= Y.c.sgsʲs.:($$j).mse * va
+            va = vertical_advection(ᶠwᵢ³ʲs, Y.c.sgsʲs.:($j).q_ice, edmfx_upwinding)
+            @. Yₜ.c.sgsʲs.:($$j).mse -= Y.c.sgsʲs.:($$j).mse * va
+            va = vertical_advection(ᶠwᵣ³ʲs, Y.c.sgsʲs.:($j).q_rai, edmfx_upwinding)
+            @. Yₜ.c.sgsʲs.:($$j).mse -= Y.c.sgsʲs.:($$j).mse * va
+            va = vertical_advection(ᶠwₛ³ʲs, Y.c.sgsʲs.:($j).q_sno, edmfx_upwinding)
+            @. Yₜ.c.sgsʲs.:($$j).mse -= Y.c.sgsʲs.:($$j).mse * va
+=#
+
+            # mse, q_tot and moisture tracers terms proportional to 1/ρ̂ ∂zρ̂
+            ᶜinv_ρ̂ = (@. lazy(divide_by_ρa(FT(1), Y.c.sgsʲs.:($$j).ρa, FT(0), Y.c.ρ, turbconv_model)))
+            ᶜ∂ρ̂∂zₗ = (@. lazy(upwind_biased_grad(-1 * Geometry.WVector(ᶜwₗʲs.:($$j)), Y.c.sgsʲs.:($$j).ρa)))
+            ᶜ∂ρ̂∂zᵢ = (@. lazy(upwind_biased_grad(-1 * Geometry.WVector(ᶜwᵢʲs.:($$j)), Y.c.sgsʲs.:($$j).ρa)))
+            ᶜ∂ρ̂∂zᵣ = (@. lazy(upwind_biased_grad(-1 * Geometry.WVector(ᶜwᵣʲs.:($$j)), Y.c.sgsʲs.:($$j).ρa)))
+            ᶜ∂ρ̂∂zₛ = (@. lazy(upwind_biased_grad(-1 * Geometry.WVector(ᶜwₛʲs.:($$j)), Y.c.sgsʲs.:($$j).ρa)))
+
+            @. Yₜ.c.sgsʲs.:($$j).q_tot -= dot(ᶜinv_ρ̂ * ᶜ∂ρ̂∂zₗ, CT3(Geometry.WVector(-1 * ᶜwₗʲs.:($$j)))) * Y.c.sgsʲs.:($$j).q_liq * (1 - Y.c.sgsʲs.:($$j).q_tot)
+            @. Yₜ.c.sgsʲs.:($$j).q_liq -= dot(ᶜinv_ρ̂ * ᶜ∂ρ̂∂zₗ, CT3(Geometry.WVector(-1 * ᶜwₗʲs.:($$j)))) * Y.c.sgsʲs.:($$j).q_liq
+            @. Yₜ.c.sgsʲs.:($$j).mse   -= dot(ᶜinv_ρ̂ * ᶜ∂ρ̂∂zₗ, CT3(Geometry.WVector(-1 * ᶜwₗʲs.:($$j)))) * (Y.c.sgsʲs.:($$j).q_liq * Y.c.sgsʲs.:($$j).mse - ᶜmseₗ)
+
+            @. Yₜ.c.sgsʲs.:($$j).q_tot -= dot(ᶜinv_ρ̂ * ᶜ∂ρ̂∂zᵢ, CT3(Geometry.WVector(-1 * ᶜwᵢʲs.:($$j)))) * Y.c.sgsʲs.:($$j).q_ice * (1 - Y.c.sgsʲs.:($$j).q_tot)
+            @. Yₜ.c.sgsʲs.:($$j).q_ice -= dot(ᶜinv_ρ̂ * ᶜ∂ρ̂∂zᵢ, CT3(Geometry.WVector(-1 * ᶜwᵢʲs.:($$j)))) * Y.c.sgsʲs.:($$j).q_ice
+            @. Yₜ.c.sgsʲs.:($$j).mse   -= dot(ᶜinv_ρ̂ * ᶜ∂ρ̂∂zᵢ, CT3(Geometry.WVector(-1 * ᶜwᵢʲs.:($$j)))) * (Y.c.sgsʲs.:($$j).q_ice * Y.c.sgsʲs.:($$j).mse - ᶜmseᵢ)
+
+            @. Yₜ.c.sgsʲs.:($$j).q_tot -= dot(ᶜinv_ρ̂ * ᶜ∂ρ̂∂zᵣ, CT3(Geometry.WVector(-1 * ᶜwᵣʲs.:($$j)))) * Y.c.sgsʲs.:($$j).q_rai * (1 - Y.c.sgsʲs.:($$j).q_tot)
+            @. Yₜ.c.sgsʲs.:($$j).q_rai -= dot(ᶜinv_ρ̂ * ᶜ∂ρ̂∂zᵣ, CT3(Geometry.WVector(-1 * ᶜwᵣʲs.:($$j)))) * Y.c.sgsʲs.:($$j).q_rai
+            @. Yₜ.c.sgsʲs.:($$j).mse   -= dot(ᶜinv_ρ̂ * ᶜ∂ρ̂∂zᵣ, CT3(Geometry.WVector(-1 * ᶜwᵣʲs.:($$j)))) * (Y.c.sgsʲs.:($$j).q_rai * Y.c.sgsʲs.:($$j).mse - ᶜmseᵣ)
+
+            @. Yₜ.c.sgsʲs.:($$j).q_tot -= dot(ᶜinv_ρ̂ * ᶜ∂ρ̂∂zₛ, CT3(Geometry.WVector(-1 * ᶜwₛʲs.:($$j)))) * Y.c.sgsʲs.:($$j).q_sno * (1 - Y.c.sgsʲs.:($$j).q_tot)
+            @. Yₜ.c.sgsʲs.:($$j).q_sno -= dot(ᶜinv_ρ̂ * ᶜ∂ρ̂∂zₛ, CT3(Geometry.WVector(-1 * ᶜwₛʲs.:($$j)))) * Y.c.sgsʲs.:($$j).q_sno
+            @. Yₜ.c.sgsʲs.:($$j).mse   -= dot(ᶜinv_ρ̂ * ᶜ∂ρ̂∂zₛ, CT3(Geometry.WVector(-1 * ᶜwₛʲs.:($$j)))) * (Y.c.sgsʲs.:($$j).q_sno * Y.c.sgsʲs.:($$j).mse - ᶜmseₛ)
+
+            # mse, q_tot and moisture tracer terms proportional to velocity gradients
+            @. Yₜ.c.sgsʲs.:($$j).q_tot -= ᶜdivᵥ(ᶠwₗ³ʲs) * Y.c.sgsʲs.:($$j).q_liq * (1 - Y.c.sgsʲs.:($$j).q_tot)
+            @. Yₜ.c.sgsʲs.:($$j).q_liq -= ᶜdivᵥ(ᶠwₗ³ʲs) * Y.c.sgsʲs.:($$j).q_liq
+            @. Yₜ.c.sgsʲs.:($$j).mse   -= ᶜdivᵥ(ᶠwₗ³ʲs) * (Y.c.sgsʲs.:($$j).q_liq * Y.c.sgsʲs.:($$j).mse - ᶜmseₗ)
+
+            @. Yₜ.c.sgsʲs.:($$j).q_tot -= ᶜdivᵥ(ᶠwᵢ³ʲs) * Y.c.sgsʲs.:($$j).q_ice * (1 - Y.c.sgsʲs.:($$j).q_tot)
+            @. Yₜ.c.sgsʲs.:($$j).q_ice -= ᶜdivᵥ(ᶠwᵢ³ʲs) * Y.c.sgsʲs.:($$j).q_ice
+            @. Yₜ.c.sgsʲs.:($$j).mse   -= ᶜdivᵥ(ᶠwᵢ³ʲs) * (Y.c.sgsʲs.:($$j).q_ice * Y.c.sgsʲs.:($$j).mse - ᶜmseᵢ)
+
+            @. Yₜ.c.sgsʲs.:($$j).q_tot -= ᶜdivᵥ(ᶠwᵣ³ʲs) * Y.c.sgsʲs.:($$j).q_rai * (1 - Y.c.sgsʲs.:($$j).q_tot)
+            @. Yₜ.c.sgsʲs.:($$j).q_rai -= ᶜdivᵥ(ᶠwᵣ³ʲs) * Y.c.sgsʲs.:($$j).q_rai
+            @. Yₜ.c.sgsʲs.:($$j).mse   -= ᶜdivᵥ(ᶠwᵣ³ʲs) * (Y.c.sgsʲs.:($$j).q_rai * Y.c.sgsʲs.:($$j).mse - ᶜmseᵣ)
+
+            @. Yₜ.c.sgsʲs.:($$j).q_tot -= ᶜdivᵥ(ᶠwₛ³ʲs) * Y.c.sgsʲs.:($$j).q_sno * (1 - Y.c.sgsʲs.:($$j).q_tot)
+            @. Yₜ.c.sgsʲs.:($$j).q_sno -= ᶜdivᵥ(ᶠwₛ³ʲs) * Y.c.sgsʲs.:($$j).q_sno
+            @. Yₜ.c.sgsʲs.:($$j).mse   -= ᶜdivᵥ(ᶠwₛ³ʲs) * (Y.c.sgsʲs.:($$j).q_sno * Y.c.sgsʲs.:($$j).mse - ᶜmseₛ)
         end
     end
 end
diff --git a/src/prognostic_equations/implicit/implicit_tendency.jl b/src/prognostic_equations/implicit/implicit_tendency.jl
@@ -70,32 +70,32 @@ end
 # expressions are less convoluted?
 
 function vertical_transport(ᶜρ, ᶠu³, ᶜχ, dt, ::Val{:none})
-    ᶜJ = Fields.local_geometry_field(ᶜρ).J
-    ᶠJ = Fields.local_geometry_field(ᶠu³).J
+    ᶜJ = Fields.local_geometry_field(axes(ᶜρ)).J
+    ᶠJ = Fields.local_geometry_field(axes(ᶠu³)).J
     return @. lazy(-(ᶜadvdivᵥ(ᶠinterp(ᶜρ * ᶜJ) / ᶠJ * ᶠu³ * ᶠinterp(ᶜχ))))
 end
 function vertical_transport(ᶜρ, ᶠu³, ᶜχ, dt, ::Val{:first_order})
-    ᶜJ = Fields.local_geometry_field(ᶜρ).J
-    ᶠJ = Fields.local_geometry_field(ᶠu³).J
+    ᶜJ = Fields.local_geometry_field(axes(ᶜρ)).J
+    ᶠJ = Fields.local_geometry_field(axes(ᶠu³)).J
     return @. lazy(-(ᶜadvdivᵥ(ᶠinterp(ᶜρ * ᶜJ) / ᶠJ * ᶠupwind1(ᶠu³, ᶜχ))))
 end
 @static if pkgversion(ClimaCore) ≥ v"0.14.22"
     function vertical_transport(ᶜρ, ᶠu³, ᶜχ, dt, ::Val{:vanleer_limiter})
-        ᶜJ = Fields.local_geometry_field(ᶜρ).J
-        ᶠJ = Fields.local_geometry_field(ᶠu³).J
+        ᶜJ = Fields.local_geometry_field(axes(ᶜρ)).J
+        ᶠJ = Fields.local_geometry_field(axes(ᶠu³)).J
         return @. lazy(
             -(ᶜadvdivᵥ(ᶠinterp(ᶜρ * ᶜJ) / ᶠJ * ᶠlin_vanleer(ᶠu³, ᶜχ, dt))),
         )
     end
 end
 function vertical_transport(ᶜρ, ᶠu³, ᶜχ, dt, ::Val{:third_order})
-    ᶜJ = Fields.local_geometry_field(ᶜρ).J
-    ᶠJ = Fields.local_geometry_field(ᶠu³).J
+    ᶜJ = Fields.local_geometry_field(axes(ᶜρ)).J
+    ᶠJ = Fields.local_geometry_field(axes(ᶠu³)).J
     return @. lazy(-(ᶜadvdivᵥ(ᶠinterp(ᶜρ * ᶜJ) / ᶠJ * ᶠupwind3(ᶠu³, ᶜχ))))
 end
 function vertical_transport(ᶜρ, ᶠu³, ᶜχ, dt, ::Val{:boris_book})
-    ᶜJ = Fields.local_geometry_field(ᶜρ).J
-    ᶠJ = Fields.local_geometry_field(ᶠu³).J
+    ᶜJ = Fields.local_geometry_field(axes(ᶜρ)).J
+    ᶠJ = Fields.local_geometry_field(axes(ᶠu³)).J
     return @. lazy(
         -(ᶜadvdivᵥ(
             ᶠinterp(ᶜρ * ᶜJ) / ᶠJ * (
@@ -109,8 +109,8 @@ function vertical_transport(ᶜρ, ᶠu³, ᶜχ, dt, ::Val{:boris_book})
     )
 end
 function vertical_transport(ᶜρ, ᶠu³, ᶜχ, dt, ::Val{:zalesak})
-    ᶜJ = Fields.local_geometry_field(ᶜρ).J
-    ᶠJ = Fields.local_geometry_field(ᶠu³).J
+    ᶜJ = Fields.local_geometry_field(axes(ᶜρ)).J
+    ᶠJ = Fields.local_geometry_field(axes(ᶠu³)).J
     return @. lazy(
         -(ᶜadvdivᵥ(
             ᶠinterp(ᶜρ * ᶜJ) / ᶠJ * (
@@ -137,8 +137,8 @@ function implicit_vertical_advection_tendency!(Yₜ, Y, p, t)
         p.atmos
     (; dt) = p
     n = n_mass_flux_subdomains(turbconv_model)
-    ᶜJ = Fields.local_geometry_field(Y.c).J
-    ᶠJ = Fields.local_geometry_field(Y.f).J
+    ᶜJ = Fields.local_geometry_field(axes(Y.c)).J
+    ᶠJ = Fields.local_geometry_field(axes(Y.f)).J
     (; ᶠgradᵥ_ᶜΦ) = p.core
     (; ᶜh_tot, ᶠu³, ᶜp) = p.precomputed
 
diff --git a/src/utils/abbreviations.jl b/src/utils/abbreviations.jl

Original file line number	Diff line number	Diff line change
`@@ -1400,6 +1400,7 @@ EDMFBoxPlots = Union{`
`1400`	`1400`	`EDMFBoxPlotsWithPrecip = Union{`
`1401`	`1401`	`Val{:prognostic_edmfx_rico_column},`
`1402`	`1402`	`Val{:prognostic_edmfx_trmm_column},`
	`1403`	`+ Val{:prognostic_edmfx_dycoms_rf02_column},`
`1403`	`1404`	`}`
`1404`	`1405`	`"""`
`1405`	`1406`	`plot_edmf_vert_profile!(grid_loc, var_group)`