P3 skeleton framework (#185)

add p3 advection case

.buildkite/pipeline.yml

@@ -135,5 +135,7 @@ steps:
         command: "julia --color=yes --project=test test/experiments/KiD_driver/KiD_driver.jl --moisture_choice=CloudyMoisture --precipitation_choice=CloudyPrecip --n_elem=128 --num_moments=7"
         artifact_paths: "test/experiments/KiD_driver/Output_CloudyMoisture_CloudyPrecip_7/figures/*"
 
-
+      - label: ":snowflake: Experiment: P3 Setup"
+        command: "julia --color=yes --project=test test/experiments/KiD_driver/KiD_driver.jl --moisture_choice=MoistureP3 --precipitation_choice=PrecipitationP3 --n_elem=24 --z_max=3000 --t_end=75 --w1=0 --dt=0.5"
+        artifact_paths: "test/experiments/KiD_driver/Output_MoistureP3_PrecipitationP3/figures/*"
 

Project.toml

@@ -30,7 +30,7 @@ Thermodynamics = "b60c26fb-14c3-4610-9d3e-2d17fe7ff00c"
 [compat]
 ClimaCore = "0.14"
 ClimaParams = "0.10.7"
-CloudMicrophysics = "0.22.1"
+CloudMicrophysics = "0.22.2"
 Cloudy = "0.5.6"
 DiffEqBase = "6"
 Distributions = "0.25"

src/Common/Common.jl

@@ -17,6 +17,7 @@ const CM1 = CM.Microphysics1M
 const CM2 = CM.Microphysics2M
 const CMAM = CM.AerosolModel
 const CMAA = CM.AerosolActivation
+const CMP3 = CM.P3Scheme
 
 include("parameters.jl")
 include("equation_types.jl")

src/Common/equation_types.jl

@@ -52,9 +52,10 @@ struct Precipitation2M{PT, ST} <: AbstractPrecipitationStyle
     rain_formation::PT
     sedimentation::ST
 end
-struct PrecipitationP3{PT, ST, P3} <: AbstractPrecipitationStyle
-    rain_formation::PT
-    sedimentation::ST
-    p3_parameters::P3
+struct PrecipitationP3{P3, CH, PT, B} <: AbstractPrecipitationStyle
+    p3_params::P3
+    Chen2022::CH
+    sb2006::PT
+    p3_boundary_condition::B
 end
 struct CloudyPrecip <: AbstractPrecipitationStyle end

src/Common/helper_functions.jl

@@ -8,6 +8,8 @@ function get_moisture_type(moisture_choice::String, toml_dict)
         moisture = NonEquilibriumMoisture(CMP.CloudLiquid(toml_dict), CMP.CloudIce(toml_dict))
     elseif moisture_choice == "CloudyMoisture"
         moisture = CloudyMoisture()
+    elseif moisture_choice == "MoistureP3"
+        moisture = MoistureP3()
     else
         error("Invalid moisture choice: $moisture_choice")
     end
@@ -22,6 +24,7 @@ function get_precipitation_type(
     toml_dict;
     rain_formation_choice::Union{Nothing, String} = nothing,
     sedimentation_choice::Union{Nothing, String} = nothing,
+    boundary::Union{Nothing, NamedTuple} = nothing,
 )
     if precipitation_choice == "NoPrecipitation"
         precip = NoPrecipitation()
@@ -77,6 +80,12 @@ function get_precipitation_type(
         end
     elseif precipitation_choice == "CloudyPrecip"
         precip = CloudyPrecip()
+    elseif precipitation_choice == "PrecipitationP3"
+        FT = Float64
+        p3_params = CMP.ParametersP3(FT)
+        Chen2022 = CMP.Chen2022VelType(FT)
+        sb2006 = CMP.SB2006(toml_dict)
+        precip = PrecipitationP3(p3_params, Chen2022, sb2006, boundary)
     else
         error("Invalid precipitation choice: $precipitation_choice")
     end

src/Common/initial_condition.jl

@@ -281,3 +281,117 @@ function cloudy_initial_condition(pdists, ip, k = 1)
 
     return merge(ip, (; moments = moments, pdists = pdists, cloudy_moments_zero))
 end
+
+function p3_initial_condition(
+    ::Type{FT},
+    kid_params,
+    thermo_params,
+    z;
+    _q_init = FT(0),
+    _N_init = FT(0),
+    _F_rim = FT(0),
+    _F_liq = FT(0),
+    _ρ_r = FT(0),
+    z_top,
+    ice_start,
+    dry = false,
+) where {FT}
+
+    # initialize water vapor profile
+    # using same setup as initial_condition_1d
+    q_vap::FT = init_profile(FT, kid_params, thermo_params, z, dry = dry).qv
+
+    # if ice_start, start with small ice bump
+    # otherwise, introduce particles solely through
+    # the boundary in advection_tendency!()
+    _z_band = z_top * FT(0.2)
+    has_ice(z) = ifelse(ice_start, z > (z_top - _z_band), false)
+    ice_bump(χ, z) = χ * cos(z - (z_top - 0.5 * _z_band))
+
+    # P3 state variables (minus B_rim -- computed below)
+    # (initialized as kg/kg before converting to kg/m3)
+    q_ice::FT = has_ice(z) ? ice_bump(_q_init, z) : FT(0)
+    N_ice_kg::FT = has_ice(z) ? ice_bump(_N_init, z) : FT(0)
+    q_liqonice::FT = _F_liq * q_ice
+    q_rim::FT = _F_rim * (q_ice - q_liqonice)
+
+    # 2M warm rain state variables
+    q_liq::FT = FT(0)
+    N_liq::FT = FT(0)
+    q_rai::FT = FT(0)
+    N_rai::FT = FT(0)
+
+    # q_tot - single p3 category + 2M categories + vapor
+    q_tot::FT = q_ice + q_liq + q_rai + q_vap
+
+    # thermodynamics:
+    # for Case 1 of Cholette et al 2019 paper
+    # use approximate temperature values from
+    # sounding: kin1d/soundings/Tsfc2.txt
+    # (path in p3 fortran github repo)
+    T::FT = -0.004 * (z - 500) + 273.15 # temperature
+    p::FT = 990 - (0.1 * z) # pressure
+    _q = TD.PhasePartition(q_tot, q_liq, q_ice)
+    ρ::FT = TD.air_density(thermo_params, T, p, _q) # total density
+    ρ_dry::FT = TD.air_density(thermo_params, T, p)
+    ts = TD.PhaseNonEquil_ρTq(thermo_params, ρ, T, _q) # thermodynamic state
+
+    # P3 scheme uses L (kg/m3) = ρ * q
+    # so we compute ρq and pass that to P3
+    # also we convert N to (1/m3)
+    # since paper init values are given
+    # in (kg/kg), (1/kg)
+    ρq_tot::FT = ρ * q_tot
+    N_ice::FT = ρ * N_ice_kg
+    ρq_ice::FT = ρ * q_ice
+    ρq_rim::FT = ρ * q_rim
+    ρq_liqonice::FT = ρ * q_liqonice
+    ρq_rai::FT = ρ * q_rai
+    ρq_liq::FT = ρ * q_liq
+    ρq_vap::FT = ρ * q_vap
+
+    # also compute B_rim from L_rim, ρ_r
+    B_rim::FT = ρq_rim / _ρ_r
+
+
+    # unused quantities:
+    q_sno::FT = FT(0)
+    ρq_sno::FT = FT(0)
+    θ_liq_ice::FT = TD.liquid_ice_pottemp(thermo_params, ts)
+    N_aer::FT = FT(0)
+    θ_dry::FT = TD.dry_pottemp(thermo_params, T, ρ)
+
+    # zero:
+    zero::FT = FT(0)
+
+    return (;
+        ρ,
+        ρ_dry,
+        T,
+        p,
+        θ_liq_ice,
+        θ_dry,
+        ρq_tot,
+        ρq_liq,
+        ρq_ice,
+        ρq_rai,
+        ρq_rim,
+        ρq_sno,
+        ρq_vap,
+        ρq_liqonice,
+        q_tot,
+        q_liq,
+        q_ice,
+        q_rai,
+        q_sno,
+        q_rim,
+        q_liqonice,
+        q_vap,
+        B_rim,
+        N_liq,
+        N_ice,
+        N_rai,
+        N_aer,
+        zero,
+    )
+end

src/Common/netcdf_io.jl

@@ -23,11 +23,21 @@ function NetCDFIO_Stats(
         :q_tot => "q_tot",
         :q_liq => "q_liq",
         :q_ice => "q_ice",
+        :q_rim => "q_rim",
+        :q_liqonice => "q_liqonice",
         :q_rai => "q_rai",
         :q_sno => "q_sno",
+        :q_vap => "q_vap",
         :N_aer => "N_aer",
         :N_liq => "N_liq",
         :N_rai => "N_rai",
+        :N_ice => "N_ice",
+        :B_rim => "B_rim",
+        :ρq_ice => "ρq_ice",
+        :ρq_rim => "ρq_rim",
+        :ρq_liqonice => "ρq_liqonice",
+        :ρq_vap => "ρq_vap",
+        :ρq_rai => "ρq_rai",
     ),
 )
     FT = Float64

src/Common/ode_utils.jl

@@ -76,11 +76,15 @@ function initialise_state(::MoistureP3, ::PrecipitationP3, initial_profiles)
         ρq_rai = initial_profiles.ρq_rai,
         ρq_ice = initial_profiles.ρq_ice,
         ρq_rim = initial_profiles.ρq_rim,
+        ρq_liqonice = initial_profiles.ρq_liqonice,
         B_rim = initial_profiles.B_rim,
         N_liq = initial_profiles.N_liq,
         N_rai = initial_profiles.N_rai,
         N_ice = initial_profiles.N_ice,
         N_aer = initial_profiles.N_aer,
+        q_vap = initial_profiles.q_vap,
+        ρq_vap = initial_profiles.ρq_vap,
+        q_rai = initial_profiles.q_rai,
     )
 end
 function initialise_state(::CloudyMoisture, ::CloudyPrecip, initial_profiles)
