Merge pull request #3959 from CliMA/he/feat-add-p3-ice-microphysics

feat: add P3 ice microphysics model

Author: haakon-e

src/cache/precipitation_precomputed_quantities.jl

@@ -5,6 +5,7 @@
 import CloudMicrophysics.MicrophysicsNonEq as CMNe
 import CloudMicrophysics.Microphysics1M as CM1
 import CloudMicrophysics.Microphysics2M as CM2
+import CloudMicrophysics.P3Scheme as CMP3
 
 import Thermodynamics as TD
 import ClimaCore.Operators as Operators
@@ -506,6 +507,68 @@ function set_precipitation_velocities!(
     return nothing
 end
 
+function set_precipitation_velocities!(
+    Y, p, ::NonEquilMoistModel, ::Microphysics2MomentP3,
+)
+    ## liquid quantities (2M warm rain)
+    (; ᶜwₗ, ᶜwᵣ, ᶜwnₗ, ᶜwnᵣ, ᶜwₜqₜ, ᶜwₕhₜ, ᶜts, ᶜu) = p.precomputed
+    (; ᶜΦ) = p.core
+
+    (; ρ, ρq_liq, ρn_liq, ρq_rai, ρn_rai) = Y.c
+    (; sb, rtv, ctv) = p.params.microphysics_2mp3_params.warm
+    thp = CAP.thermodynamics_params(p.params)
+
+    # Number- and mass weighted rain terminal velocity [m/s]
+    ᶜrai_w_terms = @. lazy(
+        CM2.rain_terminal_velocity(
+            sb, rtv,
+            max(zero(ρ), specific(ρq_rai, ρ)),
+            ρ, max(zero(ρ), ρn_rai),
+        ),
+    )
+    @. ᶜwnᵣ = getindex(ᶜrai_w_terms, 1)
+    @. ᶜwᵣ = getindex(ᶜrai_w_terms, 2)
+    # Number- and mass weighted cloud liquid terminal velocity [m/s]
+    ᶜliq_w_terms = @. lazy(
+        CM2.cloud_terminal_velocity(
+            sb.pdf_c, ctv,
+            max(zero(ρ), specific(ρq_liq, ρ)),
+            ρ, max(zero(ρ), ρn_liq),
+        ),
+    )
+    @. ᶜwnₗ = getindex(ᶜliq_w_terms, 1)
+    @. ᶜwₗ = getindex(ᶜliq_w_terms, 2)
+
+    ## Ice quantities
+    (; ρq_ice, ρn_ice, ρq_rim, ρb_rim) = Y.c
+    (; ᶜwᵢ) = p.precomputed
+    (; cold) = CAP.microphysics_2mp3_params(p.params)
+
+    # Number- and mass weighted ice terminal velocity [m/s]
+    # Calculate terminal velocities
+    (; ᶜlogλ, ᶜwnᵢ) = p.precomputed
+    use_aspect_ratio = true  # TODO: Make a config option
+    ᶜF_rim = @. lazy(ρq_rim / ρq_ice)
+    ᶜρ_rim = @. lazy(ρq_rim / ρb_rim)
+    ᶜstate_p3 = @. lazy(CMP3.P3State(cold.params,
+        max(0, ρq_ice), max(0, ρn_ice), ᶜF_rim, ᶜρ_rim,
+    ))
+    @. ᶜlogλ = CMP3.get_distribution_logλ(ᶜstate_p3)
+    args = (cold.velocity_params, ρ, ᶜstate_p3, ᶜlogλ)
+    @. ᶜwnᵢ = CMP3.ice_terminal_velocity_number_weighted(args...; use_aspect_ratio)
+    @. ᶜwᵢ = CMP3.ice_terminal_velocity_mass_weighted(args...; use_aspect_ratio)
+
+    # compute their contributions to energy and total water advection
+    @. ᶜwₜqₜ = Geometry.WVector(ᶜwₗ * ρq_liq + ᶜwᵢ * ρq_ice + ᶜwᵣ * ρq_rai) / ρ
+    @. ᶜwₕhₜ =
+        Geometry.WVector(
+            ᶜwₗ * ρq_liq * (Iₗ(thp, ᶜts) + ᶜΦ + $(Kin(ᶜwₗ, ᶜu))) +
+            ᶜwᵢ * ρq_ice * (Iᵢ(thp, ᶜts) + ᶜΦ + $(Kin(ᶜwᵢ, ᶜu))) +
+            ᶜwᵣ * ρq_rai * (Iₗ(thp, ᶜts) + ᶜΦ + $(Kin(ᶜwᵣ, ᶜu))),
+        ) / ρ
+    return nothing
+end
+
 """
     set_precipitation_cache!(Y, p, microphysics_model, turbconv_model)
 
@@ -744,6 +807,37 @@ function set_precipitation_cache!(
     return nothing
 end
 
+function set_precipitation_cache!(Y, p, ::Microphysics2MomentP3, ::Nothing)
+    ### Rainy processes (2M)
+    (; turbconv_model) = p.atmos
+    set_precipitation_cache!(Y, p, Microphysics2Moment(), turbconv_model)
+    # NOTE: the above function sets `ᶜSqᵢᵖ` to `0`. For P3, need to update `ᶜSqᵢᵖ` below!!
+
+    ### Icy processes (P3)
+    (; ᶜScoll, ᶜts, ᶜlogλ) = p.precomputed
+
+    # get thermodynamics and microphysics params
+    (; params) = p
+    params_2mp3 = CAP.microphysics_2mp3_params(params)
+    thermo_params = CAP.thermodynamics_params(params)
+
+    ᶜY_reduced = (;
+        Y.c.ρ,
+        # condensate
+        Y.c.ρq_liq, Y.c.ρn_liq, Y.c.ρq_rai, Y.c.ρn_rai,
+        # ice
+        Y.c.ρq_ice, Y.c.ρn_ice, Y.c.ρq_rim, Y.c.ρb_rim,
+    )
+
+    # compute warm precipitation sources on the grid mean (based on SB2006 2M scheme)
+    compute_cold_precipitation_sources_P3!(
+        ᶜScoll, params_2mp3, thermo_params, ᶜY_reduced, ᶜts, ᶜlogλ,
+    )
+
+    return nothing
+
+end
+
 """
     set_precipitation_surface_fluxes!(Y, p, precipitation model)
 
@@ -833,3 +927,8 @@ function set_precipitation_surface_fluxes!(
 
     return nothing
 end
+
+function set_precipitation_surface_fluxes!(Y, p, ::Microphysics2MomentP3)
+    set_precipitation_surface_fluxes!(Y, p, Microphysics2Moment())
+    # TODO: Figure out what to do for ρn_ice, ρq_rim, ρb_rim
+end

src/cache/precomputed_quantities.jl

@@ -113,7 +113,8 @@ function precomputed_quantities(Y, atmos)
             ᶜSqᵣᵖ = similar(Y.c, FT),
             ᶜSqₛᵖ = similar(Y.c, FT),
         )
-    elseif atmos.microphysics_model isa Microphysics2Moment
+    elseif atmos.microphysics_model isa Union{Microphysics2Moment, Microphysics2MomentP3}
+        # 2-moment microphysics
         precipitation_quantities = (;
             ᶜwᵣ = similar(Y.c, FT),
             ᶜwₛ = similar(Y.c, FT),
@@ -126,6 +127,23 @@ function precomputed_quantities(Y, atmos)
             ᶜSnₗᵖ = similar(Y.c, FT),
             ᶜSnᵣᵖ = similar(Y.c, FT),
         )
+        # Add additional quantities for 2M + P3
+        if atmos.microphysics_model isa Microphysics2MomentP3
+            precipitation_quantities = (;
+                # liquid quantities (2M warm rain)
+                precipitation_quantities...,
+                # ice quantities (P3)
+                ᶜwᵢ = similar(Y.c, FT),
+                ᶜwnᵢ = similar(Y.c, FT),
+                ᶜlogλ = similar(Y.c, FT),
+                ᶜScoll = similar(Y.c,
+                    @NamedTuple{
+                        ∂ₜq_c::FT, ∂ₜq_r::FT, ∂ₜN_c::FT, ∂ₜN_r::FT,
+                        ∂ₜL_rim::FT, ∂ₜL_ice::FT, ∂ₜB_rim::FT,
+                    }
+                ),
+            )
+        end
     else
         precipitation_quantities = (;)
     end

src/diagnostics/Diagnostics.jl

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

src/diagnostics/core_diagnostics.jl

@@ -738,6 +738,7 @@ function compute_pr!(
         Microphysics0Moment,
         Microphysics1Moment,
         Microphysics2Moment,
+        Microphysics2MomentP3,
     },
 )
     if isnothing(out)
