Wip on adding 2M microphysics to Atmos

src/cache/precipitation_precomputed_quantities.jl

@@ -2,20 +2,22 @@
 ##### Precomputed quantities
 #####
 import CloudMicrophysics.Microphysics1M as CM1
+import CloudMicrophysics.Microphysics2M as CM2
 
 # helper function to safely get precipitation from state
 function qₚ(ρqₚ::FT, ρ::FT) where {FT}
     return max(FT(0), ρqₚ / ρ)
 end
 
 """
-    set_precipitation_precomputed_quantities!(Y, p, t)
+    set_precipitation_precomputed_quantities!(Y, p, t, precip_model)
 
-Updates the precipitation terminal velocity stored in `p`
-for the 1-moment microphysics scheme
+Updates the precipitation terminal velocity and tracers stored in cache
 """
-function set_precipitation_precomputed_quantities!(Y, p, t)
-    @assert (p.atmos.precip_model isa Microphysics1Moment)
+function set_precipitation_precomputed_quantities!(Y, p, t, _)
+    return nothing
+end
+function set_precipitation_precomputed_quantities!(Y, p, t, ::Microphysics1Moment)
 
     (; ᶜwᵣ, ᶜwₛ, ᶜqᵣ, ᶜqₛ) = p.precomputed
 
@@ -40,3 +42,34 @@ function set_precipitation_precomputed_quantities!(Y, p, t)
     )
     return nothing
 end
+
+function set_precipitation_precomputed_quantities!(Y, p, t, ::Microphysics2Moment)
+
+    (; ᶜwᵣ, ᶜw_nᵣ, ᶜqᵣ) = p.precomputed
+
+    cmp = CAP.microphysics_precipitation_params(p.params)
+
+    # compute the precipitation specific humidities
+    @. ᶜqᵣ = qₚ(Y.c.ρq_rai, Y.c.ρ)
+    @. ᶜwᵣ = getindex(
+        CM2.rain_terminal_velocity(
+                cmp.SB2006,
+                cmp.SB2006Vel,
+                ᶜqᵣ,
+                Y.c.ρ,
+                Y.c.ρn_rai,
+        ),
+        2,
+    )
+    @. ᶜw_nᵣ = getindex(
+        CM2.rain_terminal_velocity(
+                cmp.SB2006,
+                cmp.SB2006Vel,
+                ᶜqᵣ,
+                Y.c.ρ,
+                Y.c.ρn_rai,
+        ),
+        1,
+    )
+    return nothing
+end

src/cache/precomputed_quantities.jl

@@ -146,14 +146,22 @@ function precomputed_quantities(Y, atmos)
     else
         (;)
     end
