Add 1M sedimentation terms in prognostic EDMF

trontrytel · trontrytel · commit 5b7c6138d24f · 2025-08-06T17:41:45.000-07:00
diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml
@@ -88,7 +88,7 @@ steps:
           --config_file $CONFIG_PATH/single_column_precipitation_test.yml
           --job_id single_column_precipitation_test
         artifact_paths: "single_column_precipitation_test/output_active/*"
-      
+
       - label: ":umbrella: 2-moment precipitation sanity test single column"
         command: >
           julia --color=yes --project=.buildkite .buildkite/ci_driver.jl
@@ -916,6 +916,15 @@ steps:
         agents:
           slurm_mem: 20GB
 
+      - label: ":genie: Prognostic EDMFX Dycoms RF02 in a column"
+        command: >
+          julia --color=yes --project=.buildkite .buildkite/ci_driver.jl
+          --config_file $CONFIG_PATH/prognostic_edmfx_dycoms_rf02_column.yml
+          --job_id prognostic_edmfx_dycoms_rf02_column
+        artifact_paths: "prognostic_edmfx_dycoms_rf02_column/output_active/*"
+        agents:
+          slurm_mem: 20GB
+
       - label: ":umbrella: Prognostic EDMFX Rico in a column"
         command: >
           julia --color=yes --project=.buildkite .buildkite/ci_driver.jl
@@ -1127,9 +1136,9 @@ steps:
         agents:
           slurm_gpus: 1
           slurm_mem: 16GB
-      
+
       - label: ":earth_americas: GPU: Diagnostic EDMFX ERA5 Weather Model Initial Condition"
-        command: > 
+        command: >
           julia --color=yes --project=.buildkite .buildkite/ci_driver.jl
           --config_file $COMMON_CONFIG_PATH/numerics_sphere_he6ze31.yml
           --config_file $CONFIG_PATH/diagnostic_edmfx_era5_initial_condition.yml
diff --git a/config/model_configs/prognostic_edmfx_dycoms_rf02_column.yml b/config/model_configs/prognostic_edmfx_dycoms_rf02_column.yml
@@ -0,0 +1,43 @@
+initial_condition: DYCOMS_RF02
+subsidence: DYCOMS
+scm_coriolis: DYCOMS_RF02
+rad: DYCOMS
+surface_setup: DYCOMS_RF02
+turbconv: "prognostic_edmfx"
+implicit_diffusion: false
+approximate_linear_solve_iters: 2
+edmfx_upwinding: first_order
+edmfx_entr_model: "Generalized"
+edmfx_detr_model: "Generalized"
+edmfx_sgs_mass_flux: true
+edmfx_sgs_diffusive_flux: true
+edmfx_nh_pressure: true
+edmfx_filter: true
+prognostic_tke: true
+moist: "nonequil"
+cloud_model: "quadrature_sgs"
+precip_model: "1M"
+call_cloud_diagnostics_per_stage: true
+config: "column"
+x_elem: 2
+y_elem: 2
+z_elem: 30
+z_max: 1500
+z_stretch: false
+perturb_initstate: false
+dt: 10secs
+t_end: 5hours
+dt_save_state_to_disk: 10mins
+toml: [toml/prognostic_edmfx_1M.toml]
+netcdf_interpolation_num_points: [8, 8, 30]
+diagnostics:
+  - short_name: [ts, ta, thetaa, ha, pfull, rhoa, ua, va, wa, hur, hus, cl, clw, cli, hussfc, evspsbl, pr]
+    period: 10mins
+  - short_name: [arup, waup, taup, thetaaup, haup, husup, hurup, clwup, cliup, husraup, hussnup]
+    period: 10mins
+  - short_name: [waen, taen, thetaaen, haen, husen, huren, clwen, clien, husraen, hussnen, tke]
+    period: 10mins
+  - short_name: [entr, detr, lmix, bgrad, strain, edt, evu]
+    period: 10mins
+  - short_name: [husra, hussn]
+    period: 10mins
diff --git a/config/model_configs/prognostic_edmfx_rico_column.yml b/config/model_configs/prognostic_edmfx_rico_column.yml
@@ -3,7 +3,7 @@ subsidence: "Rico"
 ls_adv: "Rico"
 surface_setup: "Rico"
 turbconv: "prognostic_edmfx"
