From cbd27bc46e05cf5889911049880df04fa6718b33 Mon Sep 17 00:00:00 2001 From: Mikhail Mints Date: Fri, 30 Jun 2023 01:12:45 -0700 Subject: [PATCH] Refactor KiD driver by wrapping in a function, add precipitation susceptibility analysis --- .buildkite/pipeline.yml | 4 +- .github/workflows/ci.yml | 2 +- Project.toml | 4 +- src/driveKiD/NetCDFIO.jl | 6 +- src/driveKiD/TimeStepping.jl | 8 +- src/driveKiD/tendency.jl | 55 +++--- test/Project.toml | 3 +- test/experiments/KiD_driver/KiD_driver.jl | 157 +----------------- .../KiD_driver/parse_commandline.jl | 54 +++--- .../KiD_driver/run_KiD_simulation.jl | 151 +++++++++++++++++ test/plotting_utils.jl | 37 ++++- .../test_precip_production.jl | 66 ++++++++ 12 files changed, 336 insertions(+), 211 deletions(-) create mode 100644 test/experiments/KiD_driver/run_KiD_simulation.jl create mode 100644 test/precip_production_tests/test_precip_production.jl diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 0f6f1fd7..4efa6dd7 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -76,8 +76,8 @@ steps: command: "julia --color=yes --project=test test/experiments/KiD_driver/KiD_driver.jl --moisture_choice=NonEquilibriumMoisture --prognostic_vars=RhodTQ --precipitation_choice=Precipitation1M" artifact_paths: "test/experiments/KiD_driver/Output_NonEquilibriumMoisture_RhodTQ_Precipitation1M_CliMA_1M/figures/*" - - label: ":crystal_ball: Experiments: equil (fixed rhod and T) + KK2000 " - 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 KK2000 --prescribed_Nd 1e8" + - label: ":crystal_ball: Experiments: equil (fixed rhod and T) + KK2000 + Float32 " + command: "julia --color=yes --project=test test/experiments/KiD_driver/KiD_driver.jl --FLOAT_TYPE=Float32 --moisture_choice=EquilibriumMoisture --prognostic_vars=RhodTQ --precipitation_choice=Precipitation1M --rain_formation_scheme_choice KK2000 --prescribed_Nd 1e8" artifact_paths: "test/experiments/KiD_driver/Output_EquilibriumMoisture_RhodTQ_Precipitation1M_KK2000/figures/*" - label: ":crystal_ball: Experiments: equil (fixed rhod and T) + B1994 " diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index a1cce55a..8b1784d0 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -16,7 +16,7 @@ jobs: fail-fast: false matrix: version: - - '1.8.1' + - '1.8.5' os: - ubuntu-latest - macOS-latest diff --git a/Project.toml b/Project.toml index fd3d58b8..eb127a5f 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "Kinematic1D" uuid = "a8a3bb05-6cc5-4342-abf6-5cc636bd2b35" authors = ["Climate Modeling Alliance"] -version = "0.3.1" +version = "0.3.2" [deps] ClimaCore = "d414da3d-4745-48bb-8d80-42e94e092884" @@ -27,7 +27,7 @@ UnPack = "3a884ed6-31ef-47d7-9d2a-63182c4928ed" [compat] ClimaCore = "0.10" ClimaCorePlots = "0.2" -CloudMicrophysics = "0.10" +CloudMicrophysics = "0.11.1" DiffEqBase = "6.75" Distributions = "0.25" EnsembleKalmanProcesses = "1.1" diff --git a/src/driveKiD/NetCDFIO.jl b/src/driveKiD/NetCDFIO.jl index e4885ca5..7cff7c9d 100644 --- a/src/driveKiD/NetCDFIO.jl +++ b/src/driveKiD/NetCDFIO.jl @@ -59,7 +59,7 @@ function NetCDFIO_Stats( return NetCDFIO_Stats(output_interval, nc_filename, output_profiles) end -function KiD_output(aux, t::Float64) +function KiD_output(aux, t::FT) where {FT <: Real} (; Stats) = aux @@ -67,9 +67,9 @@ function KiD_output(aux, t::Float64) # write simulation time to profiles and timeseries profile_t = ds.group["profiles"]["t"] - @inbounds profile_t[end + 1] = t::Float64 + @inbounds profile_t[end + 1] = t::FT ts_t = ds.group["timeseries"]["t"] - @inbounds ts_t[end + 1] = t::Float64 + @inbounds ts_t[end + 1] = t::FT # write profiles for field in (aux.moisture_variables, aux.precip_variables) diff --git a/src/driveKiD/TimeStepping.jl b/src/driveKiD/TimeStepping.jl index 2046ab11..c638d8c0 100644 --- a/src/driveKiD/TimeStepping.jl +++ b/src/driveKiD/TimeStepping.jl @@ -5,8 +5,8 @@ - simulation time `t_max`. """ -mutable struct TimeStepping - dt::Float64 - dt_io::Float64 - t_max::Float64 +mutable struct TimeStepping{FT <: Real} + dt::FT + dt_io::FT + t_max::FT end diff --git a/src/driveKiD/tendency.jl b/src/driveKiD/tendency.jl index e83f0679..9863eef9 100644 --- a/src/driveKiD/tendency.jl +++ b/src/driveKiD/tendency.jl @@ -189,8 +189,8 @@ end else error("Unrecognized rain formation scheme") end - S_qt_rain = -min(max(0, q.liq / dt), tmp) - S_qt_snow = -min(max(0, q.ice / dt), CM1.conv_q_ice_to_q_sno(microphys_params, q, ρ, T)) + S_qt_rain = -min(max(FT(0), q.liq / dt), tmp) + S_qt_snow = -min(max(FT(0), q.ice / dt), CM1.conv_q_ice_to_q_sno(microphys_params, q, ρ, T)) S_q_rai -= S_qt_rain S_q_sno -= S_qt_snow S_q_tot += S_qt_rain + S_qt_snow @@ -207,14 +207,15 @@ end end # accretion cloud water + rain - S_qr = min(max(0, q.liq / dt), tmp) + S_qr = min(max(FT(0), q.liq / dt), tmp) S_q_rai += S_qr S_q_tot -= S_qr S_q_liq -= S_qr #θ_liq_ice_tendency += S_qr / Π_m / c_pm * L_v0 # accretion cloud ice + snow - S_qs = min(max(0, q.ice / dt), CM1.accretion(microphys_params, CMT.IceType(), CMT.SnowType(), q.ice, q_sno, ρ)) + S_qs = + min(max(FT(0), q.ice / dt), CM1.accretion(microphys_params, CMT.IceType(), CMT.SnowType(), q.ice, q_sno, ρ)) S_q_sno += S_qs S_q_tot -= S_qs S_q_ice -= S_qs @@ -222,7 +223,10 @@ end # sink of cloud water via accretion cloud water + snow S_qt = - -min(max(0, q.liq / dt), CM1.accretion(microphys_params, CMT.LiquidType(), CMT.SnowType(), q.liq, q_sno, ρ)) + -min( + max(FT(0), q.liq / dt), + CM1.accretion(microphys_params, CMT.LiquidType(), CMT.SnowType(), q.liq, q_sno, ρ), + ) if T < T_fr # cloud droplets freeze to become snow) S_q_sno -= S_qt S_q_tot += S_qt @@ -238,9 +242,13 @@ end end # sink of cloud ice via accretion cloud ice - rain - S_qt = -min(max(0, q.ice / dt), CM1.accretion(microphys_params, CMT.IceType(), CMT.RainType(), q.ice, q_rai, ρ)) + S_qt = + -min( + max(FT(0), q.ice / dt), + CM1.accretion(microphys_params, CMT.IceType(), CMT.RainType(), q.ice, q_rai, ρ), + ) # sink of rain via accretion cloud ice - rain - S_qr = -min(max(0, q_rai / dt), CM1.accretion_rain_sink(microphys_params, q.ice, q_rai, ρ)) + S_qr = -min(max(FT(0), q_rai / dt), CM1.accretion_rain_sink(microphys_params, q.ice, q_rai, ρ)) S_q_tot += S_qt S_q_ice += S_qt S_q_rai += S_qr @@ -250,13 +258,13 @@ end # accretion rain - snow if T < T_fr S_qs = min( - max(0, q_rai / dt), + max(FT(0), q_rai / dt), CM1.accretion_snow_rain(microphys_params, CMT.SnowType(), CMT.RainType(), q_sno, q_rai, ρ), ) else S_qs = -min( - max(0, q_sno / dt), + max(FT(0), q_sno / dt), CM1.accretion_snow_rain(microphys_params, CMT.RainType(), CMT.SnowType(), q_rai, q_sno, ρ), ) end @@ -267,22 +275,23 @@ end if Bool(params.precip_sinks) # evaporation - S_qr = -min(max(0, q_rai / dt), -CM1.evaporation_sublimation(microphys_params, CMT.RainType(), q, q_rai, ρ, T)) + S_qr = + -min(max(FT(0), q_rai / dt), -CM1.evaporation_sublimation(microphys_params, CMT.RainType(), q, q_rai, ρ, T)) S_q_rai += S_qr S_q_tot -= S_qr S_q_vap -= S_qr # melting - S_qs = -min(max(0, q_sno / dt), CM1.snow_melt(microphys_params, q_sno, ρ, T)) + S_qs = -min(max(FT(0), q_sno / dt), CM1.snow_melt(microphys_params, q_sno, ρ, T)) S_q_rai -= S_qs S_q_sno += S_qs # deposition, sublimation tmp = CM1.evaporation_sublimation(microphys_params, CMT.SnowType(), q, q_sno, ρ, T) if tmp > 0 - S_qs = min(max(0, q_vap / dt), tmp) + S_qs = min(max(FT(0), q_vap / dt), tmp) else - S_qs = -min(max(0, q_sno / dt), -tmp) + S_qs = -min(max(FT(0), q_sno / dt), -tmp) end S_q_sno += S_qs S_q_tot -= S_qs @@ -320,43 +329,43 @@ end # autoconversion liquid to rain tmp = CM2.autoconversion(microphys_params, rf, q.liq, q_rai, ρ, N_liq) - S_qr = min(max(0, q.liq / dt), tmp.dq_rai_dt) + S_qr = min(max(FT(0), q.liq / dt), tmp.dq_rai_dt) S_q_rai += S_qr S_q_tot -= S_qr S_q_liq -= S_qr - S_Nr = min(max(0, N_liq / 2 / dt), tmp.dN_rai_dt) + S_Nr = min(max(FT(0), N_liq / 2 / dt), tmp.dN_rai_dt) S_N_rai += S_Nr S_N_liq -= 2 * S_Nr # self_collection tmp_l = CM2.liquid_self_collection(microphys_params, rf, q.liq, ρ, -S_qr) tmp_r = CM2.rain_self_collection(microphys_params, rf, q_rai, ρ, N_rai) - S_N_liq += -min(max(0, N_liq / dt), -tmp_l) - S_N_rai += -min(max(0, N_rai / dt), -tmp_r) + S_N_liq += -min(max(FT(0), N_liq / dt), -tmp_l) + S_N_rai += -min(max(FT(0), N_rai / dt), -tmp_r) # rain breakup tmp = CM2.rain_breakup(microphys_params, rf, q_rai, ρ, N_rai, tmp_r) - S_N_rai += min(max(0, N_rai / dt), tmp_r) + S_N_rai += min(max(FT(0), N_rai / dt), tmp_r) # accretion cloud water + rain tmp = CM2.accretion(microphys_params, rf, q.liq, q_rai, ρ, N_liq) - S_qr = min(max(0, q.liq / dt), tmp.dq_rai_dt) + S_qr = min(max(FT(0), q.liq / dt), tmp.dq_rai_dt) S_q_rai += S_qr S_q_tot -= S_qr S_q_liq -= S_qr - S_N_liq += -min(max(0, N_liq / dt), -tmp.dN_liq_dt) + S_N_liq += -min(max(FT(0), N_liq / dt), -tmp.dN_liq_dt) S_N_rai += FT(0) # evaporation tmp = CM2.rain_evaporation(microphys_params, rf, q, q_rai, ρ, N_rai, T) - S_Nr = -min(max(0, N_rai / dt), -tmp[1]) - S_qr = -min(max(0, q_rai / dt), -tmp[2]) + S_Nr = -min(max(FT(0), N_rai / dt), -tmp[1]) + S_qr = -min(max(FT(0), q_rai / dt), -tmp[2]) S_q_rai += S_qr S_q_tot -= S_qr S_q_vap -= S_qr S_N_rai += S_Nr - vt = CM2.rain_terminal_velocity(microphys_params, CMT.SB2006Type(), q_rai, ρ, N_rai) + vt = CM2.rain_terminal_velocity(microphys_params, rf, q_rai, ρ, N_rai) term_vel_N_rai = vt[1] term_vel_rai = vt[2] diff --git a/test/Project.toml b/test/Project.toml index 54fb1db7..3dee4683 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -7,6 +7,7 @@ CloudMicrophysics = "6a9e3e04-43cd-43ba-94b9-e8782df3c71b" Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa" DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e" EnsembleKalmanProcesses = "aa8a2aa5-91d8-4396-bcef-d4f2ec43552d" +FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41" Interpolations = "a98d9a8b-a2ab-59e6-89dd-64a1c18fca59" Kinematic1D = "a8a3bb05-6cc5-4342-abf6-5cc636bd2b35" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" @@ -27,7 +28,7 @@ UnPack = "3a884ed6-31ef-47d7-9d2a-63182c4928ed" CLIMAParameters = "0.7" ClimaCore = "0.10" ClimaCorePlots = "0.2" -CloudMicrophysics = "0.10" +CloudMicrophysics = "0.11.1" DiffEqBase = "6.75" NCDatasets = "0.12" Plots = "1.29" diff --git a/test/experiments/KiD_driver/KiD_driver.jl b/test/experiments/KiD_driver/KiD_driver.jl index 06ed4ff4..6d2d27cd 100644 --- a/test/experiments/KiD_driver/KiD_driver.jl +++ b/test/experiments/KiD_driver/KiD_driver.jl @@ -5,162 +5,21 @@ to see the list of command line arguments. """ -import NCDatasets -import OrdinaryDiffEq -import UnPack -import ArgParse - -import ClimaCore -import Thermodynamics -import CloudMicrophysics -import CLIMAParameters -import Kinematic1D - -const kid_dir = pkgdir(Kinematic1D) -include(joinpath(kid_dir, "test", "create_parameters.jl")) - -const CC = ClimaCore -const TD = Thermodynamics -const CP = CLIMAParameters -const NC = NCDatasets -const ODE = OrdinaryDiffEq -const KID = Kinematic1D -const AP = ArgParse -const CMT = CloudMicrophysics.CommonTypes - -const FT = Float64 # Get the parameter values for the simulation include("parse_commandline.jl") opts = parse_commandline() -# Equations to solve for mositure and precipitation variables -moisture_choice = opts["moisture_choice"] -prognostics_choice = opts["prognostic_vars"] -precipitation_choice = opts["precipitation_choice"] -rain_formation_choice = opts["rain_formation_scheme_choice"] +ft_choice = opts["FLOAT_TYPE"] -if moisture_choice == "EquilibriumMoisture" && prognostics_choice == "RhoThetaQ" - moisture = KID.EquilibriumMoisture_ρθq() -elseif moisture_choice == "EquilibriumMoisture" && prognostics_choice == "RhodTQ" - moisture = KID.EquilibriumMoisture_ρdTq() -elseif moisture_choice == "NonEquilibriumMoisture" && prognostics_choice == "RhoThetaQ" - moisture = KID.NonEquilibriumMoisture_ρθq() -elseif moisture_choice == "NonEquilibriumMoisture" && prognostics_choice == "RhodTQ" - moisture = KID.NonEquilibriumMoisture_ρdTq() -else - error("Invalid moisture choice: $moisture_choice or invalid prognostic variables choice: $prognostic_vars") -end -if precipitation_choice == "NoPrecipitation" - precip = KID.NoPrecipitation() -elseif precipitation_choice == "Precipitation0M" - precip = KID.Precipitation0M() -elseif precipitation_choice == "Precipitation1M" - if rain_formation_choice == "CliMA_1M" - precip = KID.Precipitation1M(KID.OneMomentRainFormation()) - elseif rain_formation_choice == "KK2000" - precip = KID.Precipitation1M(CMT.KK2000Type()) - elseif rain_formation_choice == "B1994" - precip = KID.Precipitation1M(CMT.B1994Type()) - elseif rain_formation_choice == "TC1980" - precip = KID.Precipitation1M(CMT.TC1980Type()) - elseif rain_formation_choice == "LD2004" - precip = KID.Precipitation1M(CMT.LD2004Type()) - else - error("Invalid rain formation choice: $rain_formation_choice") - end -elseif precipitation_choice == "Precipitation2M" - if rain_formation_choice == "SB2006" - precip = KID.Precipitation2M(CMT.SB2006Type()) - else - error("Invalid rain formation choice: $rain_formation_choice") - end +if ft_choice == "Float64" + const FT = Float64 +elseif ft_choice == "Float32" + const FT = Float32 else - error("Invalid precipitation choice: $precipitation_choice") -end - -# Initialize the timestepping struct -TS = KID.TimeStepping(FT(opts["dt"]), FT(opts["dt_output"]), FT(opts["t_end"])) - -# Create the coordinates -space, face_space = KID.make_function_space(FT, opts["z_min"], opts["z_max"], opts["n_elem"]) -coord = CC.Fields.coordinate_field(space) -face_coord = CC.Fields.coordinate_field(face_space) - -# Initialize the netcdf output Stats struct -output_folder = string("Output_", moisture_choice, "_", prognostics_choice, "_", precipitation_choice) -if precipitation_choice in ["Precipitation1M", "Precipitation2M"] - output_folder = output_folder * "_" * rain_formation_choice -end -if opts["qtot_flux_correction"] - output_folder = output_folder * "_wFC" + error("Invalid float type: $ft_choice") end -path = joinpath(@__DIR__, output_folder) -mkpath(path) -fname = joinpath(path, "Output.nc") -Stats = KID.NetCDFIO_Stats(fname, 1.0, vec(face_coord), vec(coord)) -# Instantiate CliMA Parameters and overwrite the defaults -toml_dict = CP.create_toml_dict(FT; dict_type = "alias") -params = create_parameter_set( - path, - toml_dict, - FT, - opts["w1"], - opts["t1"], - opts["p0"], - Int(opts["precip_sources"]), - Int(opts["precip_sinks"]), - Int(opts["qtot_flux_correction"]), - FT(opts["prescribed_Nd"]), - FT(opts["r_dry"]), - FT(opts["std_dry"]), - FT(opts["kappa"]), -) +include("run_KiD_simulation.jl") -# Solve the initial value problem for density profile -ρ_profile = KID.ρ_ivp(FT, params) -# Create the initial condition profiles -init = map(coord -> KID.init_1d_column(FT, params, ρ_profile, coord.z), coord) - -# Create state vector and apply initial condition -Y = KID.initialise_state(moisture, precip, init) - -# Create aux vector and apply initial condition -aux = KID.initialise_aux(FT, init, params, TS, Stats, face_space, moisture) - -# Output the initial condition -KID.KiD_output(aux, 0.0) - -# Define callbacks for output -callback_io = ODE.DiscreteCallback(KID.condition_io, KID.affect_io!; save_positions = (false, false)) -callbacks = ODE.CallbackSet(callback_io) - -# Collect all the tendencies into rhs function for ODE solver -# based on model choices for the solved equations -ode_rhs! = KID.make_rhs_function(moisture, precip) - -# Solve the ODE operator -problem = ODE.ODEProblem(ode_rhs!, Y, (opts["t_ini"], opts["t_end"]), aux) -solver = ODE.solve( - problem, - ODE.SSPRK33(), - dt = TS.dt, - saveat = TS.dt_io, - callback = callbacks, - progress = true, - progress_message = (dt, u, p, t) -> t, -); - -# Some basic plots -if opts["plotting_flag"] == true - - include("../../plotting_utils.jl") - - plot_folder = string("experiments/KiD_driver/", output_folder, "/figures/") - - z_centers = parent(CC.Fields.coordinate_field(space)) - plot_final_aux_profiles(z_centers, aux, output = plot_folder) - plot_animation(z_centers, solver, aux, moisture, precip, KID, output = plot_folder) - plot_timeheight(string("experiments/KiD_driver/", output_folder, "/Output.nc"), output = plot_folder) -end +run_KiD_simulation(FT, opts) diff --git a/test/experiments/KiD_driver/parse_commandline.jl b/test/experiments/KiD_driver/parse_commandline.jl index 5acdcfcc..fde37f51 100644 --- a/test/experiments/KiD_driver/parse_commandline.jl +++ b/test/experiments/KiD_driver/parse_commandline.jl @@ -2,10 +2,16 @@ Parse the command line arguments for KiD_driver.jl """ +import ArgParse as AP + function parse_commandline() s = AP.ArgParseSettings() AP.@add_arg_table! s begin + "--FLOAT_TYPE" + help = "Float type. Can be set to Float64 or Float32" + arg_type = String + default = "Float64" "--moisture_choice" help = "Mositure model choice: EquilibriumMoisture, NonEquilibriumMoisture" arg_type = String @@ -45,56 +51,56 @@ function parse_commandline() default = false "--z_min" help = "Bottom of the computational domain [m]" - arg_type = Real - default = 0.0 + arg_type = Float64 + default = Float64(0.0) "--z_max" help = "Top of the computational domain [m]" - arg_type = Real - default = 2000.0 + arg_type = Float64 + default = Float64(2000) "--n_elem" help = "Number of computational elements" arg_type = Int default = 256 "--dt" help = "Simulation time step [s]" - arg_type = Real - default = 1.0 + arg_type = Float64 + default = Float64(1) "--dt_output" help = "Output time step [s]" - arg_type = Real - default = 30.0 + arg_type = Float64 + default = Float64(30) "--t_ini" help = "Time at the beginning of the simulation [s]" - arg_type = Real - default = 0.0 + arg_type = Float64 + default = Float64(0) "--t_end" help = "Time at the end of the simulation [s]" - arg_type = Real - default = 3600.0 + arg_type = Float64 + default = Float64(3600) "--w1" help = "Maximum prescribed updraft momentum flux [m/s * kg/m3]" - arg_type = Real - default = 2.0 + arg_type = Float64 + default = Float64(2) "--t1" help = "Oscillation time of the prescribed momentum flux [s]" - arg_type = Real - default = 600.0 + arg_type = Float64 + default = Float64(600) "--p0" help = "Pressure at the surface [pa]" - arg_type = Real - default = 100000.0 + arg_type = Float64 + default = Float64(100000) "--r_dry" help = "aerosol distribution mean radius for aerosol activation calculations in 2M schemes [m]" - arg_type = Real - default = 0.04 * 1e-6 + arg_type = Float64 + default = Float64(0.04 * 1e-6) "--std_dry" help = "aerosol distribution standard deviation for aerosol activation calucaulations in 2M schemes" - arg_type = Real - default = 1.4 + arg_type = Float64 + default = Float64(1.4) "--kappa" help = "hygroscopicity of aerosols for aerosol activation calucaulations in 2M schemes" - arg_type = Real - default = 0.9 + arg_type = Float64 + default = Float64(0.9) end return AP.parse_args(s) diff --git a/test/experiments/KiD_driver/run_KiD_simulation.jl b/test/experiments/KiD_driver/run_KiD_simulation.jl new file mode 100644 index 00000000..e27a34d1 --- /dev/null +++ b/test/experiments/KiD_driver/run_KiD_simulation.jl @@ -0,0 +1,151 @@ +import NCDatasets as NC +import OrdinaryDiffEq as ODE +import UnPack +import ArgParse as AP + +import ClimaCore as CC +import Thermodynamics as TD +import CloudMicrophysics.CommonTypes as CMT +import CLIMAParameters as CP +import Kinematic1D as KID + +const kid_dir = pkgdir(KID) +include(joinpath(kid_dir, "test", "create_parameters.jl")) + +include("../../plotting_utils.jl") + + +function run_KiD_simulation(::Type{FT}, opts) where {FT} + + # Equations to solve for mositure and precipitation variables + moisture_choice = opts["moisture_choice"] + prognostics_choice = opts["prognostic_vars"] + precipitation_choice = opts["precipitation_choice"] + rain_formation_choice = opts["rain_formation_scheme_choice"] + + if moisture_choice == "EquilibriumMoisture" && prognostics_choice == "RhoThetaQ" + moisture = KID.EquilibriumMoisture_ρθq() + elseif moisture_choice == "EquilibriumMoisture" && prognostics_choice == "RhodTQ" + moisture = KID.EquilibriumMoisture_ρdTq() + elseif moisture_choice == "NonEquilibriumMoisture" && prognostics_choice == "RhoThetaQ" + moisture = KID.NonEquilibriumMoisture_ρθq() + elseif moisture_choice == "NonEquilibriumMoisture" && prognostics_choice == "RhodTQ" + moisture = KID.NonEquilibriumMoisture_ρdTq() + else + error("Invalid moisture choice: $moisture_choice or invalid prognostic variables choice: $prognostic_vars") + end + if precipitation_choice == "NoPrecipitation" + precip = KID.NoPrecipitation() + elseif precipitation_choice == "Precipitation0M" + precip = KID.Precipitation0M() + elseif precipitation_choice == "Precipitation1M" + if rain_formation_choice == "CliMA_1M" + precip = KID.Precipitation1M(KID.OneMomentRainFormation()) + elseif rain_formation_choice == "KK2000" + precip = KID.Precipitation1M(CMT.KK2000Type()) + elseif rain_formation_choice == "B1994" + precip = KID.Precipitation1M(CMT.B1994Type()) + elseif rain_formation_choice == "TC1980" + precip = KID.Precipitation1M(CMT.TC1980Type()) + elseif rain_formation_choice == "LD2004" + precip = KID.Precipitation1M(CMT.LD2004Type()) + else + error("Invalid rain formation choice: $rain_formation_choice") + end + elseif precipitation_choice == "Precipitation2M" + if rain_formation_choice == "SB2006" + precip = KID.Precipitation2M(CMT.SB2006Type()) + else + error("Invalid rain formation choice: $rain_formation_choice") + end + else + error("Invalid precipitation choice: $precipitation_choice") + end + + # Initialize the timestepping struct + TS = KID.TimeStepping(FT(opts["dt"]), FT(opts["dt_output"]), FT(opts["t_end"])) + + # Create the coordinates + space, face_space = KID.make_function_space(FT, FT(opts["z_min"]), FT(opts["z_max"]), opts["n_elem"]) + coord = CC.Fields.coordinate_field(space) + face_coord = CC.Fields.coordinate_field(face_space) + + # Initialize the netcdf output Stats struct + output_folder = string("Output_", moisture_choice, "_", prognostics_choice, "_", precipitation_choice) + if precipitation_choice in ["Precipitation1M", "Precipitation2M"] + output_folder = output_folder * "_" * rain_formation_choice + end + if opts["qtot_flux_correction"] + output_folder = output_folder * "_wFC" + end + path = joinpath(@__DIR__, output_folder) + mkpath(path) + fname = joinpath(path, "Output.nc") + Stats = KID.NetCDFIO_Stats(fname, 1.0, vec(face_coord), vec(coord)) + + # Instantiate CliMA Parameters and overwrite the defaults + toml_dict = CP.create_toml_dict(FT; dict_type = "alias") + params = create_parameter_set( + path, + toml_dict, + FT, + FT(opts["w1"]), + FT(opts["t1"]), + FT(opts["p0"]), + Int(opts["precip_sources"]), + Int(opts["precip_sinks"]), + Int(opts["qtot_flux_correction"]), + FT(opts["prescribed_Nd"]), + FT(opts["r_dry"]), + FT(opts["std_dry"]), + FT(opts["kappa"]), + ) + + # Solve the initial value problem for density profile + ρ_profile = KID.ρ_ivp(FT, params) + # Create the initial condition profiles + init = map(coord -> KID.init_1d_column(FT, params, ρ_profile, coord.z), coord) + + # Create state vector and apply initial condition + Y = KID.initialise_state(moisture, precip, init) + + # Create aux vector and apply initial condition + aux = KID.initialise_aux(FT, init, params, TS, Stats, face_space, moisture) + + # Output the initial condition + KID.KiD_output(aux, 0.0) + + # Define callbacks for output + callback_io = ODE.DiscreteCallback(KID.condition_io, KID.affect_io!; save_positions = (false, false)) + callbacks = ODE.CallbackSet(callback_io) + + # Collect all the tendencies into rhs function for ODE solver + # based on model choices for the solved equations + ode_rhs! = KID.make_rhs_function(moisture, precip) + + # Solve the ODE operator + problem = ODE.ODEProblem(ode_rhs!, Y, (FT(opts["t_ini"]), FT(opts["t_end"])), aux) + solver = ODE.solve( + problem, + ODE.SSPRK33(), + dt = TS.dt, + saveat = TS.dt_io, + callback = callbacks, + progress = true, + progress_message = (dt, u, p, t) -> t, + ) + + # Some basic plots + if opts["plotting_flag"] == true + + plot_folder = string("experiments/KiD_driver/", output_folder, "/figures/") + + z_centers = parent(CC.Fields.coordinate_field(space)) + plot_final_aux_profiles(z_centers, aux, precip, output = plot_folder) + plot_animation(z_centers, solver, aux, moisture, precip, KID, output = plot_folder) + plot_timeheight(string("experiments/KiD_driver/", output_folder, "/Output.nc"), output = plot_folder) + end + + return solver + +end diff --git a/test/plotting_utils.jl b/test/plotting_utils.jl index bd74c6b6..151357fc 100644 --- a/test/plotting_utils.jl +++ b/test/plotting_utils.jl @@ -1,6 +1,9 @@ """ Plotting utilities """ + +import CloudMicrophysics.PrecipitationSusceptibility as CMPS + ENV["GKSwstype"] = "nul" using ClimaCorePlots, Plots Plots.GRBackend() @@ -44,7 +47,7 @@ function plot_initial_profiles_comparison(KM; sdm_case = "dry") Plots.png(p, joinpath(path, fig_name)) end -function plot_final_aux_profiles(z_centers, aux; output = "output") +function plot_final_aux_profiles(z_centers, aux, precip; output = "output") path = joinpath(@__DIR__, output) mkpath(path) @@ -55,6 +58,8 @@ function plot_final_aux_profiles(z_centers, aux; output = "output") q_ice_end = parent(aux.moisture_variables.q_ice) q_rai_end = parent(aux.precip_variables.q_rai) q_sno_end = parent(aux.precip_variables.q_sno) + N_liq_end = parent(aux.precip_variables.N_liq) + ρ_end = parent(aux.moisture_variables.ρ) p1 = Plots.plot(q_tot_end .* 1e3, z_centers, xlabel = "q_tot [g/kg]", ylabel = "z [m]") p2 = Plots.plot(q_liq_end .* 1e3, z_centers, xlabel = "q_liq [g/kg]", ylabel = "z [m]") @@ -63,6 +68,33 @@ function plot_final_aux_profiles(z_centers, aux; output = "output") p5 = Plots.plot(q_rai_end .* 1e3, z_centers, xlabel = "q_rai [g/kg]", ylabel = "z [m]") p6 = Plots.plot(q_sno_end .* 1e3, z_centers, xlabel = "q_sno [g/kg]", ylabel = "z [m]") + p7 = Plots.plot(xlabel = "precipitation susceptibility", ylabel = "z [m]") + if precip isa KID.Precipitation2M + precip_sus_aut = + CMPS.precipitation_susceptibility_autoconversion.( + Ref(aux.params.microphys_params), + Ref(precip.rain_formation), + q_liq_end, + q_rai_end, + ρ_end, + N_liq_end, + ) + precip_sus_acc = + CMPS.precipitation_susceptibility_accretion.( + Ref(aux.params.microphys_params), + Ref(precip.rain_formation), + q_liq_end, + q_rai_end, + ρ_end, + N_liq_end, + ) + Plots.plot!([r.d_ln_pp_d_ln_q_liq for r in precip_sus_aut], z_centers, label = "aut, q_liq", color = :red) + Plots.plot!([r.d_ln_pp_d_ln_q_rai for r in precip_sus_aut], z_centers, label = "aut, q_rai", color = :brown) + Plots.plot!([r.d_ln_pp_d_ln_q_liq for r in precip_sus_acc], z_centers, label = "acc, q_liq", color = :blue) + Plots.plot!([r.d_ln_pp_d_ln_q_rai for r in precip_sus_acc], z_centers, label = "acc, q_rai", color = :green) + Plots.plot!(legend = :outerright) + end + p = Plots.plot( p1, p2, @@ -70,6 +102,7 @@ function plot_final_aux_profiles(z_centers, aux; output = "output") p4, p5, p6, + p7, size = (1200.0, 750.0), bottom_margin = 40.0 * Plots.PlotMeasures.px, left_margin = 80.0 * Plots.PlotMeasures.px, @@ -93,7 +126,7 @@ function plot_animation(z_centers, solver, aux, moisture, precip, KiD; output = q_liq = q_tot .* 0.0 end - if precip isa KiD.Precipitation1M + if !(precip isa Union{KiD.NoPrecipitation, KiD.Precipitation0M}) q_rai = parent(u.ρq_rai) ./ ρ .* 1e3 else q_rai = q_tot .* 0.0 diff --git a/test/precip_production_tests/test_precip_production.jl b/test/precip_production_tests/test_precip_production.jl new file mode 100644 index 00000000..e8463413 --- /dev/null +++ b/test/precip_production_tests/test_precip_production.jl @@ -0,0 +1,66 @@ +using ClimaCorePlots, Plots +Plots.GRBackend() + +using FiniteDiff + +import CloudMicrophysics.Microphysics2M as CM2 + +include("../experiments/KiD_driver/parse_commandline.jl") + +include("../experiments/KiD_driver/run_KiD_simulation.jl") + +path = joinpath(@__DIR__, "Output_test_precip_production") +mkpath(path) + +const FT = Float64 + +""" +This is a very simple metric of precipitation production that is calculated +by summing the ρq_rai at a height corresponding to the cloud base over time. +This ignores effects due to varying terminal velocity, and does not have the +correct units, but is sufficient for testing susceptibility since it is +approximately proportional to the precipitation production. +""" +function total_precipitation_production(Nd) + opts = parse_commandline() + + opts["plotting_flag"] = false + opts["precipitation_choice"] = "Precipitation2M" + opts["rain_formation_scheme_choice"] = "SB2006" + opts["prescribed_Nd"] = Nd + + cloud_base = 75 + + println(opts["prescribed_Nd"]) + + solver = run_KiD_simulation(FT, opts) + + precip = sum([parent(u.ρq_rai)[cloud_base] for u in solver.u]) + + println("Precipitation production: ", precip) + + return precip +end + +Nd_range = FT.(range(1e5, 1e8, 10)) + +# p = Plots.plot(Nd_range, total_precipitation_production.(Nd_range), yscale = :log10, +# xlabel = "Prescribed Nd", ylabel = "Precipitation Production Metric") +# Plots.png(p, joinpath(path, "precip_production.png")) + +precip_sus = + Nd_range ./ total_precipitation_production.(Nd_range) .* + FiniteDiff.finite_difference_derivative.(total_precipitation_production, Nd_range) + +p = Plots.plot( + Nd_range, + precip_sus, + xlabel = "Prescribed Nd", + ylabel = "Precipitation Susceptibility", + label = "Observed (KiD)", +) + +expected_precip_sus = fill!(similar(Nd_range), -2) +Plots.plot!(Nd_range, expected_precip_sus, label = "Expected (Glassmeier & Lohmann)") + +Plots.png(p, joinpath(path, "precip_susceptibility.png"))