From 11e7cbfffc45f058defb91f3b61e6fdca603a7b5 Mon Sep 17 00:00:00 2001
From: Julia Sloan <jsloan@caltech.edu>
Date: Tue, 24 Sep 2024 16:04:12 -0700
Subject: [PATCH] make plots for atmos, coupler separately

---
 experiments/ClimaEarth/run_amip.jl         | 22 ++++++++++++++++++----
 experiments/ClimaEarth/user_io/ci_plots.jl | 16 +++++++---------
 2 files changed, 25 insertions(+), 13 deletions(-)

diff --git a/experiments/ClimaEarth/run_amip.jl b/experiments/ClimaEarth/run_amip.jl
index 6c9d39d2eb..1f2d2df7b4 100644
--- a/experiments/ClimaEarth/run_amip.jl
+++ b/experiments/ClimaEarth/run_amip.jl
@@ -837,9 +837,23 @@ if ClimaComms.iamroot(comms_ctx)
         ## ClimaESM
         include("user_io/ci_plots.jl")
 
-        # TODO add turbulent_energy_fluxes?
-        amip_short_names = ["ta", "ua", "hus", "clw", "pr", "ts", "toa_fluxes_net", "F_turb_energy"]
-        make_ci_plots([atmos_sim.integrator.p.output_dir], dir_paths.artifacts, short_names = amip_short_names)
+        # define variable names and output directories for each diagnostic
+        amip_short_names_atmos = ["ta", "ua", "hus", "clw", "pr", "ts", "toa_fluxes_net"]
+        output_dir_atmos = atmos_sim.integrator.p.output_dir
+        amip_short_names_coupler = ["F_turb_energy"]
+        output_dir_coupler = dir_paths.artifacts
+
+        # Check if all output variables are available in the specified directories
+        all(v -> v in CAN.available_vars(output_dir_atmos), amip_short_names_atmos) && make_ci_plots(
+            output_dir_atmos,
+            joinpath(dir_paths.artifacts, "atmos"),
+            short_names = amip_short_names_atmos,
+        )
+        all(v -> v in CAN.available_vars(output_dir_coupler), amip_short_names_coupler) && make_ci_plots(
+            output_dir_coupler,
+            joinpath(dir_paths.artifacts, "coupler"),
+            short_names = amip_short_names_coupler,
+        )
 
         ## Compare against observations
         if t_end > 84600 && config_dict["output_default_diagnostics"]
@@ -954,7 +968,7 @@ if ClimaComms.iamroot(comms_ctx)
     if config_dict["ci_plots"]
         @info "Generating CI plots"
         include("user_io/ci_plots.jl")
-        make_ci_plots([atmos_sim.integrator.p.output_dir], dir_paths.artifacts)
+        make_ci_plots(atmos_sim.integrator.p.output_dir, dir_paths.artifacts)
     end
 
     ## plot all model states and coupler fields (useful for debugging)
diff --git a/experiments/ClimaEarth/user_io/ci_plots.jl b/experiments/ClimaEarth/user_io/ci_plots.jl
index c0a0a1cdf6..d2b0328acf 100644
--- a/experiments/ClimaEarth/user_io/ci_plots.jl
+++ b/experiments/ClimaEarth/user_io/ci_plots.jl
@@ -145,7 +145,7 @@ end
 
 """
     make_ci_plots(
-        output_paths::Vector{<:AbstractString},
+        output_path::AbstractString,
         plot_path::AbstractString;
         short_names::Vector{<:AbstractString} = ["mse", "lr", "edt", "ts"],
         reduction::String = "average",
@@ -156,14 +156,14 @@ to include additional diagnostics.
 The `reduction` keyword argument should be consistent with the reduction used to save the diagnostics.
 """
 function make_ci_plots(
-    output_paths::Vector{<:AbstractString},
+    output_path::AbstractString,
     plot_path::AbstractString;
     short_names::Vector{<:AbstractString} = ["mse", "lr", "edt", "ts"],
     reduction::String = "average",
 )
-    simdirs = CAN.SimDir.(output_paths)
+    simdir = CAN.SimDir(output_path)
 
-    available_periods = CAN.available_periods(simdirs[1]; short_name = short_names[1], reduction)
+    available_periods = CAN.available_periods(simdir; short_name = short_names[1], reduction)
     period = ""
     if "10d" in available_periods
         period = "10d"
@@ -174,9 +174,7 @@ function make_ci_plots(
     end
 
     # Create a CAN.OutputVar for each input field
-    vars = map_comparison(simdirs, short_names) do simdir, short_name
-        get(simdir; short_name, reduction, period)
-    end
+    vars = [get(simdir; short_name, reduction, period) for short_name in short_names]
 
     # Filter vars into 2D and 3D variable diagnostics vectors
     # 3D fields are zonally averaged platted on the lat-z plane
@@ -185,6 +183,6 @@ function make_ci_plots(
     vars_2D = filter(var -> !CAN.has_altitude(var), vars)
 
     # Generate plots and save in `plot_path`
-    make_plots_generic(output_paths, plot_path, vars_3D, time = LAST_SNAP, more_kwargs = YLINEARSCALE)
-    make_plots_generic(output_paths, plot_path, vars_2D, time = LAST_SNAP, output_name = "summary_2D")
+    make_plots_generic(output_path, plot_path, vars_3D, time = LAST_SNAP, more_kwargs = YLINEARSCALE)
+    make_plots_generic(output_path, plot_path, vars_2D, time = LAST_SNAP, output_name = "summary_2D")
 end