From 9ec80be35bda582c4f3fc901f3d14c661a773fd2 Mon Sep 17 00:00:00 2001 From: Charles Kawczynski Date: Thu, 4 Apr 2024 10:37:25 -0400 Subject: [PATCH] Add cloud diag per stage option --- .buildkite/gpu_pipeline/pipeline.yml | 1 + .buildkite/pipeline.yml | 10 +++++++++ config/default_configs/default_config.yml | 4 ++++ ...d_diag_gpu_hs_rhoe_equil_55km_nz63_0M.yaml | 18 +++++++++++++++ src/cache/cloud_fraction.jl | 22 +++++++++++++++++++ src/cache/precomputed_quantities.jl | 8 +++++-- src/callbacks/get_callbacks.jl | 8 ++++--- src/solver/model_getters.jl | 11 ++++++++++ src/solver/type_getters.jl | 3 +++ src/solver/types.jl | 2 ++ 10 files changed, 82 insertions(+), 5 deletions(-) create mode 100644 config/gpu_configs/central_cloud_diag_gpu_hs_rhoe_equil_55km_nz63_0M.yaml diff --git a/.buildkite/gpu_pipeline/pipeline.yml b/.buildkite/gpu_pipeline/pipeline.yml index b6504aa3043..c164cfa4237 100644 --- a/.buildkite/gpu_pipeline/pipeline.yml +++ b/.buildkite/gpu_pipeline/pipeline.yml @@ -110,6 +110,7 @@ steps: - mkdir -p gpu_aquaplanet_dyamond_ss_1process - > srun --cpu-bind=threads --cpus-per-task=4 + nsys profile --trace=nvtx,mpi,cuda,osrt --output=gpu_aquaplanet_dyamond_ss_1process/output_active/report julia --threads=3 --color=yes --project=examples examples/hybrid/driver.jl --config_file ${GPU_CONFIG_PATH}gpu_aquaplanet_dyamond_ss_1process.yml artifact_paths: "gpu_aquaplanet_dyamond_ss_1process/output_active/*" diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 926f05b3fe0..637cb9c9f5e 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -699,6 +699,16 @@ steps: slurm_gpus: 1 slurm_mem: 16G + - label: "GPU: GPU moist Held-Suarez cloud diagnostics per stage" + command: + - > + julia --color=yes --project=examples examples/hybrid/driver.jl + --config_file $GPU_CONFIG_PATH/central_cloud_diag_gpu_hs_rhoe_equil_55km_nz63_0M.yml + artifact_paths: "central_cloud_diag_gpu_hs_rhoe_equil_55km_nz63_0M/output_active/*" + agents: + slurm_gpus: 1 + slurm_mem: 16G + - label: ":umbrella: GPU: gpu_aquaplanet_dyamond" command: - mkdir -p gpu_aquaplanet_dyamond diff --git a/config/default_configs/default_config.yml b/config/default_configs/default_config.yml index b63e5a3fd73..176446dfb71 100644 --- a/config/default_configs/default_config.yml +++ b/config/default_configs/default_config.yml @@ -280,3 +280,7 @@ restart_file: prescribed_aerosols: help: "Which aerosols to add. List of keys from the data file (e.g., CB1, CB2)." value: [] +# TODO: remove once https://github.com/CliMA/ClimaAtmos.jl/issues/2873 is closed +call_cloud_diagnostics_per_stage: + help: "A temporary Bool for calling cloud diagnostics every stage" + value: false diff --git a/config/gpu_configs/central_cloud_diag_gpu_hs_rhoe_equil_55km_nz63_0M.yaml b/config/gpu_configs/central_cloud_diag_gpu_hs_rhoe_equil_55km_nz63_0M.yaml new file mode 100644 index 00000000000..b8e30d5ec2c --- /dev/null +++ b/config/gpu_configs/central_cloud_diag_gpu_hs_rhoe_equil_55km_nz63_0M.yaml @@ -0,0 +1,18 @@ +dt_save_state_to_disk: "Inf" +dt_save_to_sol: "Inf" +output_default_diagnostics: false +dt: "150secs" +t_end: "2days" +h_elem: 16 +z_elem: 63 +dz_bottom: 30.0 +dz_top: 3000.0 +z_max: 55000.0 +vert_diff: "true" +moist: "equil" +precip_model: "0M" +rayleigh_sponge: true +forcing: "held_suarez" +call_cloud_diagnostics_per_stage: true +job_id: "central_cloud_diag_gpu_hs_rhoe_equil_55km_nz63_0M" +toml: [toml/longrun_hs_rhoe_equil_55km_nz63_0M.toml] diff --git a/src/cache/cloud_fraction.jl b/src/cache/cloud_fraction.jl index bc6204d9ba1..c517d0c6bcf 100644 --- a/src/cache/cloud_fraction.jl +++ b/src/cache/cloud_fraction.jl @@ -33,6 +33,17 @@ NVTX.@annotate function set_cloud_fraction!( (; ᶜts, ᶜmixing_length, ᶜcloud_fraction) = p.precomputed thermo_params = CAP.thermodynamics_params(params) if isnothing(turbconv_model) + if p.atmos.call_cloud_diagnostics_per_stage isa + CallCloudDiagnosticsPerStage + (; ᶜgradᵥ_θ_virt, ᶜgradᵥ_q_tot, ᶜgradᵥ_θ_liq_ice) = p.precomputed + thermo_params = CAP.thermodynamics_params(p.params) + @. ᶜgradᵥ_θ_virt = + ᶜgradᵥ(ᶠinterp(TD.virtual_pottemp(thermo_params, ᶜts))) + @. ᶜgradᵥ_q_tot = + ᶜgradᵥ(ᶠinterp(TD.total_specific_humidity(thermo_params, ᶜts))) + @. ᶜgradᵥ_θ_liq_ice = + ᶜgradᵥ(ᶠinterp(TD.liquid_ice_pottemp(thermo_params, ᶜts))) + end compute_gm_mixing_length!(ᶜmixing_length, Y, p) end @. ᶜcloud_fraction = ifelse(TD.has_condensate(thermo_params, ᶜts), 1, 0) @@ -50,6 +61,17 @@ NVTX.@annotate function set_cloud_fraction!( (; ᶜts, ᶜp, ᶜmixing_length, ᶜcloud_fraction) = p.precomputed (; turbconv_model) = p.atmos if isnothing(turbconv_model) + if p.atmos.call_cloud_diagnostics_per_stage isa + CallCloudDiagnosticsPerStage + (; ᶜgradᵥ_θ_virt, ᶜgradᵥ_q_tot, ᶜgradᵥ_θ_liq_ice) = p.precomputed + thermo_params = CAP.thermodynamics_params(p.params) + @. ᶜgradᵥ_θ_virt = + ᶜgradᵥ(ᶠinterp(TD.virtual_pottemp(thermo_params, ᶜts))) + @. ᶜgradᵥ_q_tot = + ᶜgradᵥ(ᶠinterp(TD.total_specific_humidity(thermo_params, ᶜts))) + @. ᶜgradᵥ_θ_liq_ice = + ᶜgradᵥ(ᶠinterp(TD.liquid_ice_pottemp(thermo_params, ᶜts))) + end compute_gm_mixing_length!(ᶜmixing_length, Y, p) end diff --git a/src/cache/precomputed_quantities.jl b/src/cache/precomputed_quantities.jl index c10e25ed743..cdcfc592934 100644 --- a/src/cache/precomputed_quantities.jl +++ b/src/cache/precomputed_quantities.jl @@ -443,7 +443,9 @@ function instead of recomputing the value yourself. Otherwise, it will be difficult to ensure that the duplicated computations are consistent. """ NVTX.@annotate function set_precomputed_quantities!(Y, p, t) - (; moisture_model, turbconv_model, vert_diff, precip_model) = p.atmos + (; moisture_model, turbconv_model, vert_diff, precip_model, cloud_model) = + p.atmos + (; call_cloud_diagnostics_per_stage) = p.atmos thermo_params = CAP.thermodynamics_params(p.params) n = n_mass_flux_subdomains(turbconv_model) thermo_args = (thermo_params, moisture_model) @@ -618,7 +620,9 @@ NVTX.@annotate function set_precomputed_quantities!(Y, p, t) end # TODO - #set_cloud_fraction!(Y, p, moisture_model) + if call_cloud_diagnostics_per_stage isa CallCloudDiagnosticsPerStage + set_cloud_fraction!(Y, p, moisture_model, cloud_model) + end return nothing end diff --git a/src/callbacks/get_callbacks.jl b/src/callbacks/get_callbacks.jl index 5ac2c198963..d8cb3c9b655 100644 --- a/src/callbacks/get_callbacks.jl +++ b/src/callbacks/get_callbacks.jl @@ -272,9 +272,11 @@ function get_callbacks(config, sim_info, atmos, params, Y, p, t_start) ) end - dt_cf = FT(time_to_seconds(parsed_args["dt_cloud_fraction"])) - callbacks = - (callbacks..., call_every_dt(cloud_fraction_model_callback!, dt_cf)) + if !parsed_args["call_cloud_diagnostics_per_stage"] + dt_cf = FT(time_to_seconds(parsed_args["dt_cloud_fraction"])) + callbacks = + (callbacks..., call_every_dt(cloud_fraction_model_callback!, dt_cf)) + end if atmos.radiation_mode isa RRTMGPI.AbstractRRTMGPMode # TODO: better if-else criteria? diff --git a/src/solver/model_getters.jl b/src/solver/model_getters.jl index e8a6a765c9d..f0da55e4096 100644 --- a/src/solver/model_getters.jl +++ b/src/solver/model_getters.jl @@ -289,6 +289,17 @@ function get_forcing_type(parsed_args) end end +struct CallCloudDiagnosticsPerStage end +function get_call_cloud_diagnostics_per_stage(parsed_args) + ccdps = parsed_args["call_cloud_diagnostics_per_stage"] + @assert ccdps in (nothing, true, false) + return if ccdps in (nothing, false) + nothing + elseif ccdps == true + CallCloudDiagnosticsPerStage() + end +end + function get_subsidence_model(parsed_args, radiation_mode, FT) subsidence = parsed_args["subsidence"] subsidence == nothing && return nothing diff --git a/src/solver/type_getters.jl b/src/solver/type_getters.jl index 0c3ebf8e546..98e4d0de71a 100644 --- a/src/solver/type_getters.jl +++ b/src/solver/type_getters.jl @@ -21,6 +21,8 @@ function get_atmos(config::AtmosConfig, params) cloud_model = get_cloud_model(parsed_args) radiation_mode = get_radiation_mode(parsed_args, FT) forcing_type = get_forcing_type(parsed_args) + call_cloud_diagnostics_per_stage = + get_call_cloud_diagnostics_per_stage(parsed_args) diffuse_momentum = !(forcing_type isa HeldSuarezForcing) @@ -70,6 +72,7 @@ function get_atmos(config::AtmosConfig, params) precip_model, cloud_model, forcing_type, + call_cloud_diagnostics_per_stage, turbconv_model = get_turbconv_model(FT, parsed_args, turbconv_params), non_orographic_gravity_wave = get_non_orographic_gravity_wave_model( parsed_args, diff --git a/src/solver/types.jl b/src/solver/types.jl index f5618f017e7..3ba4df6d10f 100644 --- a/src/solver/types.jl +++ b/src/solver/types.jl @@ -327,6 +327,7 @@ Base.@kwdef struct AtmosModel{ MM, PM, CM, + CCDPS, F, S, RM, @@ -359,6 +360,7 @@ Base.@kwdef struct AtmosModel{ moisture_model::MM = nothing precip_model::PM = nothing cloud_model::CM = nothing + call_cloud_diagnostics_per_stage::CCDPS = nothing forcing_type::F = nothing subsidence::S = nothing radiation_mode::RM = nothing