Merge #102

102: Sa/add 2 m feature r=sajjadazimi a=sajjadazimi



Co-authored-by: Sajjad Azimi <[email protected]>

.buildkite/pipeline.yml

@@ -92,6 +92,14 @@ steps:
         command: "julia --color=yes --project=test test/experiments/KiD_driver/KiD_driver.jl --moisture_choice=EquilibriumMoisture --prognostic_vars=RhodTQ --precipitation_choice=Precipitation1M --rain_formation_scheme_choice LD2004 --prescribed_Nd 1e8"
         artifact_paths: "test/experiments/KiD_driver/Output_EquilibriumMoisture_RhodTQ_Precipitation1M_LD2004/figures/*"
 
+      - label: ":crystal_ball: Experiments: equil (fixed rhod and T) + SB2006 "
+        command: "julia --color=yes --project=test test/experiments/KiD_driver/KiD_driver.jl --moisture_choice=EquilibriumMoisture --prognostic_vars=RhodTQ --precipitation_choice=Precipitation2M --rain_formation_scheme_choice SB2006 --prescribed_Nd 1e8"
+        artifact_paths: "test/experiments/KiD_driver/Output_EquilibriumMoisture_RhodTQ_Precipitation2M_SB2006/figures/*"
+
+      - label: ":crystal_ball: Experiments: non equil (fixed rhod and T) + SB2006 "
+        command: "julia --color=yes --project=test test/experiments/KiD_driver/KiD_driver.jl --moisture_choice=NonEquilibriumMoisture --prognostic_vars=RhodTQ --precipitation_choice=Precipitation2M --rain_formation_scheme_choice SB2006 --prescribed_Nd 1e8"
+        artifact_paths: "test/experiments/KiD_driver/Output_NonEquilibriumMoisture_RhodTQ_Precipitation2M_SB2006/figures/*"
+
       - label: ":crystal_ball: Experiments: perfect-model calibration + EKP "
         command: "julia --color=yes --project=test test/experiments/calibrations/run_calibration.jl"
         artifact_paths: "output/*"

Project.toml

@@ -1,7 +1,7 @@
 name = "Kinematic1D"
 uuid = "a8a3bb05-6cc5-4342-abf6-5cc636bd2b35"
 authors = ["Climate Modeling Alliance"]
-version = "0.3.2"
+version = "0.3.0"
 
 [deps]
 ClimaCore = "d414da3d-4745-48bb-8d80-42e94e092884"
@@ -29,8 +29,13 @@ ClimaCore = "0.10"
 ClimaCorePlots = "0.2"
 CloudMicrophysics = "0.10"
 DiffEqBase = "6.75"
+Distributions = "0.25"
 EnsembleKalmanProcesses = "0.13"
+Interpolations = "0.14"
+JLD2 = "0.4"
 NCDatasets = "0.12"
+Optim = "1.7"
+OrdinaryDiffEq = "6.29"
 Plots = "1.29"
 ProgressLogging = "0.1"
 TerminalLoggers = "0.1"

src/Kinematic1D.jl

@@ -21,6 +21,8 @@ const CMNe = CM.MicrophysicsNonEq
 const CM0 = CM.Microphysics0M
 const CM1 = CM.Microphysics1M
 const CM2 = CM.Microphysics2M
+const CMAM = CM.AerosolModel
+const CMAA = CM.AerosolActivation
 
 # modules used only for calibrateKiD
 using Distributions

src/calibrateKiD/KiDUtils.jl

@@ -155,6 +155,12 @@ function get_variable_data_from_ODE(u, ρ::Vector{Float64}, var::String)
     elseif var == "rlr"
         _qtot = parent(u.ρq_tot) ./ ρ
         output = (parent(u.ρq_liq) .+ parent(u.ρq_rai)) ./ ρ ./ (1 .- _qtot)
+    elseif var == "Nl"
+        output = parent(u.N_liq)
+    elseif var == "Nr"
+        output = parent(u.N_rai)
+    elseif var == "Na"
+        output = parent(u.N_aer)
     elseif var == "rho"
         output = ρ
     elseif var == "rain averaged terminal velocity"
@@ -240,6 +246,12 @@ function equation_types(moisture_choice::String, precipitation_choice::String, r
         else
             error("Invalid rain formation choice: $rain_formation_choice")
         end
+    elseif precipitation_choice == "Precipitation2M"
+        if rain_formation_choice == "SB2006"
+            precip = Precipitation2M(CMT.SB2006Type())
+        else
+            error("Invalid rain formation choice: $rain_formation_choice")
+        end
     else
         error("Invalid precipitation choice: $precipitation_choice")
     end
@@ -275,6 +287,9 @@ function create_parameter_set(FT, model_settings::Dict, params_cal::Dict)
         precip_sinks = 1,
         prescribed_Nd = model_settings["Nd"],
         qtot_flux_correction = Int(model_settings["qtot_flux_correction"]),
+        r_dry = model_settings["r_dry"],
+        std_dry = model_settings["std_dry"],
+        κ = model_settings["κ"],
         microphys_params,
     )
     return param_set

