Add cloud diagnostics per stage option

.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/*"

.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

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

config/gpu_configs/central_cloud_diag_gpu_hs_rhoe_equil_55km_nz63_0M.yml

@@ -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]

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
 

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

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?

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

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,

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