From d462a7f61e95b5c2496bd78a94bff414e045a1d3 Mon Sep 17 00:00:00 2001
From: Charles Kawczynski <kawczynski.charles@gmail.com>
Date: Wed, 22 May 2024 09:18:07 -0400
Subject: [PATCH] Do not exercise diagnostics when disabled

---
 config/default_configs/default_config.yml |  3 ++
 src/solver/solve.jl                       |  2 +-
 src/solver/type_getters.jl                | 38 +++++++++++++----------
 3 files changed, 26 insertions(+), 17 deletions(-)

diff --git a/config/default_configs/default_config.yml b/config/default_configs/default_config.yml
index bf88fd2394..8ab7ec0ed5 100644
--- a/config/default_configs/default_config.yml
+++ b/config/default_configs/default_config.yml
@@ -259,6 +259,9 @@ override_τ_precip:
 output_default_diagnostics:
   help: "Output the default diagnostics associated to the selected atmospheric model"
   value: true
+enable_diagnostics:
+  help: "Set to false to fully disable the diagnostics"
+  value: true
 netcdf_output_at_levels:
   help: "Do not perform any vertical interpolation in the output NetCDF files. Instead, interpolate horizontally and output the level. "
   value: true
diff --git a/src/solver/solve.jl b/src/solver/solve.jl
index 5d53ea2528..36ba4995c7 100644
--- a/src/solver/solve.jl
+++ b/src/solver/solve.jl
@@ -103,7 +103,7 @@ function solve_atmos!(simulation)
         maxrss_str = prettymemory(maxrss())
         @info "Memory currently used (after solve!) by the process (RSS): $maxrss_str"
 
-        foreach(close, output_writers)
+        isnothing(output_writers) || foreach(close, output_writers)
     end
 end
 
diff --git a/src/solver/type_getters.jl b/src/solver/type_getters.jl
index 18478495e0..badd1612f3 100644
--- a/src/solver/type_getters.jl
+++ b/src/solver/type_getters.jl
@@ -663,17 +663,21 @@ function get_simulation(config::AtmosConfig)
     @info "get_callbacks: $s"
 
     # Initialize diagnostics
-    s = @timed_str begin
-        scheduled_diagnostics, writers = get_diagnostics(
-            config.parsed_args,
-            atmos,
-            Y,
-            p,
-            t_start,
-            sim_info.dt,
-        )
+    if config.parsed_args["enable_diagnostics"]
+        s = @timed_str begin
+            scheduled_diagnostics, writers = get_diagnostics(
+                config.parsed_args,
+                atmos,
+                Y,
+                p,
+                t_start,
+                sim_info.dt,
+            )
+        end
+        @info "initializing diagnostics: $s"
+    else
+        writers = nothing
     end
-    @info "initializing diagnostics: $s"
 
     continuous_callbacks = tuple()
     discrete_callbacks = callback
@@ -705,13 +709,15 @@ function get_simulation(config::AtmosConfig)
     end
     @info "init integrator: $s"
 
-    s = @timed_str begin
-        integrator = ClimaDiagnostics.IntegratorWithDiagnostics(
-            integrator,
-            scheduled_diagnostics,
-        )
+    if config.parsed_args["enable_diagnostics"]
+        s = @timed_str begin
+            integrator = ClimaDiagnostics.IntegratorWithDiagnostics(
+                integrator,
+                scheduled_diagnostics,
+            )
+        end
+        @info "Added diagnostics: $s"
     end
-    @info "Added diagnostics: $s"
 
     reset_graceful_exit(output_dir)