@@ -774,6 +775,7 @@ function compute_prra!(
         Microphysics0Moment,
         Microphysics1Moment,
         Microphysics2Moment,
+        Microphysics2MomentP3,
     },
 )
     if isnothing(out)
@@ -807,6 +809,7 @@ function compute_prsn!(
         Microphysics0Moment,
         Microphysics1Moment,
         Microphysics2Moment,
+        Microphysics2MomentP3,
     },
 )
     if isnothing(out)
@@ -838,7 +841,9 @@ function compute_husra!(
     state,
     cache,
     time,
-    microphysics_model::Union{Microphysics1Moment, Microphysics2Moment},
+    microphysics_model::Union{
+        Microphysics1Moment, Microphysics2Moment, Microphysics2MomentP3,
+    },
 )
     if isnothing(out)
         return state.c.ρq_rai ./ state.c.ρ
@@ -869,7 +874,9 @@ function compute_hussn!(
     state,
     cache,
     time,
-    microphysics_model::Union{Microphysics1Moment, Microphysics2Moment},
+    microphysics_model::Union{
+        Microphysics1Moment, Microphysics2Moment, Microphysics2MomentP3,
+    },
 )
     if isnothing(out)
         return state.c.ρq_sno ./ state.c.ρ
@@ -900,7 +907,7 @@ function compute_cdnc!(
     state,
     cache,
     time,
-    microphysics_model::Microphysics2Moment,
+    microphysics_model::Union{Microphysics2Moment, Microphysics2MomentP3},
 )
     if isnothing(out)
         return state.c.ρn_liq
