add 2M diagnostics

Author: sajjadazimi

src/diagnostics/core_diagnostics.jl

@@ -866,6 +866,130 @@ add_diagnostic_variable!(
     compute! = compute_hussn!,
 )
 
+compute_clwnm!(out, state, cache, time) =
+    compute_clwnm!(out, state, cache, time, cache.atmos.precip_model)
+compute_clwnm!(_, _, _, _, model::T) where {T} =
+    error_diagnostic_variable("clwnm", model)
+
+function compute_clwnm!(
+    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 = "clwnm",
+    long_name = "Number Concentration of Cloud Liquid Water",
+    standard_name = "number_concentration_of_cloud_liquid_water",
+    units = "1 m^3",
+    comments = """
+    This is calculated as the number of cloud liquid water droplets in the grid 
+    cell divided by the cell volume.
+    """,
+    compute! = compute_clwnm!,
+)
+
+compute_clinm!(out, state, cache, time) =
+    compute_clinm!(out, state, cache, time, cache.atmos.precip_model)
+compute_clinm!(_, _, _, _, model::T) where {T} =
+    error_diagnostic_variable("clinm", model)
+
+function compute_clwnm!(
+    out,
+    state,
+    cache,
+    time,
+    precip_model::Microphysics2Moment,
+)
+    if isnothing(out)
+        return state.c.ρn_ice
+    else
+        out .= state.c.ρn_ice
+    end
+end
+
+add_diagnostic_variable!(
+    short_name = "clinm",
+    long_name = "Number Concentration of Cloud Ice",
+    standard_name = "number_concentration_of_cloud_ice",
+    units = "1 m^3",
+    comments = """
+    This is calculated as the number of cloud ice chrystals in the grid 
+    cell divided by the cell volume.
+    """,
+    compute! = compute_clinm!,
+)
+
+compute_nmra!(out, state, cache, time) =
+    compute_nmra!(out, state, cache, time, cache.atmos.precip_model)
+compute_nmra!(_, _, _, _, model::T) where {T} =
+    error_diagnostic_variable("nmra", model)
+
+function compute_nmra!(
+    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 = "nmra",
+    long_name = "Number Concentration of Rain",
+    standard_name = "number_concentration_of_rain",
+    units = "1 m^3",
+    comments = """
+    This is calculated as the number of raindrops in the grid cell divided
+    by the cell volume.
+    """,
+    compute! = compute_nmra!,
+)
+
+compute_nmsn!(out, state, cache, time) =
+    compute_nmsn!(out, state, cache, time, cache.atmos.precip_model)
+compute_nmsn!(_, _, _, _, model::T) where {T} =
+    error_diagnostic_variable("nmsn", model)
+
+function compute_nmsn!(
+    out,
+    state,
+    cache,
+    time,
+    precip_model::Microphysics2Moment,
+)
+    if isnothing(out)
+        return state.c.ρn_sno
+    else
+        out .= state.c.ρn_sno
+    end
+end
+
+add_diagnostic_variable!(
+    short_name = "nmsn",
+    long_name = "Number Concentration of Snow",
+    standard_name = "number_concentration_of_snow",
+    units = "1 m^3",
+    comments = """
+    This is calculated as the number of snow flakes in the grid cell divided
+    by the cell volume.
+    """,
+    compute! = compute_nmsn!,
+)
+
 ###
 # Topography
 ###

src/diagnostics/default_diagnostics.jl

@@ -230,6 +230,19 @@ function default_diagnostics(
     return [average_func(precip_diagnostics...; output_writer, start_date)...]
 end
 
+function default_diagnostics(
+    ::Microphysics2Moment,
+    duration,
+    start_date;
+    output_writer,
+)
+    precip_diagnostics = ["husra", "hussn", "clwnm", "clinm", "nmra", "nmsn"]
+
+    average_func = frequency_averages(duration)
+
+    return [average_func(precip_diagnostics...; output_writer, start_date)...]
+end
+
 ##################
 # Radiation mode #
 ##################

src/initial_conditions/initial_conditions.jl

@@ -1200,9 +1200,13 @@ function (initial_condition::PrecipitatingColumn)(params)
     thermo_params = CAP.thermodynamics_params(params)
     p_0 = FT(101300.0)
     qᵣ = prescribed_prof(FT, 2000, 5000, 1e-6)
+    nᵣ = prescribed_prof(FT, 2000, 5000, 1e3)
     qₛ = prescribed_prof(FT, 5000, 8000, 2e-6)
+    nₛ = prescribed_prof(FT, 5000, 8000, 1e3)
     qₗ = prescribed_prof(FT, 4000, 5500, 2e-5)
+    nₗ = prescribed_prof(FT, 4000, 5500, 1e7)
     qᵢ = prescribed_prof(FT, 6000, 9000, 1e-5)
+    nᵢ = prescribed_prof(FT, 6000, 9000, 1e6)
     θ = APL.Rico_θ_liq_ice(FT)
     q_tot = APL.Rico_q_tot(FT)
     u = prescribed_prof(FT, 0, Inf, 0)
@@ -1222,7 +1226,7 @@ function (initial_condition::PrecipitatingColumn)(params)
             thermo_state = ts,
             velocity = Geometry.UVVector(u(z), v(z)),
             turbconv_state = nothing,
-            precip_state = PrecipState1M(; q_rai = qᵣ(z), q_sno = qₛ(z)),
+            precip_state = PrecipState2M(; q_rai = qᵣ(z), q_sno = qₛ(z), n_liq = nₗ(z), n_ice = nᵢ(z), n_rai = nᵣ(z), n_sno = nₛ(z)),
         )
     end
     return local_state

src/initial_conditions/local_state.jl

@@ -102,14 +102,18 @@ struct NoPrecipState{FT} <: PrecipState{FT} end
 @inline Base.getproperty(::NoPrecipState{FT}, ::Symbol) where {FT} = FT(0)
 
 """
-    PrecipState1M(; q_rai, q_sno)
+    PrecipState2M(; q_rai, q_sno)
 
-Stores the values of `ρq_rai` and `ρq_sno` for the `precip_model`.
-If no values are provided, they are set to zero.
+Stores the values of `q_rai`, `q_sno`, `n_liq`, `n_ice`, `n_rai` and `n_sno` 
+for the `precip_model`. If no values are provided, they are set to zero.
 """
-struct PrecipState1M{FT} <: PrecipState{FT}
+struct PrecipState2M{FT} <: PrecipState{FT}
     q_rai::FT
     q_sno::FT
+    n_liq::FT
+    n_ice::FT
+    n_rai::FT
+    n_sno::FT
 end
-PrecipState1M(; q_rai = 0, q_sno = 0) =
-    PrecipState1M{typeof(q_rai)}(q_rai, q_sno)
+PrecipState2M(; q_rai = 0, q_sno = 0, n_liq = 0, n_ice = 0, n_rai = 0, n_sno = 0) =
+    PrecipState2M{typeof(q_rai)}(q_rai, q_sno, n_liq, n_ice, n_rai, n_sno)