-implicit_diffusion: false
+implicit_diffusion: true
 implicit_sgs_advection: false
 approximate_linear_solve_iters: 2
 edmfx_upwinding: first_order
@@ -25,9 +25,9 @@ y_elem: 2
 z_elem: 100
 z_stretch: false
 perturb_initstate: false
-dt: "10secs"
-t_end: "6hours"
-dt_save_state_to_disk: "10mins"
+dt: "5secs"
+t_end: "180mins"
+dt_save_state_to_disk: "60mins"
 toml: [toml/prognostic_edmfx_1M.toml]
 netcdf_interpolation_num_points: [8, 8, 100]
 diagnostics:
diff --git a/post_processing/ci_plots.jl b/post_processing/ci_plots.jl
@@ -1397,6 +1397,7 @@ EDMFBoxPlots = Union{
 EDMFBoxPlotsWithPrecip = Union{
     Val{:prognostic_edmfx_rico_column},
     Val{:prognostic_edmfx_trmm_column},
+    Val{:prognostic_edmfx_dycoms_rf02_column},
 }
 
 DiagEDMFBoxPlotsWithPrecip = Union{
diff --git a/src/cache/precomputed_quantities.jl b/src/cache/precomputed_quantities.jl
@@ -138,6 +138,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
@@ -538,13 +538,73 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_precipitation
     (; ᶜSqₗᵖʲs, ᶜSqᵢᵖʲs, ᶜSqᵣᵖʲs, ᶜSqₛᵖʲs, ᶜρʲs, ᶜtsʲs) = p.precomputed
     (; ᶜSqₗᵖ⁰, ᶜSqᵢᵖ⁰, ᶜSqᵣᵖ⁰, ᶜSqₛᵖ⁰, ᶜts⁰) = 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).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)) + ᶜΦ)
