add P3 microphysics

src/cache/precipitation_precomputed_quantities.jl

@@ -4,6 +4,7 @@
 
 import CloudMicrophysics.MicrophysicsNonEq as CMNe
 import CloudMicrophysics.Microphysics1M as CM1
+import CloudMicrophysics.Microphysics2M as CM2
 
 import Thermodynamics as TD
 import ClimaCore.Operators as Operators
@@ -97,6 +98,94 @@ function set_precipitation_velocities!(
         ) / Y.c.ρ
     return nothing
 end
+function set_precipitation_velocities!(
+    Y,
+    p,
+    moisture_model::NonEquilMoistModel,
+    precip_model::Microphysics2Moment,
+)
+    (; ᶜwₗ, ᶜwᵢ, ᶜwᵣ, ᶜwₛ, ᶜwNₗ, ᶜwNᵢ, ᶜwNᵣ, ᶜwNₛ, ᶜwₜqₜ, ᶜwₕhₜ, ᶜts, ᶜu) = p.precomputed
+    (; q_liq, q_ice, q_rai, q_sno) = p.precomputed.ᶜspecific
+    (; ᶜΦ) = p.core
+
+    cm1c = 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)
+
+    # compute the precipitation terminal velocity [m/s]
+    # TODO sedimentation of snow is based on the 1M scheme
+    @. ᶜwNᵣ = getindex(CM2.rain_terminal_velocity(
+        cm2p.sb,
+        cm2p.tv,
+        q_rai,
+        Y.c.ρ,
+        Y.c.N_rai,
+    ), 1)
+    @. ᶜwᵣ = getindex(CM2.rain_terminal_velocity(
+        cm2p.sb,
+        cm2p.tv,
+        q_rai,
+        Y.c.ρ,
+        Y.c.N_rai,
+    ), 2)
+    @. ᶜwNₛ = CM1.terminal_velocity(
+        cm1p.ps,
+        cm1p.tv.snow,
+        Y.c.ρ,
+        q_sno,
+    )
+    @. ᶜwₛ = CM1.terminal_velocity(
+        cm1p.ps,
+        cm1p.tv.snow,
+        Y.c.ρ,
+        q_sno,
+    )
+    # compute sedimentation velocity for cloud condensate [m/s]
+    # TODO sedimentation velocities of cloud condensates are based 
+    # on the 1M scheme.
+    @. ᶜwNₗ = CMNe.terminal_velocity(
+        cm1c.liquid,
+        cm1c.Ch2022.rain,
+        Y.c.ρ,
+        q_liq,
+    )
+    @. ᶜwₗ = CMNe.terminal_velocity(
+        cm1c.liquid,
+        cm1c.Ch2022.rain,
+        Y.c.ρ,
+        q_liq,
+    )
+    @. ᶜwNᵢ = CMNe.terminal_velocity(
+        cm1c.ice,
+        cm1c.Ch2022.small_ice,
+        Y.c.ρ,
+        q_ice,
+    )
+    @. ᶜwᵢ = CMNe.terminal_velocity(
+        cm1c.ice,
+        cm1c.Ch2022.small_ice,
+        Y.c.ρ,
+        q_ice,
+    )
+
+    # compute their contributions to energy and total water advection
+    @. ᶜwₜqₜ =
+        Geometry.WVector(
+            ᶜwₗ * Y.c.ρq_liq +
+            ᶜwᵢ * Y.c.ρq_ice +
+            ᶜwᵣ * Y.c.ρq_rai +
+            ᶜwₛ * Y.c.ρq_sno,
+        ) / Y.c.ρ
+    @. ᶜwₕhₜ =
+        Geometry.WVector(
+            ᶜwₗ * Y.c.ρq_liq * (Iₗ(thp, ᶜts) + ᶜΦ + $(Kin(ᶜwₗ, ᶜu))) +
+            ᶜwᵢ * Y.c.ρq_ice * (Iᵢ(thp, ᶜts) + ᶜΦ + $(Kin(ᶜwᵢ, ᶜu))) +
+            ᶜwᵣ * Y.c.ρq_rai * (Iₗ(thp, ᶜts) + ᶜΦ + $(Kin(ᶜwᵣ, ᶜu))) +
+            ᶜwₛ * Y.c.ρq_sno * (Iᵢ(thp, ᶜts) + ᶜΦ + $(Kin(ᶜwₛ, ᶜu))),
+        ) / Y.c.ρ
+    return nothing
+end
 
 """
     set_precipitation_cache!(Y, p, precip_model, turbconv_model)
@@ -262,6 +351,64 @@ function set_precipitation_cache!(
     # in edmf sub-domains.
     return nothing
 end
+function set_precipitation_cache!(Y, p, ::Microphysics2Moment, _)
+    (; dt) = p
+    (; ᶜts) = p.precomputed
+    (; ᶜSqₗᵖ, ᶜSqᵢᵖ, ᶜSqᵣᵖ, ᶜSqₛᵖ) = p.precomputed
+    (; ᶜSNₗᵖ, ᶜSNᵢᵖ, ᶜSNᵣᵖ, ᶜSNₛᵖ) = p.precomputed
+
+    (; q_liq, q_rai, q_sno) = p.precomputed.ᶜspecific
+
+    ᶜSᵖ = p.scratch.ᶜtemp_scalar
+    ᶜS₂ᵖ = p.scratch.ᶜtemp_scalar_2
+
+    # get thermodynamics and microphysics params
+    (; params) = p
+    cm1p = CAP.microphysics_1m_params(params)
+    cm2p = CAP.microphysics_2m_params(params)
+    thp = CAP.thermodynamics_params(params)
+
+    # compute warm precipitation sources on the grid mean (based on SB2006 2M scheme)
+    compute_warm_precipitation_sources_2M!(
+        ᶜSᵖ,
+        ᶜS₂ᵖ,
+        ᶜSNₗᵖ,
+        ᶜSNᵣᵖ,
+        ᶜSqₗᵖ,
+        ᶜSqᵣᵖ,
+        Y.c.ρ,
+        Y.c.N_liq,
+        Y.c.N_rai,
+        q_liq,
+        q_rai,
+        ᶜts,
+        dt,
+        cm2p,
+        thp,
+    )
+
+    #TODO - implement 2M cold processes!
+
+    return nothing
+end
+function set_precipitation_cache!(
+    Y,
+    p,
+    ::Microphysics2Moment,
+    ::DiagnosticEDMFX,
+)
+    error("Not implemented yet")
+    return nothing
+end
+function set_precipitation_cache!(
+    Y,
+    p,
+    ::Microphysics2Moment,
+    ::PrognosticEDMFX,
+)
+    error("Not implemented yet")
+    return nothing
+end
 
 """
     set_precipitation_surface_fluxes!(Y, p, precipitation model)