@@ -931,7 +938,7 @@ function compute_ncra!(
     state,
     cache,
     time,
-    microphysics_model::Microphysics2Moment,
+    microphysics_model::Union{Microphysics2Moment, Microphysics2MomentP3},
 )
     if isnothing(out)
         return state.c.ρn_rai
@@ -1554,7 +1561,7 @@ function compute_rwp!(
     cache,
     time,
     moisture_model::T,
-) where {T <: Union{Microphysics1Moment, Microphysics2Moment}}
+) where {T <: Union{Microphysics1Moment, Microphysics2Moment, Microphysics2MomentP3}}
     if isnothing(out)
         out = zeros(axes(Fields.level(state.f, half)))
         rw = cache.scratch.ᶜtemp_scalar

src/initial_conditions/InitialConditions.jl

@@ -8,6 +8,7 @@ import ..NoPrecipitation
 import ..Microphysics0Moment
 import ..Microphysics1Moment
 import ..Microphysics2Moment
+import ..Microphysics2MomentP3
 import ..PrescribedSST
 import ..SlabOceanSST
 import ..ᶜinterp

src/initial_conditions/atmos_state.jl

@@ -134,6 +134,21 @@ precip_variables(ls, ::Microphysics2Moment) = (;
     ρq_sno = ls.ρ * ls.precip_state.q_sno,
 )
 
+function precip_variables(ls, ::Microphysics2MomentP3)
+    (; ρ) = ls
+    (; n_liq, n_rai, q_rai) = ls.precip_state.warm
+    (; n_ice, q_ice, q_rim, b_rim) = ls.precip_state.cold
+    # warm state
+    ls_warm = (; ρ, precip_state = ls.precip_state.warm)
+    warm_state = precip_variables(ls_warm, Microphysics2Moment())
+    # cold state
+    cold_state = (;
+        ρq_ice = ρ * q_ice, ρn_ice = ρ * n_ice,
+        ρq_rim = ρ * q_rim, ρb_rim = ρ * b_rim,
+    )
+    return (; warm_state..., cold_state...)
+end
+
 # 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
 # that there is no TKE when there is no turbconv_model.

src/parameterized_tendencies/microphysics/microphysics_wrappers.jl