-    precipitation_quantities =
-        atmos.precip_model isa Microphysics1Moment ?
+    precipitation_quantities if atmos.precip_model isa Microphysics1Moment
         (;
             ᶜwᵣ = similar(Y.c, FT),
             ᶜwₛ = similar(Y.c, FT),
             ᶜqᵣ = similar(Y.c, FT),
             ᶜqₛ = similar(Y.c, FT),
-        ) : (;)
+        )
+    elseif atmos.precip_model isa Microphysics2Moment
+        (;
+            ᶜwᵣ = similar(Y.c, FT),
+            ᶜw_nᵣ = similar(Y.c, FT),
+            ᶜqᵣ = similar(Y.c, FT),
+        )
+    else
+        (;)
+    end
     return (;
         gs_quantities...,
         sgs_quantities...,
@@ -503,9 +511,7 @@ NVTX.@annotate function set_precomputed_quantities!(Y, p, t)
     #     compute_gm_mixing_length!(ᶜmixing_length, Y, p)
     # end
 
-    if precip_model isa Microphysics1Moment
-        set_precipitation_precomputed_quantities!(Y, p, t)
-    end
+    set_precipitation_precomputed_quantities!(Y, p, t, p.atmos.precip_model)
 
     if turbconv_model isa PrognosticEDMFX
         set_prognostic_edmf_precomputed_quantities_draft_and_bc!(Y, p, ᶠuₕ³, t)

src/diagnostics/core_diagnostics.jl

@@ -648,7 +648,7 @@ function compute_husra!(
     state,
     cache,
     time,
-    precip_model::Microphysics1Moment,
+    precip_model::Union{Microphysics1Moment, Microphysics2Moment},
 )
     if isnothing(out)
         return state.c.ρq_rai ./ state.c.ρ
@@ -700,6 +700,71 @@ add_diagnostic_variable!(
     compute! = compute_hussn!,
 )
 
+###
+# Number concentrations for cloud and precipitation (3d)
+###
+compute_cdnc!(out, state, cache, time) =
+    compute_cdnc!(out, state, cache, time, cache.atmos.precip_model)
+compute_cdnc!(_, _, _, _, model::T) where {T} =
+    error_diagnostic_variable("cdnc", model)
+
+function compute_cdnc!(
+    out,
+    state,
+    cache,
+    time,
+    precip_model::Microphysics2Moment,
+)
+    if isnothing(out)
+        return state.c.ρn_liq
+    else
+        out .= state.c.ρn_liq
+    end
+end
+
+add_diagnostic_variable!(
+    short_name = "cdnc",
+    long_name = "Cloud Droplet Number Concentration in liquid water clouds",
+    standard_name = "cloud_droplet_number_concentration",
+    units = "m^-3",
+    comments = """
+    This is calculated as the number of cloud droplets in the grid cell per
+    volume of air.
+    """,
+    compute! = compute_cdnc!,
+)
+
+compute_rdnc!(out, state, cache, time) =
+    compute_rdnc!(out, state, cache, time, cache.atmos.precip_model)
+compute_rdnc!(_, _, _, _, model::T) where {T} =
+    error_diagnostic_variable("rdnc", model)
+
+function compute_rdnc!(
+    out,
+    state,
+    cache,
+    time,
+    precip_model::Microphysics2Moment,
+)
+    if isnothing(out)
+        return state.c.ρn_rai
+    else
+        out .= state.c.ρn_rai
+    end
+end
+
+add_diagnostic_variable!(
+    short_name = "rdnc",
+    long_name = "Rain Drop Number Concentration",
+    standard_name = "rain_drop_number_concentration",
+    units = "m^-3",
+    comments = """
+    This is calculated as the number of rain drops in the grid cell per
+    volume of air.
+    """,
+    compute! = compute_rdnc!,
+)
+
 ###
 # Topography
 ###

src/initial_conditions/atmos_state.jl

@@ -122,6 +122,11 @@ precip_variables(ls, ::Microphysics1Moment) = (;
     ρq_rai = ls.ρ * ls.precip_state.q_rai,
     ρq_sno = ls.ρ * ls.precip_state.q_sno,
 )
+precip_variables(ls, ::Microphysics2Moment) = (;
+    ρn_liq = ls.ρ * ls.precip_state.n_liq,
+    ρq_rai = ls.ρ * ls.precip_state.q_rai,
+    ρn_rai = ls.ρ * ls.precip_state.n_rai,
+)
 
 # 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/initial_conditions/initial_conditions.jl

@@ -426,7 +426,7 @@ function deep_atmos_baroclinic_wave_values(z, ϕ, λ, params, perturb)
     Ω = CAP.Omega(params)
     R = CAP.planet_radius(params)
 
-    # Constants from paper (See Table 1. in Ullrich et al (2014)) 
+    # Constants from paper (See Table 1. in Ullrich et al (2014))
     k = 3         # Power for temperature field
     T_e = FT(310) # Surface temperature at the equator
     T_p = FT(240) # Surface temperature at the pole
@@ -1124,6 +1124,44 @@ function (initial_condition::PrecipitatingColumn)(params)
     end
     return local_state
 end
+function (initial_condition::PrecipitatingColumn2M)(params)
+    FT = eltype(params)
+    thermo_params = CAP.thermodynamics_params(params)
+    p_0 = FT(101300.0)
+    qᵣ = prescribed_prof(FT, 2000, 5000, 1e-6)
+    qₗ = prescribed_prof(FT, 4000, 5500, 2e-5)
+    qᵢ = prescribed_prof(FT, 6000, 9000, 0)
+    nₗ = prescribed_prof(FT, 4000, 5500, 1e8)
+    nᵣ = prescribed_prof(FT, 2000, 5000, 1e6)
+    θ = APL.Rico_θ_liq_ice(FT)
+    q_tot = APL.Rico_q_tot(FT)
+    u = prescribed_prof(FT, 0, Inf, 0)
+    v = prescribed_prof(FT, 0, Inf, 0)
+    p = hydrostatic_pressure_profile(; thermo_params, p_0, θ, q_tot)
+    function local_state(local_geometry)
+        (; z) = local_geometry.coordinates
+        ts = TD.PhaseNonEquil_pθq(
+            thermo_params,
+            p(z),
+            θ(z),
+            TD.PhasePartition(q_tot(z), qₗ(z), qᵢ(z)),
+        )
+        return LocalState(;
+            params,
+            geometry = local_geometry,
+            thermo_state = ts,
+            velocity = Geometry.UVVector(u(z), v(z)),
+            turbconv_state = nothing,
+            precip_state = PrecipState2M(;
+                n_liq = nₗ(z),
+                q_rai = qᵣ(z),
+                n_rai = nᵣ(z)
+            ),
+        )
+    end
+    return local_state
+end
+
 
 """
     GCMDriven <: InitialCondition

src/initial_conditions/local_state.jl

@@ -104,7 +104,7 @@ struct NoPrecipState{FT} <: PrecipState{FT} end
 """
     PrecipState1M(; q_rai, q_sno)
 
-Stores the values of `ρq_rai` and `ρq_sno` for the `precip_model`.
+Stores the values of `q_rai` and `q_sno` for the `precip_model`.
 If no values are provided, they are set to zero.
 """
 struct PrecipState1M{FT} <: PrecipState{FT}
@@ -113,3 +113,17 @@ struct PrecipState1M{FT} <: PrecipState{FT}
 end
 PrecipState1M(; q_rai = 0, q_sno = 0) =
     PrecipState1M{typeof(q_rai)}(q_rai, q_sno)
+
+"""
+    PrecipState2M(; n_liq, q_rai, n_rai)
+
+Stores the values of `q_rai` and `n_rai` and `n_liq` for the `precip_model`.
+If no values are provided, they are set to zero.
+"""
+struct PrecipState2M{FT} <: PrecipState{FT}
+    n_liq::FT
+    q_rai::FT
+    n_rai::FT
+end
+PrecipState2M(; n_liq = 0, q_rai = 0, n_rai = 0) =
+    PrecipState2M{typeof(q_rai)}(n_liq, q_rai, n_rai)