@@ -116,7 +120,7 @@ function initialise_aux(
     if moisture isa EquilibriumMoisture
         ts = @. TD.PhaseEquil_ρθq(thermo_params, ip.ρ, ip.θ_liq_ice, ip.q_tot)
         cloud_sources = nothing
-    elseif moisture isa NonEquilibriumMoisture
+    elseif moisture isa NonEquilibriumMoisture || moisture isa MoistureP3
         q = @. TD.PhasePartition(ip.q_tot, ip.q_liq, ip.q_ice)
         ts = @. TD.PhaseNonEquil_ρθq(thermo_params, ip.ρ, ip.θ_liq_ice, q)
 
@@ -181,11 +185,20 @@ function initialise_aux(
             q_rai = ip.q_rai,
             q_ice = ip.q_ice,
             q_rim = ip.q_rim,
+            q_liqonice = ip.q_liqonice,
             B_rim = ip.B_rim,
             N_liq = ip.N_liq,
             N_rai = ip.N_rai,
             N_ice = ip.N_ice,
             N_aer = ip.N_aer,
+            ρq_tot = ip.ρq_tot,
+            ρq_liq = ip.ρq_liq,
+            ρq_rai = ip.ρq_rai,
+            ρq_ice = ip.ρq_ice,
+            ρq_rim = ip.ρq_rim,
+            ρq_liqonice = ip.ρq_liqonice,
+            q_vap = ip.q_vap,
+            ρq_vap = ip.ρq_vap,
         )
         velocities = (;
             term_vel_rai = copy(ip.zero),
@@ -199,6 +212,7 @@ function initialise_aux(
             q_rai::FT,
             q_ice::FT,
             q_rim::FT,
+            q_liqonice::FT,
             N_aer::FT,
             N_liq::FT,
             N_rai::FT,
@@ -217,11 +231,31 @@ function initialise_aux(
                 copy(ip.zero),
                 copy(ip.zero),
                 copy(ip.zero),
+                copy(ip.zero),
+            ),
+        )
+        p3_boundary_condition_eltype = @NamedTuple{
+            ice_start::Bool,
+            _magnitude::FT,
+            _q_flux::FT,
+            _N_flux::FT,
+            _F_rim::FT,
+            _F_liq::FT,
+            _ρ_r_init::FT,
+        }
+        p3_boundary_condition = @. p3_boundary_condition_eltype(
+            tuple(
+                copy(precip.p3_boundary_condition.ice_start),
+                copy(precip.p3_boundary_condition._magnitude),
+                copy(precip.p3_boundary_condition._q_flux),
+                copy(precip.p3_boundary_condition._N_flux),
+                copy(precip.p3_boundary_condition._F_rim),
+                copy(precip.p3_boundary_condition._F_liq),
+                copy(precip.p3_boundary_condition._ρ_r_init),
             ),
         )
 
         activation_sources = nothing
-        @assert moisture isa NonEquilibrumMoisture
     elseif precip isa CloudyPrecip
         microph_variables = (;
             q_tot = ip.q_tot,
@@ -241,7 +275,7 @@ function initialise_aux(
         cloudy_variables = (; nm_cloud = Val(cloudy_params.NProgMoms[1]))
         scratch = merge(scratch, (; tmp_cloudy = similar(ip.cloudy_moments_zero)))
     else
-        error("Wrong precipitation choise $precip")
+        error("Wrong precipitation choice $precip")
     end
 
     aux = (;
@@ -262,6 +296,8 @@ function initialise_aux(
 
     if precip isa CloudyPrecip
         aux = merge(aux, (; cloudy_params, cloudy_variables))
+    elseif precip isa PrecipitationP3
+        aux = merge(aux, (; p3_boundary_condition))
     end
 
     return aux