+            ) / 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:
@@ -344,7 +344,6 @@ function edmfx_sgs_vertical_advection_tendency!(
     turbconv_params = CAP.turbconv_params(params)
     α_b = CAP.pressure_normalmode_buoy_coeff1(turbconv_params)
     ᶠz = Fields.coordinate_field(Y.f).z
-    ᶜa_scalar = p.scratch.ᶜtemp_scalar
     ᶜu₃ʲ = p.scratch.ᶜtemp_C3
     ᶜKᵥʲ = p.scratch.ᶜtemp_scalar_2
     for j in 1:n
@@ -370,58 +369,127 @@ function edmfx_sgs_vertical_advection_tendency!(
     end
 
     for j in 1:n
-        @. ᶜa_scalar = draft_area(Y.c.sgsʲs.:($$j).ρa, ᶜρʲs.:($$j))
-        vtt = vertical_transport(
-            ᶜρʲs.:($j),
-            ᶠu³ʲs.:($j),
-            ᶜa_scalar,
-            dt,
-            edmfx_upwinding,
-        )
+        ᶜa = (@. lazy(draft_area(Y.c.sgsʲs.:($$j).ρa, ᶜρʲs.:($$j))))
+        edmf_upwnd = edmfx_upwinding
+
+        # Flux form vertical advection of area farction with the grid mean velocity
+        vtt = vertical_transport(ᶜρʲs.:($j), ᶠu³ʲs.:($j), ᶜa, dt, edmf_upwnd)
         @. Yₜ.c.sgsʲs.:($$j).ρa += vtt
 
-        va = vertical_advection(
-            ᶠu³ʲs.:($j),
-            Y.c.sgsʲs.:($j).mse,
-            edmfx_upwinding,
-        )
-        @. Yₜ.c.sgsʲs.:($$j).mse += va
+        # Advective form advection of mse and q_tot with the grid mean velocity
+        # TODO - make it work for multiple updrafts
+        if j > 1
+            error("Below code doesn't work for multiple updrafts")
+        end
+        sgs_q_tot_mse = (@name(c.sgsʲs.:(1).mse), @name(c.sgsʲs.:(1).q_tot))
+        MatrixFields.unrolled_foreach(sgs_q_tot_mse) do χʲ_name
+            MatrixFields.has_field(Y, χʲ_name) || return
+            ᶜχʲ = MatrixFields.get_field(Y, χʲ_name)
+            ᶜχʲₜ = MatrixFields.get_field(Yₜ, χʲ_name)
+
+            va = vertical_advection(ᶠu³ʲs.:($j), ᶜχʲ, edmf_upwnd)
+            @. ᶜχʲₜ += va
+        end
 
-        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,
-            )
-            @. Yₜ.c.sgsʲs.:($$j).q_liq += va
-            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,
+            # TODO - add contibutions to sgs mass flux from tracer sedimentation
+            # TODO - add precipitation and cloud sedimentation in implicit solver/tendency with if/else
+
+            FT = eltype(params)
+            thp = CAP.thermodynamics_params(params)
+            (; ᶜΦ) = p.core
+            (; ᶜtsʲs) = p.precomputed
+
+            # Sedimentation velocities for microphysics tracers
+            # TODO - lazify ᶜwₗʲs computation. No need to cache it.
+            sgs_microphysics_tracers = (
+                (@name(c.sgsʲs.:(1).q_liq), @name(q_liq), @name(ᶜwₗʲs.:(1))),
+                (@name(c.sgsʲs.:(1).q_ice), @name(q_ice), @name(ᶜwᵢʲs.:(1))),
+                (@name(c.sgsʲs.:(1).q_rai), @name(q_rai), @name(ᶜwᵣʲs.:(1))),
+                (@name(c.sgsʲs.:(1).q_sno), @name(q_sno), @name(ᶜwₛʲs.:(1))),
             )
-            @. Yₜ.c.sgsʲs.:($$j).q_rai += va
-            va = vertical_advection(
-                ᶠu³ʲs.:($j),
-                Y.c.sgsʲs.:($j).q_sno,
-                edmfx_upwinding,
-            )
-            @. Yₜ.c.sgsʲs.:($$j).q_sno += va
+
+            MatrixFields.unrolled_foreach(
+                sgs_microphysics_tracers,
+            ) do (qʲ_name, name, wʲ_name)
+                MatrixFields.has_field(Y, qʲ_name) || return
+
+                ᶜqʲ = MatrixFields.get_field(Y, qʲ_name)
+                ᶜqʲₜ = MatrixFields.get_field(Yₜ, qʲ_name)
+                ᶜwʲ = MatrixFields.get_field(p.precomputed, wʲ_name)
+                ᶠw³ʲ = (@. lazy(CT3(ᶠinterp(Geometry.WVector(-1 * ᶜwʲ)))))
+                ᶜw³ʲ = (@. lazy(CT3(Geometry.WVector(-1 * ᶜwʲ))))
+                ᶜaqʲ = (@. lazy(ᶜa * ᶜqʲ))
+
+                # Flux form vertical advection of rho * area with sedimentation velocities
+                # Eq (4) term (3) in the writeup
+                vtt = vertical_transport(ᶜρʲs.:($j), ᶠw³ʲ, ᶜaqʲ, dt, edmf_upwnd)
+                @. Yₜ.c.sgsʲs.:($$j).ρa += vtt
+
+                # Advective form advection of moisture tracers with the grid mean velocity
+                # Eq (2) term (1) in the writeup
+                va = vertical_advection(ᶠu³ʲs.:($j), ᶜqʲ, edmf_upwnd)
+                @. ᶜqʲₜ += va
+
+                # Advective form advection of q_tot and moisture tracers with sedimentation velocities
+                # Eq (1-2) term (2)
+                va = vertical_advection(ᶠw³ʲ, ᶜqʲ, edmf_upwnd)
+                @. Yₜ.c.sgsʲs.:($$j).q_tot += (1 - Y.c.sgsʲs.:($$j).q_tot) * va
+                @. ᶜqʲₜ += va
+                # Advective form advection of mse with sedimentation velocity
+                # Eq (3) term (2)
+                if name in (@name(q_liq), @name(q_rai))
+                    ᶜmse_li = (@. lazy(
+                        TD.internal_energy_liquid(thp, ᶜtsʲs.:($$j)) + ᶜΦ,
+                    ))
+                elseif name in (@name(q_ice), @name(q_sno))
+                    ᶜmse_li = (@. lazy(
+                        TD.internal_energy_ice(thp, ᶜtsʲs.:($$j)) + ᶜΦ,
+                    ))
+                else
+                    error("Unsupported moisture tracer variable")
+                end
+                va = vertical_advection(ᶠw³ʲ, ᶜqʲ .* ᶜmse_li, edmf_upwnd)
+                @. Yₜ.c.sgsʲs.:($$j).mse += va
+                va = vertical_advection(ᶠw³ʲ, ᶜqʲ, edmf_upwnd)
+                @. Yₜ.c.sgsʲs.:($$j).mse -= Y.c.sgsʲs.:($$j).mse * va
+
+                # mse, q_tot and moisture tracers terms proportional to 1/ρ̂ ∂zρ̂
+                # Eq (1-3) term (3)
+                ᶜinv_ρ̂ = (@. lazy(
+                    specific(
+                        FT(1),
+                        Y.c.sgsʲs.:($$j).ρa,
+                        FT(0),
+                        Y.c.ρ,
+                        turbconv_model,
+                    ),
+                ))
+                ᶜ∂ρ̂∂z = (@. lazy(
+                    upwind_biased_grad(
+                        -1 * Geometry.WVector(ᶜwʲ),
+                        Y.c.sgsʲs.:($$j).ρa,
+                    ),
+                ))
+                @. Yₜ.c.sgsʲs.:($$j).mse -=
+                    dot(ᶜinv_ρ̂ * ᶜ∂ρ̂∂z, ᶜw³ʲ) *
+                    ᶜqʲ *
+                    (ᶜmse_li - Y.c.sgsʲs.:($$j).mse)
+                @. Yₜ.c.sgsʲs.:($$j).q_tot -=
+                    dot(ᶜinv_ρ̂ * ᶜ∂ρ̂∂z, ᶜw³ʲ) *
+                    ᶜqʲ *
+                    (1 - Y.c.sgsʲs.:($$j).q_tot)
+                @. ᶜqʲₜ -= dot(ᶜinv_ρ̂ * ᶜ∂ρ̂∂z, ᶜw³ʲ) * ᶜqʲ
+
+                # mse, q_tot and moisture tracer terms proportional to velocity gradients
+                # Eq (1-3) term (4)
+                @. Yₜ.c.sgsʲs.:($$j).mse -=
+                    ᶜdivᵥ(ᶠw³ʲ) * ᶜqʲ * (ᶜmse_li - Y.c.sgsʲs.:($$j).mse)
+                @. Yₜ.c.sgsʲs.:($$j).q_tot -=
+                    ᶜdivᵥ(ᶠw³ʲ) * ᶜqʲ * (1 - Y.c.sgsʲs.:($$j).q_tot)
+                @. ᶜqʲₜ -= ᶜdivᵥ(ᶠw³ʲ) * ᶜqʲ
+            end
         end
     end
 end
diff --git a/src/prognostic_equations/implicit/implicit_tendency.jl b/src/prognostic_equations/implicit/implicit_tendency.jl
diff --git a/src/utils/abbreviations.jl b/src/utils/abbreviations.jl

Original file line number	Diff line number	Diff line change
`@@ -1397,6 +1397,7 @@ EDMFBoxPlots = Union{`
`1397`	`1397`	`EDMFBoxPlotsWithPrecip = Union{`
`1398`	`1398`	`Val{:prognostic_edmfx_rico_column},`
`1399`	`1399`	`Val{:prognostic_edmfx_trmm_column},`
	`1400`	`+ Val{:prognostic_edmfx_dycoms_rf02_column},`
`1400`	`1401`	`}`
`1401`	`1402`
`1402`	`1403`	`DiagEDMFBoxPlotsWithPrecip = Union{`