wip p3 collisions

Author: haakon-e

src/Common/equation_types.jl

@@ -18,6 +18,8 @@ export Precipitation0M
 export Precipitation1M
 export Precipitation2M
 export PrecipitationP3
+export IcePrecipitationP3
+export Precipitation2M_P3
 export CloudyPrecip
 
 abstract type AbstractStyle end
@@ -59,3 +61,12 @@ struct PrecipitationP3{P3, CH, PT, B} <: AbstractPrecipitationStyle
     p3_boundary_condition::B
 end
 struct CloudyPrecip <: AbstractPrecipitationStyle end
+# New P3 structs
+struct IcePrecipitationP3{P3, ST} <: AbstractPrecipitationStyle
+    params::P3
+    sedimentation::ST
+end
+struct Precipitation2M_P3{LP, IP} <: AbstractPrecipitationStyle
+    liq_precip::LP # The rain scheme. For now: Precipitation2M
+    ice_precip::IP # The ice scheme.  For now: IcePrecipitationP3
+end

src/Common/helper_functions.jl

@@ -86,6 +86,15 @@ function get_precipitation_type(
         Chen2022 = CMP.Chen2022VelType(FT)
         sb2006 = CMP.SB2006(toml_dict)
         precip = PrecipitationP3(p3_params, Chen2022, sb2006, boundary)
+    elseif precipitation_choice == "Precipitation2M_P3"
+        @assert sedimentation_choice == "Chen2022"
+        @assert rain_formation_choice == "SB2006"
+        params = CMP.ParametersP3(toml_dict; slope_law = :constant)
+        liq_sedimentation = CMP.Chen2022VelTypeRain(toml_dict)
+        ice_sedimentation = CMP.Chen2022VelType(toml_dict)
+        liq_precip = Precipitation2M(CMP.SB2006(toml_dict), liq_sedimentation)
+        ice_precip = IcePrecipitationP3(params, ice_sedimentation)
+        precip = Precipitation2M_P3(liq_precip, ice_precip)
     else
         error("Invalid precipitation choice: $precipitation_choice")
     end

src/Common/initial_condition.jl

@@ -162,6 +162,7 @@ function initial_condition_1d(
         ρq_ice,
         ρq_rai,
         ρq_sno,
+        ρq_rim,
         ρq_vap,
         q_tot,
         q_liq,
@@ -175,6 +176,10 @@ function initial_condition_1d(
         N_rai,
         N_aer,
         zero,
+        # P3 variables
+        F_rim = FT(0),
+        ρ_rim = FT(0),
+        logλ = FT(0),
     )
 end
 

src/Common/ode_utils.jl

@@ -47,6 +47,12 @@ function initialise_state(::MoistureP3, ::PrecipitationP3, initial_profiles)
         :N_liq, :N_rai, :N_ice, :N_aer, :q_vap, :ρq_vap, :q_rai,
     )}.(initial_profiles)
 end
+function initialise_state(::NonEquilibriumMoisture, ::Precipitation2M_P3, initial_profiles)
+    warm_state = NamedTuple{(:ρq_tot, :ρq_liq, :ρq_rai, :N_liq, :N_rai, :N_aer)}.(initial_profiles)
+    cold_state = NamedTuple{(:ρq_ice, :N_ice, :ρq_rim, :B_rim)}.(initial_profiles)
+    tmp = NamedTuple{(:ρq_sno,)}.(initial_profiles)  # only needed to use existing `Precipitation2M` code, should always be zero.
+    return merge.(warm_state, cold_state, tmp)
+end
 function initialise_state(::CloudyMoisture, ::CloudyPrecip, initial_profiles)
     return NamedTuple{(:ρq_vap, :N_aer, :moments)}.(initial_profiles)
 end
