Add warm-rain 2-moment microphysics on grid mean

Author: sajjadazimi

.buildkite/pipeline.yml

@@ -86,6 +86,13 @@ steps:
           --config_file $CONFIG_PATH/single_column_precipitation_test.yml
           --job_id single_column_precipitation_test
         artifact_paths: "single_column_precipitation_test/output_active/*"
+      
+      - label: ":umbrella: 2-moment precipitation sanity test single column"
+        command: >
+          julia --color=yes --project=.buildkite .buildkite/ci_driver.jl
+          --config_file $CONFIG_PATH/single_column_precipitation_2M_test.yml
+          --job_id single_column_precipitation_2M_test
+        artifact_paths: "single_column_precipitation_2M_test/output_active/*"
 
   - group: "Gravity wave"
     steps:
@@ -367,6 +374,20 @@ steps:
           slurm_mem: 20GB
           slurm_constraint: icelake|cascadelake|skylake|epyc
 
+      - label: ":umbrella: aquaplanet nonequil allsky monin_obukhov varying insol gravity wave (gfdl_restart) high top 2-moment"
+        command: >
+          julia --color=yes --project=.buildkite .buildkite/ci_driver.jl
+          --config_file $CONFIG_PATH/aquaplanet_nonequil_allsky_gw_res_2M.yml
+          --job_id aquaplanet_nonequil_allsky_gw_res_2M
+
+          julia --color=yes --project=.buildkite reproducibility_tests/test_mse.jl
+          --job_id aquaplanet_nonequil_allsky_gw_res_2M
+          --out_dir aquaplanet_nonequil_allsky_gw_res_2M/output_active
+        artifact_paths: "aquaplanet_nonequil_allsky_gw_res_2M/output_active/*"
+        agents:
+          slurm_mem: 20GB
+          slurm_constraint: icelake|cascadelake|skylake|epyc
+
       - label: ":computer: aquaplanet equil allsky monin_obukhov varying insol gravity wave (raw_topo) high top zonally asymmetric"
         command: >
           julia --color=yes --project=.buildkite .buildkite/ci_driver.jl

config/model_configs/aquaplanet_nonequil_allsky_gw_res_2M.yml

@@ -0,0 +1,32 @@
+z_max: 60000.0
+z_elem: 31
+dz_bottom: 50.0
+rayleigh_sponge: true
+dt: "360secs"
+t_end: "1days"
+dt_save_state_to_disk: "24hours"
+vert_diff: "DecayWithHeightDiffusion"
+implicit_diffusion: true
+approximate_linear_solve_iters: 2
+moist: "nonequil"
+precip_model: "2M"
+rad: "allskywithclear"
+aerosol_radiation: true
+prescribe_clouds_in_radiation: true
+radiation_reset_rng_seed: true
+insolation: "timevarying"
+non_orographic_gravity_wave: true
+orographic_gravity_wave: "gfdl_restart"
+surface_setup: "DefaultMoninObukhov"
+prescribe_ozone: true
+reproducibility_test: true
+prescribed_aerosols: ["SO4", "CB1", "OC1", "DST01", "SSLT01"]
+diagnostics:
+  - short_name: [edt, evu, mmrso4, mmrbcpo, mmrocpo, mmrdust, mmrss, loadss, o3, od550aer, odsc550aer]
+    reduction_time: average
+    period: 1days
+  - short_name: [hus, clw, cli, husra, hussn, cdnc, ncra]
+    reduction_time: average
+    period: 1days
+  - short_name: [reffclw, reffcli]
+    period: 6hours

config/model_configs/single_column_precipitation_2M_test.yml

@@ -0,0 +1,22 @@
+config: "column"
+initial_condition: "PrecipitatingColumn"
+surface_setup: "DefaultMoninObukhov"
+z_elem: 200
+z_max: 10000.0
+z_stretch: false
+dt: "10secs"
+t_end: "1500secs"
+dt_save_state_to_disk: "500secs"
+dt_cloud_fraction: "60secs"
+cloud_model: "grid_scale"
+moist: "nonequil"
+precip_model: "2M"
+vert_diff: "DecayWithHeightDiffusion"
+implicit_diffusion: true
+approximate_linear_solve_iters: 2
+reproducibility_test: false
+diagnostics:
+  - short_name: [hus, clw, cli, husra, hussn, cdnc, ncra, ta, wa]
+    period: 500secs
+  - short_name: [pr]
+    period: 10secs