@@ -301,7 +448,7 @@ end
 function set_precipitation_surface_fluxes!(
     Y,
     p,
-    precip_model::Microphysics1Moment,
+    precip_model::Union{Microphysics1Moment, Microphysics2Moment},
 )
     (; surface_rain_flux, surface_snow_flux) = p.precomputed
     (; col_integrated_precip_energy_tendency,) = p.conservation_check

src/cache/precomputed_quantities.jl

@@ -126,20 +126,35 @@ function precomputed_quantities(Y, atmos)
     )
     sedimentation_quantities =
         atmos.moisture_model isa NonEquilMoistModel ?
-        (; ᶜwₗ = similar(Y.c, FT), ᶜwᵢ = similar(Y.c, FT)) : (;)
+        (; ᶜwₗ = similar(Y.c, FT), ᶜwᵢ = similar(Y.c, FT), ᶜwNₗ = similar(Y.c, FT), ᶜwNᵢ = similar(Y.c, FT)) : (;)
     if atmos.precip_model isa Microphysics0Moment
         precipitation_quantities =
             (; ᶜS_ρq_tot = similar(Y.c, FT), ᶜS_ρe_tot = similar(Y.c, FT))
-    elseif atmos.precip_model isa Microphysics1Moment
-        precipitation_quantities = (;
-            ᶜwᵣ = similar(Y.c, FT),
-            ᶜwₛ = similar(Y.c, FT),
-            ᶜSqₗᵖ = similar(Y.c, FT),
-            ᶜSqᵢᵖ = similar(Y.c, FT),
-            ᶜSqᵣᵖ = similar(Y.c, FT),
-            ᶜSqₛᵖ = similar(Y.c, FT),
-        )
-    else
+        elseif atmos.precip_model isa Microphysics1Moment
+            precipitation_quantities = (;
+                ᶜwᵣ = similar(Y.c, FT),
+                ᶜwₛ = similar(Y.c, FT),
+                ᶜSqₗᵖ = similar(Y.c, FT),
+                ᶜSqᵢᵖ = similar(Y.c, FT),
+                ᶜSqᵣᵖ = similar(Y.c, FT),
+                ᶜSqₛᵖ = similar(Y.c, FT),
+            )
+        elseif atmos.precip_model isa Microphysics2Moment
+            precipitation_quantities = (;
+                ᶜwᵣ = similar(Y.c, FT),
+                ᶜwₛ = similar(Y.c, FT),
+                ᶜSqₗᵖ = similar(Y.c, FT),
+                ᶜSqᵢᵖ = similar(Y.c, FT),
+                ᶜSqᵣᵖ = similar(Y.c, FT),
+                ᶜSqₛᵖ = similar(Y.c, FT),
+                ᶜwNᵣ = similar(Y.c, FT),
+                ᶜwNₛ = similar(Y.c, FT),
+                ᶜSNₗᵖ = similar(Y.c, FT),
+                ᶜSNᵢᵖ = similar(Y.c, FT),
+                ᶜSNᵣᵖ = similar(Y.c, FT),
+                ᶜSNₛᵖ = similar(Y.c, FT),
+            )
+        else
         precipitation_quantities = (;)
     end
     precipitation_sgs_quantities =