@@ -115,11 +121,21 @@ function initialise_aux(
         velocities = zero_nt(NamedTuple{(:term_vel_rai, :term_vel_sno)})
         precip_sources = zero_nt(NamedTuple{(:q_tot, :q_liq, :q_ice, :q_rai, :q_sno)})
         activation_sources = nothing
-    elseif precip isa Precipitation2M
+    elseif precip isa Precipitation2M || precip isa Precipitation2M_P3
         microph_variables = merge.(microph_variables, NamedTuple{(:q_rai, :q_sno, :N_liq, :N_rai, :N_aer)}.(ip))
         velocities = zero_nt(NamedTuple{(:term_vel_N_rai, :term_vel_rai)})
         precip_sources = zero_nt(NamedTuple{(:q_tot, :q_liq, :q_rai, :N_aer, :N_liq, :N_rai)})
         activation_sources = zero_nt(NamedTuple{(:N_aer, :N_liq)})
+        if precip isa Precipitation2M_P3
+            # NOTE: `q_sno` from above is not used in the P3 code, but is needed to use existing `Precipitation2M` code.
+            ice_microph_variables = NamedTuple{(:F_rim, :ρ_rim, :logλ)}.(ip)  # variables needed to be precomputed 
+            ice_output_variables = NamedTuple{(:N_ice, :ρq_rim, :B_rim)}.(ip)  # variables for netCDF output
+            microph_variables = merge.(microph_variables, ice_microph_variables, ice_output_variables)
+            ice_velocities = zero_nt(NamedTuple{(:term_vel_N_ice, :term_vel_q_ice)})
+            velocities = merge.(velocities, ice_velocities)  # precompute ice terminal velocities
+            ice_precip_sources = zero_nt(NamedTuple{(:q_ice, :q_rim, :N_ice, :B_rim)})
+            precip_sources = merge.(precip_sources, ice_precip_sources)
+        end
     elseif precip isa PrecipitationP3
         microph_variables =
             NamedTuple{(

src/Common/tendency.jl

@@ -160,7 +160,7 @@ end
     @. term_vel_rai = CM1.terminal_velocity(ps.rain, ps.sedimentation.rain, ρ, q_rai)
     @. term_vel_sno = CM1.terminal_velocity(ps.snow, ps.sedimentation.snow, ρ, q_sno)
 end
-@inline function precompute_aux_precip!(ps::Precipitation2M, Y, aux)
+@inline function precompute_aux_precip!(ps::Precipitation2M, Y, aux)  # called by `make_rhs_function`
 
     FT = eltype(Y.ρq_rai)
     sb2006 = ps.rain_formation
@@ -177,6 +177,42 @@ end
     @. term_vel_N_rai = getindex(CM2.rain_terminal_velocity(sb2006, vel_scheme, q_rai, ρ, N_rai), 1)
     @. term_vel_rai = getindex(CM2.rain_terminal_velocity(sb2006, vel_scheme, q_rai, ρ, N_rai), 2)
 end
+@inline function precompute_aux_precip!(ps::IcePrecipitationP3, Y, aux)
+    # State, `Y`, includes: ρq_ice, ρq_rim, N_ice, B_rim
+
+    # update state variables
+    (; sedimentation, params) = ps
+    (; ρq_ice, ρq_rim, N_ice, B_rim) = Y
+    (; ρ) = aux.thermo_variables
+    (; F_rim, ρ_rim, logλ) = aux.microph_variables
+    (; term_vel_N_ice, term_vel_q_ice) = aux.velocities
+
+    # Calculate derived rime quantities
+    # @. F_rim = ifelse(Y.ρq_ice < eps(FT), FT(0), Y.ρq_rim / Y.ρq_ice)
+    # @. ρ_rim = ifelse(Y.B_rim < eps(FT), FT(0), Y.ρq_rim / Y.B_rim)
+
+    @. F_rim = ifelse(isnan(ρq_rim / ρq_ice), zero(ρq_rim), ρq_rim / ρq_ice)
+    @. ρ_rim = ifelse(isnan(ρq_rim / B_rim), zero(ρq_rim), ρq_rim / B_rim)
+    # Calculate distribution parameters
+    @. logλ = CMP3.get_distribution_logλ(CMP3.P3State(params, ρq_ice, N_ice, F_rim, ρ_rim))
+
+    # Calculate terminal velocities
+    use_aspect_ratio = true
+    @. term_vel_N_ice = CMP3.ice_terminal_velocity_number_weighted(
+        sedimentation, ρ, CMP3.P3State(params, ρq_ice, N_ice, F_rim, ρ_rim), logλ; use_aspect_ratio
+    )
+    @. term_vel_q_ice = CMP3.ice_terminal_velocity_mass_weighted(
+        sedimentation, ρ, CMP3.P3State(params, ρq_ice, N_ice, F_rim, ρ_rim), logλ; use_aspect_ratio
+    )
+    # args = (ps.sedimentation, ρ, ps.params, ρq_ice, N_ice, F_rim, ρ_rim, logλ)
+    # @. term_vel_N_ice = CMP3.ice_terminal_velocity_number_weighted(args...; use_aspect_ratio)
+    # @. term_vel_q_ice = CMP3.ice_terminal_velocity_mass_weighted(args...; use_aspect_ratio)
+
+end
+function precompute_aux_precip!(ps::Precipitation2M_P3, Y, aux)
+    precompute_aux_precip!(ps.liq_precip, Y, aux)
+    precompute_aux_precip!(ps.ice_precip, Y, aux)
+end
 @inline function precompute_aux_precip!(ps::PrecipitationP3, Y, aux)
 
     # update state variables
@@ -580,6 +616,13 @@ end
     end
 end
 
+@inline function precompute_aux_precip_sources!(ps::IcePrecipitationP3, Y, aux)
+    
+    # TODO: So far, implemented directly in precip_sources_tendency!
+    
+    return nothing
+end
+
 @inline function precompute_aux_precip_sources!(ps::PrecipitationP3, aux)
     # TODO [P3]
     return nothing
@@ -670,7 +713,7 @@ end
 @inline function cloud_sources_tendency!(::EquilibriumMoisture, dY, Y, aux, t) end
 @inline function cloud_sources_tendency!(ms::NonEquilibriumMoisture, dY, Y, aux, t)
 
-    precompute_aux_moisture_sources!(ms, aux)
+    precompute_aux_moisture_sources!(ms, aux)  # defined in `Common/tendency.jl`
 
     @. dY.ρq_liq += aux.thermo_variables.ρ * aux.cloud_sources.q_liq
     @. dY.ρq_ice += aux.thermo_variables.ρ * aux.cloud_sources.q_ice
@@ -740,6 +783,42 @@ end
 @inline function precip_sources_tendency!(ms::MoistureP3, ps::PrecipitationP3, dY, Y, aux, t)
     return dY
 end
+@inline function precip_sources_tendency!(ms::AbstractMoistureStyle, ps::Precipitation2M_P3, dY, Y, aux, t)
+    precip_sources_tendency!(ms, ps.liq_precip, dY, Y, aux, t)
+    # TODO [P3]
+    # precompute_aux_precip_sources!(ps, Y, aux)
+
+    #=
+    function bulk_liquid_ice_collision_sources(
+        params, logλ, L_ice, F_rim, ρ_rim,
+        psd_c, psd_r, L_c, N_c, L_r, N_r,
+        aps, tps, vel, ρₐ, T,
+    )
+    =#
+    (; ρ, ts) = aux.thermo_variables
+    (; ρq_liq, ρq_rai, ρq_ice, N_liq, N_rai, N_ice) = Y
+    (; air_params, thermo_params) = aux
+    (; logλ, F_rim, ρ_rim) = aux.microph_variables
+    (; pdf_c, pdf_r) = ps.liq_precip.rain_formation
+    (; ice_precip) = ps
+
+    coll_src = @. CMP3.bulk_liquid_ice_collision_sources(
+        ice_precip.params, logλ, ρq_ice, N_ice, F_rim, ρ_rim,
+        pdf_c, pdf_r, ρq_liq, N_liq, ρq_rai, N_rai,
+        air_params, thermo_params, (ice_precip.sedimentation,),
+        ρ, TD.air_temperature(thermo_params, ts),
+    )
+
+    @. dY.ρq_liq += ρ * coll_src.∂ₜq_c
+    @. dY.ρq_rai += ρ * coll_src.∂ₜq_r
+    @. dY.N_liq += coll_src.∂ₜN_c
+    @. dY.N_rai += coll_src.∂ₜN_r
+    @. dY.ρq_rim += coll_src.∂ₜL_rim
+    @. dY.ρq_ice += coll_src.∂ₜL_ice
+    @. dY.B_rim += coll_src.∂ₜB_rim
+
+    return dY
+end
 @inline function precip_sources_tendency!(ms::CloudyMoisture, ps::CloudyPrecip, dY, Y, aux, t)
 
     precompute_aux_precip_sources!(ps, aux)

src/K1DModel/ode_utils.jl

@@ -28,19 +28,56 @@ end
 """
 function make_rhs_function(ms::CO.AbstractMoistureStyle, ps::CO.AbstractPrecipitationStyle)
     function rhs!(dY, Y, aux, t)
-
         CO.zero_tendencies!(dY)
 
+        # defined in `K1DModel/tendency.jl`
+        # calc: aux.prescribed_velocity.ρw
         precompute_aux_prescribed_velocity!(aux, t)
+        # defined in `Common/tendency.jl`
+        # For NonEquilibriumMoisture:
+        # calc: aux.thermo_variables.{ts, ρ, p, θ_dry, θ_liq_ice}
+        # calc: aux.microph_variables.{q_tot, q_liq, q_ice}
         CO.precompute_aux_thermo!(ms, Y, aux)
-        CO.precompute_aux_precip!(ps, Y, aux)
-
-        precompute_aux_activation!(ps, dY, Y, aux, t)
-
+        # defined in `Common/tendency.jl`
+        # For Precipitation2M:
+        # calc: aux.microph_variables.{q_rai, q_sno, N_rai, N_liq}  <-- NOTE: You can assume `q_sno = ρq_sno = 0`
+        # calc: aux.velocities.{term_vel_rai, term_vel_N_rai}
+        # For IcePrecipitationP3:
+        # calc: aux.microph_variables.{ρq_ice, ρq_rim, N_ice, B_rim} <-- maybe not needed
+        # calc: aux.velocities.{term_vel_ice, term_vel_N_ice}
+        CO.precompute_aux_precip!(ps, Y, aux)  # <-- calls fn for Precipitation2M
+
+        # defined in `K1DModel/tendency.jl`
+        # calc: aux.activation_sources.{N_aer, N_liq}
+        precompute_aux_activation!(ps, dY, Y, aux, t)  # <-- calls fn for Precipitation2M
+
+        # defined in `Common/tendency.jl`
+        # calls: `precompute_aux_moisture_sources` (defined in `Common/tendency.jl`)
+        #   calc: aux.cloud_sources.{q_liq, q_ice}
+        # increments: dY.{ρq_liq, ρq_ice} += ρ * aux.cloud_sources.{q_...}
         CO.cloud_sources_tendency!(ms, dY, Y, aux, t)
-        CO.precip_sources_tendency!(ms, ps, dY, Y, aux, t)
+        # defined in `Common/tendency.jl`
+        # For Precipitation2M:
+        # calls: `precompute_aux_precip_sources` (defined in `Common/tendency.jl`)
+        #   calc: aux.precip_sources.{q_tot, q_liq, q_rai, N_liq, N_rai, N_aer}
+        # increments: dY.{ρq_tot, ρq_rai, ρq_liq, N_liq, N_rai, N_aer}
+        #               += aux.precip_sources.{...} 
+        #        and/or += aux.activation_sources.{N_liq, N_aer}
+        # For Precipitation2M_P3:
+        # calls: `precompute_aux_precip_sources` (defined in `Common/tendency.jl`)
+        # TODO: write collisions
+        CO.precip_sources_tendency!(ms, ps, dY, Y, aux, t)  # <-- calls fn for Precipitation2M + write code
 
         for eq_style in [ms, ps]
+            # defined in  `K1DModel/tendency.jl`
+            # For Precipitation2M:
+            # increments: dY.{ρq_rai, N_liq, N_rai, N_aer} += ... * aux.velocities.{term_vel_N_rai, term_vel_rai}
+            # For NonEquilibriumMoisture:
+            # increments: dY.{ρq_tot, ρq_liq, ρq_ice} += {aux.prescribed_velocity} * {Y.ρq_...}
+            # For Precipitation2M_P3
+            # - delegates to `Precipitation2M` (see above) and `IcePrecipitationP3`
+            # For IcePrecipitationP3:
+            # increments: dY.{N_ice, B_rim, ρq_rim} += ... * aux.velocities.{term_vel_N_ice, term_vel_q_ice}
             advection_tendency!(eq_style, dY, Y, aux, t)
         end
 

src/K1DModel/tendency.jl

@@ -171,6 +171,10 @@ end
     @. aux.activation_sources.N_liq = S_Nl
 end
 
+@inline function precompute_aux_activation!(ps::CO.Precipitation2M_P3, dY, Y, aux, t)
+    precompute_aux_activation!(ps.liq_precip, dY, Y, aux, t)
+end
+
 @inline function precompute_aux_activation!(::CO.CloudyPrecip, dY, Y, aux, t)
 
     (; common_params, kid_params, thermo_params, air_params, activation_params, cloudy_params) = aux
@@ -421,6 +425,38 @@ end
 
     return dY
 end
+@inline function advection_tendency!(::CO.IcePrecipitationP3, dY, Y, aux, t)
+    FT = eltype(Y.ρq_tot)
+    # Apply advection tendency to N_ice, B_rim, ρq_rim
+    (; ρw) = aux.prescribed_velocity
+    (; ρ) = aux.thermo_variables
+    (; term_vel_N_ice, term_vel_q_ice) = aux.velocities
+
+    If = CC.Operators.InterpolateC2F()
+    wvec = CC.Geometry.WVector
+    extrapolate = CC.Operators.Extrapolate()
+    zero_bc = CC.Operators.SetValue(wvec(FT(0)))
+
+    ∇_N = CC.Operators.DivergenceF2C(bottom = extrapolate, top = zero_bc)
+    ∇_q = CC.Operators.DivergenceF2C(bottom = zero_bc, top = extrapolate)
+    
+    @. dY.N_ice += -∇_N((ρw / If(ρ) - wvec(If(term_vel_N_ice))) * If(Y.N_ice))
+    @. dY.B_rim += -∇_q((ρw / If(ρ) - wvec(If(term_vel_q_ice))) * If(Y.B_rim))
+    @. dY.ρq_rim += -∇_q((ρw / If(ρ) - wvec(If(term_vel_q_ice))) * If(Y.ρq_rim))
+
+    fcc = CC.Operators.FluxCorrectionC2C(bottom = extrapolate, top = extrapolate)
+    @. dY.N_ice += fcc(ρw / If(ρ) - wvec(If(term_vel_N_ice)), Y.N_ice)
+    @. dY.B_rim += fcc(ρw / If(ρ) - wvec(If(term_vel_q_ice)), Y.B_rim)
+    @. dY.ρq_rim += fcc(ρw / If(ρ) - wvec(If(term_vel_q_ice)), Y.ρq_rim)
+
+    return dY
+end
+
+@inline function advection_tendency!(ps::CO.Precipitation2M_P3, dY, Y, aux, t)
+    advection_tendency!(ps.liq_precip, dY, Y, aux, t)
+    advection_tendency!(ps.ice_precip, dY, Y, aux, t)
+    return dY
+end
 
 @inline function advection_tendency!(::CO.PrecipitationP3, dY, Y, aux, t)
     FT = eltype(Y.ρq_tot)

test/experiments/KiD_driver/run_KiD_simulation.jl

@@ -24,7 +24,7 @@ function run_KiD_simulation(::Type{FT}, opts) where {FT}
 
     # Decide the output folder name based on options
     folder = "Output_$(moisture_choice)_$(precipitation_choice)"
-    if precipitation_choice in ["Precipitation1M", "Precipitation2M"]
+    if precipitation_choice in ["Precipitation1M", "Precipitation2M", "Precipitation2M_P3"]
         folder *= "_$(rain_formation_choice)"
         sedimentation_choice == "Chen2022" && (folder *= "_Chen2022")
     elseif precipitation_choice == "CloudyPrecip"