extend make_ci_plots for amip paperplots

This generalizes the `make_ci_plots` function to `make_plots`
which takes in variable names as an argument instead of
hardcoding them within the function

experiments/ClimaEarth/run_amip.jl

@@ -890,40 +890,17 @@ if ClimaComms.iamroot(comms_ctx)
     end
 
     ## plotting AMIP results
-    if cs.mode.name == "amip" && !isempty(cs.diagnostics)
+    if cs.mode.name == "amip"
         ## plot data that correspond to the model's last save_hdf5 call (i.e., last month)
         @info "AMIP plots"
 
         ## ClimaESM
-        include("user_io/amip_visualizer.jl")
-        post_spec = (;
-            T = (:regrid, :zonal_mean),
-            u = (:regrid, :zonal_mean),
-            q_tot = (:regrid, :zonal_mean),
-            toa_fluxes = (:regrid, :horizontal_slice),
-            precipitation_rate = (:regrid, :horizontal_slice),
-            T_sfc = (:regrid, :horizontal_slice),
-            turbulent_energy_fluxes = (:regrid, :horizontal_slice),
-            q_liq_ice = (:regrid, :zonal_mean),
-        )
+        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)
 
-        plot_spec = (;
-            T = (; clims = (190, 320), units = "K"),
-            u = (; clims = (-50, 50), units = "m/s"),
-            q_tot = (; clims = (0, 30), units = "g/kg"),
-            toa_fluxes = (; clims = (-250, 250), units = "W/m^2"),
-            precipitation_rate = (clims = (0, 1e-4), units = "kg/m^2/s"),
-            T_sfc = (clims = (225, 310), units = "K"),
-            turbulent_energy_fluxes = (; clims = (-250, 250), units = "W/m^2"),
-            q_liq_ice = (; clims = (0, 10), units = "g/kg"),
-        )
-        amip_data, fig_amip = amip_paperplots(
-            post_spec,
-            plot_spec,
-            dir_paths.output,
-            files_root = ".monthly",
-            output_dir = dir_paths.artifacts,
-        )
         # Check this because we only want monthly data for making plots
         if t_end > 84600 * 31 * 3 && config_dict["output_default_diagnostics"]
             include("leaderboard/leaderboard.jl")
@@ -936,7 +913,7 @@ if ClimaComms.iamroot(comms_ctx)
     if config_dict["ci_plots"]
         @info "Generating CI plots"
         include("user_io/ci_plots.jl")
-        make_plots(Val(:general_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)

experiments/ClimaEarth/user_io/ci_plots.jl

@@ -144,30 +144,26 @@ function map_comparison(func, simdirs, args)
 end
 
 """
-    make_plots(
-        ::Union{Val{:general_ci_plots}},
+    make_ci_plots(
         output_paths::Vector{<:AbstractString},
         plot_path::AbstractString;
+        short_names::Vector{<:AbstractString} = ["mse", "lr", "edt", "ts"],
         reduction::String = "average",
     )
 Create plots for the general CI diagnostics. The plots are saved to `plot_path`.
 This is the default plotting function for the CI diagnostics and it can be extended
 to include additional diagnostics.
 The `reduction` keyword argument should be consistent with the reduction used to save the diagnostics.
 """
-function make_plots(
-    ::Union{Val{:general_ci_plots}},
+function make_ci_plots(
     output_paths::Vector{<:AbstractString},
     plot_path::AbstractString;
+    short_names::Vector{<:AbstractString} = ["mse", "lr", "edt", "ts"],
     reduction::String = "average",
 )
     simdirs = CAN.SimDir.(output_paths)
 
-    # Default output diagnostics
-    short_names_3D = ["mse", "lr", "edt"]
-    short_names_2D = ["ts"]
-
-    available_periods = CAN.available_periods(simdirs[1]; short_name = short_names_3D[1], reduction)
+    available_periods = CAN.available_periods(simdirs[1]; short_name = short_names[1], reduction)
     period = ""
     if "10d" in available_periods
         period = "10d"
@@ -177,20 +173,18 @@ function make_plots(
         period = "12h"
     end
 
-    # Creates diagnostics vector
-    # 3D fields are zonally averaged platted onf the lat-z plane
-    # 2D fields are plotted on the lon-lat plane
-    vars_3D = map_comparison(simdirs, short_names_3D) do simdir, short_name
-        get(simdir; short_name, reduction, period) |> CAN.average_lon
-    end
-
-    available_periods = CAN.available_periods(simdirs[1]; short_name = short_names_2D[1], reduction)
-
-    vars_2D = map_comparison(simdirs, short_names_2D) do simdir, short_name
+    # 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
 
+    # Filter vars into 2D and 3D variable diagnostics vectors
+    # 3D fields are zonally averaged platted on the lat-z plane
+    # 2D fields are plotted on the lon-lat plane
+    vars_3D = map(var_3D -> CAN.average_lon(var_3D), filter(var -> CAN.has_altitude(var), vars))
+    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")
-
 end