post_processing/ci_plots.jl

@@ -621,23 +621,36 @@ function make_plots(
 end
 
 function make_plots(
-    ::Val{:single_column_precipitation_test},
+    sim_type::Union{
+        Val{:single_column_precipitation_test},
+        Val{:single_column_precipitation_2M_test},
+    },
     output_paths::Vector{<:AbstractString},
 )
     simdirs = SimDir.(output_paths)
 
     # TODO: Move this plotting code into the same framework as the other ones
     simdir = simdirs[1]
 
-    short_names = ["hus", "clw", "cli", "husra", "hussn", "ta"]
+    if sim_type isa Val{:single_column_precipitation_test}
+        short_names = ["hus", "clw", "cli", "husra", "hussn", "ta"]
+        figsize = (1200, 600)
+        pr_row = 3
+    else
+        short_names =
+            ["hus", "clw", "cli", "husra", "hussn", "ta", "cdnc", "ncra"]
+        figsize = (1200, 800)
+        pr_row = 4
+    end
+
     vars = [
         slice(get(simdir; short_name), x = 0.0, y = 0.0) for
         short_name in short_names
     ]
 
     # We first prepare the axes with all the nice labels with ClimaAnalysis, then we use
     # CairoMakie to add the additional lines.
-    fig = CairoMakie.Figure(; size = (1200, 600))
+    fig = CairoMakie.Figure(; size = figsize)
 
     p_loc = [1, 1]
 
@@ -677,7 +690,7 @@ function make_plots(
     viz.line_plot1D!(
         fig,
         slice(surface_precip, x = 0.0, y = 0.0);
-        p_loc = [3, 1:3],
+        p_loc = [pr_row, 1:3],
     )
 
     file_path = joinpath(output_paths[1], "summary.pdf")
@@ -1130,6 +1143,7 @@ AquaplanetPlots = Union{
     Val{:edonly_edmfx_aquaplanet},
     Val{:mpi_sphere_aquaplanet_rhoe_equil_clearsky},
     Val{:aquaplanet_nonequil_allsky_gw_res},
+    Val{:aquaplanet_nonequil_allsky_gw_res_2M},
     Val{:rcemipii_sphere_diagnostic_edmfx},
     Val{:longrun_aquaplanet_allsky_0M},
     Val{:longrun_aquaplanet_allsky_diagedmf_0M},
@@ -1141,11 +1155,17 @@ AquaplanetPlots = Union{
     Val{:amip_target_edonly},
 }
 
-function make_plots(::AquaplanetPlots, output_paths::Vector{<:AbstractString})
+function make_plots(
+    sim_type::AquaplanetPlots,
+    output_paths::Vector{<:AbstractString},
+)
     simdirs = SimDir.(output_paths)
 
     reduction = "average"
-    short_names_3D = ["ua", "ta", "hus"]
+    short_names_3D =
+        sim_type isa Val{:aquaplanet_nonequil_allsky_gw_res_2M} ?
+        ["ua", "ta", "hus", "clw", "cli", "husra", "hussn", "cdnc", "ncra"] :
+        ["ua", "ta", "hus"]
     short_names_2D = [
         "rsdt",
         "rsds",

reproducibility_tests/ref_counter.jl

@@ -1,4 +1,4 @@
-243
+244
 
 # **README**
 #
@@ -20,6 +20,9 @@
 
 
 #=
+244
+- Add a new test case: aquaplanet with 2M microphysics
+
 243
 - Consolidate eddy diffusivity logic in eddy_diffusion_closures.jl and mass flux
   logic in mass_flux_closures.jl, update buoyancy calculations to be consistent throughout

src/cache/precipitation_precomputed_quantities.jl

@@ -4,6 +4,7 @@
 
 import CloudMicrophysics.MicrophysicsNonEq as CMNe
 import CloudMicrophysics.Microphysics1M as CM1
+import CloudMicrophysics.Microphysics2M as CM2
 
 import Thermodynamics as TD
 import ClimaCore.Operators as Operators
@@ -97,6 +98,59 @@ function set_precipitation_velocities!(
         ) / Y.c.ρ
     return nothing
 end
+function set_precipitation_velocities!(
+    Y,
+    p,
+    moisture_model::NonEquilMoistModel,
+    precip_model::Microphysics2Moment,
+)
+    (; ᶜwₗ, ᶜwᵢ, ᶜwᵣ, ᶜwₛ, ᶜwnₗ, ᶜwnᵣ, ᶜwₜqₜ, ᶜwₕhₜ, ᶜts, ᶜu) = p.precomputed
+    (; q_liq, q_ice, q_rai, q_sno) = p.precomputed.ᶜspecific
+    (; ᶜΦ) = p.core
+
+    cm1c = CAP.microphysics_cloud_params(p.params)
+    cm1p = CAP.microphysics_1m_params(p.params)
+    cm2p = CAP.microphysics_2m_params(p.params)
+    thp = CAP.thermodynamics_params(p.params)
+
+    # compute the precipitation terminal velocity [m/s]
+    # TODO sedimentation of snow is based on the 1M scheme
+    @. ᶜwnᵣ = getindex(
+        CM2.rain_terminal_velocity(cm2p.sb, cm2p.tv, q_rai, Y.c.ρ, Y.c.ρn_rai),
+        1,
+    )
+    @. ᶜwᵣ = getindex(
+        CM2.rain_terminal_velocity(cm2p.sb, cm2p.tv, q_rai, Y.c.ρ, Y.c.ρn_rai),
+        2,
+    )
+    @. ᶜwₛ = CM1.terminal_velocity(cm1p.ps, cm1p.tv.snow, Y.c.ρ, q_sno)
+    # compute sedimentation velocity for cloud condensate [m/s]
+    # TODO sedimentation velocities of cloud condensates are based 
+    # on the 1M scheme. Sedimentation velocity of cloud number concentration
+    # is equal to that of the mass.
+    @. ᶜwnₗ =
+        CMNe.terminal_velocity(cm1c.liquid, cm1c.Ch2022.rain, Y.c.ρ, q_liq)
+    @. ᶜwₗ = CMNe.terminal_velocity(cm1c.liquid, cm1c.Ch2022.rain, Y.c.ρ, q_liq)
+    @. ᶜwᵢ =
+        CMNe.terminal_velocity(cm1c.ice, cm1c.Ch2022.small_ice, Y.c.ρ, q_ice)
+
+    # compute their contributions to energy and total water advection
+    @. ᶜwₜqₜ =
+        Geometry.WVector(
+            ᶜwₗ * Y.c.ρq_liq +
+            ᶜwᵢ * Y.c.ρq_ice +
+            ᶜwᵣ * Y.c.ρq_rai +
+            ᶜwₛ * Y.c.ρq_sno,
+        ) / Y.c.ρ
+    @. ᶜwₕhₜ =
+        Geometry.WVector(
+            ᶜwₗ * Y.c.ρq_liq * (Iₗ(thp, ᶜts) + ᶜΦ + $(Kin(ᶜwₗ, ᶜu))) +
+            ᶜwᵢ * Y.c.ρq_ice * (Iᵢ(thp, ᶜts) + ᶜΦ + $(Kin(ᶜwᵢ, ᶜu))) +
+            ᶜwᵣ * Y.c.ρq_rai * (Iₗ(thp, ᶜts) + ᶜΦ + $(Kin(ᶜwᵣ, ᶜu))) +
+            ᶜwₛ * Y.c.ρq_sno * (Iᵢ(thp, ᶜts) + ᶜΦ + $(Kin(ᶜwₛ, ᶜu))),
+        ) / Y.c.ρ
+    return nothing
+end
 
 """
     set_precipitation_cache!(Y, p, precip_model, turbconv_model)
@@ -262,6 +316,65 @@ function set_precipitation_cache!(
     # in edmf sub-domains.
     return nothing
 end
+function set_precipitation_cache!(Y, p, ::Microphysics2Moment, _)
+    (; dt) = p
+    (; ᶜts) = p.precomputed
+    (; ᶜSqₗᵖ, ᶜSqᵢᵖ, ᶜSqᵣᵖ, ᶜSqₛᵖ) = p.precomputed
+    (; ᶜSnₗᵖ, ᶜSnᵣᵖ) = p.precomputed
+
+    (; q_liq, q_rai, n_liq, n_rai) = p.precomputed.ᶜspecific
+
+    ᶜSᵖ = p.scratch.ᶜtemp_scalar
+    ᶜS₂ᵖ = p.scratch.ᶜtemp_scalar_2
+
+    # get thermodynamics and microphysics params
+    (; params) = p
+    cmp = CAP.microphysics_2m_params(params)
+    thp = CAP.thermodynamics_params(params)
+
+    # compute warm precipitation sources on the grid mean (based on SB2006 2M scheme)
+    compute_warm_precipitation_sources_2M!(
+        ᶜSᵖ,
+        ᶜS₂ᵖ,
+        ᶜSnₗᵖ,
+        ᶜSnᵣᵖ,
+        ᶜSqₗᵖ,
+        ᶜSqᵣᵖ,
+        Y.c.ρ,
+        n_liq,
+        n_rai,
+        q_liq,
+        q_rai,
+        ᶜts,
+        dt,
+        cmp,
+        thp,
+    )
+
+    #TODO - implement 2M cold processes!
+    @. ᶜSqᵢᵖ = 0
+    @. ᶜSqₛᵖ = 0
+
+    return nothing
+end
+function set_precipitation_cache!(
+    Y,
+    p,
+    ::Microphysics2Moment,
+    ::DiagnosticEDMFX,
+)
+    error("Not implemented yet")
+    return nothing
+end
+function set_precipitation_cache!(
+    Y,
+    p,
+    ::Microphysics2Moment,
+    ::PrognosticEDMFX,
+)
+    error("Not implemented yet")
+    return nothing
+end
 
 """
     set_precipitation_surface_fluxes!(Y, p, precipitation model)
@@ -301,7 +414,7 @@ end
 function set_precipitation_surface_fluxes!(
     Y,
     p,
-    precip_model::Microphysics1Moment,
+    precip_model::Union{Microphysics1Moment, Microphysics2Moment},
 )
     (; surface_rain_flux, surface_snow_flux) = p.precomputed
     (; col_integrated_precip_energy_tendency,) = p.conservation_check

src/cache/precomputed_quantities.jl

@@ -139,6 +139,19 @@ function precomputed_quantities(Y, atmos)
             ᶜSqᵣᵖ = similar(Y.c, FT),
             ᶜSqₛᵖ = similar(Y.c, FT),
         )
+    elseif atmos.precip_model isa Microphysics2Moment
+        precipitation_quantities = (;
+            ᶜwᵣ = similar(Y.c, FT),
+            ᶜwₛ = similar(Y.c, FT),
+            ᶜSqₗᵖ = similar(Y.c, FT),
+            ᶜSqᵢᵖ = similar(Y.c, FT),
+            ᶜSqᵣᵖ = similar(Y.c, FT),
+            ᶜSqₛᵖ = similar(Y.c, FT),
+            ᶜwnₗ = similar(Y.c, FT),
+            ᶜwnᵣ = similar(Y.c, FT),
+            ᶜSnₗᵖ = similar(Y.c, FT),
+            ᶜSnᵣᵖ = similar(Y.c, FT),
+        )
     else
         precipitation_quantities = (;)
     end

src/diagnostics/Diagnostics.jl

@@ -31,6 +31,7 @@ import ..NonEquilMoistModel
 import ..NoPrecipitation
 import ..Microphysics0Moment
 import ..Microphysics1Moment
+import ..Microphysics2Moment
 
 # radiation
 import ClimaAtmos.RRTMGPInterface as RRTMGPI