Wip on adding cloud condensate sedimentation

src/cache/precomputed_quantities.jl

@@ -148,6 +148,12 @@ function precomputed_quantities(Y, atmos)
     else
         (;)
     end
+    sedimentation_quantities =
+        atmos.moisture_model isa NonEquilMoistModel ?
+        (;
+            ᶜwₗ = similar(Y.c, FT),
+            ᶜwᵢ = similar(Y.c, FT),
+        ) : (;)
     precipitation_quantities =
         atmos.precip_model isa Microphysics1Moment ?
         (;
@@ -162,6 +168,7 @@ function precomputed_quantities(Y, atmos)
         advective_sgs_quantities...,
         diagnostic_sgs_quantities...,
         vert_diff_quantities...,
+        sedimentation_quantities...,
         precipitation_quantities...,
         cloud_diagnostics_tuple,
     )

src/cache/sedimentation_precomputed_quantities.jl

@@ -0,0 +1,38 @@
+#####
+##### Precomputed quantities
+#####
+import CloudMicrophysics.MicrophysicsNonEq as CMNe
+
+# TODO - duplicated with precip, should be moved to some common helper module
+# helper function to safely get precipitation from state
+function q_cc(ρq::FT, ρ::FT) where {FT}
+    return max(FT(0), ρq / ρ)
+end
+
+"""
+    set_sedimentation_precomputed_quantities!(Y, p, t)
+
+Updates the sedimentation terminal velocity stored in `p`
+for the non-equilibrium microphysics scheme
+"""
+function set_sedimentation_precomputed_quantities!(Y, p, t)
+    @assert (p.atmos.moisture_model isa MicrophysicsNonEq)
+
+    (; ᶜwₗ, ᶜwᵢ) = p.precomputed
+    cmc = CAP.microphysics_cloud_params(p.params)
+
+    # compute the precipitation terminal velocity [m/s]
+    @. ᶜwₗ = CMNe.terminal_velocity(
+        cmc.liquid,
+        cmc.Ch2022.rain,
+        Y.c.ρ,
+        q_cc(Y.c.ρq_liq, Y.c.ρ),
+    )
+    @. ᶜwᵢ = CMNe.terminal_velocity(
+        cmc.ice,
+        cmc.Ch2022.small_ice,
+        Y.c.ρ,
+        q_cc(Y.c.ρq_ice, Y.c.ρ),
+    )
+    return nothing
+end

src/parameters/create_parameters.jl

@@ -96,6 +96,7 @@ function create_parameter_set(config::AtmosConfig)
         (;
             liquid = CM.Parameters.CloudLiquid(toml_dict),
             ice = CM.Parameters.CloudIce(toml_dict),
+            Ch2022 = CM.Parameters.Chen2022VelType(toml_dict),
         )
     else
         nothing

src/prognostic_equations/implicit/implicit_solver.jl

@@ -402,6 +402,10 @@ NVTX.@annotate function Wfact!(A, Y, p, dtγ, t)
         p.precomputed.ᶜK,
         p.precomputed.ᶜts,
         p.precomputed.ᶜp,
+        (
+            p.atmos.moisture_model isa NonEquilMoistModel ?
+            (; p.precomputed.ᶜwₗ, p.precomputed.ᶜwᵢ) : (;)
+        )...,
         (
             p.atmos.precip_model isa Microphysics1Moment ?
             (; p.precomputed.ᶜwᵣ, p.precomputed.ᶜwₛ) : (;)
@@ -679,8 +683,12 @@ function update_implicit_equation_jacobian!(A, Y, p, dtγ)
         end
 
         ᶠlg = Fields.local_geometry_field(Y.f)
-        precip_info =
-            ((@name(c.ρq_rai), @name(ᶜwᵣ)), (@name(c.ρq_sno), @name(ᶜwₛ)))
+        precip_info = (
+            (@name(c.ρq_liq), @name(ᶜwₗ)),
+            (@name(c.ρq_ice), @name(ᶜwᵢ)),
+            (@name(c.ρq_rai), @name(ᶜwᵣ)),
+            (@name(c.ρq_sno), @name(ᶜwₛ)),
+        )
         MatrixFields.unrolled_foreach(precip_info) do (ρqₚ_name, wₚ_name)
             MatrixFields.has_field(Y, ρqₚ_name) || return
             ∂ᶜρqₚ_err_∂ᶜρqₚ = matrix[ρqₚ_name, ρqₚ_name]

src/prognostic_equations/implicit/implicit_tendency.jl

@@ -146,11 +146,23 @@ function implicit_vertical_advection_tendency!(Yₜ, Y, p, t)
         )
     end
 
+    # Advection of cloud condensate and precipitation with the mean flow
+    # is done with other passive tracers in the explicit tendency.
+    # Here we add the advection with cloud condensate and precipitation
+    # terminal velocity using downward biasing
+    # and free outflow bottom boundary condition
+    if moisture_model isa NonEquilMoistModel
+        (; ᶜwₗ, ᶜwᵢ) = p.precomputed
+        @. Yₜ.c.ρq_liq -= ᶜprecipdivᵥ(
+            ᶠwinterp(ᶜJ, Y.c.ρ) *
+            ᶠright_bias(Geometry.WVector(-(ᶜwₗ)) * ᶜspecific.q_liq),
+        )
+        @. Yₜ.c.ρq_ice -= ᶜprecipdivᵥ(
+            ᶠwinterp(ᶜJ, Y.c.ρ) *
+            ᶠright_bias(Geometry.WVector(-(ᶜwᵢ)) * ᶜspecific.q_ice),
+        )
+    end
     if precip_model isa Microphysics1Moment
-        # Advection of precipitation with the mean flow
-        # is done with other passive tracers in the explicit tendency.
-        # Here we add the advection with precipitation terminal velocity
-        # using downward biasing and free outflow bottom boundary condition
         (; ᶜwᵣ, ᶜwₛ) = p.precomputed
         @. Yₜ.c.ρq_rai -= ᶜprecipdivᵥ(
             ᶠwinterp(ᶜJ, Y.c.ρ) *