src/calibrateKiD/ReferenceModels.jl

@@ -143,6 +143,12 @@ function get_single_obs_field(
             _data = (_data_pysdm["qc"] .+ _data_pysdm["qr"]) .* 1e-3
         elseif var == "rtr"
             _data = _data_pysdm["qv"] .+ (_data_pysdm["qc"] .+ _data_pysdm["qr"]) .* 1e-3
+        elseif var == "Nl"
+            _data = _data_pysdm["nc"]
+        elseif var == "Nr"
+            _data = _data_pysdm["nr"]
+        elseif var == "Na"
+            _data = _data_pysdm["na"]
         else
             _data = _data_pysdm[var]
         end

src/driveKiD/EquationTypes.jl

@@ -19,6 +19,7 @@ export NonEquilibriumMoisture_ρdTq
 export NoPrecipitation
 export Precipitation0M
 export Precipitation1M
+export Precipitation2M
 
 export OneMomentRainFormation
 
@@ -42,5 +43,8 @@ struct Precipitation0M <: AbstractPrecipitationStyle end
 struct Precipitation1M{PT} <: AbstractPrecipitationStyle
     rain_formation::PT
 end
+struct Precipitation2M{PT} <: AbstractPrecipitationStyle
+    rain_formation::PT
+end
 
 struct OneMomentRainFormation <: AbstractRainFormationStyle end

src/driveKiD/KiD_model.jl

@@ -79,6 +79,26 @@ function initialise_state(::NonEquilibriumMoisture, ::Precipitation1M, initial_p
         ρq_sno = initial_profiles.ρq_sno,
     )
 end
+function initialise_state(::EquilibriumMoisture, ::Precipitation2M, initial_profiles)
+    return CC.Fields.FieldVector(;
+        ρq_tot = initial_profiles.ρq_tot,
+        ρq_rai = initial_profiles.ρq_rai,
+        N_liq = initial_profiles.N_liq,
+        N_rai = initial_profiles.N_rai,
+        N_aer = initial_profiles.N_aer,
+    )
+end
+function initialise_state(::NonEquilibriumMoisture, ::Precipitation2M, initial_profiles)
+    return CC.Fields.FieldVector(;
+        ρq_tot = initial_profiles.ρq_tot,
+        ρq_liq = initial_profiles.ρq_liq,
+        ρq_ice = initial_profiles.ρq_ice,
+        ρq_rai = initial_profiles.ρq_rai,
+        N_liq = initial_profiles.N_liq,
+        N_rai = initial_profiles.N_rai,
+        N_aer = initial_profiles.N_aer,
+    )
+end
 
 """
    Interface to ODE solver. It initializes the auxiliary state.
@@ -119,16 +139,28 @@ function initialise_aux(FT, ip, params, TS, Stats, face_space, moisture)
             q_ice = ip.q_ice,
             ts = ts,
         ),
-        precip_variables = CC.Fields.FieldVector(; q_rai = ip.q_rai, q_sno = ip.q_sno),
+        precip_variables = CC.Fields.FieldVector(;
+            q_rai = ip.q_rai,
+            q_sno = ip.q_sno,
+            N_liq = ip.N_liq,
+            N_rai = ip.N_rai,
+        ),
+        aerosol_variables = CC.Fields.FieldVector(; N_aer = ip.N_aer, N_aer_0 = ip.N_aer_0, S_N_aer = ip.S_Na),
         moisture_sources = CC.Fields.FieldVector(; S_q_liq = ip.S_ql_moisture, S_q_ice = ip.S_qi_moisture),
         precip_sources = CC.Fields.FieldVector(;
             S_q_tot = ip.S_qt_precip,
             S_q_liq = ip.S_ql_precip,
             S_q_ice = ip.S_qi_precip,
             S_q_rai = ip.S_qr_precip,
             S_q_sno = ip.S_qs_precip,
+            S_N_liq = ip.S_Nl_precip,
+            S_N_rai = ip.S_Nr_precip,
+        ),
+        precip_velocities = CC.Fields.FieldVector(;
+            term_vel_rai = term_vel_rai,
+            term_vel_sno = term_vel_sno,
+            term_vel_N_rai = term_vel_rai,
         ),
-        precip_velocities = CC.Fields.FieldVector(; term_vel_rai = term_vel_rai, term_vel_sno = term_vel_sno),
         prescribed_velocity = CC.Fields.FieldVector(; ρw = ρw, ρw0 = ρw0),
         params = params,
         q_surf = q_surf,

src/driveKiD/Parameters.jl

@@ -22,6 +22,9 @@ Base.@kwdef struct KinematicParameters{FT, MP} <: AKP
     precip_sinks::Int
     prescribed_Nd::FT
     qtot_flux_correction::Int
+    r_dry::FT
+    std_dry::FT
+    κ::FT
     microphys_params::MP
 end
 thermodynamics_params(ps::AKP) = CM.Parameters.thermodynamics_params(ps.microphys_params)
@@ -33,6 +36,10 @@ t1(ps::AKP) = ps.t1
 precip_sources(ps::AKP) = ps.precip_sources
 precip_sinks(ps::AKP) = ps.precip_sinks
 prescribed_Nd(ps::AKP) = ps.prescribed_Nd
+#aerosol parameters for 2M scheme
+r_dry(ps::AKP) = ps.r_dry
+std_dry(ps::AKP) = ps.std_dry
+κ(ps::AKP) = ps.κ
 
 # Forward parameters to Thermodynamics
 const TDPS = TD.Parameters.ThermodynamicsParameters

src/driveKiD/helper_functions.jl

@@ -4,3 +4,40 @@
 @inline function ρw_helper(t, w1, t1)
     return t < t1 ? w1 * sin(pi * t / t1) : 0.0
 end
+
+"""
+    Returns the number of new activated aerosol particles and updates aerosol number density
+"""
+@inline function aerosol_activation_helper(params, q_tot, q_liq, N_aer, N_aer_0, T, p, ρ, ρw, dt)
+
+    microphys_params = KP.microphysics_params(params)
+    thermo_params = CM.Parameters.thermodynamics_params(microphys_params)
+
+    FT = eltype(q_tot)
+    S_Nl::FT = FT(0)
+    S_Na::FT = FT(0)
+
+    q = TD.PhasePartition(q_tot, q_liq, FT(0))
+    S::FT = TD.supersaturation(thermo_params, q, ρ, T, TD.Liquid())
+
+    if (S < FT(0))
+        return (; S_Nl, S_Na)
+    end
+
+    r_dry = KP.r_dry(params)
+    std_dry = KP.std_dry(params)
+    κ = KP.κ(params)
+    w = ρw / ρ
+
+    aerosol_distribution = CMAM.AerosolDistribution((CMAM.Mode_κ(r_dry, std_dry, N_aer_0, FT(1), FT(1), FT(0), κ, 1),))
+    N_act = CMAA.N_activated_per_mode(microphys_params, aerosol_distribution, T, p, w, q)[1]
+
+    if isnan(N_act)
+        return (; S_Nl, S_Na)
+    end
+
+    S_Nl = max(0, N_act - (N_aer_0 - N_aer)) / dt
+    S_Na = -S_Nl
+
+    return (; S_Nl, S_Na)
+end

src/driveKiD/initial_condition.jl

@@ -111,6 +111,7 @@ function init_1d_column(::Type{FT}, params, ρ_profile, z; dry = false) where {F
 
     ρ::FT = ρ_profile(z)
     ρ_dry::FT = SDM_ρ_dry_of_ρ(ρ, q_vap)
+    ρ_SDP = 1.225
 
     # assuming no cloud condensate in the initial profile
     θ_liq_ice::FT = θ_std # TODO - compute this based on TS
@@ -132,6 +133,10 @@ function init_1d_column(::Type{FT}, params, ρ_profile, z; dry = false) where {F
     q_sno::FT = FT(0.0)
     ρq_rai::FT = q_rai * ρ
     ρq_sno::FT = q_sno * ρ
+    N_liq::FT = FT(0)
+    N_rai::FT = FT(0)
+    N_aer_0::FT = Parameters.prescribed_Nd(params) * ρ_dry / ρ_SDP
+    N_aer::FT = N_aer_0
 
     S_ql_moisture::FT = FT(0.0)
     S_qi_moisture::FT = FT(0.0)
@@ -141,6 +146,9 @@ function init_1d_column(::Type{FT}, params, ρ_profile, z; dry = false) where {F
     S_qi_precip::FT = FT(0.0)
     S_qr_precip::FT = FT(0.0)
     S_qs_precip::FT = FT(0.0)
+    S_Nl_precip::FT = FT(0)
+    S_Nr_precip::FT = FT(0)
+    S_Na::FT = FT(0)
 
     return (;
         ρ,
@@ -159,12 +167,19 @@ function init_1d_column(::Type{FT}, params, ρ_profile, z; dry = false) where {F
         q_ice,
         q_rai,
         q_sno,
+        N_liq,
+        N_rai,
+        N_aer,
+        N_aer_0,
         S_ql_moisture,
         S_qi_moisture,
         S_qt_precip,
         S_ql_precip,
         S_qi_precip,
         S_qr_precip,
         S_qs_precip,
+        S_Nl_precip,
+        S_Nr_precip,
+        S_Na,
     )
 end