CliMA · sajjadazimi · Aug 13, 2025 · Aug 11, 2025
diff --git a/.buildkite/Manifest-v1.11.toml b/.buildkite/Manifest-v1.11.toml
@@ -462,9 +462,9 @@ version = "0.1.13"
 
 [[deps.CloudMicrophysics]]
 deps = ["ClimaParams", "DocStringExtensions", "ForwardDiff", "LazyArtifacts", "LogExpFunctions", "RootSolvers", "SpecialFunctions", "StaticArrays", "Thermodynamics"]
-git-tree-sha1 = "755a88a92c3281975eec4da19dbbec8839c607aa"
+git-tree-sha1 = "3a938de9cf21b603108aa273d973997fbd8c882d"
 uuid = "6a9e3e04-43cd-43ba-94b9-e8782df3c71b"
-version = "0.26.2"
+version = "0.26.4"
 
     [deps.CloudMicrophysics.extensions]
     EmulatorModelsExt = ["DataFrames", "MLJ"]

diff --git a/Project.toml b/Project.toml
@@ -50,7 +50,7 @@ ClimaInterpolations = "0.1.0"
 ClimaParams = "0.11.1"
 ClimaTimeSteppers = "0.8.2"
 ClimaUtilities = "0.1.23"
-CloudMicrophysics = "0.26.2"
+CloudMicrophysics = "0.26.4"
 Dates = "1"
 ForwardDiff = "1"
 Insolation = "0.9.5"

diff --git a/config/model_configs/prognostic_edmfx_rico_column_2M.yml b/config/model_configs/prognostic_edmfx_rico_column_2M.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]
 prescribed_aerosols: ["SO4", "CB1", "OC1", "DST01", "SSLT01", "SSLT02", "SSLT03", "SSLT04", "SSLT05"]