@@ -751,3 +751,30 @@ function compute_warm_precipitation_sources_2M!(
     @. Snᵣᵖ += Sᵖ + S₂ᵖ
 
 end
+
+function compute_cold_precipitation_sources_P3!(
+    ᶜScoll,         # NamedTuple-valued Field with P3 liquid-ice collision sources
+    params_2mp3,    # Parameters for 2M and P3 schemes, see `get_microphysics_2m_p3_parameters`
+    thermo_params,  # An instance of `Thermodynamics.Parameters.ThermodynamicsParameters`
+    ᶜY_reduced,     # A reduced set of prognostic variables needed for P3 sources
+    ᶜts,            # Thermodynamic state
+    ᶜlogλ,          # Logarithm of the P3 distribution slope parameter
+)
+
+    (; warm, cold) = params_2mp3
+    (; ρ, ρq_liq, ρn_liq, ρq_rai, ρn_rai, ρq_ice, ρn_ice, ρq_rim, ρb_rim) = ᶜY_reduced
+
+    ᶜF_rim = @. lazy(ρq_rim / ρq_ice)
+    ᶜρ_rim = @. lazy(ρq_rim / ρb_rim)
+
+    @. ᶜScoll = CMP3.bulk_liquid_ice_collision_sources(cold.params, ᶜlogλ,
+        ρq_ice, max(0, ρn_ice),
+        ᶜF_rim, ᶜρ_rim,
+        warm.sb.pdf_c, warm.sb.pdf_r,
+        ρq_liq, ρn_liq, ρq_rai, ρn_rai,
+        warm.aps, thermo_params, (cold.velocity_params,),
+        ρ, Tₐ(thermo_params, ᶜts),
+    )
+
+    return nothing
+end

src/parameterized_tendencies/microphysics/precipitation.jl

@@ -259,3 +259,33 @@ function precipitation_tendency!(
         @. Yₜ.c.ρn_rai += Y.c.sgsʲs.:($$j).ρa * ᶜSnᵣᵖʲs.:($$j)
     end
 end
+
+####
+#### 2-moment warm microphysics with P3 scheme
+####
+
+function precipitation_tendency!(Yₜ, Y, p, t,
+    ne::NonEquilMoistModel, ::Microphysics2MomentP3, ::Nothing,
+)
+    (; ᶜSqₗᵖ, ᶜSqᵢᵖ, ᶜSqᵣᵖ, ᶜSqₛᵖ, ᶜScoll) = p.precomputed
+    (; ᶜSnₗᵖ, ᶜSnᵣᵖ) = p.precomputed
+
+    ## Update grid mean tendencies
+    # 2 moment scheme (warm)
+    precipitation_tendency!(Yₜ, Y, p, t, ne, Microphysics2Moment(), nothing)
+
+    # P3 scheme (cold)
+    # @. Yₜ.c.ρq_ice += Y.c.ρ * ᶜSqᵢᵖ  # updated in the above function
+    # @. Yₜ.c.ρn_ice += Y.c.ρ * ᶜSnᵢᵖ  # TODO: Additional sources for `ρn_ice`
+
+    # Collisions
+    @. Yₜ.c.ρq_liq += Y.c.ρ * ᶜScoll.∂ₜq_c
+    @. Yₜ.c.ρq_rai += Y.c.ρ * ᶜScoll.∂ₜq_r
+    @. Yₜ.c.ρn_liq += ᶜScoll.∂ₜN_c
+    @. Yₜ.c.ρn_rai += ᶜScoll.∂ₜN_r
+    @. Yₜ.c.ρq_rim += ᶜScoll.∂ₜL_rim
+    @. Yₜ.c.ρq_ice += ᶜScoll.∂ₜL_ice
+    @. Yₜ.c.ρb_rim += ᶜScoll.∂ₜB_rim
+
+    return nothing
+end

src/parameters/Parameters.jl

@@ -69,6 +69,7 @@ Base.@kwdef struct ClimaAtmosParameters{
     MP0M,
     MP1M,
     MP2M,
+    MP2MP3,
     SFP,
     TCP,
     STP,
@@ -82,6 +83,7 @@ Base.@kwdef struct ClimaAtmosParameters{
     microphysics_0m_params::MP0M
     microphysics_1m_params::MP1M
     microphysics_2m_params::MP2M
+    microphysics_2mp3_params::MP2MP3
     surface_fluxes_params::SFP
     turbconv_params::TCP
     surface_temp_params::STP