src/diagnostics/Diagnostics.jl

@@ -28,6 +28,7 @@ import ..NonEquilMoistModel
 import ..NoPrecipitation
 import ..Microphysics0Moment
 import ..Microphysics1Moment
+import ..Microphysics2Moment
 
 # radiation
 import ClimaAtmos.RRTMGPInterface as RRTMGPI

src/diagnostics/core_diagnostics.jl

@@ -707,6 +707,7 @@ function compute_pr!(
         NoPrecipitation,
         Microphysics0Moment,
         Microphysics1Moment,
+        Microphysics2Moment,
     },
 )
     if isnothing(out)
@@ -742,6 +743,7 @@ function compute_prra!(
         NoPrecipitation,
         Microphysics0Moment,
         Microphysics1Moment,
+        Microphysics2Moment,
     },
 )
     if isnothing(out)
@@ -774,6 +776,7 @@ function compute_prsn!(
         NoPrecipitation,
         Microphysics0Moment,
         Microphysics1Moment,
+        Microphysics2Moment,
     },
 )
     if isnothing(out)
@@ -805,7 +808,10 @@ function compute_husra!(
     state,
     cache,
     time,
-    precip_model::Microphysics1Moment,
+    precip_model::Union{
+        Microphysics1Moment,
+        Microphysics2Moment,
+    },
 )
     if isnothing(out)
         return state.c.ρq_rai ./ state.c.ρ
@@ -836,7 +842,10 @@ function compute_hussn!(
     state,
     cache,
     time,
-    precip_model::Microphysics1Moment,
+    precip_model::Union{
+        Microphysics1Moment,
+        Microphysics2Moment,
+    },
 )
     if isnothing(out)
         return state.c.ρq_sno ./ state.c.ρ

src/initial_conditions/InitialConditions.jl

@@ -7,6 +7,7 @@ import ..NonEquilMoistModel
 import ..NoPrecipitation
 import ..Microphysics0Moment
 import ..Microphysics1Moment