diff --git a/src/cache/precipitation_precomputed_quantities.jl b/src/cache/precipitation_precomputed_quantities.jl
@@ -107,7 +107,7 @@ function set_precipitation_velocities!(
     (; ᶜwₗ, ᶜwᵢ, ᶜwᵣ, ᶜwₛ, ᶜwnₗ, ᶜwnᵣ, ᶜwₜqₜ, ᶜwₕhₜ, ᶜts, ᶜu) = p.precomputed
     (; ᶜΦ) = p.core
 
-    cm1c = CAP.microphysics_cloud_params(p.params)
+    cmc = CAP.microphysics_cloud_params(p.params)
     cm1p = CAP.microphysics_1m_params(p.params)
     cm2p = CAP.microphysics_2m_params(p.params)
     thp = CAP.thermodynamics_params(p.params)
@@ -117,50 +117,56 @@ function set_precipitation_velocities!(
     @. ᶜwnᵣ = getindex(
         CM2.rain_terminal_velocity(
             cm2p.sb,
-            cm2p.tv,
-            specific(Y.c.ρq_rai, Y.c.ρ),
+            cm2p.rtv,
+            max(zero(Y.c.ρ), specific(Y.c.ρq_rai, Y.c.ρ)),
             Y.c.ρ,
-            Y.c.ρn_rai,
+            max(zero(Y.c.ρ), Y.c.ρn_rai),
         ),
         1,
     )
     @. ᶜwᵣ = getindex(
         CM2.rain_terminal_velocity(
             cm2p.sb,
-            cm2p.tv,
-            specific(Y.c.ρq_rai, Y.c.ρ),
+            cm2p.rtv,
+            max(zero(Y.c.ρ), specific(Y.c.ρq_rai, Y.c.ρ)),
             Y.c.ρ,
-            Y.c.ρn_rai,
+            max(zero(Y.c.ρ), Y.c.ρn_rai),
         ),
         2,
     )
     @. ᶜwₛ = CM1.terminal_velocity(
         cm1p.ps,
         cm1p.tv.snow,
         Y.c.ρ,
-        specific(Y.c.ρq_sno, Y.c.ρ),
+        max(zero(Y.c.ρ), specific(Y.c.ρq_sno, Y.c.ρ)),
     )
     # compute sedimentation velocity for cloud condensate [m/s]
-    # TODO sedimentation velocities of cloud condensates are based
-    # on the 1M scheme. Sedimentation velocity of cloud number concentration
-    # is equal to that of the mass.
-    @. ᶜwnₗ = CMNe.terminal_velocity(
-        cm1c.liquid,
-        cm1c.Ch2022.rain,
-        Y.c.ρ,
-        specific(Y.c.ρq_liq, Y.c.ρ),
+    # TODO sedimentation of ice is based on the 1M scheme
+    @. ᶜwnₗ = getindex(
+        CM2.cloud_terminal_velocity(
+            cm2p.sb.pdf_c,
+            cm2p.ctv,
+            max(zero(Y.c.ρ), specific(Y.c.ρq_liq, Y.c.ρ)),
+            Y.c.ρ,
+            max(zero(Y.c.ρ), Y.c.ρn_liq),
+        ),
+        1,
     )
-    @. ᶜwₗ = CMNe.terminal_velocity(
-        cm1c.liquid,
-        cm1c.Ch2022.rain,
-        Y.c.ρ,
-        specific(Y.c.ρq_liq, Y.c.ρ),
+    @. ᶜwₗ = getindex(
+        CM2.cloud_terminal_velocity(
+            cm2p.sb.pdf_c,
+            cm2p.ctv,
+            max(zero(Y.c.ρ), specific(Y.c.ρq_liq, Y.c.ρ)),
+            Y.c.ρ,
+            max(zero(Y.c.ρ), Y.c.ρn_liq),
+        ),
+        2,
     )
     @. ᶜwᵢ = CMNe.terminal_velocity(
-        cm1c.ice,
-        cm1c.Ch2022.small_ice,
+        cmc.ice,
+        cmc.Ch2022.small_ice,
         Y.c.ρ,
-        specific(Y.c.ρq_ice, Y.c.ρ),
+        max(zero(Y.c.ρ), specific(Y.c.ρq_ice, Y.c.ρ)),
     )
 
     # compute their contributions to energy and total water advection

diff --git a/src/cache/precomputed_quantities.jl b/src/cache/precomputed_quantities.jl
@@ -157,6 +157,14 @@ function precomputed_quantities(Y, atmos)
             ᶜSqₛᵖʲs = similar(Y.c, NTuple{n, FT}),
             ᶜSnₗᵖʲs = similar(Y.c, NTuple{n, FT}),
             ᶜSnᵣᵖʲ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}),
+            ᶜ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
@@ -769,8 +769,10 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_precipitation
 )
 
     (; params, dt) = p
+    (; ᶜΦ,) = p.core
     thp = CAP.thermodynamics_params(params)
-    cmp = CAP.microphysics_2m_params(params)
+    cm1p = CAP.microphysics_1m_params(p.params)
+    cm2p = CAP.microphysics_2m_params(p.params)
     cmc = CAP.microphysics_cloud_params(params)
 
     (;
@@ -786,6 +788,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_precipitation
     ) = p.precomputed
     (; ᶜSqₗᵖ⁰, ᶜSqᵢᵖ⁰, ᶜSqᵣᵖ⁰, ᶜSqₛᵖ⁰, ᶜSnₗᵖ⁰, ᶜSnᵣᵖ⁰, ᶜts⁰, ᶜu⁰) =
         p.precomputed
+    (; ᶜwₗʲs, ᶜwᵢʲs, ᶜwᵣʲs, ᶜwₛʲs, ᶜwₙₗʲs, ᶜwₙᵣʲs, ᶜwₜʲs, ᶜwₕʲs) = p.precomputed
 
     ᶜSᵖ = p.scratch.ᶜtemp_scalar
     ᶜS₂ᵖ = p.scratch.ᶜtemp_scalar_2
@@ -801,7 +804,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_precipitation
             seasalt_mean_radius,
             sulfate_num,
             p.tracers.prescribed_aerosols_field,
-            cmp.aerosol,
+            cm2p.aerosol,
         )
     else
         @. seasalt_num = 0
@@ -811,7 +814,95 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_precipitation
 
     # Compute sources
     n = n_mass_flux_subdomains(p.atmos.turbconv_model)
+    FT = eltype(params)
     for j in 1:n
