From 691d4dfda56f438df6ac5fe860ed11bd5f451415 Mon Sep 17 00:00:00 2001
From: Julia Sloan <jsloan@caltech.edu>
Date: Wed, 24 Apr 2024 16:48:25 -0700
Subject: [PATCH] make diagnostics optional

---
 config/benchmark_configs/amip_diagedmf.yml    |  1 +
 .../benchmark_configs/gpu_amip_diagedmf.yml   |  1 +
 experiments/AMIP/cli_options.jl               |  4 ++
 experiments/AMIP/coupler_driver.jl            | 57 ++++++++++---------
 4 files changed, 37 insertions(+), 26 deletions(-)

diff --git a/config/benchmark_configs/amip_diagedmf.yml b/config/benchmark_configs/amip_diagedmf.yml
index df8ab111b5..961aac02b8 100644
--- a/config/benchmark_configs/amip_diagedmf.yml
+++ b/config/benchmark_configs/amip_diagedmf.yml
@@ -13,3 +13,4 @@ run_name: "amip_diagedmf"
 start_date: "19790301"
 t_end: "12hours"
 turb_flux_partition: "CombinedStateFluxes"
+use_coupler_diagnostics: false
diff --git a/config/benchmark_configs/gpu_amip_diagedmf.yml b/config/benchmark_configs/gpu_amip_diagedmf.yml
index 3e1eef626a..d353dbb8db 100644
--- a/config/benchmark_configs/gpu_amip_diagedmf.yml
+++ b/config/benchmark_configs/gpu_amip_diagedmf.yml
@@ -13,3 +13,4 @@ run_name: "gpu_amip_diagedmf"
 start_date: "19790301"
 t_end: "12hours"
 turb_flux_partition: "CombinedStateFluxes"
+use_coupler_diagnostics: false
diff --git a/experiments/AMIP/cli_options.jl b/experiments/AMIP/cli_options.jl
index 8a7ef7b397..a533f6abd5 100644
--- a/experiments/AMIP/cli_options.jl
+++ b/experiments/AMIP/cli_options.jl
@@ -78,6 +78,10 @@ function argparse_settings()
         help = "Device type to use [`auto` (default) `CPUSingleThreaded`, `CPUMultiThreaded`, `CUDADevice`]"
         arg_type = String
         default = "auto"
+        "--use_coupler_diagnostics"
+        help = "Boolean flag indicating whether to compute and output coupler diagnostics [`true` (default), `false`]"
+        arg_type = Bool
+        default = true
         # ClimaAtmos specific
         "--surface_setup"
         help = "Triggers ClimaAtmos into the coupled mode [`PrescribedSurface` (default)]" # retained here for standalone Atmos benchmarks
diff --git a/experiments/AMIP/coupler_driver.jl b/experiments/AMIP/coupler_driver.jl
index aba26201b7..f97fce4de6 100644
--- a/experiments/AMIP/coupler_driver.jl
+++ b/experiments/AMIP/coupler_driver.jl
@@ -128,6 +128,7 @@ restart_dir = config_dict["restart_dir"]
 restart_t = Int(config_dict["restart_t"])
 evolving_ocean = config_dict["evolving_ocean"]
 dt_rad = config_dict["dt_rad"]
+use_coupler_diagnostics = config_dict["use_coupler_diagnostics"]
 
 #=
 ## Setup Communication Context
@@ -497,27 +498,30 @@ User can write custom diagnostics in the `user_diagnostics.jl`.
 Note, this will be replaced by the diagnostics framework currently in ClimaAtmos, once it is abstracted
 into a more general package, so we can use it to save fields from surface models.
 =#
+if use_coupler_diagnostics
+    monthly_3d_diags = Diagnostics.init_diagnostics(
+        (:T, :u, :q_tot, :q_liq_ice),
+        atmos_sim.domain.center_space;
+        save = TimeManager.Monthly(),
+        operations = (; accumulate = Diagnostics.TimeMean([Int(0)])),
+        output_dir = COUPLER_OUTPUT_DIR,
+        name_tag = "monthly_mean_3d_",
+    )
 
-monthly_3d_diags = Diagnostics.init_diagnostics(
-    (:T, :u, :q_tot, :q_liq_ice),
-    atmos_sim.domain.center_space;
-    save = TimeManager.Monthly(),
-    operations = (; accumulate = Diagnostics.TimeMean([Int(0)])),
-    output_dir = COUPLER_OUTPUT_DIR,
-    name_tag = "monthly_mean_3d_",
-)
-
-monthly_2d_diags = Diagnostics.init_diagnostics(
-    (:precipitation_rate, :toa_fluxes, :T_sfc, :tubulent_energy_fluxes),
-    boundary_space;
-    save = TimeManager.Monthly(),
-    operations = (; accumulate = Diagnostics.TimeMean([Int(0)])),
-    output_dir = COUPLER_OUTPUT_DIR,
-    name_tag = "monthly_mean_2d_",
-)
+    monthly_2d_diags = Diagnostics.init_diagnostics(
+        (:precipitation_rate, :toa_fluxes, :T_sfc, :tubulent_energy_fluxes),
+        boundary_space;
+        save = TimeManager.Monthly(),
+        operations = (; accumulate = Diagnostics.TimeMean([Int(0)])),
+        output_dir = COUPLER_OUTPUT_DIR,
+        name_tag = "monthly_mean_2d_",
+    )
 
-diagnostics = (monthly_3d_diags, monthly_2d_diags)
-Utilities.show_memory_usage(comms_ctx, Dict(diagnostics => "diagnostics"))
+    diagnostics = (monthly_3d_diags, monthly_2d_diags)
+    Utilities.show_memory_usage(comms_ctx, Dict(diagnostics => "diagnostics"))
+else
+    diagnostics = ()
+end
 
 #=
 ## Initialize Conservation Checks
@@ -743,13 +747,14 @@ function solve_coupler!(cs)
             CO2_current = BCReader.interpolate_midmonth_to_daily(cs.dates.date[1], cs.mode.CO2_info)
             Interfacer.update_field!(atmos_sim, Val(:co2), CO2_current)
 
-            ## calculate and accumulate diagnostics at each timestep
-            ClimaComms.barrier(comms_ctx)
-            Diagnostics.accumulate_diagnostics!(cs)
-
-            ## save and reset monthly averages
-            Diagnostics.save_diagnostics(cs)
+            ## calculate and accumulate diagnostics at each timestep, if we're using diagnostics in this run
+            if !isempty(cs.diagnostics)
+                ClimaComms.barrier(comms_ctx)
+                Diagnostics.accumulate_diagnostics!(cs)
 
+                ## save and reset monthly averages
+                Diagnostics.save_diagnostics(cs)
+            end
         end
 
         ## compute global energy
@@ -860,7 +865,7 @@ if ClimaComms.iamroot(comms_ctx)
     end
 
     ## plotting AMIP results
-    if cs.mode.name == "amip"
+    if cs.mode.name == "amip" && !isempty(cs.diagnostics)
         @info "AMIP plots"
 
         ## ClimaESM