+import ..Microphysics2Moment
 import ..PrescribedSurfaceTemperature
 import ..PrognosticSurfaceTemperature
 import ..ᶜinterp

src/initial_conditions/atmos_state.jl

@@ -127,6 +127,32 @@ precip_variables(ls, ::Microphysics1Moment) = (;
     ρq_rai = ls.ρ * ls.precip_state.q_rai,
     ρq_sno = ls.ρ * ls.precip_state.q_sno,
 )
+precip_variables(ls, ::Microphysics2Moment) = (;
+    ρq_rai = ls.ρ * ls.precip_state.q_rai,
+    ρq_sno = ls.ρ * ls.precip_state.q_sno,
+    N_rai = ls.precip_state.N_rai,
+    N_sno = ls.precip_state.N_sno,
+    N_liq = ls.precip_state.N_liq,
+    N_ice = ls.precip_state.N_ice,
+)
+
+"""
+    precip_variables(ls, ::MicrophysicsP3)
+
+Define prognostic variables for P3 microphysics scheme.
+"""
+precip_variables(ls, ::MicrophysicsP3) = (;
+    # Warm rain variables (from 2M)
+    ρq_rai = ls.ρ * ls.precip_state.q_rai,
+    N_rai = ls.precip_state.N_rai,
+    N_liq = ls.precip_state.N_liq,
+    # P3 ice variables
+    ρq_ice = ls.ρ * ls.precip_state.q_ice,
+    N_ice = ls.precip_state.N_ice,  # Rimed mass mixing ratio
+    L_ice = ls.precip_state.L_ice,  # Ice mass mixing ratio
+    B_rim = ls.precip_state.B_rim,  # Rimed volume mixing ratio
+    L_rim = ls.precip_state.L_rim,  # Rimed mass mixing ratio
+)
 
 # We can use paper-based cases for LES type configurations (no TKE)
 # or SGS type configurations (initial TKE needed), so we do not need to assert

src/parameterized_tendencies/microphysics/cloud_condensate.jl

@@ -16,7 +16,7 @@ function cloud_condensate_tendency!(
     ::Union{NoPrecipitation, Microphysics0Moment},
 )
     error(
-        "NonEquilMoistModel can only be run with Microphysics1Moment precipitation",
+        "NonEquilMoistModel can only be run with Microphysics1Moment or Microphysics2Moment precipitation",
     )
 end
 
@@ -37,3 +37,24 @@ function cloud_condensate_tendency!(
     @. Yₜ.c.ρq_liq += Y.c.ρ * cloud_sources(cmc.liquid, thp, ᶜts, q_rai, dt)
     @. Yₜ.c.ρq_ice += Y.c.ρ * cloud_sources(cmc.ice, thp, ᶜts, q_sno, dt)
 end
+
+function cloud_condensate_tendency!(
+    Yₜ,
+    Y,
+    p,
+    ::NonEquilMoistModel,
+    ::Microphysics2Moment,
+)
+    (; ᶜts) = p.precomputed
+    (; params, dt) = p
+    (; q_rai, q_sno) = p.precomputed.ᶜspecific
+    FT = eltype(params)
+    thp = CAP.thermodynamics_params(params)
+    cmc = CAP.microphysics_cloud_params(params)
+
+    @. Yₜ.c.ρq_liq += Y.c.ρ * cloud_sources(cmc.liquid, thp, ᶜts, q_rai, dt)
+    @. Yₜ.c.ρq_ice += Y.c.ρ * cloud_sources(cmc.ice, thp, ᶜts, q_sno, dt)
+
+    @. Yₜ.c.N_liq += aerosol_activation_sources(cmc.liquid, thp, ᶜts, Y.c.ρ, q_rai, Y.c.N_ice, cmc.N_cloud_liquid_droplets, dt)
+    @. Yₜ.c.N_ice += aerosol_activation_sources(cmc.ice, thp, ᶜts, Y.c.ρ, q_sno, Y.c.N_ice, cmc.N_cloud_liquid_droplets, dt)
+end