+
+        # compute terminal velocity for precipitation
+        # TODO sedimentation of snow is based on the 1M scheme
+        @. ᶜwₙᵣʲs.:($$j) = getindex(
+            CM2.rain_terminal_velocity(
+                cm2p.sb,
+                cm2p.rtv,
+                max(zero(Y.c.ρ), Y.c.sgsʲs.:($$j).q_rai),
+                ᶜρʲs.:($$j),
+                max(zero(Y.c.ρ), Y.c.sgsʲs.:($$j).n_rai),
+            ),
+            1,
+        )
+        @. ᶜwᵣʲs.:($$j) = getindex(
+            CM2.rain_terminal_velocity(
+                cm2p.sb,
+                cm2p.rtv,
+                max(zero(Y.c.ρ), Y.c.sgsʲs.:($$j).q_rai),
+                ᶜρʲs.:($$j),
+                max(zero(Y.c.ρ), Y.c.sgsʲs.:($$j).n_rai),
+            ),
+            2,
+        )
+        @. ᶜwₛʲs.:($$j) = CM1.terminal_velocity(
+            cm1p.ps,
+            cm1p.tv.snow,
+            ᶜρʲs.:($$j),
+            max(zero(Y.c.ρ), Y.c.sgsʲs.:($$j).q_sno),
+        )
+        # compute sedimentation velocity for cloud condensate [m/s]
+        # TODO sedimentation of ice is based on the 1M scheme
+        @. ᶜwₙₗʲs.:($$j) = getindex(
+            CM2.cloud_terminal_velocity(
+                cm2p.sb.pdf_c,
+                cm2p.ctv,
+                max(zero(Y.c.ρ), Y.c.sgsʲs.:($$j).q_liq),
+                ᶜρʲs.:($$j),
+                max(zero(Y.c.ρ), ᶜρʲs.:($$j) * Y.c.sgsʲs.:($$j).n_liq),
+            ),
+            1,
+        )
+        @. ᶜwₗʲs.:($$j) = getindex(
+            CM2.cloud_terminal_velocity(
+                cm2p.sb.pdf_c,
+                cm2p.ctv,
+                max(zero(Y.c.ρ), Y.c.sgsʲs.:($$j).q_liq),
+                ᶜρʲs.:($$j),
+                max(zero(Y.c.ρ), ᶜρʲs.:($$j) * Y.c.sgsʲs.:($$j).n_liq),
+            ),
+            2,
+        )
+        @. ᶜ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_warm_precipitation_sources_2M!(
             ᶜSᵖ,
@@ -830,7 +921,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_precipitation
             Y.c.sgsʲs.:($j).q_sno,
             ᶜtsʲs.:($j),
             dt,
-            cmp,
+            cm2p,
             thp,
         )
         @. ᶜSqᵢᵖʲs.:($$j) = 0
@@ -870,7 +961,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_precipitation
             Y.c.sgsʲs.:($$j).n_liq + Y.c.sgsʲs.:($$j).n_rai,
             ᶜρʲs.:($$j),
             max(0, w_component.(Geometry.WVector.(ᶜuʲs.:($$j)))),
-            (cmp,),
+            (cm2p,),
             thp,
             ᶜtsʲs.:($$j),
             dt,
@@ -903,7 +994,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_precipitation
         ᶜq_sno⁰,
         ᶜts⁰,
         dt,
-        cmp,
+        cm2p,
         thp,
     )
     @. ᶜSqᵢᵖ⁰ = 0
@@ -943,7 +1034,7 @@ NVTX.@annotate function set_prognostic_edmf_precomputed_quantities_precipitation
         ᶜn_liq⁰ + ᶜn_rai⁰,
         ᶜρ⁰,
         w_component.(Geometry.WVector.(ᶜu⁰)),
-        (cmp,),
+        (cm2p,),
         thp,
         ᶜts⁰,
         dt,

diff --git a/src/parameters/create_parameters.jl b/src/parameters/create_parameters.jl
@@ -163,7 +163,8 @@ microphys_2m_parameters(::Type{FT}) where {FT <: AbstractFloat} =
 microphys_2m_parameters(toml_dict::CP.AbstractTOMLDict) = (;
     sb = CM.Parameters.SB2006(toml_dict),
     aps = CM.Parameters.AirProperties(toml_dict),
-    tv = CM.Parameters.SB2006VelType(toml_dict),
+    ctv = CM.Parameters.StokesRegimeVelType(toml_dict),
+    rtv = CM.Parameters.SB2006VelType(toml_dict),
     liquid = CM.Parameters.CloudLiquid(toml_dict),
     ice = CM.Parameters.CloudIce(toml_dict),
     arg = CM.Parameters.AerosolActivationParameters(toml_dict),

diff --git a/src/prognostic_equations/advection.jl b/src/prognostic_equations/advection.jl
@@ -404,8 +404,10 @@ function edmfx_sgs_vertical_advection_tendency!(
             @. ᶜχʲₜ += va
         end
 
-        if p.atmos.moisture_model isa NonEquilMoistModel &&
-           p.atmos.microphysics_model isa Microphysics1Moment
+        if p.atmos.moisture_model isa NonEquilMoistModel && (
+            p.atmos.microphysics_model isa Microphysics1Moment ||
+            p.atmos.microphysics_model isa Microphysics2Moment
+        )
             # TODO - add contibutions to sgs mass flux from tracer sedimentation
             # TODO - add precipitation and cloud sedimentation in implicit solver/tendency with if/else
 
@@ -506,16 +508,55 @@ function edmfx_sgs_vertical_advection_tendency!(
         end
         if p.atmos.moisture_model isa NonEquilMoistModel &&
            p.atmos.microphysics_model isa Microphysics2Moment
-            # TODO add sedimentation for number concentrations
 
-            sgs_microphysics_tracers =
-                (:q_liq, :q_ice, :q_rai, :q_sno, :n_liq, :n_rai)
+            FT = eltype(params)
+
+            # Sedimentation velocities for microphysics number concentrations
+            # (or any tracers that does not directly participate in variations of q_tot and mse)
+            sgs_microphysics_tracers = (
+                (@name(c.sgsʲs.:(1).n_liq), @name(ᶜwₙₗʲs.:(1))),
+                (@name(c.sgsʲs.:(1).n_rai), @name(ᶜwₙᵣʲs.:(1))),
+            )
+
+            MatrixFields.unrolled_foreach(
+                sgs_microphysics_tracers,
+            ) do (χʲ_name, χʲ_name)
+                MatrixFields.has_field(Y, χʲ_name) || return
 
-            for χ_name in sgs_microphysics_tracers
-                ᶜχʲ = getproperty(Y.c.sgsʲs.:($j), χ_name)
-                va = vertical_advection(ᶠu³ʲs.:($j), ᶜχʲ, edmfx_upwinding)
-                ᶜχʲₜ = getproperty(Yₜ.c.sgsʲs.:($j), χ_name)
+                ᶜχʲ = MatrixFields.get_field(Y, χʲ_name)
+                ᶜχʲₜ = MatrixFields.get_field(Yₜ, χʲ_name)
+                ᶜwʲ = MatrixFields.get_field(p.precomputed, wʲ_name)
+                ᶠw³ʲ = (@. lazy(CT3(ᶠinterp(Geometry.WVector(-1 * ᶜwʲ)))))
+                ᶜw³ʲ = (@. lazy(CT3(Geometry.WVector(-1 * ᶜwʲ))))
+
+                # Advective form advection of moisture tracers with the grid mean velocity
+                va = vertical_advection(ᶠu³ʲs.:($j), ᶜχʲ, edmf_upwnd)
+                @. ᶜχʲₜ += va
+
+                # Advective form advection of moisture tracers with sedimentation velocities
+                va = vertical_advection(ᶠw³ʲ, ᶜχʲ, edmf_upwnd)
                 @. ᶜχʲₜ += va
+
+                # moisture tracers terms proportional to 1/ρ̂ ∂zρ̂
+                ᶜ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,
+                    ),
+                ))
+                @. ᶜχʲₜ -= dot(ᶜinv_ρ̂ * ᶜ∂ρ̂∂z, ᶜw³ʲ) * ᶜχʲ
+
+                # moisture tracer term proportional to velocity gradients
+                @. ᶜχʲₜ -= ᶜdivᵥ(ᶠw³ʲ) * ᶜχʲ
             end
 
         end